From 7f17f18dc694e5386fc3b0f9826fb378d2b2fe98 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 19:09:05 +0200 Subject: [PATCH 001/115] Test isolated native Mastra memory replay --- packages/mastra/package.json | 1 + packages/mastra/test/helpers/memory-agent.ts | 332 ++++++++++++++ .../mastra/test/native-memory-replay.test.ts | 410 ++++++++++++++++++ pnpm-lock.yaml | 130 +++++- 4 files changed, 869 insertions(+), 4 deletions(-) create mode 100644 packages/mastra/test/helpers/memory-agent.ts create mode 100644 packages/mastra/test/native-memory-replay.test.ts diff --git a/packages/mastra/package.json b/packages/mastra/package.json index 65eeeca8b..a6cf10ea0 100644 --- a/packages/mastra/package.json +++ b/packages/mastra/package.json @@ -45,6 +45,7 @@ }, "devDependencies": { "@mastra/core": "1.67.0", + "@mastra/memory": "1.30.0", "zod": "3.25.76" }, "scripts": { diff --git a/packages/mastra/test/helpers/memory-agent.ts b/packages/mastra/test/helpers/memory-agent.ts new file mode 100644 index 000000000..e97e3ab1a --- /dev/null +++ b/packages/mastra/test/helpers/memory-agent.ts @@ -0,0 +1,332 @@ +import { Agent } from "@mastra/core/agent"; +import type { InputProcessor } from "@mastra/core/processors"; +import { RequestContext } from "@mastra/core/request-context"; +import { InMemoryStore } from "@mastra/core/storage"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { createTool } from "@mastra/core/tools"; +import { LocalSkillSource, Workspace } from "@mastra/core/workspace"; +import { Memory } from "@mastra/memory"; +import { z } from "zod/v4"; + +export const THREAD = "historical-thread"; +export const RESOURCE = "historical-resource"; +export const FILE_URL = "https://files.invalid/report.pdf"; +export const FILE_BYTES = new Uint8Array([37, 80, 68, 70, 45, 49, 0, 255]); +export type ModelCall = Parameters[0]; +type StreamPart = + Awaited< + ReturnType + >["stream"] extends ReadableStream + ? T + : never; + +export function textStream( + text: string, +): Awaited> { + return streamParts([ + { type: "text-start", id: "text" }, + { type: "text-delta", id: "text", delta: text }, + { type: "text-end", id: "text" }, + ]); +} + +export function streamParts( + parts: StreamPart[], + finishReason: "stop" | "tool-calls" = "stop", +): Awaited> { + return { + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: "stream-start", warnings: [] }); + for (const part of parts) controller.enqueue(part); + controller.enqueue({ + type: "finish", + finishReason, + usage: { inputTokens: 5, outputTokens: 5, totalTokens: 10 }, + }); + controller.close(); + }, + }), + }; +} + +export function memoryModel( + kind: "observer" | "reflector", + wait?: () => Promise, +): { calls: ModelCall[]; model: MastraLanguageModelV2Mock } { + const calls: ModelCall[] = []; + return { + calls, + model: new MastraLanguageModelV2Mock({ + modelId: kind, + provider: "fixture", + doStream: async (options) => { + calls.push(options); + await wait?.(); + const observation = + kind === "observer" + ? `OBSERVED_REPLAY: ${"The user changed the preference to replay-green. ".repeat(15)}` + : "REFLECTED_REPLAY: user prefers replay-green."; + return textStream( + `\n${observation}\n\nContinue.`, + ); + }, + }), + }; +} + +export function createMemoryRuntime( + options: { observerWait?: () => Promise; messageTokens?: number } = {}, +) { + const store = new InMemoryStore(); + const domain = store.stores.memory; + if (!domain) throw new Error("Missing native memory domain"); + const observer = memoryModel("observer", options.observerWait); + const reflector = memoryModel("reflector"); + const memory = new Memory({ + storage: store, + options: { + lastMessages: 20, + semanticRecall: false, + workingMemory: { + enabled: true, + scope: "thread", + schema: z.object({ + preference: z.string(), + notes: z.string().optional(), + }), + }, + observationalMemory: { + scope: "thread", + observation: { + model: observer.model, + messageTokens: options.messageTokens ?? 300, + bufferTokens: 0.2, + bufferActivation: 1, + blockAfter: 1.1, + bufferOnIdle: true, + }, + reflection: { + model: reflector.model, + observationTokens: 100, + bufferActivation: 0.5, + blockAfter: 1.1, + }, + }, + }, + }); + return { store, domain, memory, observer, reflector }; +} +export type MemoryRuntime = ReturnType; + +export async function seedMemory(runtime: MemoryRuntime) { + await runtime.memory.createThread({ + threadId: THREAD, + resourceId: RESOURCE, + title: "Historical conversation", + }); + await runtime.domain.saveResource({ + resource: { + id: RESOURCE, + createdAt: new Date(0), + updatedAt: new Date(0), + metadata: { tenant: "fixture" }, + }, + }); + await runtime.memory.updateWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + workingMemory: JSON.stringify({ preference: "historical-blue" }), + }); + await runtime.memory.saveMessages({ + messages: [ + { + id: "historical-message", + role: "user", + content: { + format: 2, + parts: [{ type: "text", text: "HISTORICAL_MESSAGE: remember blue." }], + }, + createdAt: new Date("2026-01-01T00:00:00Z"), + threadId: THREAD, + resourceId: RESOURCE, + }, + ], + }); + const record = await runtime.domain.initializeObservationalMemory({ + threadId: THREAD, + resourceId: RESOURCE, + scope: "thread", + config: { fixture: true }, + }); + await runtime.domain.updateActiveObservations({ + id: record.id, + observations: "HISTORICAL_OBSERVATION: likes blue.", + tokenCount: 10, + lastObservedAt: new Date("2025-12-31T00:00:00Z"), + observedMessageIds: [], + }); +} + +// The caller must hold exclusive application-level access across settled() and +// these reads. settled() joins this instance's jobs; it is not a storage lock. +export async function snapshotMemory( + runtime: MemoryRuntime, + exclusive: boolean, +) { + if (!exclusive) throw new Error("Exclusive source-thread ownership required"); + await runtime.memory.settled(); + const thread = await runtime.domain.getThreadById({ threadId: THREAD }); + const resource = await runtime.domain.getResourceById({ + resourceId: RESOURCE, + }); + const { messages } = await runtime.domain.listMessages({ + threadId: THREAD, + perPage: false, + }); + const records = await runtime.domain.getObservationalMemoryHistory( + THREAD, + RESOURCE, + ); + if ( + records.some( + (record) => + record.isObserving || + record.isReflecting || + record.isBufferingObservation || + record.isBufferingReflection, + ) + ) { + throw new Error("Unjoined observational-memory work"); + } + return structuredClone({ thread, resource, messages, records }); +} + +export async function restoreMemory( + snapshot: Awaited>, + runtime: MemoryRuntime, +) { + const copy = structuredClone(snapshot); + if (copy.thread) + await runtime.domain.saveThread({ thread: structuredClone(copy.thread) }); + if (copy.resource) + await runtime.domain.saveResource({ resource: copy.resource }); + await runtime.domain.saveMessages({ messages: copy.messages }); + // saveMessages advances thread.updatedAt; restore the historical record last. + if (copy.thread) + await runtime.domain.saveThread({ thread: structuredClone(copy.thread) }); + for (const record of copy.records) + await runtime.domain.insertObservationalMemoryRecord(record); +} + +export function makeMemoryAgent(options: { + runtime: MemoryRuntime; + model: MastraLanguageModelV2Mock; + skillsPath: string; + resolveFile: (url: string) => Promise; + processorRuns: string[]; + configurationCalls?: string[]; + processors?: InputProcessor[]; +}) { + const context = new RequestContext<{ + applicationInstructions: string; + temperature: number; + }>(); + context.set( + "applicationInstructions", + "APPLICATION_INSTRUCTIONS: answer using memory.", + ); + context.set("temperature", 0.25); + const workspace = new Workspace({ + skills: ["skills"], + skillSource: new LocalSkillSource({ basePath: options.skillsPath }), + }); + const processor: InputProcessor = { + id: "bound-file-resolver", + async processInput({ messages }) { + options.processorRuns.push(FILE_URL); + const bytes = await options.resolveFile(FILE_URL); + return messages.map((message) => ({ + ...message, + content: { + ...message.content, + parts: message.content.parts.map((part) => + part.type === "file" && part.data === FILE_URL + ? { + ...part, + data: `data:application/pdf;base64,${Buffer.from(bytes).toString("base64")}`, + } + : part, + ), + }, + })); + }, + }; + const agent = new Agent({ + id: "native-memory-proof", + name: "Native memory proof", + instructions: ({ requestContext }) => { + options.configurationCalls?.push("instructions"); + return String(requestContext.get("applicationInstructions")); + }, + model: () => { + options.configurationCalls?.push("model"); + return options.model; + }, + defaultOptions: ({ requestContext }) => { + options.configurationCalls?.push("defaultOptions"); + return { + maxSteps: 5, + modelSettings: { + temperature: Number(requestContext.get("temperature")), + }, + }; + }, + memory: () => options.runtime.memory, + tools: { + readConversationEvidence: createTool({ + id: "readConversationEvidence", + description: "Read recorded conversation evidence", + inputSchema: z.object({}), + execute: async () => { + await options.runtime.memory.settled(); + return { + evidence: "The replay user now prefers green. ".repeat(400), + }; + }, + }), + }, + workspace, + inputProcessors: [processor, ...(options.processors ?? [])], + }); + return { agent, context, workspace }; +} + +export async function consumeMemoryAgent( + fixture: ReturnType, + text = "Use the attachment.", +) { + const output = await fixture.agent.stream( + [ + { + role: "user", + content: [ + { type: "text", text }, + { + type: "file", + data: new URL(FILE_URL), + mimeType: "application/pdf", + }, + ], + }, + ], + { + memory: { thread: THREAD, resource: RESOURCE }, + requestContext: fixture.context, + system: "EXTRA_SYSTEM_CONTEXT: preserve this contribution.", + }, + ); + let result = ""; + for await (const chunk of output.textStream) result += chunk; + return { output, text: result }; +} diff --git a/packages/mastra/test/native-memory-replay.test.ts b/packages/mastra/test/native-memory-replay.test.ts new file mode 100644 index 000000000..ad1eb22b4 --- /dev/null +++ b/packages/mastra/test/native-memory-replay.test.ts @@ -0,0 +1,410 @@ +import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { APICallError } from "ai"; +import { afterEach, beforeEach, expect, it, vi } from "vitest"; +import { + consumeMemoryAgent, + createMemoryRuntime, + FILE_BYTES, + FILE_URL, + type ModelCall, + makeMemoryAgent, + RESOURCE, + restoreMemory, + seedMemory, + snapshotMemory, + streamParts, + THREAD, + textStream, +} from "./helpers/memory-agent.js"; + +let skillsPath: string; +beforeEach(async () => { + skillsPath = await mkdtemp(join(tmpdir(), "mastra-memory-proof-")); + await mkdir(join(skillsPath, "skills", "triage"), { recursive: true }); + await writeFile( + join(skillsPath, "skills", "triage", "SKILL.md"), + "---\nname: triage\ndescription: HISTORICAL_SKILL classify support issues.\n---\nKeep the original evidence.\n", + ); +}); +afterEach(async () => { + vi.restoreAllMocks(); + await rm(skillsPath, { recursive: true, force: true }); +}); + +it("restores historical native memory into a separate store and reruns the file processor", async () => { + const production = createMemoryRuntime({ messageTokens: 6000 }); + await seedMemory(production); + const historical = await snapshotMemory(production, true); + const baselineCalls: ModelCall[] = []; + const baselineProcessors: string[] = []; + const configurationCalls: string[] = []; + const capturedFiles = new Map(); + let liveFile: Uint8Array | undefined = FILE_BYTES; + const liveResolver = vi.fn(async (url: string) => { + if (!liveFile) throw new Error("Original file is unavailable"); + capturedFiles.set(url, new Uint8Array(liveFile)); + return new Uint8Array(liveFile); + }); + const model = (calls: ModelCall[]) => + new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async (options) => { + calls.push(options); + return textStream("done"); + }, + }); + const baseline = makeMemoryAgent({ + runtime: production, + model: model(baselineCalls), + skillsPath, + resolveFile: liveResolver, + processorRuns: baselineProcessors, + configurationCalls, + }); + expect((await consumeMemoryAgent(baseline)).text).toBe("done"); + await production.memory.settled(); + liveFile = undefined; + await production.memory.updateWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + workingMemory: '{"preference":"production-red"}', + }); + const record = await production.domain.getObservationalMemory( + THREAD, + RESOURCE, + ); + if (!record) throw new Error("Missing observational memory"); + await production.domain.updateActiveObservations({ + id: record.id, + observations: "PRODUCTION_RED", + tokenCount: 10, + lastObservedAt: new Date(), + }); + + const replay = createMemoryRuntime({ messageTokens: 6000 }); + await restoreMemory(historical, replay); + expect(await snapshotMemory(replay, true)).toEqual(historical); + const denied: string[] = []; + for (const name of Object.getOwnPropertyNames( + Object.getPrototypeOf(production.domain), + )) { + if (name === "constructor") continue; + const descriptor = Object.getOwnPropertyDescriptor( + Object.getPrototypeOf(production.domain), + name, + ); + if (typeof descriptor?.value === "function") { + vi.spyOn(production.domain, name as "listMessages").mockImplementation( + () => { + denied.push(name); + throw new Error(`Production access: ${name}`); + }, + ); + } + } + const replayCalls: ModelCall[] = []; + const replayProcessors: string[] = []; + const fixture = makeMemoryAgent({ + runtime: replay, + model: model(replayCalls), + skillsPath, + resolveFile: async (url) => { + const bytes = capturedFiles.get(url); + if (!bytes) throw new Error("Missing captured file"); + return new Uint8Array(bytes); + }, + processorRuns: replayProcessors, + }); + const nativeStream = vi.spyOn(fixture.agent, "stream"); + const result = await consumeMemoryAgent(fixture); + expect(result.text).toBe("done"); + expect(result.output).toBe(await nativeStream.mock.results[0]?.value); + expect(await fixture.agent.getMemory()).toBe(replay.memory); + expect(await fixture.agent.getMemory()).toBe(replay.memory); + await replay.memory.settled(); + expect(denied).toEqual([]); + expect(liveResolver).toHaveBeenCalledExactlyOnceWith(FILE_URL); + expect(baselineProcessors).toHaveLength(1); + expect(replayProcessors).toHaveLength(1); + expect(configurationCalls).toEqual( + expect.arrayContaining(["instructions", "model", "defaultOptions"]), + ); + expect(baselineCalls[0]?.temperature).toBe(0.25); + const prompt = JSON.stringify(replayCalls[0]?.prompt); + for (const marker of [ + "historical-blue", + "HISTORICAL_OBSERVATION", + "HISTORICAL_SKILL", + "EXTRA_SYSTEM_CONTEXT", + "APPLICATION_INSTRUCTIONS", + "HISTORICAL_MESSAGE", + ]) + expect(prompt).toContain(marker); + expect(prompt).not.toContain("production-red"); + expect(prompt).not.toContain("PRODUCTION_RED"); + expect(replayCalls[0]?.tools?.map((tool) => tool.name)).toContain( + "updateWorkingMemory", + ); + expect( + replayCalls[0]?.tools?.some((tool) => tool.name.includes("skill")), + ).toBe(true); + const files = (calls: ModelCall[]) => + calls[0]?.prompt.flatMap((message) => + message.role !== "system" + ? message.content.flatMap((part) => + part.type === "file" + ? [{ data: part.data, mediaType: part.mediaType }] + : [], + ) + : [], + ); + expect(files(replayCalls)).toEqual(files(baselineCalls)); + expect(files(replayCalls)?.[0]).toEqual({ + data: Buffer.from(FILE_BYTES).toString("base64"), + mediaType: "application/pdf", + }); +}, 20000); + +it("runs native working-memory updates, observation and reflection inside an isolated turn", async () => { + const production = createMemoryRuntime({ messageTokens: 600 }); + await seedMemory(production); + const replay = createMemoryRuntime({ messageTokens: 600 }); + await restoreMemory(await snapshotMemory(production, true), replay); + const calls: ModelCall[] = []; + const model = new MastraLanguageModelV2Mock({ + modelId: "evolving-actor", + provider: "fixture", + doStream: async (options) => { + calls.push(options); + return calls.length === 1 + ? streamParts( + [ + { + type: "tool-call", + toolCallId: "update-1", + toolName: "updateWorkingMemory", + input: JSON.stringify({ + memory: { + preference: "replay-green", + notes: "A new preference", + }, + }), + }, + ], + "tool-calls", + ) + : calls.length === 2 + ? streamParts( + [ + { + type: "tool-call", + toolCallId: "evidence-1", + toolName: "readConversationEvidence", + input: "{}", + }, + ], + "tool-calls", + ) + : textStream("evolved"); + }, + }); + const fixture = makeMemoryAgent({ + runtime: replay, + model, + skillsPath, + resolveFile: async () => FILE_BYTES, + processorRuns: [], + }); + expect((await consumeMemoryAgent(fixture)).text).toBe("evolved"); + await replay.memory.settled(); + expect(calls).toHaveLength(3); + expect(replay.observer.calls.length).toBeGreaterThan(0); + expect(replay.reflector.calls.length).toBeGreaterThan(0); + expect( + await replay.memory.getWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + }), + ).toContain("replay-green"); + expect(JSON.stringify(calls[2]?.prompt)).toContain("REFLECTED_REPLAY"); + const secondSystem = calls[1]?.prompt.filter( + (message) => message.role === "system", + ); + // Native working-memory injection runs once. The subsequent tool result + // conveys the update; changing that behavior would cease to be native replay. + expect(JSON.stringify(secondSystem)).toContain("historical-blue"); + expect(JSON.stringify(secondSystem)).not.toContain("replay-green"); + expect( + JSON.stringify( + calls[1]?.prompt.filter((message) => message.role === "tool"), + ), + ).toContain("replay-green"); + expect( + await production.memory.getWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + }), + ).toContain("historical-blue"); + const evolved = await snapshotMemory(replay, true); + expect( + evolved.records.some((record) => record.originType === "reflection"), + ).toBe(true); + const restored = createMemoryRuntime({ messageTokens: 600 }); + await restoreMemory(evolved, restored); + expect(await snapshotMemory(restored, true)).toEqual(evolved); +}, 20000); + +it("captures both retry attempts at public doStream after late native processor changes", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const received: ModelCall[] = []; + const attempts: { id: number; options: ModelCall; failed: boolean }[] = []; + const model = new MastraLanguageModelV2Mock({ + modelId: "retry-actor", + provider: "fixture", + doStream: async (options) => { + received.push(options); + if (received.length === 1) + throw new APICallError({ + message: "Retryable fixture failure", + url: "https://model.invalid", + requestBodyValues: {}, + statusCode: 503, + isRetryable: true, + }); + return textStream("retried"); + }, + }); + const instrumented = new Proxy(model, { + get(target, key) { + if (key === "doStream") + return async (options: ModelCall) => { + const attempt = { id: attempts.length, options, failed: false }; + attempts.push(attempt); + try { + return await target.doStream(options); + } catch (error) { + attempt.failed = true; + throw error; + } + }; + const value = Reflect.get(target, key, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + const fixture = makeMemoryAgent({ + runtime, + model: instrumented, + skillsPath, + resolveFile: async () => FILE_BYTES, + processorRuns: [], + processors: [ + { + id: "late-settings-and-prompt", + processInputStep: () => ({ + modelSettings: { + temperature: 0.75, + maxOutputTokens: 91, + maxRetries: 1, + }, + }), + processLLMRequest: ({ prompt }) => ({ + prompt: [ + ...prompt, + { role: "system", content: "LATE_PROMPT_TRANSFORMATION" }, + ], + }), + }, + ], + }); + expect((await consumeMemoryAgent(fixture)).text).toBe("retried"); + await runtime.memory.settled(); + expect(attempts.map(({ id, failed }) => ({ id, failed }))).toEqual([ + { id: 0, failed: true }, + { id: 1, failed: false }, + ]); + expect(received).toHaveLength(2); + for (const [index, attempt] of attempts.entries()) { + expect(attempt.options).toStrictEqual(received[index]); + expect(attempt.options).toMatchObject({ + temperature: 0.75, + maxOutputTokens: 91, + }); + expect(JSON.stringify(attempt.options.prompt)).toContain( + "LATE_PROMPT_TRANSFORMATION", + ); + } +}, 20000); + +it("joins owned background observation without discarding persisted buffers", async () => { + let releaseObserver!: () => void; + let observerStarted!: () => void; + const blocked = new Promise((resolve) => { + releaseObserver = resolve; + }); + const started = new Promise((resolve) => { + observerStarted = resolve; + }); + const runtime = createMemoryRuntime({ + messageTokens: 10000, + observerWait: async () => { + observerStarted(); + await blocked; + }, + }); + await seedMemory(runtime); + const fixture = makeMemoryAgent({ + runtime, + skillsPath, + model: new MastraLanguageModelV2Mock({ + doStream: async () => textStream("done"), + }), + resolveFile: async () => FILE_BYTES, + processorRuns: [], + }); + try { + expect((await consumeMemoryAgent(fixture)).text).toBe("done"); + await started; + let settled = false; + const snapshot = snapshotMemory(runtime, true).then((value) => { + settled = true; + return value; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(settled).toBe(false); + releaseObserver(); + const historical = await snapshot; + expect( + historical.records[0]?.bufferedObservationChunks?.length, + ).toBeGreaterThan(0); + expect(historical.records[0]?.isBufferingObservation).toBe(false); + const replay = createMemoryRuntime({ messageTokens: 10000 }); + await restoreMemory(historical, replay); + expect(await snapshotMemory(replay, true)).toEqual(historical); + } finally { + releaseObserver(); + await runtime.memory.settled(); + } +}, 20000); + +it("rejects a source without exclusive access or with work this Memory cannot join", async () => { + const runtime = createMemoryRuntime(); + await seedMemory(runtime); + const read = vi.spyOn(runtime.domain, "listMessages"); + await expect(snapshotMemory(runtime, false)).rejects.toThrow( + "Exclusive source-thread ownership required", + ); + expect(read).not.toHaveBeenCalled(); + const record = await runtime.domain.getObservationalMemory(THREAD, RESOURCE); + if (!record) throw new Error("Missing observational memory"); + await runtime.domain.setBufferingObservationFlag(record.id, true); + // No promise on this Memory instance corresponds to the persisted flag. + await runtime.memory.settled(); + await expect(snapshotMemory(runtime, true)).rejects.toThrow( + "Unjoined observational-memory work", + ); +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index cbee10dcc..3037fe930 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -101,6 +101,9 @@ importers: '@mastra/core': specifier: 1.67.0 version: 1.67.0(ai@7.0.107(zod@3.25.76))(express@5.2.1)(zod@3.25.76) + '@mastra/memory': + specifier: 1.30.0 + version: 1.30.0(@mastra/core@1.67.0(ai@7.0.107(zod@3.25.76))(express@5.2.1)(zod@3.25.76)) zod: specifier: 3.25.76 version: 3.25.76 @@ -294,6 +297,12 @@ packages: peerDependencies: zod: ^3.25.0 || ^4.0.0 + '@mastra/memory@1.30.0': + resolution: {integrity: sha512-rzl1tQLorwvHU6Ayi9T0ZiEiE/zKJGgPIpxmnZXTTJ9mn+Bs5FQ7tB4DHPidHNy1SVHRUwg3HRtqMX+nAeREaQ==} + engines: {node: '>=22.13.0'} + peerDependencies: + '@mastra/core': '>=1.4.1-0 <2.0.0-0' + '@mastra/schema-compat@1.3.10': resolution: {integrity: sha512-ZrjvutWy3QJoJROA5PeSFdeFwDgyfGN1sOOObFirex9h9c/R4ReqRllCojNHSN0cKsKxkRlH6zBbWVTiybamUw==} engines: {node: '>=22.13.0'} @@ -553,6 +562,9 @@ packages: resolution: {integrity: sha512-Izi8RQcffqCeNVgFigKli1ssklIbpHnCYc6AknXGYoB6grJqyeby7jv12JUQgmTAnIDnbck1uxksT4dzN3PWBA==} engines: {node: '>=12'} + async-mutex@0.5.0: + resolution: {integrity: sha512-1A94B18jkJ3DYq284ohPxoXbfTA5HsQ7/Mf4DEhcyLx3Bz27Rh59iScbB6EPiP+B+joue6YCxcMXSbFC1tZKwA==} + bail@2.0.2: resolution: {integrity: sha512-0xO6mYd7JB2YesxDKplafRpsiOzPt9V02ddPCLbY1xYGPOX24NTyN50qnUxgCPcSoYMhKpAuBTjQoRZCAkUDRw==} @@ -637,6 +649,22 @@ packages: resolution: {integrity: sha512-uV2QOWP2nWzsy2aMp8aRibhi9dlzF5Hgh5SHaB9OiTGEyDTiJJyx0uy51QXdyWbtAHNua4XJzUKca3OzKUd3vA==} engines: {node: '>= 8'} + debug@2.6.9: + resolution: {integrity: sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==} + peerDependencies: + supports-color: '*' + peerDependenciesMeta: + supports-color: + optional: true + + debug@3.2.7: + resolution: {integrity: sha512-CFjzYYAi4ThfiQvizrFQevTTXHtnCqWfe7x1AhgEscTz6ZbLbfoLRLPugTQyBth6f8ZERVUSyWHFD/7Wu4t1XQ==} + peerDependencies: + supports-color: '*' + peerDependenciesMeta: + supports-color: + optional: true + debug@4.4.3: resolution: {integrity: sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==} engines: {node: '>=6.0'} @@ -664,6 +692,10 @@ packages: devlop@1.1.0: resolution: {integrity: sha512-RWmIqhcFf1lRYBvNmr7qTNuyCt/7/ns2jbpp1+PalgE/rDQcBT0fioSMUpJ93irlUhC5hrg4cYqe6U+0ImW0rA==} + diff@8.0.4: + resolution: {integrity: sha512-DPi0FmjiSU5EvQV0++GFDOJ9ASQUVFh5kD+OzOnYdi7n3Wpm9hWWGfB/O2blfHcMVTL5WkQXSnRiK9makhrcnw==} + engines: {node: '>=0.3.1'} + dotenv@17.4.2: resolution: {integrity: sha512-nI4U3TottKAcAD9LLud4Cb7b2QztQMUEfHbvhTH09bqXTxnSie8WnjPALV/WMCrJZ6UV/qHJ6L03OqO3LcdYZw==} engines: {node: '>=12'} @@ -818,6 +850,10 @@ packages: resolution: {integrity: sha512-eKCa6bwnJhvxj14kZk5NCPc6Hb6BdsU9DZcOnmQKSnO1VKrfV0zCvtttPZUsBvjmNDn8rpcJfpwSYnHBjc95MQ==} engines: {node: '>=18.18.0'} + iconv-lite@0.4.24: + resolution: {integrity: sha512-v3MXnZAcvnywkTUEZomIActle7RXXeedOR31wwl7VlyoXO4Qi9arvSenNQWne1TcRwhCL1HwLI21bEqdpj8/rA==} + engines: {node: '>=0.10.0'} + iconv-lite@0.7.3: resolution: {integrity: sha512-IKXpvIzjnC9XTAUbVBcMfGS0EPaIXtW6v+zr+RRp+hqULEpo0owZax6wyRwPOJbWbzjYspQwusTsfVr0ifh4uQ==} engines: {node: '>=0.10.0'} @@ -972,6 +1008,9 @@ packages: resolution: {integrity: sha512-WkUDrojuJs0xkgGf2udWxa3yGBRxPtxUkB79i6aCZLRgc7PM8fZe9TosfPDcvEpQZbuFASnHYmRLBLUbmLOIIA==} engines: {node: '>= 12.0.0'} + lodash.merge@4.6.2: + resolution: {integrity: sha512-0KpjqXRVvrYyCsX1swR/XTK0va6VQkQM6MNo7PqW77ByjAhoARA8EfrP1N4+KlKj8YS0ZUCtRT/YUuhyYDujIQ==} + longest-streak@3.1.0: resolution: {integrity: sha512-9Ri+o0JYgehTaVBBDoMqIl8GXtbWg711O3srftcHhZ0dqnETqLaoIK0x17fUw9rFSlK/0NlsKe0Ahhyl5pXE2g==} @@ -1126,6 +1165,9 @@ packages: resolution: {integrity: sha512-7o1wEA2RyMP7Iu7GNba9vc0RWWGACJOCZBJX2GJWip0ikV+wcOsgVuY9uE8CPiyQhkGFSlhuSkZPavN7u1c2Fw==} engines: {node: '>=10'} + ms@2.0.0: + resolution: {integrity: sha512-Tpp60P6IUJDTuOq/5Z8cdskzJujfwqfOTkrwIwj7IRISpnkJnT6SyJ4PCPnGMoFjC9ddhal5KVIYtAt97ix05A==} + ms@2.1.3: resolution: {integrity: sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==} @@ -1134,6 +1176,11 @@ packages: engines: {node: ^10 || ^12 || ^13.7 || ^14 || >=15.0.1} hasBin: true + needle@2.9.1: + resolution: {integrity: sha512-6R9fqJ5Zcmf+uYaFgdIHmLwNldn5HbK8L5ybn7Uz+ylX/rnOsSp1AHcvQSrCaFN+qNM1wpymHqD7mVasEOlHGQ==} + engines: {node: '>= 4.4.x'} + hasBin: true + negotiator@1.1.0: resolution: {integrity: sha512-NMPBRMJgiQHjbd8phG3Vebdx4kZ1H121rbl5IkMqeOsahptB9BKo/d7oJ3zTXqTgagn2bWlNSXkh0QUGM31RYg==} engines: {node: '>=18'} @@ -1225,6 +1272,9 @@ packages: resolution: {integrity: sha512-gjVS5hOP+M3wMm5nmNOucbIrqudzs9v/57bWRHQWLYklXqoXKrVfYW2W9+glfGsqtPgpiz5WwyEEB+ksXIx3gQ==} engines: {node: '>=18'} + probe-image-size@7.4.0: + resolution: {integrity: sha512-cdEprVtZxV+awMde9X+4jILBFYh4CARxVrQaMl4wY4YcPWbul9jntXrIW95NInBDyJwcVUP3U0T6yukN8rMBaQ==} + proxy-addr@2.0.8: resolution: {integrity: sha512-5nnx0yGyVUcY6t9RnWcARWtwT9F1D8O9rt08htPvnd49W1IgZtmLkhu9WfMzQj1cFxjHIO6connUNVW5k7AVyQ==} engines: {node: '>= 0.10'} @@ -1273,6 +1323,10 @@ packages: safer-buffer@2.1.2: resolution: {integrity: sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==} + sax@1.6.1: + resolution: {integrity: sha512-42tBVwLWnaQvW5zc4HbZrTuWccECCZfBi92FDuwtqxasH+JbPB3/FOKb1m222K42R4WxuxzzMsTswfzgtSu64Q==} + engines: {node: '>=11.0.0'} + section-matter@1.0.0: resolution: {integrity: sha512-vfD3pmTzGpufjScBh50YHKzEu2lxBWhVEHsNGoEXmCmn2hKGfeNLYMzCJpe8cD7gqX7TJluOVpBkAequ6dgMmA==} engines: {node: '>=4'} @@ -1336,6 +1390,9 @@ packages: std-env@4.2.0: resolution: {integrity: sha512-oCUKSupKTHX53EyjDtuZQ64pjLJ6yYCtpmEw0goYxtjG9KpbRe8KAsl2tBUGU9DyMcJ0RwJ8GqJAFzMXcXW1Rw==} + stream-parser@0.3.1: + resolution: {integrity: sha512-bJ/HgKq41nlKvlhccD5kaCr/P+Hu0wPNKPJOH7en+YrJu/9EgqUF+88w5Jb6KNcjOFMhfX4B2asfeAtIGuHObQ==} + strip-bom-string@1.0.0: resolution: {integrity: sha512-uCC2VHvQRYu+lMh4My/sFNmF2klFymLX1wHJeXnbEJERpV/ZsVuonzerjfrGpIGF7LBVa1O7i9kjiWvJiFck8g==} engines: {node: '>=0.10.0'} @@ -1834,12 +1891,33 @@ snapshots: - supports-color - utf-8-validate + '@mastra/memory@1.30.0(@mastra/core@1.67.0(ai@7.0.107(zod@3.25.76))(express@5.2.1)(zod@3.25.76))': + dependencies: + '@mastra/core': 1.67.0(ai@7.0.107(zod@3.25.76))(express@5.2.1)(zod@3.25.76) + '@mastra/schema-compat': 1.3.10(zod@4.4.3) + async-mutex: 0.5.0 + diff: 8.0.4 + json-schema: 0.4.0 + lru-cache: 11.5.2 + probe-image-size: 7.4.0 + tokenx: 1.3.0 + xxhash-wasm: 1.1.0 + zod: 4.4.3 + transitivePeerDependencies: + - supports-color + '@mastra/schema-compat@1.3.10(zod@3.25.76)': dependencies: json-schema-to-zod: 2.8.1 zod: 3.25.76 zod-from-json-schema: 0.5.6 + '@mastra/schema-compat@1.3.10(zod@4.4.3)': + dependencies: + json-schema-to-zod: 2.8.1 + zod: 4.4.3 + zod-from-json-schema: 0.5.6 + '@mastra/schema-compat@1.3.8(zod@4.4.3)': dependencies: json-schema-to-zod: 2.8.1 @@ -2076,6 +2154,10 @@ snapshots: assertion-error@2.0.1: {} + async-mutex@0.5.0: + dependencies: + tslib: 2.8.1 + bail@2.0.2: {} balanced-match@1.0.2: {} @@ -2179,6 +2261,14 @@ snapshots: shebang-command: 2.0.0 which: 2.0.2 + debug@2.6.9: + dependencies: + ms: 2.0.0 + + debug@3.2.7: + dependencies: + ms: 2.1.3 + debug@4.4.3(supports-color@10.2.2): dependencies: ms: 2.1.3 @@ -2200,6 +2290,8 @@ snapshots: dependencies: dequal: 2.0.3 + diff@8.0.4: {} + dotenv@17.4.2: {} dunder-proto@1.0.1: @@ -2402,6 +2494,10 @@ snapshots: human-signals@8.0.1: {} + iconv-lite@0.4.24: + dependencies: + safer-buffer: 2.1.2 + iconv-lite@0.7.3: dependencies: safer-buffer: 2.1.2 @@ -2506,6 +2602,8 @@ snapshots: lightningcss-win32-arm64-msvc: 1.33.0 lightningcss-win32-x64-msvc: 1.33.0 + lodash.merge@4.6.2: {} + longest-streak@3.1.0: {} lru-cache@11.5.2: {} @@ -2830,10 +2928,20 @@ snapshots: dependencies: brace-expansion: 2.1.4 + ms@2.0.0: {} + ms@2.1.3: {} nanoid@3.3.19: {} + needle@2.9.1: + dependencies: + debug: 3.2.7 + iconv-lite: 0.4.24 + sax: 1.6.1 + transitivePeerDependencies: + - supports-color + negotiator@1.1.0: dependencies: content-type: 2.1.0 @@ -2915,6 +3023,14 @@ snapshots: dependencies: parse-ms: 4.0.0 + probe-image-size@7.4.0: + dependencies: + lodash.merge: 4.6.2 + needle: 2.9.1 + stream-parser: 0.3.1 + transitivePeerDependencies: + - supports-color + proxy-addr@2.0.8: dependencies: forwarded: 0.2.0 @@ -3002,8 +3118,9 @@ snapshots: - supports-color optional: true - safer-buffer@2.1.2: - optional: true + safer-buffer@2.1.2: {} + + sax@1.6.1: {} section-matter@1.0.0: dependencies: @@ -3093,6 +3210,12 @@ snapshots: std-env@4.2.0: {} + stream-parser@0.3.1: + dependencies: + debug: 2.6.9 + transitivePeerDependencies: + - supports-color + strip-bom-string@1.0.0: {} strip-final-newline@4.0.0: {} @@ -3117,8 +3240,7 @@ snapshots: trough@2.2.0: {} - tslib@2.8.1: - optional: true + tslib@2.8.1: {} type-fest@4.41.0: {} From b04612b121284cb15e89b2050a7ed8b852d9d7e5 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 19:20:33 +0200 Subject: [PATCH 002/115] Add versioned Mastra memory snapshots --- packages/mastra/src/index.ts | 27 + packages/mastra/src/memory-binding.ts | 321 ++++++++++ packages/mastra/src/memory-snapshot.ts | 578 +++++++++++++++++++ packages/mastra/test/memory-binding.test.ts | 217 +++++++ packages/mastra/test/memory-snapshot.test.ts | 223 +++++++ 5 files changed, 1366 insertions(+) create mode 100644 packages/mastra/src/memory-binding.ts create mode 100644 packages/mastra/src/memory-snapshot.ts create mode 100644 packages/mastra/test/memory-binding.test.ts create mode 100644 packages/mastra/test/memory-snapshot.test.ts diff --git a/packages/mastra/src/index.ts b/packages/mastra/src/index.ts index 07ffc713c..a809a245b 100644 --- a/packages/mastra/src/index.ts +++ b/packages/mastra/src/index.ts @@ -1,4 +1,31 @@ export { KitaruAgent } from "./agent.js"; +export type { + MastraExclusiveMemoryAccess, + MastraMemoryCaptureBinding, + MastraMemoryCaptureOptions, + MastraMemoryMutation, + MastraMemorySelector, +} from "./memory-binding.js"; +export { + createMemoryCaptureBinding, + createProcessLocalMemoryAccess, +} from "./memory-binding.js"; +export type { + MastraFileManifestEntry, + MastraMemoryReplayEnvelope, + MastraMemoryReplayInput, + MastraMemorySnapshot, + MastraRecordedFile, +} from "./memory-snapshot.js"; +export { + createMemoryReplayEnvelope, + decodeMemoryReplayEnvelope, + decodeMemoryValue, + encodeMemoryValue, + MEMORY_REPLAY_KEY, + restoreMemoryReplayEnvelope, + validateMemorySnapshot, +} from "./memory-snapshot.js"; export type { MastraEvaluatorOptions, RunnableMastraScorer, diff --git a/packages/mastra/src/memory-binding.ts b/packages/mastra/src/memory-binding.ts new file mode 100644 index 000000000..ba8d38cd0 --- /dev/null +++ b/packages/mastra/src/memory-binding.ts @@ -0,0 +1,321 @@ +import type { MemoryStorage } from "@mastra/core/storage"; +import { type JsonValue, toRecorderJson } from "@zenml-io/kitaru"; +import { recordedToolPayloadConversion } from "@zenml-io/kitaru/adapter"; +import { + decodeMemoryValue, + encodeMemoryValue, + type MastraMemorySnapshot, + validateMemorySnapshot, +} from "./memory-snapshot.js"; + +export interface MastraMemorySelector { + threadId: string; + resourceId: string; +} + +/** All writers must participate in this application's lease, including other processes. */ +export interface MastraExclusiveMemoryAccess { + acquire( + selector: MastraMemorySelector, + onConflict?: () => void, + ): Promise<() => Promise>; +} + +/** + * A real process-local lease. Use only when every writer runs in this process + * and shares this instance; distributed writers require a distributed lease. + */ +export function createProcessLocalMemoryAccess(): MastraExclusiveMemoryAccess { + const activeThreads = new Map void) | undefined>(); + return { + async acquire({ threadId }, onConflict) { + if (activeThreads.has(threadId)) { + activeThreads.get(threadId)?.(); + throw new Error("Exclusive source-thread ownership is unavailable."); + } + activeThreads.set(threadId, onConflict); + let released = false; + return async () => { + if (!released) activeThreads.delete(threadId); + released = true; + }; + }, + }; +} + +export interface MastraMemoryMutation { + id: string; + invocationId: string; + revision: number; + complete: boolean; + method: string; + arguments: JsonValue; + result: JsonValue; + requestId?: string; +} + +export interface MastraMemoryCaptureOptions extends MastraMemorySelector { + invocationId: string; + domain: MemoryStorage; + exclusiveAccess: MastraExclusiveMemoryAccess; + recordMutation: (event: MastraMemoryMutation) => Promise; + getRequestId?: () => string | undefined; + onIncomplete?: (reason: string) => void; +} + +export interface MastraMemoryCaptureBinding { + /** Supply this public domain to the invocation's native Memory storage. */ + domain: MemoryStorage; + readonly revision: number; + readonly incompleteReasons: readonly string[]; + captureInitial(memory: { + settled(): Promise; + }): Promise; + markIncomplete(reason: string): void; + drain(): Promise; + release(): Promise; +} + +// The pinned public MemoryStorage mutation inventory. Delegation binds `this` to +// the original domain, so a native method's own helper calls record only once. +const MUTATIONS = new Set([ + "saveThread", + "updateThread", + "patchThread", + "deleteThread", + "saveMessages", + "updateMessages", + "deleteMessages", + "copyThread", + "cloneThread", + "updateThreadResourceId", + "saveResource", + "updateResource", + "initializeObservationalMemory", + "updateActiveObservations", + "updateBufferedObservations", + "swapBufferedToActive", + "createReflectionGeneration", + "updateBufferedReflection", + "swapBufferedReflectionToActive", + "setReflectingFlag", + "setObservingFlag", + "setBufferingObservationFlag", + "setBufferingReflectionFlag", + "insertObservationalMemoryRecord", + "clearObservationalMemory", + "setPendingMessageTokens", + "updateObservationalMemoryConfig", +]); + +/** Capture one native invocation without changing the shared source domain or Agent. */ +export function createMemoryCaptureBinding( + options: MastraMemoryCaptureOptions, +): MastraMemoryCaptureBinding { + let revision = 0; + let started = false; + let capturing = false; + let released = false; + let releaseLease: (() => Promise) | undefined; + let mutations = Promise.resolve(); + let evidence = Promise.resolve(); + const reasons: string[] = []; + const methods = new Map(); + + function markIncomplete(reason: string): void { + if (reasons.includes(reason)) return; + reasons.push(reason); + try { + options.onIncomplete?.(reason); + } catch { + /* Diagnostics must not affect native calls. */ + } + } + + const domain = new Proxy(options.domain, { + get(target, property) { + const value: unknown = Reflect.get(target, property, target); + if (typeof value !== "function") return value; + if (methods.has(property)) return methods.get(property); + if (!MUTATIONS.has(property as keyof MemoryStorage)) { + const bound = value.bind(target); + methods.set(property, bound); + return bound; + } + const bound = (...args: unknown[]): Promise => { + let encodedArguments: JsonValue = null; + let complete = true; + let requestId: string | undefined; + try { + encodedArguments = encodeMemoryValue(args); + requestId = options.getRequestId?.(); + } catch { + complete = false; + markIncomplete( + "Memory mutation arguments or request attribution could not be recorded safely.", + ); + } + if (!started || released) + markIncomplete( + "Memory mutation occurred outside the owned invocation lifecycle.", + ); + const duringCapture = capturing; + const result = mutations.then(async () => { + let output: unknown; + try { + output = await Reflect.apply(value, target, args); + } catch (error) { + markIncomplete("Native memory storage mutation failed."); + throw error; + } + // Joined work from a previous turn belongs to the initial snapshot. + if (duringCapture) return output; + revision += 1; + let encodedResult: JsonValue = null; + try { + encodedResult = encodeMemoryValue(output); + } catch { + complete = false; + markIncomplete( + "Memory mutation result could not be recorded safely.", + ); + } + let event: MastraMemoryMutation = { + id: `${options.invocationId}:memory:${revision}`, + invocationId: options.invocationId, + revision, + complete, + method: String(property), + arguments: encodedArguments, + result: encodedResult, + ...(requestId === undefined ? {} : { requestId }), + }; + try { + toRecorderJson(event); + if ( + recordedToolPayloadConversion(event, "Mastra memory mutation") + .lossy + ) + throw new Error("Lossy event"); + } catch { + markIncomplete( + "Memory mutation evidence exceeds replay payload bounds or contains credentials.", + ); + event = { + ...event, + complete: false, + arguments: null, + result: null, + }; + } + evidence = evidence.then(async () => { + try { + await options.recordMutation(event); + } catch { + markIncomplete("Memory mutation evidence persistence failed."); + } + }); + return output; + }); + mutations = result.then( + () => {}, + () => {}, + ); + return result; + }; + methods.set(property, bound); + return bound; + }, + }); + + async function drain(): Promise { + // Evidence can grow while a storage operation settles; follow both tails. + while (true) { + const currentMutations = mutations; + await currentMutations; + const currentEvidence = evidence; + await currentEvidence; + if (currentMutations === mutations && currentEvidence === evidence) + return; + } + } + + return { + domain, + get revision() { + return revision; + }, + get incompleteReasons() { + return [...reasons]; + }, + markIncomplete, + async captureInitial(memory) { + if (started || released) { + markIncomplete( + "Initial memory capture may run only once per invocation.", + ); + return undefined; + } + started = true; + capturing = true; + try { + try { + releaseLease = await options.exclusiveAccess.acquire(options, () => + markIncomplete( + "Exclusive source-thread ownership was invalidated by an overlapping invocation.", + ), + ); + } catch { + markIncomplete("Exclusive source-thread ownership is unavailable."); + return undefined; + } + await memory.settled(); + await mutations; + const thread = await options.domain.getThreadById({ + threadId: options.threadId, + }); + const resource = await options.domain.getResourceById({ + resourceId: options.resourceId, + }); + const { messages } = await options.domain.listMessages({ + threadId: options.threadId, + perPage: false, + }); + const records = await options.domain.getObservationalMemoryHistory( + options.threadId, + options.resourceId, + ); + const snapshot = { + threadId: options.threadId, + resourceId: options.resourceId, + thread, + resource, + messages, + records, + }; + // Copy through the explicit codec: no storage-owned objects or Dates escape. + const copy = decodeMemoryValue(encodeMemoryValue(snapshot)); + validateMemorySnapshot(copy); + return reasons.length === 0 ? copy : undefined; + } catch { + markIncomplete( + "Initial memory capture failed: unsupported, altered, or Unjoined observational-memory state.", + ); + return undefined; + } finally { + capturing = false; + } + }, + drain, + async release() { + if (released) return; + await drain(); + released = true; + try { + await releaseLease?.(); + } catch { + markIncomplete("Exclusive source-thread lease release failed."); + } + }, + }; +} diff --git a/packages/mastra/src/memory-snapshot.ts b/packages/mastra/src/memory-snapshot.ts new file mode 100644 index 000000000..b7f6d5b70 --- /dev/null +++ b/packages/mastra/src/memory-snapshot.ts @@ -0,0 +1,578 @@ +import { createHash } from "node:crypto"; +import type { MastraDBMessage, StorageThreadType } from "@mastra/core/memory"; +import type { + ObservationalMemoryRecord, + StorageResourceType, +} from "@mastra/core/storage"; +import { type JsonValue, toRecorderJson } from "@zenml-io/kitaru"; +import { + MAX_RECORDED_PAYLOAD_CHARS, + recordedToolPayloadConversion, +} from "@zenml-io/kitaru/adapter"; + +export const MEMORY_REPLAY_KEY = "mastra_memory_replay"; +const CODEC_KEY = "$mastra"; + +export interface MastraMemorySnapshot { + threadId: string; + resourceId: string; + thread: StorageThreadType | null; + resource: StorageResourceType | null; + messages: MastraDBMessage[]; + records: ObservationalMemoryRecord[]; +} + +export interface MastraRecordedFile { + url: string; + mediaType: string; + bytes: Uint8Array; +} + +export interface MastraFileManifestEntry { + [key: string]: JsonValue; + url: string; + mediaType: string; + base64: string; + length: number; + sha256: string; +} + +export interface MastraMemoryReplayInput { + invocationId: string; + rawInput: unknown; + initialSnapshot: MastraMemorySnapshot; + /** Materialized options only. Models and schemas need explicit JSON representations. */ + configuration: Record; + requestContext: Record; + files: MastraRecordedFile[]; +} + +export interface MastraMemoryReplayEnvelope { + [key: string]: JsonValue; + version: 2; + complete: boolean; + reasons: string[]; + invocationId: string; + rawInput: JsonValue; + initialSnapshot: JsonValue; + configuration: JsonValue; + requestContext: JsonValue; + files: MastraFileManifestEntry[]; +} + +class MemoryReplayError extends Error {} + +function unsupported(reason: string): Error { + return new MemoryReplayError(`Unsupported Mastra memory replay: ${reason}`); +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function requireValue(condition: unknown, reason: string): asserts condition { + if (!condition) throw unsupported(reason); +} + +function checkBudget(value: JsonValue): void { + try { + toRecorderJson(value); + } catch { + throw unsupported( + "Memory value exceeds the replay JSON depth/item limits.", + ); + } +} + +function hash(bytes: Uint8Array): string { + return createHash("sha256").update(bytes).digest("hex"); +} + +function binary(bytes: Uint8Array): { + base64: string; + length: number; + sha256: string; +} { + requireValue( + bytes.byteLength <= MAX_RECORDED_PAYLOAD_CHARS, + "Binary content exceeds the replay payload limit.", + ); + return { + base64: Buffer.from(bytes).toString("base64"), + length: bytes.byteLength, + sha256: hash(bytes), + }; +} + +function readBinary(value: Record): Uint8Array { + requireValue( + typeof value.base64 === "string" && + typeof value.length === "number" && + Number.isSafeInteger(value.length) && + value.length >= 0 && + typeof value.sha256 === "string", + "Malformed binary content.", + ); + const bytes = new Uint8Array(Buffer.from(value.base64, "base64")); + requireValue( + Buffer.from(bytes).toString("base64") === value.base64 && + bytes.length === value.length && + hash(bytes) === value.sha256, + "Corrupt binary content hash, length, or encoding.", + ); + return bytes; +} + +function validateUrl(value: string): URL { + const url = new URL(value); + requireValue( + !url.username && !url.password, + "URL credentials are not replayable.", + ); + const query = recordedToolPayloadConversion( + Object.fromEntries(url.searchParams), + "Mastra file URL", + ); + requireValue(!query.lossy, "URL query credentials are not replayable."); + return url; +} + +/** Encode the few non-JSON values in native memory without losing their types. */ +export function encodeMemoryValue(value: unknown): JsonValue { + let items = 0; + const active = new Set(); + function visit(current: unknown, depth: number): JsonValue { + requireValue( + ++items <= 10_000 && depth < 64, + "Memory value exceeds the replay depth/item limits.", + ); + if (current === undefined) return { [CODEC_KEY]: "undefined" }; + if (current === null || typeof current === "boolean") return current; + if (typeof current === "string") { + requireValue( + current.length <= MAX_RECORDED_PAYLOAD_CHARS, + "Memory value exceeds the replay payload limit.", + ); + return current; + } + if (typeof current === "number") { + requireValue(Number.isFinite(current), "Non-finite memory number."); + return current; + } + requireValue( + typeof current === "object", + "Unsupported memory value; functions and live dependencies need explicit representations.", + ); + if (current instanceof Date) { + requireValue(Number.isFinite(current.getTime()), "Invalid memory Date."); + return { [CODEC_KEY]: "date", value: current.toISOString() }; + } + if (current instanceof URL) { + validateUrl(current.href); + return { [CODEC_KEY]: "url", value: current.href }; + } + if (current instanceof Uint8Array) + return { [CODEC_KEY]: "bytes", ...binary(current) }; + requireValue(!active.has(current), "Circular memory value."); + active.add(current); + try { + if (Array.isArray(current)) + return current.map((item) => visit(item, depth + 1)); + requireValue( + Object.getPrototypeOf(current) === Object.prototype || + Object.getPrototypeOf(current) === null, + "Unsupported memory object; use explicit JSON configuration.", + ); + requireValue( + !Object.hasOwn(current, CODEC_KEY) && + Reflect.ownKeys(current).every((key) => typeof key === "string"), + "Reserved or symbolic memory key.", + ); + const result: Record = Object.create(null); + for (const [key, descriptor] of Object.entries( + Object.getOwnPropertyDescriptors(current), + )) { + requireValue( + descriptor.enumerable && "value" in descriptor, + "Accessor or hidden memory properties are unsupported.", + ); + result[key] = visit(descriptor.value, depth + 1); + } + return result; + } finally { + active.delete(current); + } + } + const encoded = visit(value, 0); + const converted = recordedToolPayloadConversion( + encoded, + "Mastra memory replay", + ); + requireValue( + !converted.lossy, + "Memory content was altered by credential protection or replay payload bounds.", + ); + checkBudget(converted.value); + return converted.value; +} + +/** Decode an already bounded value, rejecting ambiguous or damaged codec records. */ +export function decodeMemoryValue(value: JsonValue): unknown { + const converted = recordedToolPayloadConversion( + value, + "Mastra memory replay", + ); + requireValue( + !converted.lossy, + "Memory content was altered by credential protection or replay payload bounds.", + ); + checkBudget(converted.value); + function visit(current: JsonValue): unknown { + if (Array.isArray(current)) return current.map(visit); + if (!isRecord(current)) return current; + if (Object.hasOwn(current, CODEC_KEY)) { + const kind = current[CODEC_KEY]; + if (kind === "undefined" && Object.keys(current).length === 1) + return undefined; + if ( + (kind === "date" || kind === "url") && + Object.keys(current).length === 2 && + typeof current.value === "string" + ) { + if (kind === "url") return validateUrl(current.value); + const date = new Date(current.value); + requireValue( + Number.isFinite(date.getTime()) && + date.toISOString() === current.value, + "Malformed memory Date.", + ); + return date; + } + if (kind === "bytes" && Object.keys(current).length === 4) + return readBinary(current); + throw unsupported("Malformed memory codec tag."); + } + return Object.fromEntries( + Object.entries(current).map(([key, item]) => [ + key, + visit(item as JsonValue), + ]), + ); + } + return visit(converted.value); +} + +/** Validate the complete native state before an isolated store receives any writes. */ +export function validateMemorySnapshot( + value: unknown, +): asserts value is MastraMemorySnapshot { + requireValue( + isRecord(value) && + typeof value.threadId === "string" && + value.threadId.length > 0 && + typeof value.resourceId === "string" && + value.resourceId.length > 0 && + Array.isArray(value.messages) && + Array.isArray(value.records), + "Malformed initial memory snapshot.", + ); + const dates = (record: Record) => + record.createdAt instanceof Date && record.updatedAt instanceof Date; + requireValue( + value.thread === null || + (isRecord(value.thread) && + value.thread.id === value.threadId && + value.thread.resourceId === value.resourceId && + dates(value.thread)), + "Malformed or mismatched thread record.", + ); + requireValue( + value.resource === null || + (isRecord(value.resource) && + value.resource.id === value.resourceId && + dates(value.resource)), + "Malformed or mismatched resource record.", + ); + requireValue( + value.thread !== null || + (value.messages.length === 0 && value.records.length === 0), + "Orphaned memory state.", + ); + const messageIds = new Set(); + for (const message of value.messages) { + requireValue( + isRecord(message) && + typeof message.id === "string" && + !messageIds.has(message.id) && + message.threadId === value.threadId && + (message.resourceId === undefined || + message.resourceId === value.resourceId) && + message.createdAt instanceof Date && + ["system", "user", "assistant", "tool"].includes( + String(message.role), + ) && + isRecord(message.content) && + message.content.format === 2 && + Array.isArray(message.content.parts), + "Malformed or mismatched stored message.", + ); + messageIds.add(message.id); + } + const recordIds = new Set(); + for (const record of value.records) { + requireValue( + isRecord(record) && + typeof record.id === "string" && + !recordIds.has(record.id) && + record.scope === "thread" && + record.threadId === value.threadId && + record.resourceId === value.resourceId && + dates(record) && + typeof record.activeObservations === "string" && + isRecord(record.config), + "Malformed or unsupported observational-memory record.", + ); + recordIds.add(record.id); + for (const key of [ + "generationCount", + "totalTokensObserved", + "observationTokenCount", + "pendingMessageTokens", + "lastBufferedAtTokens", + ]) + requireValue( + typeof record[key] === "number" && + Number.isFinite(record[key]) && + record[key] >= 0, + "Malformed observational-memory counter.", + ); + for (const key of [ + "isObserving", + "isReflecting", + "isBufferingObservation", + "isBufferingReflection", + ]) + requireValue( + record[key] === false, + "Unjoined observational-memory work or missing work flag.", + ); + requireValue( + record.lastBufferedAtTime === null || + record.lastBufferedAtTime instanceof Date, + "Malformed observational-memory buffer cursor.", + ); + requireValue( + record.lastObservedAt === undefined || + record.lastObservedAt instanceof Date, + "Malformed observational-memory observation cursor.", + ); + requireValue( + record.originType === "initial" || record.originType === "reflection", + "Malformed observational-memory generation origin.", + ); + for (const key of [ + "bufferedObservations", + "bufferedReflection", + "observedTimezone", + ]) + requireValue( + record[key] === undefined || typeof record[key] === "string", + "Malformed observational-memory text.", + ); + for (const key of [ + "bufferedObservationTokens", + "bufferedReflectionTokens", + "bufferedReflectionInputTokens", + "reflectedObservationLineCount", + ]) + requireValue( + record[key] === undefined || + (typeof record[key] === "number" && + Number.isFinite(record[key]) && + record[key] >= 0), + "Malformed observational-memory buffer counter.", + ); + for (const key of ["observedMessageIds", "bufferedMessageIds"]) + requireValue( + record[key] === undefined || + (Array.isArray(record[key]) && + record[key].every((id) => typeof id === "string")), + "Malformed observational-memory message identities.", + ); + const chunks = record.bufferedObservationChunks; + requireValue( + chunks === undefined || Array.isArray(chunks), + "Malformed observation buffer.", + ); + if (Array.isArray(chunks)) + for (const chunk of chunks) { + requireValue( + isRecord(chunk) && + typeof chunk.id === "string" && + typeof chunk.cycleId === "string" && + typeof chunk.observations === "string" && + chunk.createdAt instanceof Date && + chunk.lastObservedAt instanceof Date && + typeof chunk.tokenCount === "number" && + Number.isFinite(chunk.tokenCount) && + chunk.tokenCount >= 0 && + typeof chunk.messageTokens === "number" && + Number.isFinite(chunk.messageTokens) && + chunk.messageTokens >= 0 && + Array.isArray(chunk.messageIds) && + chunk.messageIds.every((id) => typeof id === "string"), + "Malformed observation buffer chunk.", + ); + } + } +} + +function validateConfiguration( + configuration: unknown, +): asserts configuration is Record { + requireValue(isRecord(configuration), "Malformed resolved configuration."); + const memory = configuration.memoryConfig ?? configuration.memory; + if (!isRecord(memory)) return; + requireValue( + memory.semanticRecall === undefined || memory.semanticRecall === false, + "Semantic recall is outside isolated memory replay scope.", + ); + for (const key of ["workingMemory", "observationalMemory"]) { + const feature = memory[key]; + if (isRecord(feature) && feature.enabled !== false) + requireValue( + feature.scope === "thread", + "Only explicitly thread-scoped memory is replayable.", + ); + } +} + +/** Build safe diagnostic evidence even when complete replay prerequisites are unavailable. */ +export function createMemoryReplayEnvelope( + input: MastraMemoryReplayInput, +): MastraMemoryReplayEnvelope { + const incomplete = (reason: string): MastraMemoryReplayEnvelope => ({ + version: 2, + complete: false, + reasons: [reason], + invocationId: "", + rawInput: null, + initialSnapshot: null, + configuration: null, + requestContext: null, + files: [], + }); + try { + validateMemorySnapshot(input.initialSnapshot); + validateConfiguration(input.configuration); + const envelope: MastraMemoryReplayEnvelope = { + version: 2, + complete: true, + reasons: [], + invocationId: input.invocationId, + rawInput: encodeMemoryValue(input.rawInput), + initialSnapshot: encodeMemoryValue(input.initialSnapshot), + configuration: encodeMemoryValue(input.configuration), + requestContext: encodeMemoryValue(input.requestContext), + files: input.files.map((file) => ({ + url: file.url, + mediaType: file.mediaType, + ...binary(file.bytes), + })), + }; + // The combined envelope, including encoded bytes and metadata, shares one budget. + const converted = recordedToolPayloadConversion( + envelope, + "Mastra memory replay envelope", + ); + requireValue( + !converted.lossy, + "Envelope exceeds replay limits or contains credentials.", + ); + decodeMemoryReplayEnvelope(converted.value); + return converted.value as unknown as MastraMemoryReplayEnvelope; + } catch (error) { + return incomplete( + error instanceof MemoryReplayError + ? error.message + : "Memory replay prerequisites could not be captured safely.", + ); + } +} + +export function decodeMemoryReplayEnvelope( + input: unknown, +): MastraMemoryReplayInput { + const converted = recordedToolPayloadConversion( + input, + "Mastra memory replay envelope", + ); + requireValue( + !converted.lossy, + "Envelope exceeds replay limits or contains credentials.", + ); + checkBudget(converted.value); + const value = converted.value; + requireValue( + isRecord(value) && + value.version === 2 && + value.complete === true && + Array.isArray(value.reasons) && + value.reasons.length === 0 && + typeof value.invocationId === "string" && + value.invocationId.length > 0 && + Array.isArray(value.files), + "Missing, incomplete, or unknown version of memory replay envelope.", + ); + for (const key of [ + "rawInput", + "initialSnapshot", + "configuration", + "requestContext", + ]) + requireValue( + Object.hasOwn(value, key), + "Missing memory replay prerequisite.", + ); + const initialSnapshot = decodeMemoryValue(value.initialSnapshot as JsonValue); + validateMemorySnapshot(initialSnapshot); + const configuration = decodeMemoryValue(value.configuration as JsonValue); + validateConfiguration(configuration); + const requestContext = decodeMemoryValue(value.requestContext as JsonValue); + requireValue(isRecord(requestContext), "Malformed recorded request context."); + const urls = new Set(); + const files = value.files.map((file) => { + requireValue( + isRecord(file) && + typeof file.url === "string" && + !urls.has(file.url) && + typeof file.mediaType === "string" && + file.mediaType.length > 0, + "Malformed or duplicate recorded file.", + ); + validateUrl(file.url); + urls.add(file.url); + return { + url: file.url, + mediaType: file.mediaType, + bytes: readBinary(file), + }; + }); + return { + invocationId: value.invocationId, + rawInput: decodeMemoryValue(value.rawInput as JsonValue), + initialSnapshot, + configuration, + requestContext, + files, + }; +} + +/** Return undefined for legacy inputs; a present but invalid v2 envelope always rejects. */ +export function restoreMemoryReplayEnvelope( + input: unknown, +): MastraMemoryReplayInput | undefined { + if (!isRecord(input) || !Object.hasOwn(input, MEMORY_REPLAY_KEY)) + return undefined; + return decodeMemoryReplayEnvelope(input[MEMORY_REPLAY_KEY]); +} diff --git a/packages/mastra/test/memory-binding.test.ts b/packages/mastra/test/memory-binding.test.ts new file mode 100644 index 000000000..dabbd3a86 --- /dev/null +++ b/packages/mastra/test/memory-binding.test.ts @@ -0,0 +1,217 @@ +import { expect, it, vi } from "vitest"; +import { + createMemoryCaptureBinding, + createProcessLocalMemoryAccess, +} from "../src/memory-binding.js"; +import { + createMemoryRuntime, + RESOURCE, + seedMemory, + THREAD, +} from "./helpers/memory-agent.js"; + +function required(value: T | undefined | null): T { + if (value === undefined || value === null) + throw new Error("Missing fixture value"); + return value; +} + +async function fixture( + invocationId = "invocation-1", + access = createProcessLocalMemoryAccess(), +) { + const runtime = createMemoryRuntime(); + await seedMemory(runtime); + const recordMutation = vi.fn(async (_event: unknown) => {}); + const binding = createMemoryCaptureBinding({ + invocationId, + domain: runtime.domain, + threadId: THREAD, + resourceId: RESOURCE, + exclusiveAccess: access, + recordMutation, + }); + return { runtime, binding, recordMutation }; +} + +it("captures initial memory under a real lease, records ordered mutations, and releases", async () => { + const { runtime, binding, recordMutation } = await fixture(); + const initial = await binding.captureInitial(runtime.memory); + expect(initial?.thread?.id).toBe(THREAD); + await binding.domain.updateThread({ + id: THREAD, + metadata: { workingMemory: "new" }, + }); + const record = await binding.domain.getObservationalMemory(THREAD, RESOURCE); + await binding.domain.setPendingMessageTokens(required(record).id, 12); + await binding.drain(); + expect(binding.revision).toBe(2); + expect(recordMutation.mock.calls.map(([event]) => event)).toMatchObject([ + { id: "invocation-1:memory:1", revision: 1, method: "updateThread" }, + { + id: "invocation-1:memory:2", + revision: 2, + method: "setPendingMessageTokens", + }, + ]); + expect(binding.incompleteReasons).toEqual([]); + expect(initial?.thread?.metadata?.workingMemory).not.toBe("new"); + await binding.release(); +}); + +it("rejects shared-thread overlap while allowing independent source threads", async () => { + const access = createProcessLocalMemoryAccess(); + const one = await fixture("one", access); + const two = await fixture("two", access); + await one.binding.captureInitial(one.runtime.memory); + expect(await two.binding.captureInitial(two.runtime.memory)).toBeUndefined(); + expect(two.binding.incompleteReasons.join()).toMatch(/exclusive/i); + await one.binding.release(); + await two.binding.release(); +}); + +it("preserves native mutation results despite evidence persistence failure", async () => { + const { runtime, binding, recordMutation } = await fixture(); + await binding.captureInitial(runtime.memory); + recordMutation.mockRejectedValue(new Error("sink failed")); + const result = await binding.domain.updateThread({ + id: THREAD, + title: "still succeeds", + }); + expect(result.title).toBe("still succeeds"); + await binding.drain(); + expect(binding.incompleteReasons.join()).toMatch(/persistence/); + await binding.release(); +}); + +it("marks failed capture or unjoined work incomplete without throwing", async () => { + const { runtime, binding } = await fixture(); + const record = await runtime.domain.getObservationalMemory(THREAD, RESOURCE); + await runtime.domain.setObservingFlag(required(record).id, true); + expect(await binding.captureInitial(runtime.memory)).toBeUndefined(); + expect(binding.incompleteReasons.join()).toMatch(/Unjoined/); + await binding.release(); +}); + +it("invalidates the first recording when a conflicting invocation cannot get its lease", async () => { + const access = createProcessLocalMemoryAccess(); + const one = await fixture("one", access); + const two = await fixture("two", access); + await one.binding.captureInitial(one.runtime.memory); + await two.binding.captureInitial(two.runtime.memory); + expect(one.binding.incompleteReasons.join()).toMatch(/overlapping/); + await one.binding.release(); + await two.binding.release(); +}); + +it("does not return a coherent initial snapshot after overlap during capture", async () => { + const access = createProcessLocalMemoryAccess(); + const one = await fixture("one", access); + const two = await fixture("two", access); + const settled = async () => { + await two.binding.captureInitial(two.runtime.memory); + }; + expect(await one.binding.captureInitial({ settled })).toBeUndefined(); + await one.binding.release(); + await two.binding.release(); +}); + +it("serializes overlapping native mutations and preserves original storage errors", async () => { + const { runtime, binding, recordMutation } = await fixture(); + await binding.captureInitial(runtime.memory); + let unblock: (() => void) | undefined; + const wait = new Promise((resolve) => { + unblock = resolve; + }); + const native = runtime.domain.updateThread.bind(runtime.domain); + const calls: string[] = []; + const spy = vi + .spyOn(runtime.domain, "updateThread") + .mockImplementation(async (args) => { + calls.push(required(args.title)); + if (args.title === "first") await wait; + return native(args); + }); + const first = binding.domain.updateThread({ id: THREAD, title: "first" }); + const second = binding.domain.updateThread({ id: THREAD, title: "second" }); + await Promise.resolve(); + expect(calls).toEqual(["first"]); + required(unblock)(); + await Promise.all([first, second]); + await binding.drain(); + expect(calls).toEqual(["first", "second"]); + expect(recordMutation.mock.calls).toHaveLength(2); + spy.mockRestore(); + const fault = new Error("native write failed"); + vi.spyOn(runtime.domain, "saveResource").mockRejectedValueOnce(fault); + await expect( + binding.domain.saveResource({ + resource: { id: RESOURCE, createdAt: new Date(), updatedAt: new Date() }, + }), + ).rejects.toBe(fault); + expect(binding.revision).toBe(2); + expect(binding.incompleteReasons.join()).toMatch(/Native memory storage/); + await binding.release(); +}); + +it("records OM flags, buffers, config, activation and working memory with stable request identity", async () => { + const { runtime, binding, recordMutation } = await fixture(); + await binding.captureInitial(runtime.memory); + const record = await binding.domain.getObservationalMemory(THREAD, RESOURCE); + await binding.domain.setBufferingObservationFlag( + required(record).id, + true, + 10, + ); + await binding.domain.updateBufferedObservations({ + id: required(record).id, + chunk: { + cycleId: "cycle", + observations: "buffer", + tokenCount: 3, + messageIds: ["historical-message"], + messageTokens: 20, + lastObservedAt: new Date(20), + }, + lastBufferedAtTime: new Date(21), + }); + await binding.domain.setBufferingObservationFlag(required(record).id, false); + await binding.domain.updateObservationalMemoryConfig({ + id: required(record).id, + config: { observation: { messageTokens: 20 } }, + }); + await binding.domain.updateActiveObservations({ + id: required(record).id, + observations: "changed", + tokenCount: 3, + lastObservedAt: new Date(20), + }); + await binding.domain.updateThread({ + id: THREAD, + metadata: { workingMemory: "changed" }, + }); + await binding.drain(); + expect( + recordMutation.mock.calls.map( + ([event]) => (event as { revision: number }).revision, + ), + ).toEqual([1, 2, 3, 4, 5, 6]); + await binding.release(); +}); + +it("marks credential-altered mutation evidence incomplete but keeps native arguments and results", async () => { + const { runtime, binding, recordMutation } = await fixture(); + await binding.captureInitial(runtime.memory); + const result = await binding.domain.updateThread({ + id: THREAD, + metadata: { apiKey: "private-value" }, + }); + expect(result.metadata?.apiKey).toBe("private-value"); + await binding.drain(); + expect(binding.incompleteReasons.length).toBeGreaterThan(0); + expect(JSON.stringify(recordMutation.mock.calls)).not.toContain( + "private-value", + ); + expect(recordMutation.mock.calls[0]?.[0]).toMatchObject({ complete: false }); + await binding.release(); +}); diff --git a/packages/mastra/test/memory-snapshot.test.ts b/packages/mastra/test/memory-snapshot.test.ts new file mode 100644 index 000000000..2bb069e91 --- /dev/null +++ b/packages/mastra/test/memory-snapshot.test.ts @@ -0,0 +1,223 @@ +import { expect, it } from "vitest"; +import { + createContextInput, + restoreConversationContext, +} from "../src/conversation-context.js"; +import { + createMemoryReplayEnvelope, + decodeMemoryReplayEnvelope, + decodeMemoryValue, + encodeMemoryValue, +} from "../src/memory-snapshot.js"; +import { + createMemoryRuntime, + FILE_BYTES, + FILE_URL, + seedMemory, + snapshotMemory, +} from "./helpers/memory-agent.js"; + +function required(value: T | undefined | null): T { + if (value === undefined || value === null) + throw new Error("Missing fixture value"); + return value; +} + +async function fixture() { + const runtime = createMemoryRuntime(); + await seedMemory(runtime); + return { + invocationId: "invocation-1", + rawInput: [ + { role: "user", content: [{ type: "file", data: new URL(FILE_URL) }] }, + ], + initialSnapshot: { + ...(await snapshotMemory(runtime, true)), + threadId: "historical-thread", + resourceId: "historical-resource", + }, + configuration: { + instructions: "Answer", + model: { provider: "fixture", modelId: "actor" }, + memory: { + semanticRecall: false, + workingMemory: { scope: "thread", schema: { type: "object" } }, + observationalMemory: { scope: "thread" }, + }, + }, + requestContext: { locale: "en" }, + files: [{ url: FILE_URL, mediaType: "application/pdf", bytes: FILE_BYTES }], + }; +} + +it("round-trips historical memory Dates, undefined fields, URL input and lossless binary", async () => { + const input = await fixture(); + required(input.initialSnapshot.records[0]).lastBufferedAtTime = new Date(123); + required(input.initialSnapshot.records[0]).bufferedReflection = + "pending reflection"; + required(input.initialSnapshot.records[0]).bufferedMessageIds = [ + "historical-message", + ]; + const envelope = createMemoryReplayEnvelope(input); + expect(envelope.complete).toBe(true); + const restored = decodeMemoryReplayEnvelope( + JSON.parse(JSON.stringify(envelope)), + ); + expect(restored).toEqual(input); + expect(restored.initialSnapshot.messages[0]?.createdAt).toBeInstanceOf(Date); + expect(decodeMemoryValue(encodeMemoryValue(FILE_BYTES))).toEqual(FILE_BYTES); +}); + +it.each(["version", "hash", "missing", "inflight"])( + "rejects invalid replay prerequisites: %s", + async (kind) => { + const envelope = createMemoryReplayEnvelope(await fixture()); + if (kind === "version") envelope.version = 3 as 2; + if (kind === "hash") required(envelope.files[0]).sha256 = "0".repeat(64); + if (kind === "missing") envelope.initialSnapshot = {}; + if (kind === "inflight") { + const input = await fixture(); + required(input.initialSnapshot.records[0]).isObserving = true; + Object.assign(envelope, createMemoryReplayEnvelope(input)); + } + expect(() => decodeMemoryReplayEnvelope(envelope)).toThrow( + /Unsupported Mastra memory replay/, + ); + }, +); + +it.each([ + { apiKey: "private-value" }, + { callback: () => "live" }, + { value: "a".repeat(1_048_576) }, + { value: new Map([["key", "value"]]) }, + { value: Array.from({ length: 10_001 }, () => 1) }, +])( + "marks altered or oversized state incomplete without exposing credentials", + async (configuration) => { + const envelope = createMemoryReplayEnvelope({ + ...(await fixture()), + configuration, + }); + expect(envelope.complete).toBe(false); + expect(JSON.stringify(envelope)).not.toContain("private-value"); + expect(() => decodeMemoryReplayEnvelope(envelope)).toThrow(); + }, +); + +it("rejects reserved codec tags and cyclic objects instead of accepting ambiguous data", () => { + const cycle: Record = {}; + cycle.self = cycle; + expect(() => encodeMemoryValue(cycle)).toThrow(); + expect(() => + encodeMemoryValue({ $mastra: "date", value: "2026-01-01" }), + ).toThrow(); +}); + +it("keeps the version-1 recalled conversation contract", () => { + const messages = [{ role: "user", content: "hello" }]; + expect( + restoreConversationContext(createContextInput(messages, messages)), + ).toEqual(messages); +}); + +it("round-trips complete buffered chunks and every historical OM generation", async () => { + const input = await fixture(); + required(input.initialSnapshot.records[0]).bufferedObservationChunks = [ + { + id: "chunk-1", + cycleId: "cycle-1", + observations: "buffered", + tokenCount: 3, + messageIds: ["historical-message"], + messageTokens: 7, + lastObservedAt: new Date(100), + createdAt: new Date(110), + extractedValues: { next: "wait" }, + }, + ]; + input.initialSnapshot.records.push({ + ...required(input.initialSnapshot.records[0]), + id: "older-generation", + generationCount: 1, + originType: "reflection", + }); + expect(decodeMemoryReplayEnvelope(createMemoryReplayEnvelope(input))).toEqual( + input, + ); +}); + +it("round-trips an empty initial conversation", async () => { + const input = await fixture(); + input.initialSnapshot = { + threadId: "new", + resourceId: "new-resource", + thread: null, + resource: null, + messages: [], + records: [], + }; + expect( + decodeMemoryReplayEnvelope(createMemoryReplayEnvelope(input)) + .initialSnapshot, + ).toEqual(input.initialSnapshot); +}); + +it("counts aggregate envelope items, depth, and binary expansion against the shared budget", async () => { + const input = await fixture(); + expect( + createMemoryReplayEnvelope({ + ...input, + configuration: { values: Array.from({ length: 6000 }, () => 1) }, + requestContext: { values: Array.from({ length: 6000 }, () => 1) }, + }).complete, + ).toBe(false); + const deep = Array.from({ length: 65 }).reduce( + (value) => ({ value }), + null, + ); + expect(() => encodeMemoryValue(deep)).toThrow(/depth/); + const withFile = await fixture(); + required(withFile.files[0]).bytes = new Uint8Array(800_000); + expect(createMemoryReplayEnvelope(withFile).complete).toBe(false); +}); + +it.each(["resourceScope", "semanticRecall", "buffer", "ids", "date", "url"])( + "rejects malformed or out-of-scope state: %s", + async (kind) => { + const input = await fixture(); + if (kind === "resourceScope") + required(input.initialSnapshot.records[0]).scope = "resource"; + if (kind === "semanticRecall") + input.configuration.memory.semanticRecall = true; + if (kind === "buffer") + required(input.initialSnapshot.records[0]).bufferedObservationChunks = [ + {} as never, + ]; + if (kind === "ids") + required(input.initialSnapshot.messages[0]).threadId = "different-thread"; + if (kind === "date") + required(input.initialSnapshot.messages[0]).createdAt = new Date( + Number.NaN, + ); + if (kind === "url") + required(required(input.rawInput[0]).content[0]).data = new URL( + "https://files.invalid/file?apiKey=private-value", + ); + const envelope = createMemoryReplayEnvelope(input); + expect(envelope.complete).toBe(false); + expect(JSON.stringify(envelope)).not.toContain("private-value"); + }, +); + +it("rejects changed file lengths, noncanonical base64, and malformed date tags", async () => { + const input = await fixture(); + for (const patch of [{ length: 100 }, { base64: "???" }]) { + const envelope = createMemoryReplayEnvelope(input); + Object.assign(required(envelope.files[0]), patch); + expect(() => decodeMemoryReplayEnvelope(envelope)).toThrow(/binary/); + } + expect(() => + decodeMemoryValue({ $mastra: "date", value: "2026-01-01" }), + ).toThrow(/Date/); +}); From fa0bba8e305aa80b0e5c87c73c2b80ca27e4b576 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 19:21:47 +0200 Subject: [PATCH 003/115] Reject mutations during initial memory capture --- packages/mastra/src/memory-binding.ts | 11 +++++++++++ packages/mastra/test/memory-binding.test.ts | 16 ++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/packages/mastra/src/memory-binding.ts b/packages/mastra/src/memory-binding.ts index ba8d38cd0..5780e89bf 100644 --- a/packages/mastra/src/memory-binding.ts +++ b/packages/mastra/src/memory-binding.ts @@ -79,6 +79,8 @@ export interface MastraMemoryCaptureBinding { // The pinned public MemoryStorage mutation inventory. Delegation binds `this` to // the original domain, so a native method's own helper calls record only once. const MUTATIONS = new Set([ + "dangerouslyClearAll", + "prune", "saveThread", "updateThread", "patchThread", @@ -115,6 +117,7 @@ export function createMemoryCaptureBinding( let revision = 0; let started = false; let capturing = false; + let readingSnapshot = false; let released = false; let releaseLease: (() => Promise) | undefined; let mutations = Promise.resolve(); @@ -159,6 +162,12 @@ export function createMemoryCaptureBinding( markIncomplete( "Memory mutation occurred outside the owned invocation lifecycle.", ); + if (readingSnapshot) + markIncomplete("Memory mutation overlapped initial snapshot reads."); + if (property === "dangerouslyClearAll" || property === "prune") + markIncomplete( + "Storage-wide mutation is outside the captured thread scope.", + ); const duringCapture = capturing; const result = mutations.then(async () => { let output: unknown; @@ -271,6 +280,7 @@ export function createMemoryCaptureBinding( } await memory.settled(); await mutations; + readingSnapshot = true; const thread = await options.domain.getThreadById({ threadId: options.threadId, }); @@ -304,6 +314,7 @@ export function createMemoryCaptureBinding( return undefined; } finally { capturing = false; + readingSnapshot = false; } }, drain, diff --git a/packages/mastra/test/memory-binding.test.ts b/packages/mastra/test/memory-binding.test.ts index dabbd3a86..856707e36 100644 --- a/packages/mastra/test/memory-binding.test.ts +++ b/packages/mastra/test/memory-binding.test.ts @@ -30,6 +30,7 @@ async function fixture( resourceId: RESOURCE, exclusiveAccess: access, recordMutation, + getRequestId: () => "request-1", }); return { runtime, binding, recordMutation }; } @@ -215,3 +216,18 @@ it("marks credential-altered mutation evidence incomplete but keeps native argum expect(recordMutation.mock.calls[0]?.[0]).toMatchObject({ complete: false }); await binding.release(); }); + +it("rejects same-invocation writes that interleave with initial snapshot reads", async () => { + const { runtime, binding } = await fixture(); + const native = runtime.domain.getThreadById.bind(runtime.domain); + vi.spyOn(runtime.domain, "getThreadById").mockImplementationOnce( + async (args) => { + const thread = await native(args); + await binding.domain.updateThread({ id: THREAD, title: "interleaved" }); + return thread; + }, + ); + expect(await binding.captureInitial(runtime.memory)).toBeUndefined(); + expect(binding.incompleteReasons.join()).toMatch(/overlapped initial/); + await binding.release(); +}); From 2517ace8ded5abbacfae39443023293ecb3d2137 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 19:27:55 +0200 Subject: [PATCH 004/115] Restore historical Mastra memory in isolated stores --- packages/mastra/package.json | 8 +- packages/mastra/src/memory-replay.ts | 221 +++++++++++++++++ packages/mastra/test/memory-replay.test.ts | 261 +++++++++++++++++++++ 3 files changed, 489 insertions(+), 1 deletion(-) create mode 100644 packages/mastra/src/memory-replay.ts create mode 100644 packages/mastra/test/memory-replay.test.ts diff --git a/packages/mastra/package.json b/packages/mastra/package.json index a6cf10ea0..d1a080801 100644 --- a/packages/mastra/package.json +++ b/packages/mastra/package.json @@ -41,7 +41,13 @@ "@zenml-io/kitaru": "workspace:0.5.0" }, "peerDependencies": { - "@mastra/core": ">=1.51.0 <1.68.0" + "@mastra/core": ">=1.51.0 <1.68.0", + "@mastra/memory": "1.30.0" + }, + "peerDependenciesMeta": { + "@mastra/memory": { + "optional": true + } }, "devDependencies": { "@mastra/core": "1.67.0", diff --git a/packages/mastra/src/memory-replay.ts b/packages/mastra/src/memory-replay.ts new file mode 100644 index 000000000..5ddccc8ca --- /dev/null +++ b/packages/mastra/src/memory-replay.ts @@ -0,0 +1,221 @@ +import { createRequire } from "node:module"; +import type { MastraModelConfig } from "@mastra/core/llm"; +import type { MemoryConfigInternal } from "@mastra/core/memory"; +import { + standardSchemaToJSONSchema, + toStandardSchema, +} from "@mastra/core/schema"; +import { InMemoryStore, MastraCompositeStore } from "@mastra/core/storage"; +import { + createMemoryCaptureBinding, + createProcessLocalMemoryAccess, + type MastraMemoryCaptureOptions, +} from "./memory-binding.js"; +import { + decodeMemoryValue, + encodeMemoryValue, + type MastraMemorySnapshot, + validateMemorySnapshot, +} from "./memory-snapshot.js"; + +function record(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function unsupported(message: string): never { + throw new Error(`Unsupported Mastra memory replay: ${message}`); +} + +/** Require the dependency pair exercised by the native memory proof. */ +export function assertMemoryReplayVersions(): void { + const require = createRequire(import.meta.url); + for (const [name, version] of [ + ["@mastra/core", "1.67.0"], + ["@mastra/memory", "1.30.0"], + ]) { + const metadata: unknown = require(`${name}/package.json`); + if (!record(metadata) || metadata.version !== version) + unsupported(`requires ${name}@${version}.`); + } +} + +/** Save a model identity, never the provider client or its credentials. */ +export function getMemoryModelId(model: unknown): string { + if (typeof model === "string" && model.length > 0) return model; + if ( + record(model) && + typeof model.modelId === "string" && + typeof model.provider === "string" + ) + return `${model.provider}/${model.modelId}`; + if (record(model) && typeof model.id === "string") return model.id; + return unsupported("Memory models require a static model identity."); +} + +function checkConfiguration(config: Record): void { + const allowed = new Set([ + "readOnly", + "lastMessages", + "semanticRecall", + "workingMemory", + "observationalMemory", + "generateTitle", + "filterIncompleteToolCalls", + ]); + for (const key of Object.keys(config)) + if (!allowed.has(key)) + unsupported(`Memory option '${key}' is not supported.`); + if (config.semanticRecall !== undefined && config.semanticRecall !== false) + unsupported("Semantic recall requires external state."); + if (config.generateTitle !== undefined && config.generateTitle !== false) + unsupported("Automatic title generation is not supported."); + for (const key of ["workingMemory", "observationalMemory"]) { + const feature = config[key]; + if (feature === undefined || feature === false) continue; + if (!record(feature)) + unsupported(`${key} requires explicit thread-scoped configuration.`); + if (feature.enabled !== false && feature.scope !== "thread") + unsupported(`${key} must use thread scope.`); + } +} + +/** Convert the supported native schema and OM models to self-contained configuration. */ +export function serializeMemoryConfiguration( + config: MemoryConfigInternal, +): Record { + checkConfiguration(config); + const copy: Record = { ...config }; + if (config.workingMemory?.schema) { + copy.workingMemory = { + ...config.workingMemory, + schema: standardSchemaToJSONSchema( + toStandardSchema(config.workingMemory.schema), + ), + }; + } + if (record(config.observationalMemory)) { + const om = { ...config.observationalMemory }; + if (om.model !== undefined) om.model = getMemoryModelId(om.model); + for (const name of ["observation", "reflection"]) { + if (record(om[name])) { + const phase = { ...om[name] }; + if (phase.model !== undefined) + phase.model = getMemoryModelId(phase.model); + om[name] = phase; + } + } + copy.observationalMemory = om; + } + return decodeMemoryValue(encodeMemoryValue(copy)) as Record; +} + +export async function restoreMemoryConfiguration( + configuration: Record, + resolveModel: (id: string) => Promise | MastraModelConfig, +): Promise { + const copy = decodeMemoryValue(encodeMemoryValue(configuration)); + if (!record(copy)) return unsupported("Missing native memory configuration."); + checkConfiguration(copy); + if (record(copy.observationalMemory)) { + const om = copy.observationalMemory; + if (om.model !== undefined) { + if (typeof om.model !== "string") + unsupported("Invalid observer model identity."); + om.model = await resolveModel(om.model); + } + for (const name of ["observation", "reflection"]) { + const phase = om[name]; + if (record(phase) && phase.model !== undefined) { + if (typeof phase.model !== "string") + unsupported("Invalid memory model identity."); + phase.model = await resolveModel(phase.model); + } + } + } + return copy as MemoryConfigInternal; +} + +export interface IsolatedMemoryReplayOptions { + invocationId: string; + initialSnapshot: MastraMemorySnapshot; + configuration: Record; + resolveModel: (id: string) => Promise | MastraModelConfig; + recordMutation: MastraMemoryCaptureOptions["recordMutation"]; + getRequestId?: MastraMemoryCaptureOptions["getRequestId"]; + onIncomplete?: MastraMemoryCaptureOptions["onIncomplete"]; +} + +/** Restore historical state into a fresh store; no production store is accepted. */ +export async function createIsolatedMemoryReplay( + options: IsolatedMemoryReplayOptions, +) { + assertMemoryReplayVersions(); + validateMemorySnapshot(options.initialSnapshot); + const snapshot = decodeMemoryValue( + encodeMemoryValue(options.initialSnapshot), + ) as MastraMemorySnapshot; + const configuration = await restoreMemoryConfiguration( + options.configuration, + options.resolveModel, + ); + const { Memory } = await import("@mastra/memory"); + const store = new InMemoryStore(); + const domain = store.stores.memory; + if (!domain) return unsupported("Native in-memory storage is unavailable."); + try { + if (snapshot.thread) + await domain.saveThread({ thread: structuredClone(snapshot.thread) }); + if (snapshot.resource) + await domain.saveResource({ resource: snapshot.resource }); + if (snapshot.messages.length) + await domain.saveMessages({ messages: snapshot.messages }); + // Saving messages updates thread metadata, including updatedAt. + if (snapshot.thread) + await domain.saveThread({ thread: structuredClone(snapshot.thread) }); + for (const value of snapshot.records) + await domain.insertObservationalMemoryRecord(value); + const binding = createMemoryCaptureBinding({ + invocationId: options.invocationId, + threadId: snapshot.threadId, + resourceId: snapshot.resourceId, + domain, + exclusiveAccess: createProcessLocalMemoryAccess(), + recordMutation: options.recordMutation, + getRequestId: options.getRequestId, + onIncomplete: options.onIncomplete, + }); + const storage = new MastraCompositeStore({ + id: `kitaru-replay-${options.invocationId}`, + domains: { memory: binding.domain }, + }); + const memory = new Memory({ storage, options: configuration }); + const initialSnapshot = await binding.captureInitial(memory); + if (!initialSnapshot) { + await binding.release(); + return unsupported( + "Restored memory did not produce a coherent initial snapshot.", + ); + } + let finished: Promise | undefined; + return { + memory, + binding, + initialSnapshot, + finish(): Promise { + finished ??= (async () => { + try { + await memory.settled(); + await binding.drain(); + } finally { + await binding.release(); + await store.close(); + } + })(); + return finished; + }, + }; + } catch (error) { + await store.close(); + throw error; + } +} diff --git a/packages/mastra/test/memory-replay.test.ts b/packages/mastra/test/memory-replay.test.ts new file mode 100644 index 000000000..34cf2102e --- /dev/null +++ b/packages/mastra/test/memory-replay.test.ts @@ -0,0 +1,261 @@ +import { Agent } from "@mastra/core/agent"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { createTool } from "@mastra/core/tools"; +import { expect, it } from "vitest"; +import { z } from "zod/v4"; +import type { MastraMemoryMutation } from "../src/memory-binding.js"; +import { + createIsolatedMemoryReplay, + serializeMemoryConfiguration, +} from "../src/memory-replay.js"; +import { + createMemoryRuntime, + type ModelCall, + RESOURCE, + seedMemory, + snapshotMemory, + streamParts, + THREAD, + textStream, +} from "./helpers/memory-agent.js"; + +async function baseline() { + const source = createMemoryRuntime(); + await seedMemory(source); + const initialSnapshot = { + ...(await snapshotMemory(source, true)), + threadId: THREAD, + resourceId: RESOURCE, + }; + const configuration = serializeMemoryConfiguration( + source.memory.getMergedThreadConfig(), + ); + return { source, initialSnapshot, configuration }; +} + +it("restores historical memory with original timestamps into independent native stores", async () => { + const { source, initialSnapshot, configuration } = await baseline(); + await source.domain.updateThread({ + id: THREAD, + metadata: { workingMemory: "production advanced" }, + }); + const options = { + initialSnapshot, + configuration, + resolveModel: async (id: string) => + id.endsWith("observer") ? source.observer.model : source.reflector.model, + recordMutation: async () => {}, + }; + const one = await createIsolatedMemoryReplay({ + ...options, + invocationId: "one", + }); + const two = await createIsolatedMemoryReplay({ + ...options, + invocationId: "two", + }); + expect(one.initialSnapshot).toEqual(initialSnapshot); + await one.memory.updateWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + workingMemory: "replay changed", + }); + expect( + await two.memory.getWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + }), + ).toContain("historical-blue"); + expect( + (await source.domain.getThreadById({ threadId: THREAD }))?.metadata + ?.workingMemory, + ).toBe("production advanced"); + await one.finish(); + await two.finish(); +}); + +it("serializes native schema and model identities without serializing live model objects", async () => { + const { configuration } = await baseline(); + expect(JSON.stringify(configuration)).toContain('"type":"object"'); + expect(JSON.stringify(configuration)).toContain("fixture/observer"); + expect(JSON.stringify(configuration)).not.toContain("doStream"); +}); + +it("rejects unsupported memory dependencies and unjoined work before model resolution", async () => { + expect(() => serializeMemoryConfiguration({ semanticRecall: true })).toThrow( + /semantic/i, + ); + expect(() => + serializeMemoryConfiguration({ + workingMemory: { enabled: true, scope: "resource" }, + }), + ).toThrow(/thread/i); + const { initialSnapshot, configuration } = await baseline(); + const firstRecord = initialSnapshot.records[0]; + if (!firstRecord) throw new Error("Missing native memory fixture"); + firstRecord.isObserving = true; + let resolved = false; + await expect( + createIsolatedMemoryReplay({ + invocationId: "bad", + initialSnapshot, + configuration, + resolveModel: async () => { + resolved = true; + throw new Error("must not resolve"); + }, + recordMutation: async () => {}, + }), + ).rejects.toThrow(/Unjoined/); + expect(resolved).toBe(false); +}); + +it("runs native observation, reflection and working-memory changes with ordered evidence", async () => { + const source = createMemoryRuntime({ messageTokens: 600 }); + await seedMemory(source); + const changes: MastraMemoryMutation[] = []; + const runtime = await createIsolatedMemoryReplay({ + invocationId: "evolving", + initialSnapshot: { + ...(await snapshotMemory(source, true)), + threadId: THREAD, + resourceId: RESOURCE, + }, + configuration: serializeMemoryConfiguration( + source.memory.getMergedThreadConfig(), + ), + resolveModel: (id) => + id.endsWith("observer") ? source.observer.model : source.reflector.model, + recordMutation: async (event) => { + changes.push(event); + }, + }); + const calls: ModelCall[] = []; + const model = new MastraLanguageModelV2Mock({ + provider: "fixture", + modelId: "actor", + doStream: async (options) => { + calls.push(options); + if (calls.length > 2) return textStream("done"); + return streamParts( + [ + { + type: "tool-call", + toolCallId: `call-${calls.length}`, + toolName: calls.length === 1 ? "updateWorkingMemory" : "evidence", + input: + calls.length === 1 + ? JSON.stringify({ memory: { preference: "replay-green" } }) + : "{}", + }, + ], + "tool-calls", + ); + }, + }); + const agent = new Agent({ + id: "isolated", + name: "Isolated", + instructions: "Answer from memory", + model, + memory: runtime.memory, + tools: { + evidence: createTool({ + id: "evidence", + description: "Read evidence", + inputSchema: z.object({}), + execute: async () => { + await runtime.memory.settled(); + return "The user now prefers replay-green. ".repeat(400); + }, + }), + }, + }); + const result = await agent.stream("Remember my new preference.", { + maxSteps: 5, + memory: { thread: THREAD, resource: RESOURCE }, + }); + for await (const _ of result.textStream) { + /* Consume the native stream. */ + } + await runtime.finish(); + expect(source.observer.calls.length).toBeGreaterThan(0); + expect(source.reflector.calls.length).toBeGreaterThan(0); + expect(JSON.stringify(calls.at(-1)?.prompt)).toContain("REFLECTED_REPLAY"); + expect( + changes.some((event) => event.method === "createReflectionGeneration"), + ).toBe(true); + expect(changes.map((event) => event.revision)).toEqual( + changes.map((_, index) => index + 1), + ); + expect(runtime.binding.incompleteReasons).toEqual([]); + expect( + await source.memory.getWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + }), + ).toContain("historical-blue"); +}, 20000); + +it("waits for native background observation and evidence before releasing replay state", async () => { + let release!: () => void; + let signalStarted!: () => void; + const blocked = new Promise((resolve) => { + release = resolve; + }); + const started = new Promise((resolve) => { + signalStarted = resolve; + }); + const source = createMemoryRuntime({ + messageTokens: 10000, + observerWait: async () => { + signalStarted(); + await blocked; + }, + }); + await seedMemory(source); + const runtime = await createIsolatedMemoryReplay({ + invocationId: "background", + initialSnapshot: { + ...(await snapshotMemory(source, true)), + threadId: THREAD, + resourceId: RESOURCE, + }, + configuration: serializeMemoryConfiguration( + source.memory.getMergedThreadConfig(), + ), + resolveModel: (id) => + id.endsWith("observer") ? source.observer.model : source.reflector.model, + recordMutation: async () => {}, + }); + try { + const agent = new Agent({ + id: "background", + name: "Background", + instructions: "Answer", + model: new MastraLanguageModelV2Mock({ + doStream: async () => textStream("done"), + }), + memory: runtime.memory, + }); + const result = await agent.stream("Remember this preference.", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + for await (const _ of result.textStream) { + /* Consume native stream. */ + } + await started; + let finished = false; + const completion = runtime.finish().then(() => { + finished = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(finished).toBe(false); + release(); + await completion; + expect(runtime.binding.incompleteReasons).toEqual([]); + } finally { + release(); + await runtime.finish(); + } +}, 20000); From e2b93bb0c06c204822c771c268ef4dc0ec4f395e Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 19:28:44 +0200 Subject: [PATCH 005/115] Capture effective Mastra provider requests --- packages/mastra/src/request-capture.ts | 315 ++++++++++++++++ packages/mastra/src/step-recorder.ts | 21 +- packages/mastra/test/request-capture.test.ts | 367 +++++++++++++++++++ 3 files changed, 697 insertions(+), 6 deletions(-) create mode 100644 packages/mastra/src/request-capture.ts create mode 100644 packages/mastra/test/request-capture.test.ts diff --git a/packages/mastra/src/request-capture.ts b/packages/mastra/src/request-capture.ts new file mode 100644 index 000000000..d44ebb878 --- /dev/null +++ b/packages/mastra/src/request-capture.ts @@ -0,0 +1,315 @@ +import type { MessageList } from "@mastra/core/agent/message-list"; +import type { JsonValue } from "@zenml-io/kitaru"; +import { + boundedRecorderConversion, + type RecordingLimits, + recordedToolPayloadConversion, +} from "@zenml-io/kitaru/adapter"; +import { encodeMemoryValue } from "./memory-snapshot.js"; + +export interface RequestEvidence { + externalId: string; + invocationId: string; + stepNumber: number; + attemptNumber: number; + memoryRevision: number | null; + method: "doStream" | "doGenerate"; + modelId: string; + provider: string; + startedAt: string; + inputs: JsonValue; + modelSettings: Record; + provenance: JsonValue; + complete: boolean; + reasons: string[]; +} + +export interface RequestCaptureOptions { + invocationId: string; + getMemoryRevision: () => number; + recordingLimits?: RecordingLimits; + onFailedAttempt?: ( + evidence: RequestEvidence, + error: unknown, + ) => void | Promise; + onCaptureError?: (error: unknown) => void; +} + +export interface RequestStepContext { + stepNumber: number; + messageList?: MessageList; + applicationInstructions?: unknown; + extraContext?: unknown; +} + +interface PublicModel { + specificationVersion: string; + modelId: string; + provider: string; +} + +const SETTINGS = [ + "maxOutputTokens", + "temperature", + "topP", + "topK", + "presencePenalty", + "frequencyPenalty", + "stopSequences", + "seed", + "responseFormat", + "reasoning", + "providerOptions", +] as const; + +function asRecord(value: unknown): Record { + return typeof value === "object" && value !== null + ? (value as Record) + : {}; +} + +/** Capture the final provider arguments without consuming or replacing its output stream. */ +export function createRequestCapture(options: RequestCaptureOptions) { + let context: { + stepNumber: number; + provenance: JsonValue; + reasons: string[]; + } = { + stepNumber: 0, + provenance: null, + reasons: ["Step provenance was not supplied."], + }; + const attempts = new Map(); + const unfinished = new Map< + string, + { evidence: RequestEvidence; returned: boolean } + >(); + const proxies = new WeakMap(); + const writes = new Set>(); + + function report(error: unknown): void { + try { + options.onCaptureError?.(error); + } catch { + /* Diagnostics cannot change native execution. */ + } + } + + function convert( + value: unknown, + label: string, + reasons: string[], + ): JsonValue { + try { + const encoded = encodeMemoryValue(value); + // Provider options can contain custom transport headers whose keys are + // not recognizable credential names. Do not persist that transport bag. + function containsTransport(current: JsonValue): boolean { + if (current === null || typeof current !== "object") return false; + if (Array.isArray(current)) return current.some(containsTransport); + return Object.entries(current).some( + ([key, item]) => + /^(headers|abortsignal)$/i.test(key) || containsTransport(item), + ); + } + if (containsTransport(encoded)) { + reasons.push(`${label} contains transport metadata.`); + return null; + } + if (options.recordingLimits === undefined) return encoded; + const converted = boundedRecorderConversion( + encoded, + label, + options.recordingLimits, + ); + if (converted.lossy) + reasons.push( + `${label} exceeded recording limits or required redaction.`, + ); + return converted.value; + } catch (error) { + reasons.push(`${label} could not be recorded losslessly.`); + report(error); + return null; + } + } + + function beginStep(step: RequestStepContext): void { + const reasons: string[] = []; + let provenance: JsonValue = null; + try { + const list = step.messageList; + const sources = list?.makeMessageSourceChecker(); + provenance = convert( + { + applicationInstructions: step.applicationInstructions, + extraContext: step.extraContext, + systemMessages: list?.serializeForSpan().systemMessages ?? [], + messages: + list?.get.all.db().map((message) => ({ + id: message.id, + source: sources?.getSource(message) ?? null, + role: message.role, + content: message.content, + })) ?? [], + }, + "Prompt provenance", + reasons, + ); + } catch (error) { + reasons.push("Prompt provenance could not be read."); + report(error); + } + context = { stepNumber: step.stepNumber, provenance, reasons }; + } + + function capture( + model: PublicModel, + method: RequestEvidence["method"], + args: unknown, + ): RequestEvidence { + const reasons = [...context.reasons]; + const attemptNumber = (attempts.get(context.stepNumber) ?? 0) + 1; + attempts.set(context.stepNumber, attemptNumber); + let memoryRevision: number | null = null; + try { + memoryRevision = options.getMemoryRevision(); + } catch (error) { + reasons.push("Memory revision could not be read."); + report(error); + } + const request = asRecord(args); + let inputs: JsonValue = null; + const modelSettings: Record = {}; + try { + inputs = convert( + { + prompt: request.prompt, + tools: request.tools, + toolChoice: request.toolChoice, + }, + "Effective model request", + reasons, + ); + for (const key of SETTINGS) { + if (request[key] !== undefined) + modelSettings[key] = convert( + request[key], + `Model setting ${key}`, + reasons, + ); + } + } catch (error) { + reasons.push("Effective model arguments could not be read."); + report(error); + } + const combined = recordedToolPayloadConversion( + { inputs, modelSettings, provenance: context.provenance }, + "Effective request evidence", + ); + if (combined.lossy) { + reasons.push("Combined request evidence exceeded recording limits."); + inputs = null; + } + return { + externalId: globalThis.crypto.randomUUID(), + invocationId: options.invocationId, + stepNumber: context.stepNumber, + attemptNumber, + memoryRevision, + method, + modelId: model.modelId, + provider: model.provider, + startedAt: new Date().toISOString(), + inputs, + modelSettings: combined.lossy ? {} : modelSettings, + provenance: combined.lossy ? null : context.provenance, + complete: reasons.length === 0, + reasons, + }; + } + + function instrumentModel(model: T): T { + if (!["v2", "v3", "v4"].includes(model.specificationVersion)) + throw new TypeError("Request capture requires a v2, v3, or v4 model"); + const existing = proxies.get(model); + if (existing) return existing as T; + const proxy = new Proxy(model, { + get(target, key) { + const value = Reflect.get(target, key, target); + if ( + (key === "doStream" || key === "doGenerate") && + typeof value === "function" + ) { + return async (...args: unknown[]) => { + const evidence = capture(target, key, args[0]); + const attempt = { evidence, returned: false }; + unfinished.set(evidence.externalId, attempt); + try { + const result = await Reflect.apply(value, target, args); + attempt.returned = true; + return result; + } catch (error) { + unfinished.delete(evidence.externalId); + // Queue the sink separately: retries and provider failures retain native behavior. + const write = Promise.resolve() + .then(() => options.onFailedAttempt?.(evidence, error)) + .catch(report); + writes.add(write); + void write.finally(() => writes.delete(write)); + throw error; + } + }; + } + return typeof value === "function" ? value.bind(target) : value; + }, + }); + proxies.set(model, proxy); + proxies.set(proxy, proxy); + return proxy; + } + + function takeSuccessful(stepNumber?: number): RequestEvidence | undefined { + for (const [id, attempt] of unfinished) { + if ( + attempt.returned && + (stepNumber === undefined || attempt.evidence.stepNumber === stepNumber) + ) { + unfinished.delete(id); + return attempt.evidence; + } + } + return undefined; + } + + /** Retain requests whose step never completed, including aborts and abandoned output streams. */ + function flushUnfinished(): RequestEvidence[] { + const evidence = [...unfinished.values()].map( + (attempt) => attempt.evidence, + ); + unfinished.clear(); + return evidence; + } + + async function drain(): Promise { + while (writes.size > 0) await Promise.all(writes); + } + + return { beginStep, instrumentModel, takeSuccessful, flushUnfinished, drain }; +} + +/** Attributes shared by normal step completion and failed or unfinished attempt records. */ +export function requestEvidenceAttributes( + evidence: RequestEvidence, +): Record { + return { + invocation_id: evidence.invocationId, + step_number: evidence.stepNumber, + attempt_number: evidence.attemptNumber, + memory_revision: evidence.memoryRevision, + request_method: evidence.method, + request_complete: evidence.complete, + request_incomplete_reasons: evidence.reasons, + prompt_provenance: evidence.provenance, + }; +} diff --git a/packages/mastra/src/step-recorder.ts b/packages/mastra/src/step-recorder.ts index 7c6cdb9e8..0b20fe14f 100644 --- a/packages/mastra/src/step-recorder.ts +++ b/packages/mastra/src/step-recorder.ts @@ -10,7 +10,10 @@ import { recordNormalizedStep, resolveCost, } from "@zenml-io/kitaru/adapter"; - +import { + type RequestEvidence, + requestEvidenceAttributes, +} from "./request-capture.js"; import type { KitaruCostCalculator, PublicModelIdentity } from "./types.js"; export type RecordedStep = LLMStepResult & { @@ -174,6 +177,7 @@ export async function recordStep( step: RecordedStep, costCalculator?: KitaruCostCalculator, limits?: RecordingLimits, + requestEvidence?: RequestEvidence, ): Promise { const calls = step.toolCalls.flatMap((item) => { const call = toolCallPayload(item); @@ -224,11 +228,13 @@ export async function recordStep( }; }); const failed = step.finishReason === "error" || step.tripwire !== undefined; - const servedModelId = step.response?.modelId ?? step.model?.modelId; + const servedModelId = + step.response?.modelId ?? step.model?.modelId ?? requestEvidence?.modelId; + const provider = step.model?.provider ?? requestEvidence?.provider; const tokens = usageTokens(step.usage); const cost = await resolveCost(costCalculator, { model: servedModelId ?? "", - provider: step.model?.provider ?? "", + provider: provider ?? "", requestedModelId: state.requestedModelId, tokens, }); @@ -236,6 +242,7 @@ export async function recordStep( await recordNormalizedStep(state, { attributes: { cost: cost.attribute, + ...(requestEvidence ? requestEvidenceAttributes(requestEvidence) : {}), ...(isRecord(step.providerMetadata) ? { provider_metadata: projectRecordedMetadata(step.providerMetadata) } : {}), @@ -244,12 +251,14 @@ export async function recordStep( error: failed ? errorMessage(step.error ?? step.tripwire?.reason, "Model step failed") : undefined, - externalId: step.response?.id, + externalId: requestEvidence?.externalId ?? step.response?.id, failed, - inputs: null, + inputs: requestEvidence?.inputs ?? null, model: servedModelId, + modelSettings: requestEvidence?.modelSettings, + startedAt: requestEvidence?.startedAt, outputs: stepOutputs(step, tools), - provider: step.model?.provider, + provider, tokens, tools, }); diff --git a/packages/mastra/test/request-capture.test.ts b/packages/mastra/test/request-capture.test.ts new file mode 100644 index 000000000..eac072d10 --- /dev/null +++ b/packages/mastra/test/request-capture.test.ts @@ -0,0 +1,367 @@ +import { Agent } from "@mastra/core/agent"; +import { MessageList } from "@mastra/core/agent/message-list"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { KitaruClient } from "@zenml-io/kitaru"; +import { RunRecorder } from "@zenml-io/kitaru/adapter"; +import { APICallError } from "ai"; +import { afterEach, expect, it, vi } from "vitest"; +import { decodeMemoryValue } from "../src/memory-snapshot.js"; +import { createRequestCapture } from "../src/request-capture.js"; +import { recordStep } from "../src/step-recorder.js"; +import { type ModelCall, textStream } from "./helpers/memory-agent.js"; +import { AGENT_ID, installTestApi, textStep } from "./helpers.js"; + +function required(value: T | undefined): T { + if (value === undefined) throw new Error("Expected recorded evidence"); + return value; +} + +afterEach(() => { + vi.unstubAllGlobals(); + vi.restoreAllMocks(); +}); + +it.each(["v2", "v3", "v4"])( + "preserves %s model receiver, arguments, result and unread stream", + async (version) => { + const result = textStream("native"); + const received: unknown[] = []; + class Model { + #id = "private-model"; + specificationVersion = version; + provider = "fixture"; + get modelId() { + return this.#id; + } + async doStream(args: unknown) { + received.push(args); + return result; + } + async doGenerate(args: unknown) { + received.push(args); + return { text: this.#id }; + } + } + const capture = createRequestCapture({ + invocationId: "invocation", + getMemoryRevision: () => 4, + }); + capture.beginStep({ stepNumber: 2 }); + const model = capture.instrumentModel(new Model()); + const args = { + prompt: [ + { + role: "user", + content: [ + { + type: "file", + data: new Uint8Array([1, 0, 255]), + mediaType: "image/png", + }, + ], + }, + ], + tools: [ + { type: "function", name: "local", inputSchema: { type: "object" } }, + ], + toolChoice: { type: "auto" }, + temperature: 0.4, + headers: { authorization: "SECRET" }, + abortSignal: new AbortController().signal, + }; + expect(await model.doStream(args)).toBe(result); + expect(result.stream.locked).toBe(false); + expect(received[0]).toBe(args); + const evidence = required(capture.takeSuccessful(2)); + expect(evidence).toMatchObject({ + invocationId: "invocation", + stepNumber: 2, + attemptNumber: 1, + memoryRevision: 4, + modelId: "private-model", + complete: true, + }); + expect(decodeMemoryValue(evidence.inputs)).toEqual({ + prompt: args.prompt, + tools: args.tools, + toolChoice: args.toolChoice, + }); + expect(evidence.modelSettings).toEqual({ temperature: 0.4 }); + expect(JSON.stringify(evidence)).not.toContain("SECRET"); + expect(capture.takeSuccessful(2)).toBeUndefined(); + expect(await model.doGenerate({ prompt: [] })).toEqual({ + text: "private-model", + }); + expect(capture.takeSuccessful(2)).toMatchObject({ + attemptNumber: 2, + method: "doGenerate", + }); + }, +); + +it("captures retries after late prompt and settings changes without changing the provider arguments", async () => { + const received: ModelCall[] = []; + const failed = vi.fn(); + const capture = createRequestCapture({ + invocationId: "retry", + getMemoryRevision: () => 3, + onFailedAttempt: failed, + }); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async (args) => { + received.push(args); + if (received.length === 1) + throw new APICallError({ + message: "Retry", + url: "https://fixture.invalid", + requestBodyValues: {}, + statusCode: 503, + isRetryable: true, + }); + return textStream("done"); + }, + }); + const agent = new Agent({ + id: "capture", + name: "capture", + instructions: "APP", + model: capture.instrumentModel(model), + inputProcessors: [ + { + id: "late", + processInputStep: ({ stepNumber, messageList }) => { + capture.beginStep({ stepNumber, messageList }); + return { modelSettings: { temperature: 0.7, maxRetries: 1 } }; + }, + processLLMRequest: ({ prompt }) => ({ + prompt: [...prompt, { role: "system", content: "LATE" }], + }), + }, + ], + }); + const output = await agent.stream("question"); + await output.consumeStream(); + expect(await output.text).toBe("done"); + await capture.drain(); + expect(failed).toHaveBeenCalledTimes(1); + const first = required(failed.mock.calls[0])?.[0]; + const second = required(capture.takeSuccessful()); + expect(first.externalId).not.toBe(second.externalId); + expect([first.attemptNumber, second.attemptNumber]).toEqual([1, 2]); + for (const [index, evidence] of [first, second].entries()) { + expect(decodeMemoryValue(evidence.inputs)).toEqual({ + prompt: required(received[index]).prompt, + tools: required(received[index]).tools, + toolChoice: required(received[index]).toolChoice, + }); + expect(evidence.modelSettings.temperature).toBe(0.7); + expect(JSON.stringify(evidence.inputs)).toContain("LATE"); + } +}, 10000); + +it("uses public tags and message sources for provenance", () => { + const messageList = new MessageList(); + messageList.addSystem("application", "application"); + messageList.addSystem("remembered", "memory"); + messageList.add([{ role: "user", content: "old" }], "memory"); + messageList.add([{ role: "user", content: "extra" }], "context"); + const capture = createRequestCapture({ + invocationId: "sources", + getMemoryRevision: () => 0, + }); + capture.beginStep({ stepNumber: 0, messageList }); + const model = capture.instrumentModel({ + specificationVersion: "v2", + modelId: "actor", + provider: "fixture", + doGenerate: async () => ({}), + }); + return model.doGenerate().then(() => { + const evidence = required(capture.takeSuccessful()); + expect(evidence.provenance).toMatchObject({ + systemMessages: [{ tag: "application" }, { tag: "memory" }], + messages: [{ source: "memory" }, { source: "context" }], + }); + }); +}); + +it("contains capture and telemetry failures while preserving native errors", async () => { + const error = new Error("provider failed"); + const onCaptureError = vi.fn(() => { + throw new Error("diagnostic failure"); + }); + const capture = createRequestCapture({ + invocationId: "failure", + getMemoryRevision: () => { + throw new Error("revision failed"); + }, + onFailedAttempt: async () => { + throw new Error("write failed"); + }, + onCaptureError, + }); + capture.beginStep({ stepNumber: 0 }); + const model = capture.instrumentModel({ + specificationVersion: "v2", + modelId: "actor", + provider: "fixture", + doStream: async (_args: unknown) => { + throw error; + }, + }); + await expect(model.doStream({ prompt: [] })).rejects.toBe(error); + await expect(capture.drain()).resolves.toBeUndefined(); + expect(onCaptureError).toHaveBeenCalled(); +}); + +it("marks unsupported or bounded evidence incomplete without failing native calls", async () => { + const capture = createRequestCapture({ + invocationId: "bounded", + getMemoryRevision: () => 0, + recordingLimits: { maxStringChars: 10 }, + }); + capture.beginStep({ stepNumber: 0 }); + const model = capture.instrumentModel({ + specificationVersion: "v2", + modelId: "actor", + provider: "fixture", + doGenerate: async (_args: unknown) => "native", + }); + await expect( + model.doGenerate({ + prompt: [{ role: "system", content: "a".repeat(100) }], + tools: [], + toolChoice: { type: "auto" }, + }), + ).resolves.toBe("native"); + expect(capture.takeSuccessful()).toMatchObject({ complete: false }); +}); + +it("retains unfinished calls and captures independently replaced models", async () => { + const capture = createRequestCapture({ + invocationId: "unfinished", + getMemoryRevision: () => 1, + }); + for (const stepNumber of [0, 1]) { + capture.beginStep({ stepNumber }); + const original = { + specificationVersion: "v2", + modelId: `model-${stepNumber}`, + provider: "fixture", + doStream: async () => textStream("unread"), + }; + const model = capture.instrumentModel(original); + expect(capture.instrumentModel(model)).toBe(model); + await model.doStream(); + } + expect(capture.flushUnfinished()).toMatchObject([ + { stepNumber: 0, modelId: "model-0", attemptNumber: 1 }, + { stepNumber: 1, modelId: "model-1", attemptNumber: 1 }, + ]); + expect(capture.flushUnfinished()).toEqual([]); + expect(capture.takeSuccessful()).toBeUndefined(); +}); + +it("excludes nested transport metadata and credentials from provider settings", async () => { + const capture = createRequestCapture({ + invocationId: "redaction", + getMemoryRevision: () => 0, + }); + capture.beginStep({ stepNumber: 0 }); + const model = capture.instrumentModel({ + specificationVersion: "v3", + modelId: "actor", + provider: "fixture", + doGenerate: async (_args: unknown) => "native", + }); + await model.doGenerate({ + prompt: [], + providerOptions: { + vendor: { headers: { "x-custom-access": "SECRET_HEADER" } }, + }, + }); + const evidence = required(capture.takeSuccessful()); + expect(evidence.complete).toBe(false); + expect(JSON.stringify(evidence)).not.toContain("SECRET_HEADER"); + await model.doGenerate({ + prompt: [], + providerOptions: { vendor: { apiKey: "SECRET_KEY" } }, + }); + const credentials = required(capture.takeSuccessful()); + expect(credentials.complete).toBe(false); + expect(JSON.stringify(credentials)).not.toContain("SECRET_KEY"); +}); + +it("keeps ordinary multi-kilobyte prompts complete within replay payload bounds", async () => { + const capture = createRequestCapture({ + invocationId: "long", + getMemoryRevision: () => 0, + }); + capture.beginStep({ stepNumber: 0 }); + const model = capture.instrumentModel({ + specificationVersion: "v2", + modelId: "actor", + provider: "fixture", + doGenerate: async (_args: unknown) => "native", + }); + const prompt = [{ role: "system", content: "x".repeat(10000) }]; + await model.doGenerate({ prompt }); + const evidence = required(capture.takeSuccessful()); + expect(evidence.complete).toBe(true); + expect(decodeMemoryValue(evidence.inputs)).toMatchObject({ prompt }); +}); + +it("enriches the ordinary LLM node without duplicating nodes or changing legacy input", async () => { + const api = installTestApi(); + const recorder = await RunRecorder.create({ + adapterVersion: "test", + agentId: AGENT_ID, + client: new KitaruClient({ apiUrl: "https://api.example" }), + effectiveInput: "input", + framework: "mastra", + requestedModelId: "requested", + }); + const capture = createRequestCapture({ + invocationId: "enrich", + getMemoryRevision: () => 7, + }); + capture.beginStep({ stepNumber: 0 }); + await capture + .instrumentModel({ + specificationVersion: "v2", + provider: "fixture", + modelId: "actor", + doGenerate: async (_args: unknown) => ({}), + }) + .doGenerate({ prompt: [{ role: "user", content: [] }], temperature: 0.6 }); + const evidence = required(capture.takeSuccessful()); + await recordStep( + recorder.state, + textStep("enriched"), + undefined, + undefined, + evidence, + ); + await recordStep(recorder.state, textStep("legacy")); + const nodes = api + .nodeBatches() + .flat() + .filter((node) => node.node_type === "llm_call"); + expect(nodes).toHaveLength(2); + expect(nodes[0]).toMatchObject({ + external_id: evidence.externalId, + inputs: evidence.inputs, + model_params: { temperature: 0.6 }, + attributes: { + invocation_id: "enrich", + memory_revision: 7, + attempt_number: 1, + }, + }); + expect(nodes[1]).toMatchObject({ + external_id: "response-legacy", + inputs: null, + }); +}); From bdfd2b984b4bbf8853291cf238e6300168815b40 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 19:39:16 +0200 Subject: [PATCH 006/115] Pin file and skill content for Mastra replay --- packages/mastra/src/stateful-files.ts | 53 ++++++++ packages/mastra/src/stateful-workspace.ts | 117 ++++++++++++++++++ packages/mastra/test/stateful-files.test.ts | 25 ++++ .../mastra/test/stateful-workspace.test.ts | 36 ++++++ 4 files changed, 231 insertions(+) create mode 100644 packages/mastra/src/stateful-files.ts create mode 100644 packages/mastra/src/stateful-workspace.ts create mode 100644 packages/mastra/test/stateful-files.test.ts create mode 100644 packages/mastra/test/stateful-workspace.test.ts diff --git a/packages/mastra/src/stateful-files.ts b/packages/mastra/src/stateful-files.ts new file mode 100644 index 000000000..176cd5f1b --- /dev/null +++ b/packages/mastra/src/stateful-files.ts @@ -0,0 +1,53 @@ +import type { MastraRecordedFile } from "./memory-snapshot.js"; + +export interface ResolvedMemoryFile { + bytes: Uint8Array; + mediaType: string; +} + +export type MemoryFileResolver = (url: string) => Promise; + +/** Bind processors to the invocation's declared content, with no live fallback. */ +export function restoreCapturedFiles(recorded: readonly MastraRecordedFile[]) { + const files = recorded.map((file) => ({ + ...file, + bytes: new Uint8Array(file.bytes), + })); + const lookup = new Map(files.map((file) => [file.url, file])); + if (lookup.size !== files.length) + throw new Error("Duplicate recorded file URL"); + return { + files, + resolveFile: async (url: string): Promise => { + const file = lookup.get(url); + if (!file) + throw new Error( + "Unsupported Mastra memory replay: file URL was not recorded.", + ); + return { bytes: new Uint8Array(file.bytes), mediaType: file.mediaType }; + }, + }; +} + +/** Fetch each declared file before recording the immutable session input. */ +export async function createCapturedFiles( + urls: readonly string[], + resolveFile: MemoryFileResolver, +) { + const files: MastraRecordedFile[] = []; + for (const url of new Set(urls)) { + const resolved = await resolveFile(url); + if ( + !(resolved.bytes instanceof Uint8Array) || + typeof resolved.mediaType !== "string" || + !resolved.mediaType + ) + throw new TypeError("File resolver must return bytes and mediaType"); + files.push({ + url, + bytes: new Uint8Array(resolved.bytes), + mediaType: resolved.mediaType, + }); + } + return restoreCapturedFiles(files); +} diff --git a/packages/mastra/src/stateful-workspace.ts b/packages/mastra/src/stateful-workspace.ts new file mode 100644 index 000000000..345274e88 --- /dev/null +++ b/packages/mastra/src/stateful-workspace.ts @@ -0,0 +1,117 @@ +import { createHash } from "node:crypto"; +import { lstat, readdir, readFile } from "node:fs/promises"; +import { join, posix } from "node:path"; +import { + type SkillSource, + type SkillSourceEntry, + Workspace, +} from "@mastra/core/workspace"; +import { MAX_RECORDED_PAYLOAD_CHARS } from "@zenml-io/kitaru/adapter"; + +export interface SkillsManifest { + files: { path: string; length: number; sha256: string }[]; +} + +/** Read and pin an artifact's skills. Native tools subsequently read only these bytes. */ +export async function loadSkillsWorkspace( + skillsDirectory: string, + expectedManifest?: SkillsManifest, +) { + const files = new Map(); + const directories = new Map(); + let size = 0; + async function visit(relative: string): Promise { + const path = join(skillsDirectory, relative); + const info = await lstat(path); + if (info.isSymbolicLink()) + throw new Error("Unsupported Mastra skills symlink"); + if (info.isDirectory()) { + const entries = await readdir(path, { withFileTypes: true }); + directories.set( + relative, + entries.map((entry) => ({ + name: entry.name, + type: entry.isDirectory() ? "directory" : "file", + })), + ); + for (const entry of entries.sort((a, b) => + a.name < b.name ? -1 : a.name > b.name ? 1 : 0, + )) + await visit(posix.join(relative, entry.name)); + } else if (info.isFile()) { + if (size + info.size > MAX_RECORDED_PAYLOAD_CHARS) + throw new Error("Skills content exceeds the supported replay limit"); + const content = await readFile(path); + size += content.length; + files.set(relative, content); + } else throw new Error("Unsupported Mastra skills file type"); + if ( + files.size + directories.size > 10_000 || + size > MAX_RECORDED_PAYLOAD_CHARS + ) + throw new Error("Skills content exceeds the supported replay limit"); + } + await visit("."); + const manifest: SkillsManifest = { + files: [...files].map(([path, content]) => ({ + path, + length: content.length, + sha256: createHash("sha256").update(content).digest("hex"), + })), + }; + if ( + expectedManifest && + JSON.stringify(expectedManifest) !== JSON.stringify(manifest) + ) + throw new Error( + "Unsupported Mastra memory replay: skills artifact changed.", + ); + function normalize(path: string): string { + const normalized = posix.normalize(path.replaceAll("\\", "/")); + if ( + posix.isAbsolute(normalized) || + normalized === ".." || + normalized.startsWith("../") + ) + throw new Error("Skill path is outside the pinned artifact"); + return normalized; + } + const source: SkillSource = { + async exists(path) { + const key = normalize(path); + return files.has(key) || directories.has(key); + }, + async stat(path) { + const key = normalize(path); + const directory = directories.has(key); + const file = files.get(key); + if (!directory && !file) + throw new Error("Skill path is absent from pinned artifact"); + return { + name: posix.basename(key), + type: directory ? "directory" : "file", + size: file?.length ?? 0, + createdAt: new Date(0), + modifiedAt: new Date(0), + }; + }, + async readFile(path) { + const file = files.get(normalize(path)); + if (!file) throw new Error("Skill file is absent from pinned artifact"); + return Buffer.from(file); + }, + async readdir(path) { + const entries = directories.get(normalize(path)); + if (!entries) + throw new Error("Skill directory is absent from pinned artifact"); + return entries.map((entry) => ({ ...entry })); + }, + async realpath(path) { + return normalize(path); + }, + }; + return { + manifest, + workspace: new Workspace({ skills: ["."], skillSource: source }), + }; +} diff --git a/packages/mastra/test/stateful-files.test.ts b/packages/mastra/test/stateful-files.test.ts new file mode 100644 index 000000000..aeaf8238b --- /dev/null +++ b/packages/mastra/test/stateful-files.test.ts @@ -0,0 +1,25 @@ +import { expect, it, vi } from "vitest"; +import { + createCapturedFiles, + restoreCapturedFiles, +} from "../src/stateful-files.js"; + +it("preloads once and supplies independent historical bytes without network fallback", async () => { + const resolver = vi.fn(async () => ({ + bytes: new Uint8Array([0, 255, 1]), + mediaType: "application/pdf", + })); + const captured = await createCapturedFiles( + ["https://files.invalid/a", "https://files.invalid/a"], + resolver, + ); + const replay = restoreCapturedFiles(captured.files); + (await captured.resolveFile("https://files.invalid/a")).bytes[0] = 9; + expect((await replay.resolveFile("https://files.invalid/a")).bytes).toEqual( + new Uint8Array([0, 255, 1]), + ); + await expect( + replay.resolveFile("https://files.invalid/missing"), + ).rejects.toThrow(/not recorded/); + expect(resolver).toHaveBeenCalledTimes(1); +}); diff --git a/packages/mastra/test/stateful-workspace.test.ts b/packages/mastra/test/stateful-workspace.test.ts new file mode 100644 index 000000000..087c239bc --- /dev/null +++ b/packages/mastra/test/stateful-workspace.test.ts @@ -0,0 +1,36 @@ +import { mkdir, mkdtemp, rm, symlink, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { expect, it } from "vitest"; +import { loadSkillsWorkspace } from "../src/stateful-workspace.js"; + +it("hashes skills by content and refuses changed content or symlinks", async () => { + const root = await mkdtemp(join(tmpdir(), "kitaru-memory-skills-")); + try { + await mkdir(join(root, "triage")); + const path = join(root, "triage", "SKILL.md"); + await writeFile( + path, + "---\nname: triage\ndescription: Support triage\n---\nHistorical instructions.\n", + ); + const baseline = await loadSkillsWorkspace(root); + expect(baseline.manifest.files[0]?.path).toBe("triage/SKILL.md"); + expect(await baseline.workspace.skills?.list()).toEqual( + expect.arrayContaining([expect.objectContaining({ name: "triage" })]), + ); + expect( + (await loadSkillsWorkspace(root, baseline.manifest)).manifest, + ).toEqual(baseline.manifest); + await writeFile(path, "changed"); + expect( + JSON.stringify(await baseline.workspace.skills?.get("triage")), + ).toContain("Historical instructions."); + await expect(loadSkillsWorkspace(root, baseline.manifest)).rejects.toThrow( + /changed/, + ); + await symlink(path, join(root, "linked.md")); + await expect(loadSkillsWorkspace(root)).rejects.toThrow(/symlink/); + } finally { + await rm(root, { recursive: true, force: true }); + } +}); From b7bc1d26c069a7bd5bfbfb955ef48ae713c90473 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 19:43:48 +0200 Subject: [PATCH 007/115] Link memory changes to provider attempts --- packages/mastra/src/memory-replay.ts | 11 ++++++----- packages/mastra/src/request-capture.ts | 13 ++++++++++++- packages/mastra/src/stateful-workspace.ts | 7 ++----- packages/mastra/test/request-capture.test.ts | 6 ++++++ 4 files changed, 26 insertions(+), 11 deletions(-) diff --git a/packages/mastra/src/memory-replay.ts b/packages/mastra/src/memory-replay.ts index 5ddccc8ca..1cebcc662 100644 --- a/packages/mastra/src/memory-replay.ts +++ b/packages/mastra/src/memory-replay.ts @@ -1,11 +1,6 @@ import { createRequire } from "node:module"; import type { MastraModelConfig } from "@mastra/core/llm"; import type { MemoryConfigInternal } from "@mastra/core/memory"; -import { - standardSchemaToJSONSchema, - toStandardSchema, -} from "@mastra/core/schema"; -import { InMemoryStore, MastraCompositeStore } from "@mastra/core/storage"; import { createMemoryCaptureBinding, createProcessLocalMemoryAccess, @@ -86,6 +81,9 @@ export function serializeMemoryConfiguration( checkConfiguration(config); const copy: Record = { ...config }; if (config.workingMemory?.schema) { + const { standardSchemaToJSONSchema, toStandardSchema } = createRequire( + import.meta.url, + )("@mastra/core/schema") as typeof import("@mastra/core/schema"); copy.workingMemory = { ...config.workingMemory, schema: standardSchemaToJSONSchema( @@ -159,6 +157,9 @@ export async function createIsolatedMemoryReplay( options.resolveModel, ); const { Memory } = await import("@mastra/memory"); + const { InMemoryStore, MastraCompositeStore } = await import( + "@mastra/core/storage" + ); const store = new InMemoryStore(); const domain = store.stores.memory; if (!domain) return unsupported("Native in-memory storage is unavailable."); diff --git a/packages/mastra/src/request-capture.ts b/packages/mastra/src/request-capture.ts index d44ebb878..3eef999e7 100644 --- a/packages/mastra/src/request-capture.ts +++ b/packages/mastra/src/request-capture.ts @@ -70,6 +70,7 @@ function asRecord(value: unknown): Record { /** Capture the final provider arguments without consuming or replacing its output stream. */ export function createRequestCapture(options: RequestCaptureOptions) { + let currentRequestId: string | undefined; let context: { stepNumber: number; provenance: JsonValue; @@ -243,6 +244,7 @@ export function createRequestCapture(options: RequestCaptureOptions) { ) { return async (...args: unknown[]) => { const evidence = capture(target, key, args[0]); + currentRequestId = evidence.externalId; const attempt = { evidence, returned: false }; unfinished.set(evidence.externalId, attempt); try { @@ -295,7 +297,16 @@ export function createRequestCapture(options: RequestCaptureOptions) { while (writes.size > 0) await Promise.all(writes); } - return { beginStep, instrumentModel, takeSuccessful, flushUnfinished, drain }; + return { + beginStep, + instrumentModel, + takeSuccessful, + flushUnfinished, + drain, + get currentRequestId() { + return currentRequestId; + }, + }; } /** Attributes shared by normal step completion and failed or unfinished attempt records. */ diff --git a/packages/mastra/src/stateful-workspace.ts b/packages/mastra/src/stateful-workspace.ts index 345274e88..7940941e6 100644 --- a/packages/mastra/src/stateful-workspace.ts +++ b/packages/mastra/src/stateful-workspace.ts @@ -1,11 +1,7 @@ import { createHash } from "node:crypto"; import { lstat, readdir, readFile } from "node:fs/promises"; import { join, posix } from "node:path"; -import { - type SkillSource, - type SkillSourceEntry, - Workspace, -} from "@mastra/core/workspace"; +import type { SkillSource, SkillSourceEntry } from "@mastra/core/workspace"; import { MAX_RECORDED_PAYLOAD_CHARS } from "@zenml-io/kitaru/adapter"; export interface SkillsManifest { @@ -17,6 +13,7 @@ export async function loadSkillsWorkspace( skillsDirectory: string, expectedManifest?: SkillsManifest, ) { + const { Workspace } = await import("@mastra/core/workspace"); const files = new Map(); const directories = new Map(); let size = 0; diff --git a/packages/mastra/test/request-capture.test.ts b/packages/mastra/test/request-capture.test.ts index eac072d10..261e93ef1 100644 --- a/packages/mastra/test/request-capture.test.ts +++ b/packages/mastra/test/request-capture.test.ts @@ -101,6 +101,7 @@ it.each(["v2", "v3", "v4"])( it("captures retries after late prompt and settings changes without changing the provider arguments", async () => { const received: ModelCall[] = []; + const requestIds: (string | undefined)[] = []; const failed = vi.fn(); const capture = createRequestCapture({ invocationId: "retry", @@ -112,6 +113,7 @@ it("captures retries after late prompt and settings changes without changing the provider: "fixture", doStream: async (args) => { received.push(args); + requestIds.push(capture.currentRequestId); if (received.length === 1) throw new APICallError({ message: "Retry", @@ -143,6 +145,9 @@ it("captures retries after late prompt and settings changes without changing the }); const output = await agent.stream("question"); await output.consumeStream(); + expect(requestIds).toHaveLength(2); + expect(typeof requestIds[0]).toBe("string"); + expect(requestIds[0]).not.toBe(requestIds[1]); expect(await output.text).toBe("done"); await capture.drain(); expect(failed).toHaveBeenCalledTimes(1); @@ -151,6 +156,7 @@ it("captures retries after late prompt and settings changes without changing the expect(first.externalId).not.toBe(second.externalId); expect([first.attemptNumber, second.attemptNumber]).toEqual([1, 2]); for (const [index, evidence] of [first, second].entries()) { + expect(requestIds[index]).toBe(evidence.externalId); expect(decodeMemoryValue(evidence.inputs)).toEqual({ prompt: required(received[index]).prompt, tools: required(received[index]).tools, From 68cc5135527acc73f99e06a231ae78a96b2c9b29 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 19:49:31 +0200 Subject: [PATCH 008/115] Enable isolated Mastra memory stream replay --- packages/mastra/src/index.ts | 7 + packages/mastra/src/replay-guards.ts | 2 +- packages/mastra/src/stateful-agent.ts | 572 ++++++++++++++++ packages/mastra/src/stateful-tools.ts | 137 ++++ packages/mastra/src/stream-recording.ts | 69 +- packages/mastra/src/tool-policies.ts | 3 +- packages/mastra/test/processor-replay.test.ts | 431 ++++++++++++ .../mastra/test/stateful-overrides.test.ts | 614 ++++++++++++++++++ 8 files changed, 1814 insertions(+), 21 deletions(-) create mode 100644 packages/mastra/src/stateful-agent.ts create mode 100644 packages/mastra/src/stateful-tools.ts create mode 100644 packages/mastra/test/processor-replay.test.ts create mode 100644 packages/mastra/test/stateful-overrides.test.ts diff --git a/packages/mastra/src/index.ts b/packages/mastra/src/index.ts index a809a245b..13d831dcf 100644 --- a/packages/mastra/src/index.ts +++ b/packages/mastra/src/index.ts @@ -31,6 +31,13 @@ export type { RunnableMastraScorer, } from "./scorers.js"; export { createMastraEvaluator } from "./scorers.js"; +export type { + MemoryReplayAgentBindings, + MemoryReplayAgentFactory, + MemoryReplayAgentOptions, +} from "./stateful-agent.js"; + +export { createMemoryReplayAgent } from "./stateful-agent.js"; export type { ConfiguredAfterToolCall, ConfiguredBeforeToolCall, diff --git a/packages/mastra/src/replay-guards.ts b/packages/mastra/src/replay-guards.ts index 89ed5da93..34c008b24 100644 --- a/packages/mastra/src/replay-guards.ts +++ b/packages/mastra/src/replay-guards.ts @@ -58,7 +58,7 @@ function formatMastraToolName(toolName: string): string { return formatted.slice(0, 63); } -function assertStableToolName(toolName: string): void { +export function assertStableToolName(toolName: string): void { const runtimeName = formatMastraToolName(toolName); if (runtimeName !== toolName) { throw new ToolPolicyError( diff --git a/packages/mastra/src/stateful-agent.ts b/packages/mastra/src/stateful-agent.ts new file mode 100644 index 000000000..e8c3189b5 --- /dev/null +++ b/packages/mastra/src/stateful-agent.ts @@ -0,0 +1,572 @@ +import { createRequire } from "node:module"; +import { Agent, type AgentConfig } from "@mastra/core/agent"; +import type { MastraModelConfig } from "@mastra/core/llm"; +import type { Mastra } from "@mastra/core/mastra"; +import type { MemoryConfigInternal } from "@mastra/core/memory"; +import type { InputProcessor } from "@mastra/core/processors"; +import { RequestContext } from "@mastra/core/request-context"; +import type { MemoryStorage } from "@mastra/core/storage"; +import type { Memory } from "@mastra/memory"; +import { + type JsonValue, + KitaruClient, + type SessionNodeCreateRequest, +} from "@zenml-io/kitaru"; +import { + type AdapterRunState, + normalizeRecordingLimits, + parseModelSettings, + ROOT_NODE_EXTERNAL_ID, + resolveReplayContext, +} from "@zenml-io/kitaru/adapter"; +import { + createMemoryCaptureBinding, + type MastraExclusiveMemoryAccess, + type MastraMemoryCaptureBinding, + type MastraMemoryMutation, +} from "./memory-binding.js"; +import { + assertMemoryReplayVersions, + createIsolatedMemoryReplay, + getMemoryModelId, + serializeMemoryConfiguration, +} from "./memory-replay.js"; +import { + createMemoryReplayEnvelope, + decodeMemoryValue, + encodeMemoryValue, + type MastraMemorySnapshot, + MEMORY_REPLAY_KEY, + restoreMemoryReplayEnvelope, +} from "./memory-snapshot.js"; +import { assertStableToolName } from "./replay-guards.js"; +import { + createRequestCapture, + type RequestEvidence, + requestEvidenceAttributes, +} from "./request-capture.js"; +import { createCapturedFiles, restoreCapturedFiles } from "./stateful-files.js"; +import { + bindMemoryToolIdentity, + createStatefulToolProcessors, +} from "./stateful-tools.js"; +import { loadSkillsWorkspace } from "./stateful-workspace.js"; +import { + StatefulRecordingError, + streamWithRecording, +} from "./stream-recording.js"; +import type { KitaruAgentOptions, RuntimeStreamOptions } from "./types.js"; + +export interface MemoryReplayAgentOptions extends KitaruAgentOptions { + /** Registry context passed to baseline dynamic configuration resolvers. */ + mastra?: Mastra; + /** Called only for new recordings. All writers must share exclusiveAccess. */ + sourceMemory(): + | Promise<{ + settled(): Promise; + domain: MemoryStorage; + configuration: MemoryConfigInternal; + exclusiveAccess: MastraExclusiveMemoryAccess; + }> + | { + settled(): Promise; + domain: MemoryStorage; + configuration: MemoryConfigInternal; + exclusiveAccess: MastraExclusiveMemoryAccess; + }; + /** Return only approved replay-relevant JSON context. Credentials are forbidden. */ + captureRequestContext?(context: RequestContext): Record; + files?: readonly string[]; + resolveFile?: ( + url: string, + ) => Promise<{ bytes: Uint8Array; mediaType: string }>; + skillsDirectory?: string; + resolveModel: (id: string) => MastraModelConfig | Promise; +} + +export interface MemoryReplayAgentBindings { + memory: Memory; + resolveFile(url: string): Promise<{ bytes: Uint8Array; mediaType: string }>; + workspace?: Awaited>["workspace"]; +} + +export type MemoryReplayAgentFactory = ( + bindings: MemoryReplayAgentBindings, +) => AgentConfig | Promise; + +function record(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} +function requireRecord(value: unknown, label: string): Record { + if (!record(value)) + throw new Error(`Unsupported Mastra memory replay: missing ${label}.`); + return value; +} +function getSelector(options: RuntimeStreamOptions) { + const memory = requireRecord( + options.memory, + "memory thread/resource selectors", + ); + const threadId = + typeof memory.thread === "string" + ? memory.thread + : record(memory.thread) + ? memory.thread.id + : undefined; + if (typeof threadId !== "string" || typeof memory.resource !== "string") + throw new Error( + "Memory replay requires explicit memory.thread and memory.resource strings.", + ); + return { threadId, resourceId: memory.resource }; +} +function assertSupportedConfiguration( + config: AgentConfig, + options: RuntimeStreamOptions, +): void { + for (const name of [ + "agents", + "workflows", + "voice", + "browser", + "backgroundTasks", + "editor", + "defaultGenerateOptionsLegacy", + "defaultStreamOptionsLegacy", + "outputProcessors", + "errorProcessors", + "hooks", + ]) { + if ((config as unknown as Record)[name] !== undefined) + throw new Error(`Unsupported memory replay configuration '${name}'.`); + } + for (const name of [ + "inputProcessors", + "outputProcessors", + "hooks", + "toolsets", + "clientTools", + "prepareStep", + "instructions", + "model", + "requestContext", + "abortSignal", + "onFinish", + "onError", + "onAbort", + "onStepFinish", + ]) { + if (options[name] !== undefined) + throw new Error(`Unsupported serialized memory replay option '${name}'.`); + } + for (const name of Object.keys(config.tools ?? {})) + assertStableToolName(name); + for (const name of ["experimental_sandbox", "delegation", "backgroundTasks"]) + if (options[name] !== undefined) + throw new Error(`Unsupported memory replay option ${name}.`); + if ( + typeof config.tools === "function" || + typeof config.inputProcessors === "function" || + typeof config.workspace === "function" + ) + throw new Error( + "Memory replay requires static tools, processors and supplied workspace bindings.", + ); + if ( + config.inputProcessors?.some( + (processor) => + !record(processor) || + "loadTools" in processor || + "createRun" in processor, + ) + ) + throw new Error( + "Memory replay processors must use the supplied dependencies and ordinary processor methods.", + ); +} + +/** Construct each streamed invocation with historical configuration and isolated replay memory. */ +export function createMemoryReplayAgent( + factory: MemoryReplayAgentFactory, + supplied: MemoryReplayAgentOptions, +): Pick { + const options = { + ...supplied, + recordingLimits: normalizeRecordingLimits(supplied.recordingLimits), + }; + const client = new KitaruClient({ + apiKey: options.apiKey, + apiUrl: options.apiUrl, + timeoutMs: options.timeoutMs, + }); + async function stream( + rawInput: unknown, + callerOptions: RuntimeStreamOptions = {}, + ): Promise { + assertMemoryReplayVersions(); + const { resolveModelConfig } = await import("@mastra/core/llm"); + const { MastraCompositeStore } = await import("@mastra/core/storage"); + const startedAt = new Date().toISOString(); + const invocationId = globalThis.crypto.randomUUID(); + const replay = await resolveReplayContext({ + allowedReplayModels: options.allowedReplayModels, + callerInput: encodeMemoryValue(rawInput), + client, + requestedModelId: options.requestedModelId, + }); + const historical = restoreMemoryReplayEnvelope(replay.effectiveInput); + const invocationInput = + historical?.rawInput ?? + decodeMemoryValue(replay.effectiveRuntimeInput as JsonValue); + if (Boolean(replay.spec) !== Boolean(historical)) + throw new Error( + "Memory replay requires a complete version-2 recorded invocation and an active Kitaru replay.", + ); + if (replay.override?.prompt != null) + throw new Error( + "Memory replay supports system_prompt overrides; replacing raw invocation input requires a new recording.", + ); + const abort = new AbortController(); + let state: AdapterRunState | undefined; + let requestCapture: ReturnType | undefined; + const getState = () => { + if (!state) throw new Error("Memory recorder has not initialized."); + return state; + }; + const writeNode = async (node: SessionNodeCreateRequest) => { + const active = getState(); + await active.enqueueStep(async () => { + await active.client.upsertSessionNodes(active.sessionId, { + nodes: [node], + }); + }); + }; + const recordMutation = async (event: MastraMemoryMutation) => + writeNode({ + external_id: event.id, + parent_external_id: ROOT_NODE_EXTERNAL_ID, + node_type: "span", + name: "memory_mutation", + status: event.complete ? "completed" : "failed", + inputs: event.arguments, + outputs: event.result, + attributes: { + invocation_id: event.invocationId, + memory_revision: event.revision, + memory_method: event.method, + request_id: event.requestId ?? null, + evidence_complete: event.complete, + }, + }); + let runtime: { + memory: Memory; + binding: MastraMemoryCaptureBinding; + initialSnapshot: MastraMemorySnapshot | undefined; + finish(): Promise; + }; + if (historical) { + runtime = await createIsolatedMemoryReplay({ + invocationId, + initialSnapshot: historical.initialSnapshot, + configuration: requireRecord( + historical.configuration.memoryConfig, + "memory configuration", + ), + resolveModel: options.resolveModel, + recordMutation, + getRequestId: () => requestCapture?.currentRequestId, + }); + } else { + const source = await options.sourceMemory(); + const selector = getSelector(callerOptions); + const binding = createMemoryCaptureBinding({ + invocationId, + ...selector, + domain: source.domain, + exclusiveAccess: source.exclusiveAccess, + recordMutation, + getRequestId: () => requestCapture?.currentRequestId, + }); + const { Memory } = await import("@mastra/memory"); + const memory = new Memory({ + storage: new MastraCompositeStore({ + id: `kitaru-baseline-${invocationId}`, + domains: { memory: binding.domain }, + }), + options: source.configuration, + }); + const initialSnapshot = await binding.captureInitial(source); + let finished: Promise | undefined; + runtime = { + memory, + binding, + initialSnapshot, + finish() { + finished ??= (async () => { + try { + await memory.settled(); + await binding.drain(); + } finally { + await binding.release(); + } + })(); + return finished; + }, + }; + } + try { + const files = historical + ? restoreCapturedFiles(historical.files) + : await createCapturedFiles( + options.files ?? [], + options.resolveFile ?? + (async () => { + throw new Error("Missing controlled file resolver."); + }), + ); + const workspace = options.skillsDirectory + ? await loadSkillsWorkspace( + options.skillsDirectory, + historical?.configuration.workspaceManifest as Parameters< + typeof loadSkillsWorkspace + >[1], + ) + : undefined; + if (historical?.configuration.workspaceManifest && !workspace) + throw new Error("Recorded skills workspace is missing."); + const owned = bindMemoryToolIdentity(runtime.memory); + const config = await factory({ + memory: owned.memory, + resolveFile: files.resolveFile, + workspace: workspace?.workspace, + }); + if (config.memory !== undefined && config.memory !== owned.memory) + throw new Error("Agent factory must use its supplied Memory instance."); + if ( + config.workspace !== undefined && + config.workspace !== workspace?.workspace + ) + throw new Error( + "Agent factory must use its supplied pinned workspace.", + ); + const liveContext = callerOptions.requestContext ?? new RequestContext(); + const recordedContext = + historical?.requestContext ?? + options.captureRequestContext?.(liveContext) ?? + Object.fromEntries(liveContext.entries()); + // Validate before any dynamic resolver can observe an unrecordable value. + const safeContext = requireRecord( + decodeMemoryValue(encodeMemoryValue(recordedContext)), + "request context", + ); + const requestContext = new RequestContext(); + for (const [key, value] of Object.entries(safeContext)) + requestContext.set(key, value); + const dynamic = { requestContext, mastra: options.mastra }; + const instructions = historical + ? historical.configuration.instructions + : typeof config.instructions === "function" + ? await config.instructions(dynamic) + : config.instructions; + const modelConfiguration = historical + ? await options.resolveModel( + replay.replacementModelId ?? + String(historical.configuration.modelId), + ) + : typeof config.model === "function" + ? await config.model(dynamic) + : config.model; + if (Array.isArray(modelConfiguration)) + throw new Error( + "Model fallback arrays are outside memory replay support.", + ); + const nativeModel = await resolveModelConfig( + modelConfiguration, + requestContext, + ); + const resolvedDefaults = historical + ? historical.configuration.defaultOptions + : typeof config.defaultOptions === "function" + ? await config.defaultOptions(dynamic) + : (config.defaultOptions ?? {}); + const defaults = requireRecord(resolvedDefaults, "default options"); + const { deepMerge } = await import("@mastra/core/utils"); + const callerData = { ...callerOptions }; + delete callerData.requestContext; + delete callerData.abortSignal; + for (const callback of ["onFinish", "onError", "onAbort", "onStepFinish"]) + delete callerData[callback]; + const effective = historical + ? requireRecord( + historical.configuration.runOptions, + "invocation options", + ) + : deepMerge(defaults, callerData); + assertSupportedConfiguration(config, effective); + if (record(effective.memory) && effective.memory.options !== undefined) + throw new Error( + "Per-call memory.options are unsupported. Set the complete memory configuration in sourceMemory instead.", + ); + const overrideSettings = parseModelSettings( + replay.override?.model_params, + ); + if (overrideSettings) + effective.modelSettings = { + ...(record(effective.modelSettings) ? effective.modelSettings : {}), + ...overrideSettings, + }; + const applicationInstructions = + replay.override?.system_prompt ?? instructions; + const configuration = { + instructions: applicationInstructions, + modelId: getMemoryModelId(modelConfiguration), + defaultOptions: defaults, + runOptions: effective, + memoryConfig: + historical?.configuration.memoryConfig ?? + serializeMemoryConfiguration(runtime.memory.getMergedThreadConfig()), + ...(workspace ? { workspaceManifest: workspace.manifest } : {}), + }; + const envelope = createMemoryReplayEnvelope({ + invocationId, + rawInput: invocationInput, + initialSnapshot: runtime.initialSnapshot as MastraMemorySnapshot, + configuration, + requestContext: safeContext, + files: files.files, + }); + if (!envelope.complete && historical) + throw new Error(envelope.reasons.join(" ")); + const writeAttempt = (evidence: RequestEvidence, error: unknown) => + writeNode({ + external_id: evidence.externalId, + parent_external_id: ROOT_NODE_EXTERNAL_ID, + node_type: "llm_call", + name: "model_request", + status: "failed", + error: error instanceof Error ? error.name : "Model request failed", + inputs: evidence.inputs, + outputs: null, + model: evidence.modelId, + model_params: evidence.modelSettings, + started_at: evidence.startedAt, + ended_at: new Date().toISOString(), + attributes: requestEvidenceAttributes(evidence), + }); + const capture = createRequestCapture({ + invocationId, + getMemoryRevision: () => runtime.binding.revision, + onFailedAttempt: writeAttempt, + onCaptureError: () => + runtime.binding.markIncomplete( + "Actor request evidence was incomplete.", + ), + }); + requestCapture = capture; + const policy = createStatefulToolProcessors({ + tokens: owned.tokens, + getState, + abort(reason) { + state?.storeFailure(reason); + abort.abort(reason); + }, + adapter: options, + }); + const requestProcessor: InputProcessor = { + id: "kitaru-effective-request", + async processInputStep(args) { + capture.beginStep({ + stepNumber: args.stepNumber, + messageList: args.messageList, + applicationInstructions, + extraContext: effective.system ?? null, + }); + const resolved = await resolveModelConfig(args.model, requestContext); + return { + model: capture.instrumentModel(resolved) as typeof args.model, + }; + }, + }; + const agent = new Agent({ + ...config, + memory: owned.memory, + workspace: workspace?.workspace, + instructions: applicationInstructions as AgentConfig["instructions"], + model: capture.instrumentModel(nativeModel) as MastraModelConfig, + defaultOptions: {}, + inputProcessors: [ + policy.first, + ...((config.inputProcessors as InputProcessor[]) ?? []), + policy.last, + requestProcessor, + ], + }); + const runtimeOptions: RuntimeStreamOptions = { + ...effective, + onFinish: callerOptions.onFinish, + onError: callerOptions.onError, + onAbort: callerOptions.onAbort, + onStepFinish: callerOptions.onStepFinish, + requestContext, + abortSignal: callerOptions.abortSignal + ? AbortSignal.any([callerOptions.abortSignal, abort.signal]) + : abort.signal, + }; + const version = createRequire(import.meta.url)("../package.json") as { + version: string; + }; + return await streamWithRecording({ + adapterVersion: version.version, + agent: agent as unknown as Parameters< + typeof streamWithRecording + >[0]["agent"], + callerMessages: invocationInput, + callerOptions: runtimeOptions, + client, + options, + replayInput: replay.effectiveInput, + replay, + requestedModelId: + replay.replacementModelId ?? String(configuration.modelId), + sessionName: options.sessionName, + startedAt, + stateful: { + input: { [MEMORY_REPLAY_KEY]: envelope }, + initialize(value) { + state = value; + }, + takeRequest() { + const evidence = capture.takeSuccessful(); + if (evidence && !evidence.complete) + runtime.binding.markIncomplete( + "Actor request evidence was incomplete.", + ); + return evidence; + }, + async finish() { + await runtime.finish(); + await capture.drain(); + for (const pending of capture.flushUnfinished()) + await writeAttempt( + pending, + new Error("Unfinished model attempt"), + ); + if (runtime.binding.incompleteReasons.length) { + const message = runtime.binding.incompleteReasons.join(" "); + if ( + runtime.binding.incompleteReasons.includes( + "Native memory storage mutation failed.", + ) + ) + throw new Error(message); + throw new StatefulRecordingError(message); + } + }, + }, + }); + } catch (error) { + await runtime.finish(); + throw error; + } + } + return { stream: stream as Agent["stream"] }; +} diff --git a/packages/mastra/src/stateful-tools.ts b/packages/mastra/src/stateful-tools.ts new file mode 100644 index 000000000..be6f27274 --- /dev/null +++ b/packages/mastra/src/stateful-tools.ts @@ -0,0 +1,137 @@ +import type { InputProcessor } from "@mastra/core/processors"; +import type { Memory } from "@mastra/memory"; +import { + type AdapterRunState, + assertInterceptableTool, + assertSupportedToolPolicy, +} from "@zenml-io/kitaru/adapter"; +import { assertStableToolName } from "./replay-guards.js"; +import { createToolHooks } from "./tool-policies.js"; +import type { KitaruAgentOptions } from "./types.js"; + +function record(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +/** Brand only tools returned by this invocation's native Memory, before native conversion. */ +export function bindMemoryToolIdentity(memory: Memory) { + const tokens = new Set(); + const copies = new WeakMap(); + const bound = new Proxy(memory, { + get(target, key) { + if (key === "listTools") + return (...args: Parameters) => + Object.fromEntries( + Object.entries(target.listTools(...args)).map(([name, tool]) => { + let copy = copies.get(tool); + if (!copy) { + const id = `kitaru-memory-${globalThis.crypto.randomUUID()}`; + tokens.add(id); + copy = { ...tool, id }; + copies.set(tool, copy); + } + return [name, copy]; + }), + ); + const value = Reflect.get(target, key, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + return { memory: bound, tokens }; +} + +/** Enforce policies on the final executable inventory, including processor-added tools. */ +export function createStatefulToolProcessors(options: { + tokens: ReadonlySet; + getState(): AdapterRunState; + abort(reason: unknown): void; + adapter: KitaruAgentOptions; +}) { + const trusted = new WeakSet<(...args: never[]) => unknown>(); + const wrappersByName = new Map< + string, + WeakMap< + (...args: never[]) => unknown, + (...args: unknown[]) => Promise + > + >(); + let inspected = false; + const first: InputProcessor = { + id: "kitaru-memory-tool-identity", + processInputStep({ tools }) { + if (inspected) return; + inspected = true; + for (const tool of Object.values(tools ?? {})) { + if ( + record(tool) && + typeof tool.id === "string" && + options.tokens.has(tool.id) && + typeof tool.execute === "function" + ) + trusted.add(tool.execute as (...args: never[]) => unknown); + } + }, + }; + const last: InputProcessor = { + id: "kitaru-final-tool-policy", + processInputStep({ tools }) { + const output: Record = {}; + for (const [name, tool] of Object.entries(tools ?? {})) { + if ( + !record(tool) || + typeof tool.execute !== "function" || + (tool.requireApproval !== undefined && + tool.requireApproval !== false) || + tool.hasSuspendSchema === true || + tool.type === "provider-defined" + ) + throw new Error( + `Unsupported replay tool '${name}': an ordinary interceptable executor is required.`, + ); + assertStableToolName(name); + assertInterceptableTool(name, true); + const state = options.getState(); + if (state.spec) assertSupportedToolPolicy(state.spec, name); + const wrappers = wrappersByName.get(name) ?? new WeakMap(); + wrappersByName.set(name, wrappers); + const execute = tool.execute as (...args: never[]) => unknown; + let wrapper = wrappers.get(execute); + if (!wrapper) { + const isMemory = trusted.has(execute); + wrapper = async (input: unknown, context: unknown) => { + const hooks = createToolHooks({ + state: options.getState(), + abortReplay: options.abort, + trustedMemoryTool: isMemory, + configuredBeforeToolCall: + options.adapter.configuredBeforeToolCall, + configuredAfterToolCall: options.adapter.configuredAfterToolCall, + limits: options.adapter.recordingLimits, + }); + const event = { toolName: name, input, context, metadata: {} }; + const before = await hooks.beforeToolCall?.(event); + if (before?.proceed === false) return before.output; + let result: unknown; + try { + result = await Reflect.apply(execute, tool, [input, context]); + } catch (error) { + await hooks.afterToolCall?.({ + ...event, + output: undefined, + error, + }); + throw error; + } + await hooks.afterToolCall?.({ ...event, output: result }); + return result; + }; + wrappers.set(execute, wrapper); + wrappers.set(wrapper, wrapper); + } + output[name] = { ...tool, execute: wrapper }; + } + return { tools: output }; + }, + }; + return { first, last }; +} diff --git a/packages/mastra/src/stream-recording.ts b/packages/mastra/src/stream-recording.ts index f8530185a..0b3e9291f 100644 --- a/packages/mastra/src/stream-recording.ts +++ b/packages/mastra/src/stream-recording.ts @@ -1,6 +1,7 @@ import { createRequire } from "node:module"; import type { JsonValue, KitaruClient } from "@zenml-io/kitaru"; import { + type AdapterRunState, parseModelSettings, type ReplayContext, type RunRecorder, @@ -21,6 +22,7 @@ import { assertReplayToolCoverage, stripLiveMemoryOptions, } from "./replay-guards.js"; +import type { RequestEvidence } from "./request-capture.js"; import { type RecordedStep, recordStep } from "./step-recorder.js"; import { createToolHooks } from "./tool-policies.js"; import type { @@ -37,7 +39,17 @@ type StreamAgent = { listConfiguredInputProcessors?: (requestContext?: unknown) => unknown; }; +export class StatefulRecordingError extends Error {} + +export interface StatefulStreamRecording { + input: JsonValue; + initialize(state: AdapterRunState): void; + takeRequest(): RequestEvidence | undefined; + finish(): Promise; +} + interface StreamRecordingOptions { + stateful?: StatefulStreamRecording; adapterVersion: string; agent: StreamAgent; callerMessages: unknown; @@ -196,6 +208,7 @@ class StreamLifecycle { constructor( readonly recorder: RunRecorder, readonly options: KitaruAgentOptions, + readonly stateful?: StatefulStreamRecording, ) {} async record(step: RecordedStep): Promise { @@ -207,6 +220,7 @@ class StreamLifecycle { step, this.options.costCalculator, this.options.recordingLimits, + this.stateful?.takeRequest(), ), ); this.#stepTail = write.catch((error: unknown) => { @@ -273,6 +287,13 @@ class StreamLifecycle { private async finalize(complete: boolean, result?: unknown): Promise { this.#finalizerPromise ??= (async () => { await this.#stepTail; + try { + await this.stateful?.finish(); + } catch (error) { + if (error instanceof StatefulRecordingError) + this.requestRecordingFailure("complete", error); + else this.requestFailure(error); + } if (this.#recordingError !== undefined) { await this.cleanup(this.#recordingError.error, "recording"); return; @@ -331,6 +352,7 @@ class StreamLifecycle { } export async function streamWithRecording({ + stateful, adapterVersion, agent, callerMessages, @@ -355,8 +377,10 @@ export async function streamWithRecording({ : {}; const { deepMerge } = await import("@mastra/core/utils"); const effective = deepMerge(defaults, callerOptions) as RuntimeStreamOptions; - const needsContext = hasMemoryOptions(effective); - const contextMessages = restoreConversationContext(replayInput); + const needsContext = !stateful && hasMemoryOptions(effective); + const contextMessages = stateful + ? undefined + : restoreConversationContext(replayInput); if (contextMessages && !replay.spec) { throw new Error( "A recorded Mastra conversation context can only be restored through a Kitaru replay. Start a replay for this session to keep live memory isolated.", @@ -380,7 +404,7 @@ export async function streamWithRecording({ effective.context = []; } let replayAbortController: AbortController | undefined; - if (replay.replacementModelId !== undefined) { + if (!stateful && replay.replacementModelId !== undefined) { if (!options.resolveModel) { throw new Error( `Cannot resolve replacement model '${replay.replacementModelId}' without resolveModel`, @@ -395,6 +419,7 @@ export async function streamWithRecording({ effective.model = resolved; } if ( + !stateful && replay.override?.system_prompt !== undefined && replay.override.system_prompt !== null ) { @@ -419,13 +444,15 @@ export async function streamWithRecording({ replayAbortController.signal, ]) : replayAbortController.signal; - stripLiveMemoryOptions(effective); - await assertReplayToolCoverage({ - agent, - methodType: "stream", - runtimeOptions: effective, - spec: replay.spec, - }); + if (!stateful) { + stripLiveMemoryOptions(effective); + await assertReplayToolCoverage({ + agent, + methodType: "stream", + runtimeOptions: effective, + spec: replay.spec, + }); + } effective.toolCallConcurrency = 1; } const processors = @@ -434,6 +461,7 @@ export async function streamWithRecording({ ? await agent.listConfiguredInputProcessors(effective.requestContext) : undefined); if ( + !stateful && processors != null && (!Array.isArray(processors) || processors.length > 0) ) { @@ -444,9 +472,10 @@ export async function streamWithRecording({ await assertSupportedOptions(agent, effective); let recordedInput = - needsContext && !replay.spec + stateful?.input ?? + (needsContext && !replay.spec ? createContextInput(replayInput) - : replayInput; + : replayInput); let lifecycle: StreamLifecycle | undefined; let initializePromise: Promise | undefined; const initialize = (): Promise => { @@ -473,7 +502,8 @@ export async function streamWithRecording({ await recorder.fail(error).catch(() => undefined); throw error; } - lifecycle = new StreamLifecycle(recorder, options); + stateful?.initialize(recorder.state); + lifecycle = new StreamLifecycle(recorder, options, stateful); return lifecycle; })(); return initializePromise; @@ -592,12 +622,13 @@ export async function streamWithRecording({ limits: options.recordingLimits, state: (await initialize()).recorder.state, }); - effective.hooks = { - beforeToolCall: async (event) => - (await getToolHooks()).beforeToolCall?.(event), - afterToolCall: async (event) => - (await getToolHooks()).afterToolCall?.(event), - }; + if (!stateful) + effective.hooks = { + beforeToolCall: async (event) => + (await getToolHooks()).beforeToolCall?.(event), + afterToolCall: async (event) => + (await getToolHooks()).afterToolCall?.(event), + }; try { return await agent.stream(effectiveMessages, effective); diff --git a/packages/mastra/src/tool-policies.ts b/packages/mastra/src/tool-policies.ts index 9940c3943..7db6e7f59 100644 --- a/packages/mastra/src/tool-policies.ts +++ b/packages/mastra/src/tool-policies.ts @@ -21,6 +21,7 @@ import type { } from "./types.js"; interface ToolHookOptions { + trustedMemoryTool?: boolean; abortReplay?: (reason: unknown) => void; callerHooks?: ToolHooks; configuredAfterToolCall?: ConfiguredAfterToolCall; @@ -107,7 +108,7 @@ export function createToolHooks(options: ToolHookOptions): ToolHooks { `tool '${hookContext.toolName}' input`, limits, ); - if (state.spec) { + if (state.spec && !options.trustedMemoryTool) { const decision = await decideToolCall(state, { callId, inputs: converted.value, diff --git a/packages/mastra/test/processor-replay.test.ts b/packages/mastra/test/processor-replay.test.ts new file mode 100644 index 000000000..2b903ef65 --- /dev/null +++ b/packages/mastra/test/processor-replay.test.ts @@ -0,0 +1,431 @@ +import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { InputProcessor } from "@mastra/core/processors"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { createTool } from "@mastra/core/tools"; +import { afterEach, expect, it, vi } from "vitest"; +import { z } from "zod/v4"; +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, + MEMORY_REPLAY_KEY, +} from "../src/index.js"; +import { + createMemoryRuntime, + FILE_URL, + RESOURCE, + seedMemory, + streamParts, + THREAD, + textStream, +} from "./helpers/memory-agent.js"; +import { + AGENT_ID, + installTestApi, + ORIGINAL_SESSION_ID, + REPLAY_ID, +} from "./helpers.js"; + +afterEach(() => { + vi.unstubAllEnvs(); + vi.unstubAllGlobals(); + vi.restoreAllMocks(); +}); + +it("runs the native file processor with historical bytes, skills and complete large request evidence", async () => { + const directory = await mkdtemp(join(tmpdir(), "kitaru-stateful-files-")); + await mkdir(join(directory, "triage")); + await writeFile( + join(directory, "triage", "SKILL.md"), + "---\nname: triage\ndescription: HISTORICAL_SKILL.\n---\nUse historical knowledge.\n", + ); + const runtime = createMemoryRuntime({ messageTokens: 100000 }); + await seedMemory(runtime); + const nativeFetch = globalThis.fetch; + const api = installTestApi({ + replaySpec: { + id: REPLAY_ID, + baseline_session_id: ORIGINAL_SESSION_ID, + status: "pending", + override: { system_prompt: "Changed application instruction" }, + tool_policy: { + default: { type: "history", on_miss: "fail", scope: "baseline" }, + tools: {}, + }, + }, + }); + const apiFetch = globalThis.fetch; + vi.stubGlobal("fetch", (( + input: Parameters[0], + init: Parameters[1], + ) => + String(input).startsWith("data:") + ? nativeFetch(input, init) + : apiFetch(input, init)) as typeof fetch); + const requests: unknown[] = []; + const bytes = new Uint8Array(40000).fill(65); + const fetchFile = vi.fn(async () => ({ + bytes, + mediaType: "application/pdf", + })); + const processFile = vi.fn(); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async (args) => { + requests.push(args); + return textStream("done"); + }, + }); + const adapter = createMemoryReplayAgent( + ({ memory, workspace, resolveFile }) => ({ + id: "files", + name: "Files", + instructions: "Original application instruction", + model, + memory, + workspace, + inputProcessors: [ + { + id: "file-content", + async processInput({ messages }) { + processFile(); + const content = await resolveFile(FILE_URL); + return messages.map((message) => ({ + ...message, + content: { + ...message.content, + parts: message.content.parts.map((part) => + part.type === "file" + ? { + ...part, + data: Buffer.from(content.bytes).toString("base64"), + } + : part, + ), + }, + })); + }, + }, + ], + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: async (id) => + id.includes("observer") + ? runtime.observer.model + : id.includes("reflector") + ? runtime.reflector.model + : model, + files: [FILE_URL], + resolveFile: fetchFile, + skillsDirectory: directory, + }, + ); + try { + const baseline = await adapter.stream( + [ + { + role: "user", + content: [ + { type: "text", text: "Please read" }, + { + type: "file", + data: new URL(FILE_URL), + mimeType: "application/pdf", + }, + ], + }, + ], + { + memory: { thread: THREAD, resource: RESOURCE }, + system: "Extra context", + }, + ); + await baseline.consumeStream(); + const input = api.calls.find( + (call) => call.path === "/api/v1/sessions" && call.method === "POST", + )?.body?.inputs; + expect( + (input as Record)[MEMORY_REPLAY_KEY] + ?.complete, + ).toBe(true); + expect( + api.calls.filter((call) => call.method === "PATCH").at(-1)?.body?.status, + ).toBe("completed"); + fetchFile.mockRejectedValue(new Error("Original file unavailable")); + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(input)); + const replay = await adapter.stream("ignored"); + await replay.consumeStream(); + expect(fetchFile).toHaveBeenCalledTimes(1); + expect(processFile).toHaveBeenCalledTimes(2); + const json = JSON.stringify(requests[1]); + expect(json).toContain("HISTORICAL_SKILL"); + expect(json).toContain("historical-blue"); + expect(json).toContain("Extra context"); + expect(json).toContain("Changed application instruction"); + const modelNodes = api + .nodeBatches() + .flat() + .filter((node) => node.node_type === "llm_call"); + expect(modelNodes).toHaveLength(2); + expect( + modelNodes.every( + (node) => + (node.attributes as Record).request_complete === + true, + ), + ).toBe(true); + expect(JSON.stringify(modelNodes[1]?.inputs).length).toBeGreaterThan(40000); + } finally { + await runtime.store.close(); + await rm(directory, { recursive: true, force: true }); + } +}); + +it.each(["late", "memory-name-spoof", "copied-memory-id"])( + "applies history failure to a %s processor tool before it can execute", + async (kind) => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const api = installTestApi({ + replaySpec: { + id: REPLAY_ID, + baseline_session_id: ORIGINAL_SESSION_ID, + status: "pending", + override: null, + tool_policy: { + default: { type: "history", on_miss: "fail", scope: "baseline" }, + tools: {}, + }, + }, + }); + let replaying = false; + const execute = vi.fn(async () => ({ sideEffect: true })); + const toolName = kind === "late" ? "lateTool" : "updateWorkingMemory"; + const processor: InputProcessor = { + id: "late-tool", + processInputStep({ tools }) { + return { + tools: { + ...tools, + [toolName]: + kind === "copied-memory-id" + ? { + ...(tools?.updateWorkingMemory as Record), + execute, + } + : createTool({ + id: toolName, + description: "An external tool", + inputSchema: z.object({}), + execute, + }), + }, + }; + }, + }; + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => + replaying + ? streamParts( + [ + { + type: "tool-call", + toolName, + toolCallId: "external-call", + input: "{}", + }, + ], + "tool-calls", + ) + : textStream("baseline"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "late", + name: "Late", + instructions: "Use a tool", + model, + memory, + inputProcessors: [processor], + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: async (id) => + id.includes("observer") + ? runtime.observer.model + : id.includes("reflector") + ? runtime.reflector.model + : model, + }, + ); + const baseline = await adapter.stream("record", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await baseline.consumeStream(); + const input = api.calls.find( + (call) => call.path === "/api/v1/sessions" && call.method === "POST", + )?.body?.inputs; + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(input)); + replaying = true; + const result = await adapter.stream("ignored"); + await result.consumeStream(); + await vi.waitFor(() => + expect( + api.calls.filter((call) => call.method === "PATCH").at(-1)?.body + ?.status, + ).toBe("failed"), + ); + expect(execute).not.toHaveBeenCalled(); + expect( + api.calls.filter((call) => call.path.endsWith("tool-lookup")), + ).toHaveLength(1); + await runtime.store.close(); + }, +); + +it("keeps policies separate when processor tools share the same executor", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const staticPolicy = (result: string) => ({ + type: "static", + on_miss: "fail", + cases: [{ match: null, match_mode: "exact", result }], + }); + const api = installTestApi({ + replaySpec: { + id: REPLAY_ID, + baseline_session_id: ORIGINAL_SESSION_ID, + status: "pending", + override: null, + tool_policy: { + default: { type: "history", on_miss: "fail", scope: "baseline" }, + tools: { + firstAlias: staticPolicy("FIRST"), + secondAlias: staticPolicy("SECOND"), + }, + }, + }, + }); + let replaying = false; + let step = 0; + const execute = vi.fn(async () => "live"); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => + replaying && step++ === 0 + ? streamParts( + [ + { + type: "tool-call", + toolName: "firstAlias", + toolCallId: "first", + input: "{}", + }, + { + type: "tool-call", + toolName: "secondAlias", + toolCallId: "second", + input: "{}", + }, + ], + "tool-calls", + ) + : textStream("done"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "aliases", + name: "Aliases", + instructions: "Use tools", + model, + memory, + inputProcessors: [ + { + id: "aliases", + processInputStep({ tools }) { + return { + tools: { + ...tools, + firstAlias: createTool({ + id: "firstAlias", + description: "First", + inputSchema: z.object({}), + execute, + }), + secondAlias: createTool({ + id: "secondAlias", + description: "Second", + inputSchema: z.object({}), + execute, + }), + }, + }; + }, + }, + ], + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: async (id) => + id.includes("observer") + ? runtime.observer.model + : id.includes("reflector") + ? runtime.reflector.model + : model, + }, + ); + const baseline = await adapter.stream("record", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await baseline.consumeStream(); + const input = api.calls.find( + (call) => call.path === "/api/v1/sessions" && call.method === "POST", + )?.body?.inputs; + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(input)); + replaying = true; + const replay = await adapter.stream("ignored"); + await replay.consumeStream(); + const tools = api + .nodeBatches(api.sessionIds[1]) + .flat() + .filter((node) => node.node_type === "tool_call"); + expect(tools.map((node) => [node.name, node.outputs])).toEqual([ + ["firstAlias", "FIRST"], + ["secondAlias", "SECOND"], + ]); + expect(execute).not.toHaveBeenCalled(); + await runtime.store.close(); +}); diff --git a/packages/mastra/test/stateful-overrides.test.ts b/packages/mastra/test/stateful-overrides.test.ts new file mode 100644 index 000000000..30d95afd0 --- /dev/null +++ b/packages/mastra/test/stateful-overrides.test.ts @@ -0,0 +1,614 @@ +import { Agent } from "@mastra/core/agent"; +import type { InputProcessor } from "@mastra/core/processors"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { afterEach, expect, it, vi } from "vitest"; +import { + createMemoryRuntime, + RESOURCE, + seedMemory, + streamParts, + THREAD, + textStream, +} from "./helpers/memory-agent.js"; + +it("retains invocation memory tool identity through public native conversion", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const marker = "kitaru-owned-memory-identity"; + const memory = new Proxy(runtime.memory, { + get(target, key) { + if (key === "listTools") + return (...args: Parameters) => + Object.fromEntries( + Object.entries(target.listTools(...args)).map(([name, tool]) => [ + name, + { ...tool, id: marker }, + ]), + ); + const value = Reflect.get(target, key, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + const identities: unknown[] = []; + const executions: string[] = []; + const initial: InputProcessor = { + id: "identity-first", + processInputStep({ tools }) { + identities.push((tools?.updateWorkingMemory as { id?: string })?.id); + }, + }; + const final: InputProcessor = { + id: "identity-last", + processInputStep({ tools }) { + const wrapped = Object.fromEntries( + Object.entries(tools ?? {}).map(([name, tool]) => { + const native = tool as { execute?: (...args: unknown[]) => unknown }; + return [ + name, + { + ...native, + execute: async (...args: unknown[]) => { + executions.push(name); + return native.execute?.(...args); + }, + }, + ]; + }), + ); + return { tools: wrapped }; + }, + }; + let calls = 0; + const agent = new Agent({ + id: "marker-proof", + name: "Marker proof", + instructions: "Update memory", + memory, + inputProcessors: [initial, final], + model: new MastraLanguageModelV2Mock({ + doStream: async () => + ++calls === 1 + ? streamParts( + [ + { + type: "tool-call", + toolCallId: "memory-call", + toolName: "updateWorkingMemory", + input: JSON.stringify({ memory: { preference: "green" } }), + }, + ], + "tool-calls", + ) + : textStream("done"), + }), + }); + const result = await agent.stream("Green please", { + memory: { thread: THREAD, resource: RESOURCE }, + maxSteps: 3, + }); + await result.consumeStream(); + await runtime.memory.settled(); + expect(identities).toEqual([marker, marker]); + expect(executions).toEqual(["updateWorkingMemory"]); + expect( + await runtime.memory.getWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + }), + ).toContain("green"); + await runtime.store.close(); +}); + +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, + MEMORY_REPLAY_KEY, +} from "../src/index.js"; +import { + AGENT_ID, + installTestApi, + ORIGINAL_SESSION_ID, + REPLAY_ID, +} from "./helpers.js"; + +afterEach(() => { + vi.unstubAllEnvs(); + vi.unstubAllGlobals(); + vi.restoreAllMocks(); +}); + +it("records and replays native evolving memory without re-resolving live configuration", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + let replaying = false; + const dynamicCalls: string[] = []; + const requests: unknown[] = []; + let calls = 0; + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async (args) => { + requests.push(args); + return ++calls % 2 === 1 + ? streamParts( + [ + { + type: "tool-call", + toolCallId: `call-${calls}`, + toolName: "updateWorkingMemory", + input: JSON.stringify({ + memory: { + preference: replaying ? "replay-green" : "baseline-red", + }, + }), + }, + ], + "tool-calls", + ) + : textStream("done"); + }, + }); + const source = vi.fn(() => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + })); + const baselineApi = installTestApi({ + replaySpec: { + id: REPLAY_ID, + baseline_session_id: ORIGINAL_SESSION_ID, + status: "pending", + override: { system_prompt: "New instructions" }, + tool_policy: { + default: { type: "history", scope: "baseline", on_miss: "fail" }, + tools: {}, + }, + }, + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "stateful", + name: "Stateful", + memory, + instructions: () => { + dynamicCalls.push("instructions"); + return "Original instructions"; + }, + model: () => { + dynamicCalls.push("model"); + return model; + }, + defaultOptions: () => { + dynamicCalls.push("defaults"); + return { maxSteps: 3 }; + }, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + apiKey: "fixture", + requestedModelId: "fixture/actor", + sourceMemory: source, + resolveModel: async (id) => + id.includes("observer") + ? runtime.observer.model + : id.includes("reflector") + ? runtime.reflector.model + : model, + }, + ); + const baseline = await adapter.stream("Green please", { + memory: { thread: THREAD, resource: RESOURCE }, + system: "Extra context. Original instructions", + }); + await baseline.consumeStream(); + expect(dynamicCalls).toEqual(["instructions", "model", "defaults"]); + const recorded = baselineApi.calls.find( + (call) => call.method === "POST" && call.path === "/api/v1/sessions", + )?.body?.inputs; + expect(recorded).toHaveProperty(MEMORY_REPLAY_KEY); + expect( + baselineApi.calls.filter((call) => call.method === "PATCH").at(-1)?.body + ?.status, + ).toBe("completed"); + const baselineNodes = baselineApi.nodeBatches().flat(); + expect(baselineNodes.some((node) => node.name === "memory_mutation")).toBe( + true, + ); + expect( + baselineNodes + .filter((node) => node.node_type === "llm_call") + .every((node) => node.inputs), + ).toBe(true); + await runtime.memory.updateWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + workingMemory: JSON.stringify({ preference: "production-today" }), + }); + replaying = true; + source.mockImplementation(() => { + throw new Error("Production source used during replay"); + }); + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(recorded)); + const replayApi = baselineApi; + const replay = await adapter.stream("ignored", { + memory: { thread: "today", resource: "today" }, + }); + await replay.consumeStream(); + expect(dynamicCalls).toEqual(["instructions", "model", "defaults"]); + expect(source).toHaveBeenCalledTimes(1); + expect(JSON.stringify(requests[2])).toContain("historical-blue"); + expect(JSON.stringify(requests[2])).toContain("New instructions"); + expect( + JSON.stringify(requests[2]).match(/Original instructions/g), + ).toHaveLength(1); + expect(JSON.stringify(requests[2])).toContain("Extra context"); + expect(JSON.stringify(requests[3])).toContain("replay-green"); + expect( + replayApi.calls.filter((call) => call.method === "PATCH").at(-1)?.body + ?.status, + ).toBe("completed"); + expect( + replayApi.calls.some((call) => call.path.endsWith("tool-lookup")), + ).toBe(false); + expect( + await runtime.memory.getWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + }), + ).toContain("production-today"); + await runtime.store.close(); +}); + +it("waits for owned native observation before reporting completion and releases its lease", async () => { + let release!: () => void; + let started!: () => void; + const blocked = new Promise((resolve) => { + release = resolve; + }); + const observing = new Promise((resolve) => { + started = resolve; + }); + const runtime = createMemoryRuntime({ + messageTokens: 10000, + observerWait: async () => { + started(); + await blocked; + }, + }); + await seedMemory(runtime); + const api = installTestApi(); + const lease = createProcessLocalMemoryAccess(); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("done"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "settled", + name: "Settled", + instructions: "Answer", + model, + memory, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: lease, + }), + resolveModel: () => model, + }, + ); + try { + const result = await adapter.stream("Observe this message", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + const consuming = result.consumeStream(); + await observing; + await new Promise((resolve) => setImmediate(resolve)); + expect( + api.calls.some( + (call) => call.method === "PATCH" && call.body?.status === "completed", + ), + ).toBe(false); + release(); + await consuming; + expect( + api.calls.filter((call) => call.method === "PATCH").at(-1)?.body?.status, + ).toBe("completed"); + const releaseLease = await lease.acquire({ + threadId: THREAD, + resourceId: RESOURCE, + }); + await releaseLease(); + expect(runtime.observer.calls.length).toBeGreaterThan(0); + expect( + api + .nodeBatches() + .flat() + .some( + (node) => + node.name === "memory_mutation" && + (node.attributes as Record).memory_method === + "updateBufferedObservations", + ), + ).toBe(true); + } finally { + release(); + await runtime.memory.settled(); + await runtime.store.close(); + } +}); + +it("keeps baseline output when initial snapshot evidence fails and rejects its incomplete replay", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const api = installTestApi(); + vi.spyOn(runtime.domain, "getResourceById").mockRejectedValueOnce( + new Error("Snapshot read failed"), + ); + const source = vi.fn(() => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + })); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("native output"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "partial", + name: "Partial", + instructions: "Answer", + memory, + model, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: source, + resolveModel: () => model, + }, + ); + const output = await adapter.stream("Hello", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await output.consumeStream(); + expect(await output.text).toBe("native output"); + const input = api.calls.find( + (call) => call.method === "POST" && call.path === "/api/v1/sessions", + )?.body?.inputs as Record; + expect(input[MEMORY_REPLAY_KEY]?.complete).toBe(false); + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(input)); + await expect(adapter.stream("ignored")).rejects.toThrow(/incomplete/i); + expect(source).toHaveBeenCalledTimes(1); + await runtime.store.close(); +}); + +it("rebuilds observation/reflection on the replay trajectory with an overridden actor model", async () => { + const runtime = createMemoryRuntime({ messageTokens: 600 }); + await seedMemory(runtime); + const api = installTestApi({ + replaySpec: { + id: REPLAY_ID, + baseline_session_id: ORIGINAL_SESSION_ID, + status: "pending", + override: { model: "fixture/replacement" }, + tool_policy: { default: { type: "passthrough" }, tools: {} }, + }, + }); + const requests: unknown[] = []; + let calls = 0; + const replacement = new MastraLanguageModelV2Mock({ + modelId: "replacement", + provider: "fixture", + doStream: async (args) => { + requests.push(args); + calls++; + return calls === 1 + ? streamParts( + [ + { + type: "tool-call", + toolCallId: "change-memory", + toolName: "updateWorkingMemory", + input: JSON.stringify({ + memory: { preference: "replay-green" }, + }), + }, + ], + "tool-calls", + ) + : calls === 2 + ? streamParts( + [ + { + type: "tool-call", + toolCallId: "evidence", + toolName: "readEvidence", + input: "{}", + }, + ], + "tool-calls", + ) + : textStream("evolved"); + }, + }); + const original = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("original"), + }); + const { createTool } = await import("@mastra/core/tools"); + const { z } = await import("zod/v4"); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "evolving", + name: "Evolving", + instructions: "Use evidence", + model: original, + memory, + defaultOptions: { maxSteps: 5 }, + tools: { + readEvidence: createTool({ + id: "readEvidence", + description: "Read conversation evidence", + inputSchema: z.object({}), + execute: async () => { + await memory.settled(); + return { + evidence: "The replay user now prefers green. ".repeat(400), + }; + }, + }), + }, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + allowedReplayModels: ["fixture/replacement"], + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: async (id) => + id.includes("observer") + ? runtime.observer.model + : id.includes("reflector") + ? runtime.reflector.model + : id === "fixture/replacement" + ? replacement + : original, + }, + ); + const baseline = await adapter.stream("Baseline turn", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await baseline.consumeStream(); + const input = api.calls.find( + (call) => call.method === "POST" && call.path === "/api/v1/sessions", + )?.body?.inputs; + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(input)); + const nativeStream = Agent.prototype.stream; + let nativeResult: unknown; + vi.spyOn(Agent.prototype, "stream").mockImplementation(async function ( + this: Agent, + ...args + ) { + nativeResult = await Reflect.apply(nativeStream, this, args); + return nativeResult as Awaited>; + }); + const output = await adapter.stream("ignored"); + expect(output).toBe(nativeResult); + await output.consumeStream(); + expect(await output.text).toBe("evolved"); + expect(requests).toHaveLength(3); + expect(runtime.observer.calls.length).toBeGreaterThan(0); + expect(runtime.reflector.calls.length).toBeGreaterThan(0); + expect(JSON.stringify(requests[2])).toContain("REFLECTED_REPLAY"); + const resultNodes = api.nodeBatches(api.sessionIds[1]).flat(); + expect( + resultNodes + .filter((node) => node.node_type === "llm_call") + .every((node) => node.model === "replacement"), + ).toBe(true); + expect( + resultNodes.some( + (node) => + node.name === "memory_mutation" && + (node.attributes as Record).memory_method === + "createReflectionGeneration", + ), + ).toBe(true); + expect( + api.calls.filter((call) => call.method === "PATCH").at(-1)?.body?.status, + ).toBe("completed"); + expect( + await runtime.memory.getWorkingMemory({ + threadId: THREAD, + resourceId: RESOURCE, + }), + ).toContain("historical-blue"); + await runtime.store.close(); +}); + +it("releases the source lease on setup failure and cancellation", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const api = installTestApi(); + const lease = createProcessLocalMemoryAccess(); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("unused"), + }); + let failSetup = true; + const adapter = createMemoryReplayAgent( + ({ memory }) => { + if (failSetup) throw new Error("Factory failed"); + return { + id: "cancelled", + name: "Cancelled", + instructions: "Answer", + model, + memory, + }; + }, + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: lease, + }), + resolveModel: () => model, + }, + ); + await expect( + adapter.stream("Hello", { memory: { thread: THREAD, resource: RESOURCE } }), + ).rejects.toThrow("Factory failed"); + const firstRelease = await lease.acquire({ + threadId: THREAD, + resourceId: RESOURCE, + }); + await firstRelease(); + failSetup = false; + const abort = new AbortController(); + abort.abort(new Error("Cancelled")); + const output = await adapter.stream("Hello", { + memory: { thread: THREAD, resource: RESOURCE }, + abortSignal: abort.signal, + }); + await output.consumeStream(); + await vi.waitFor(() => + expect( + api.calls.filter((call) => call.method === "PATCH").at(-1)?.body?.status, + ).toBe("failed"), + ); + const release = await lease.acquire({ + threadId: THREAD, + resourceId: RESOURCE, + }); + await release(); + await runtime.store.close(); +}); From c8c0e12f7bdac389e8270a69f4acdd498070cce5 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 19:53:27 +0200 Subject: [PATCH 009/115] Simplify memory source and tool hook setup --- packages/mastra/src/stateful-agent.ts | 21 ++++++++------------- packages/mastra/src/stateful-tools.ts | 17 ++++++++--------- 2 files changed, 16 insertions(+), 22 deletions(-) diff --git a/packages/mastra/src/stateful-agent.ts b/packages/mastra/src/stateful-agent.ts index e8c3189b5..fdf5e173a 100644 --- a/packages/mastra/src/stateful-agent.ts +++ b/packages/mastra/src/stateful-agent.ts @@ -57,23 +57,18 @@ import { } from "./stream-recording.js"; import type { KitaruAgentOptions, RuntimeStreamOptions } from "./types.js"; +interface MastraMemorySource { + settled(): Promise; + domain: MemoryStorage; + configuration: MemoryConfigInternal; + exclusiveAccess: MastraExclusiveMemoryAccess; +} + export interface MemoryReplayAgentOptions extends KitaruAgentOptions { /** Registry context passed to baseline dynamic configuration resolvers. */ mastra?: Mastra; /** Called only for new recordings. All writers must share exclusiveAccess. */ - sourceMemory(): - | Promise<{ - settled(): Promise; - domain: MemoryStorage; - configuration: MemoryConfigInternal; - exclusiveAccess: MastraExclusiveMemoryAccess; - }> - | { - settled(): Promise; - domain: MemoryStorage; - configuration: MemoryConfigInternal; - exclusiveAccess: MastraExclusiveMemoryAccess; - }; + sourceMemory(): MastraMemorySource | Promise; /** Return only approved replay-relevant JSON context. Credentials are forbidden. */ captureRequestContext?(context: RequestContext): Record; files?: readonly string[]; diff --git a/packages/mastra/src/stateful-tools.ts b/packages/mastra/src/stateful-tools.ts index be6f27274..7fe71b783 100644 --- a/packages/mastra/src/stateful-tools.ts +++ b/packages/mastra/src/stateful-tools.ts @@ -98,16 +98,15 @@ export function createStatefulToolProcessors(options: { let wrapper = wrappers.get(execute); if (!wrapper) { const isMemory = trusted.has(execute); + const hooks = createToolHooks({ + state, + abortReplay: options.abort, + trustedMemoryTool: isMemory, + configuredBeforeToolCall: options.adapter.configuredBeforeToolCall, + configuredAfterToolCall: options.adapter.configuredAfterToolCall, + limits: options.adapter.recordingLimits, + }); wrapper = async (input: unknown, context: unknown) => { - const hooks = createToolHooks({ - state: options.getState(), - abortReplay: options.abort, - trustedMemoryTool: isMemory, - configuredBeforeToolCall: - options.adapter.configuredBeforeToolCall, - configuredAfterToolCall: options.adapter.configuredAfterToolCall, - limits: options.adapter.recordingLimits, - }); const event = { toolName: name, input, context, metadata: {} }; const before = await hooks.beforeToolCall?.(event); if (before?.proceed === false) return before.output; From 96d7b1fd373cae3dd6f6254f800134804b019d44 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 19:55:03 +0200 Subject: [PATCH 010/115] Identify extra context in memory request evidence --- packages/mastra/src/stateful-agent.ts | 5 ++++- packages/mastra/test/stateful-overrides.test.ts | 9 ++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/packages/mastra/src/stateful-agent.ts b/packages/mastra/src/stateful-agent.ts index fdf5e173a..5b1de65ec 100644 --- a/packages/mastra/src/stateful-agent.ts +++ b/packages/mastra/src/stateful-agent.ts @@ -473,7 +473,10 @@ export function createMemoryReplayAgent( stepNumber: args.stepNumber, messageList: args.messageList, applicationInstructions, - extraContext: effective.system ?? null, + extraContext: { + system: effective.system ?? null, + context: effective.context ?? [], + }, }); const resolved = await resolveModelConfig(args.model, requestContext); return { diff --git a/packages/mastra/test/stateful-overrides.test.ts b/packages/mastra/test/stateful-overrides.test.ts index 30d95afd0..923e693e0 100644 --- a/packages/mastra/test/stateful-overrides.test.ts +++ b/packages/mastra/test/stateful-overrides.test.ts @@ -200,7 +200,9 @@ it("records and replays native evolving memory without re-resolving live configu ); const baseline = await adapter.stream("Green please", { memory: { thread: THREAD, resource: RESOURCE }, - system: "Extra context. Original instructions", + context: [ + { role: "system", content: "Extra context. Original instructions" }, + ], }); await baseline.consumeStream(); expect(dynamicCalls).toEqual(["instructions", "model", "defaults"]); @@ -221,6 +223,11 @@ it("records and replays native evolving memory without re-resolving live configu .filter((node) => node.node_type === "llm_call") .every((node) => node.inputs), ).toBe(true); + expect( + baselineNodes.find((node) => node.node_type === "llm_call")?.attributes, + ).toHaveProperty("prompt_provenance.extraContext.context", [ + { role: "system", content: "Extra context. Original instructions" }, + ]); await runtime.memory.updateWorkingMemory({ threadId: THREAD, resourceId: RESOURCE, From e56e22cb6202a06f7497c4a6b3e2dc3e5d4f98b0 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 20:03:32 +0200 Subject: [PATCH 011/115] Accept SDK job objects in MCP activity pages --- src/kitaru/mcp/models/common.py | 3 ++ tests/mcp/test_handlers_protocol.py | 44 +++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/src/kitaru/mcp/models/common.py b/src/kitaru/mcp/models/common.py index 1144e7656..974ac7693 100644 --- a/src/kitaru/mcp/models/common.py +++ b/src/kitaru/mcp/models/common.py @@ -130,6 +130,9 @@ class RegistryReadResult(ToolResult): # field descriptions in the discovery schema. The literal mirrors JobKind while # avoiding a separate enum definition in this already budget-constrained union. class _MCPJob(JobResponse): + # Paginated SDK reads contain JobResponse instances, not dictionaries. + model_config = ConfigDict(from_attributes=True) + id: uuid.UUID kind: Literal["session_run", "import", "evaluation", "replay"] diff --git a/tests/mcp/test_handlers_protocol.py b/tests/mcp/test_handlers_protocol.py index 96eaf029f..8b7c8cdcc 100644 --- a/tests/mcp/test_handlers_protocol.py +++ b/tests/mcp/test_handlers_protocol.py @@ -24,6 +24,7 @@ ImportResponse, ) from kitaru.api_models.v1.investigation import InvestigationSessionResponse +from kitaru.api_models.v1.job import JobResponse from kitaru.api_models.v1.session import SessionDetailResponse, TokenUsage from kitaru.api_models.v1.session_node import SessionNodeListParams, SessionNodeResponse from kitaru.api_models.v1.tag import TagResponse @@ -334,6 +335,49 @@ async def test_activity_returns_exactly_one_page_and_preserves_cursor() -> None: assert result.page.has_more is True +@pytest.mark.parametrize("operation", ["list", "list_children"]) +async def test_activity_job_pages_accept_sdk_responses(operation: str) -> None: + now = datetime.now(UTC) + job = JobResponse( + id=uuid.uuid4(), + owner_id=uuid.uuid4(), + kind="replay", + status="completed", + created=now, + updated=now, + ) + + async def list_jobs(*_args: object) -> Page[JobResponse]: + return Page(items=[job], next_cursor="next-job") + + client = SimpleNamespace( + jobs=SimpleNamespace(list=list_jobs), + experiment_runs=SimpleNamespace(list_jobs=list_jobs), + ) + server, context = build_server_context(client) + request = {"operation": operation, "size": 1} + if operation == "list": + request["kind"] = "job" + else: + request.update(kind="experiment_run_jobs", parent_id=str(uuid.uuid4())) + result = await server.call_tool( + "kitaru_activity_read", {"request": request}, context + ) + + assert isinstance(result, CallToolResult) + assert result.is_error is False + assert result.structured_content is not None + data = result.structured_content["data"] + assert data["items"][0]["id"] == str(job.id) + assert data["items"][0]["kind"] == "replay" + assert data["items"][0]["status"] == "completed" + assert data["page"] == { + "size": 1, + "next_cursor": "next-job", + "has_more": True, + } + + async def test_activity_import_get_and_list_use_the_imports_resource() -> None: client = FakeClient() state = _get_state(client) From 1f6e19ee0e7b6a9c0cae3b9f4c0455acaac2130d Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 20:05:07 +0200 Subject: [PATCH 012/115] Prove authenticated headless Mastra memory replay --- devtools/check_mastra_memory_replay.mjs | 273 ++++++++++ devtools/check_mastra_memory_replay.py | 672 ++++++++++++++++++++++++ 2 files changed, 945 insertions(+) create mode 100644 devtools/check_mastra_memory_replay.mjs create mode 100644 devtools/check_mastra_memory_replay.py diff --git a/devtools/check_mastra_memory_replay.mjs b/devtools/check_mastra_memory_replay.mjs new file mode 100644 index 000000000..a35ec87cd --- /dev/null +++ b/devtools/check_mastra_memory_replay.mjs @@ -0,0 +1,273 @@ +/** Provider-free agent entrypoint for check_mastra_memory_replay.py. */ +import assert from "node:assert/strict"; +import { readFile, writeFile } from "node:fs/promises"; +import { createRequire } from "node:module"; +import { join } from "node:path"; +import { pathToFileURL } from "node:url"; +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, +} from "../packages/mastra/dist/index.js"; + +const require = createRequire( + new URL("../packages/mastra/package.json", import.meta.url), +); +const nativeImport = (name) => + import(pathToFileURL(require.resolve(name).replace(/\.cjs$/, ".js")).href); +const [{ InMemoryStore }, { Memory }, { z }] = await Promise.all([ + nativeImport("@mastra/core/storage"), + nativeImport("@mastra/memory"), + nativeImport("zod/v4"), +]); +const directory = process.env.CHECK_DIRECTORY; +const replayId = process.env.KITARU_REPLAY_ID; +const thread = "historical-thread"; +const resource = "historical-resource"; +const fileUrl = "https://files.invalid/historical.pdf"; +const bytes = new Uint8Array(40000).fill(65); +const report = { + task_id: process.env.KITARU_TASK_ID, + replay_id: replayId ?? null, + task_inputs_in_environment: process.env.KITARU_TASK_INPUTS !== undefined, + source_calls: 0, + file_calls: 0, + dynamic_calls: 0, + processor_calls: 0, + actor_calls: 0, +}; +assert(process.env.KITARU_API_TOKEN, "Expected the ordinary worker task token"); +assert.equal( + process.env.KITARU_API_KEY, + undefined, + "Administrative key leaked to task", +); +const denied = await fetch( + `${process.env.KITARU_API_URL}/api/v1/blobs/${process.env.CHECK_UNRELATED_BLOB}/content`, + { + headers: { Authorization: `Bearer ${process.env.KITARU_API_TOKEN}` }, + }, +); +assert.equal(denied.status, 403, "Task unexpectedly read an unrelated blob"); +report.unrelated_blob_status = denied.status; +const stream = (parts, finishReason = "stop") => ({ + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: "stream-start", warnings: [] }); + for (const part of parts) controller.enqueue(part); + controller.enqueue({ + type: "finish", + finishReason, + usage: { inputTokens: 5, outputTokens: 5, totalTokens: 10 }, + }); + controller.close(); + }, + }), +}); +const text = (value) => + stream([ + { type: "text-start", id: "text" }, + { type: "text-delta", id: "text", delta: value }, + { type: "text-end", id: "text" }, + ]); +const actor = (id) => ({ + specificationVersion: "v2", + supportedUrls: {}, + doGenerate: async () => { + throw new Error("Unexpected generate call"); + }, + provider: "fixture", + modelId: id, + doStream: async (args) => { + report.actor_calls++; + const prompt = JSON.stringify(args.prompt); + assert(prompt.includes("historical-blue")); + assert(prompt.includes("Extra context")); + assert(prompt.includes("HISTORICAL_SKILL")); + assert(!prompt.includes("production-today")); + if (replayId) { + assert(prompt.includes("Replay application instruction")); + assert(!prompt.includes("Original application instruction")); + assert.equal(id, "replacement"); + } + if (report.actor_calls === 1) + return stream( + [ + { + type: "tool-call", + toolCallId: "change-preference", + toolName: "updateWorkingMemory", + input: JSON.stringify({ + memory: { + preference: replayId ? "replay-green" : "baseline-red", + }, + }), + }, + ], + "tool-calls", + ); + assert(prompt.includes(replayId ? "replay-green" : "baseline-red")); + return text("done"); + }, +}); +const models = { actor: actor("actor"), replacement: actor("replacement") }; +const store = new InMemoryStore(); +const memory = new Memory({ + storage: store, + options: { + lastMessages: 10, + semanticRecall: false, + workingMemory: { + enabled: true, + scope: "thread", + schema: z.object({ preference: z.string() }), + }, + }, +}); +const domain = store.stores.memory; +assert(domain); +await memory.createThread({ + threadId: thread, + resourceId: resource, + title: "Historical thread", +}); +const currentProduction = JSON.parse( + await readFile(join(directory, "production.json"), "utf8"), +); +await memory.updateWorkingMemory({ + threadId: thread, + resourceId: resource, + workingMemory: JSON.stringify(currentProduction), +}); +const sourceDomain = new Proxy(domain, { + get(target, key) { + const value = Reflect.get(target, key, target); + if (typeof value !== "function") return value; + return (...args) => { + if (replayId) + throw new Error( + `Production memory accessed during replay: ${String(key)}`, + ); + return Reflect.apply(value, target, args); + }; + }, +}); +const dynamic = (value) => () => { + if (replayId) + throw new Error("Live dynamic configuration was resolved during replay"); + report.dynamic_calls++; + return value; +}; +const agent = createMemoryReplayAgent( + ({ memory: isolated, resolveFile, workspace }) => ({ + id: "headless-memory", + name: "Headless memory", + memory: isolated, + workspace, + instructions: dynamic("Original application instruction"), + model: dynamic(models.actor), + defaultOptions: dynamic({ maxSteps: 3 }), + inputProcessors: [ + { + id: "historical-file", + async processInput({ messages }) { + report.processor_calls++; + const file = await resolveFile(fileUrl); + assert.deepEqual(file.bytes, bytes); + return messages.map((message) => ({ + ...message, + content: { + ...message.content, + parts: message.content.parts.map((part) => + part.type === "file" + ? { + ...part, + data: Buffer.from(file.bytes).toString("base64"), + } + : part, + ), + }, + })); + }, + }, + ], + }), + { + agentId: process.env.CHECK_AGENT_ID, + requestedModelId: "fixture/actor", + allowedReplayModels: ["fixture/replacement"], + sourceMemory: () => { + report.source_calls++; + if (replayId) + throw new Error("Production source requested during replay"); + return { + domain: sourceDomain, + configuration: memory.getMergedThreadConfig(), + settled: () => memory.settled(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }; + }, + resolveModel: (id) => { + const model = models[id.split("/").at(-1)]; + assert(model, `Unexpected model ${id}`); + return model; + }, + files: [fileUrl], + resolveFile: async () => { + report.file_calls++; + if (replayId) throw new Error("Live file requested during replay"); + return { bytes, mediaType: "application/pdf" }; + }, + skillsDirectory: join(directory, "skills"), + }, +); +try { + const result = await agent.stream( + [ + { + role: "user", + content: [ + { type: "text", text: "Update the preference using the attachment." }, + { type: "file", data: new URL(fileUrl), mimeType: "application/pdf" }, + ], + }, + ], + { + memory: { thread, resource }, + context: [{ role: "system", content: "Extra context" }], + }, + ); + await result.consumeStream(); + assert.equal(await result.text, "done"); + assert.equal(report.processor_calls, 1); + assert.equal(report.actor_calls, 2); + assert.equal(report.dynamic_calls, replayId ? 0 : 3); + assert.equal(report.source_calls, replayId ? 0 : 1); + assert.equal(report.file_calls, replayId ? 0 : 1); + if (replayId) { + assert.equal(report.task_inputs_in_environment, false); + assert.deepEqual( + JSON.parse(await readFile(join(directory, "production.json"), "utf8")), + { preference: "production-today" }, + ); + assert( + ( + await memory.getWorkingMemory({ + threadId: thread, + resourceId: resource, + }) + ).includes("production-today"), + ); + } + report.result = "passed"; +} catch (error) { + report.result = "failed"; + report.error = String(error); + throw error; +} finally { + await memory.settled(); + await store.close(); + await writeFile( + join(directory, `${report.task_id}.json`), + JSON.stringify(report, null, 2), + ); +} diff --git a/devtools/check_mastra_memory_replay.py b/devtools/check_mastra_memory_replay.py new file mode 100644 index 000000000..2f8d0f618 --- /dev/null +++ b/devtools/check_mastra_memory_replay.py @@ -0,0 +1,672 @@ +"""Prove SDK, CLI and MCP memory replay on an authenticated disposable stack. + +Run ``pnpm run build:packages`` first, then +``uv run python devtools/check_mastra_memory_replay.py``. +No model provider is called. Only this invocation's server and database are removed. +""" + +import argparse +import asyncio +import copy +import json +import os +import shutil +import subprocess +import sys +import tempfile +import uuid +from pathlib import Path +from typing import Any + +import asyncpg +from mcp.client.session import ClientSession +from mcp.client.stdio import StdioServerParameters, stdio_client +from stack import ( + DB_HOST, + DB_PORT, + DB_PWD, + DB_USER, + bootstrap_api_key, + create_database, + drop_database, + ensure_postgres, + get_free_port, + start_server, + wait_for_health, +) + +from kitaru.api_models.v1.agent import AgentCreateRequest +from kitaru.api_models.v1.agent_version import AgentVersionCreateRequest, RunSpec +from kitaru.api_models.v1.evaluator import ( + EvaluatorCreateRequest, + EvaluatorVersionCreateRequest, +) +from kitaru.api_models.v1.job import JobResponse, JobStatus +from kitaru.api_models.v1.plugin import EvaluatorConfig, ScriptPluginSource +from kitaru.api_models.v1.replay import ReplayCreateRequest +from kitaru.api_models.v1.replay_config import ( + ReplayOverride, + ToolPolicy, +) +from kitaru.api_models.v1.session import ( + SessionCreateRequest, + SessionOrigin, + SessionStatus, +) +from kitaru.api_models.v1.session_node import SessionNodeListParams +from kitaru.api_models.v1.session_run import SessionRunCreateRequest +from kitaru.client.api_client import KitaruAPIClient +from kitaru.worker import Worker, WorkerConfig + +ARTIFACT = Path(__file__).with_suffix(".mjs").resolve() +OVERRIDE = { + "system_prompt": "Replay application instruction", + "model": "fixture/replacement", +} +POLICY = { + "default": {"type": "history", "scope": "baseline", "on_miss": "fail"}, + "tools": {}, +} + + +async def await_job( + client: KitaruAPIClient, job_id: uuid.UUID, label: str, timeout: float +) -> JobResponse: + """Wait for one owned job to reach a terminal state.""" + deadline = asyncio.get_running_loop().time() + timeout + while True: + job = await client.jobs.get(job_id) + if job.status in {JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELED}: + return job + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError(f"{label} did not complete within {timeout}s") + await asyncio.sleep(0.2) + + +async def call_mcp(session: ClientSession, tool: str, request: dict[str, Any]) -> Any: + """Call a public MCP tool and require its structured success envelope.""" + result = await session.call_tool(tool, {"request": request}) + assert not result.is_error, result + envelope = result.structured_content + assert envelope and envelope["ok"], envelope + return envelope["data"] + + +async def run_cli(*args: str) -> dict[str, Any]: + """Run the installed CLI without blocking the worker's event loop.""" + process = await asyncio.create_subprocess_exec( + str(Path(sys.executable).with_name("kitaru")), + *args, + "--output", + "json", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), 60) + except BaseException: + if process.returncode is None: + process.kill() + await process.wait() + raise + assert process.returncode == 0, stderr.decode() + stdout.decode() + return json.loads(stdout) + + +async def inspect_session( + client: KitaruAPIClient, session_id: uuid.UUID +) -> dict[str, Any]: + """Require complete, hydrated request evidence and ordered memory changes.""" + session = await client.sessions.get(session_id) + assert session.status == SessionStatus.COMPLETED, session.error + nodes = [ + node + async for node in client.sessions.iter_nodes( + session_id, SessionNodeListParams(include_payloads=True, size=1) + ) + ] + model_nodes = [node for node in nodes if node.node_type == "llm_call"] + mutations = [node for node in nodes if node.name == "memory_mutation"] + assert len(model_nodes) == 2, [(node.name, node.node_type) for node in nodes] + assert mutations, "Memory changes were not recorded" + revisions = [node.attributes["memory_revision"] for node in mutations] + assert revisions == sorted(set(revisions)), revisions + assert all(node.attributes["evidence_complete"] for node in mutations) + request_ids = {node.external_id for node in model_nodes} + assert any(node.attributes["request_id"] in request_ids for node in mutations) + assert all(node.attributes["request_complete"] for node in model_nodes) + assert all(len(json.dumps(node.inputs)) > 40000 for node in model_nodes) + first_request = json.dumps(model_nodes[0].inputs) + assert "historical-blue" in first_request + assert "production-today" not in first_request + assert "Replay application instruction" in first_request + assert "Extra context" in first_request + assert "HISTORICAL_SKILL" in first_request + assert "replay-green" in json.dumps(model_nodes[1].inputs) + assert all(node.model == "replacement" for node in model_nodes) + return { + "session_id": str(session_id), + "task_id": str(session.task_id), + "node_ids": [str(node.id) for node in nodes], + "model_nodes": len(model_nodes), + "memory_nodes": len(mutations), + } + + +async def finish_replay( + client: KitaruAPIClient, replay_id: uuid.UUID +) -> dict[str, Any]: + """Wait for both replay and evaluator completion, then inspect the result.""" + replay = await client.replays.get(replay_id) + assert replay.job_id + job = await await_job(client, replay.job_id, "memory replay", 120) + if job.status != JobStatus.COMPLETED: + tasks = await client.jobs.list_tasks(job.id) + raise AssertionError( + [(task.kind, task.status, task.error) for task in tasks.items] + ) + replay = await client.replays.get(replay_id) + assert replay.status == "completed", replay + assert replay.result_session_id + return { + "replay_id": str(replay.id), + **await inspect_session(client, replay.result_session_id), + } + + +async def check_mcp( + client: KitaruAPIClient, + url: str, + directory: Path, + agent_id: uuid.UUID, + version_id: uuid.UUID, + evaluator_id: uuid.UUID, + baseline_id: uuid.UUID, +) -> dict[str, Any]: + """Create an experiment and inspect every result page over real MCP stdio.""" + mcp_console = str(Path(sys.executable).with_name("kitaru-mcp")) + with (directory / "mcp.log").open("w") as mcp_log: + parameters = StdioServerParameters( + command=mcp_console, + args=["--server", url, "--mode", "standard"], + env=dict(os.environ), + ) + async with ( + stdio_client(parameters, errlog=mcp_log) as (reader, writer), + ClientSession(reader, writer) as mcp, + ): + await mcp.initialize() + cohort = await call_mcp( + mcp, + "kitaru_cohorts_manage", + { + "operation": "create", + "agent_id": str(agent_id), + "name": "memory-cohort", + }, + ) + cohort_version = await call_mcp( + mcp, + "kitaru_cohorts_manage", + { + "operation": "create_version", + "cohort_id": cohort["id"], + "add_session_ids": [str(baseline_id)], + }, + ) + experiment = await call_mcp( + mcp, + "kitaru_experiments_manage", + { + "operation": "create", + "agent_id": str(agent_id), + "name": "memory-experiment", + "override": OVERRIDE, + "tool_policy": POLICY, + "evaluators": [{"evaluator_id": str(evaluator_id), "version": 1}], + }, + ) + started = await call_mcp( + mcp, + "kitaru_workflow_start", + { + "operation": "experiment_run", + "experiment_id": experiment["id"], + "cohort_version_id": cohort_version["id"], + "agent_version_id": str(version_id), + "baseline_evaluation_mode": "none", + }, + ) + run_id = started["result"]["id"] + deadline = asyncio.get_running_loop().time() + 120 + while True: + run = await call_mcp( + mcp, + "kitaru_activity_read", + { + "operation": "get", + "kind": "experiment_run", + "id": run_id, + }, + ) + if run["status"] in {"completed", "failed", "canceled"}: + break + assert asyncio.get_running_loop().time() < deadline, ( + "MCP experiment timed out" + ) + await asyncio.sleep(0.2) + assert run["status"] == "completed", run + children = await call_mcp( + mcp, + "kitaru_activity_read", + { + "operation": "list_children", + "kind": "experiment_run_jobs", + "parent_id": run_id, + "size": 1, + }, + ) + assert len(children["items"]) == 1 + replay_page = await call_mcp( + mcp, + "kitaru_activity_read", + { + "operation": "list", + "kind": "replay", + "size": 1, + "filter": { + "field": "experiment_run_id", + "op": "eq", + "value": run_id, + }, + }, + ) + assert len(replay_page["items"]) == 1 + mcp_replay = replay_page["items"][0] + assert mcp_replay["status"] == "completed" + mcp_result = await finish_replay(client, uuid.UUID(mcp_replay["id"])) + mcp_session = await call_mcp( + mcp, + "kitaru_activity_read", + { + "operation": "get", + "kind": "session", + "id": mcp_result["session_id"], + }, + ) + assert mcp_session["status"] == "completed" + assert mcp_session["inputs"]["mastra_memory_replay"]["complete"] + mcp_nodes = [] + cursor = None + while True: + page = await call_mcp( + mcp, + "kitaru_activity_read", + { + "operation": "list_children", + "kind": "session_nodes", + "parent_id": mcp_result["session_id"], + "size": 1, + "cursor": cursor, + "include_payloads": True, + }, + ) + mcp_nodes.extend(page["items"]) + if not page["page"]["has_more"]: + break + cursor = page["page"]["next_cursor"] + assert cursor + assert {item["id"] for item in mcp_nodes} == set(mcp_result["node_ids"]) + assert all( + item["inputs"] and item["attributes"]["request_complete"] + for item in mcp_nodes + if item["node_type"] == "llm_call" + ) + assert any( + item["name"] == "memory_mutation" + and item["attributes"]["memory_revision"] > 0 + for item in mcp_nodes + ) + parameters = StdioServerParameters( + command=mcp_console, + args=["--server", url, "--mode", "read-only"], + env=dict(os.environ), + ) + async with ( + stdio_client(parameters, errlog=mcp_log) as (reader, writer), + ClientSession(reader, writer) as mcp, + ): + await mcp.initialize() + names = {tool.name for tool in (await mcp.list_tools()).tools} + assert "kitaru_workflow_start" not in names + assert "kitaru_experiments_manage" not in names + read = await call_mcp( + mcp, + "kitaru_activity_read", + { + "operation": "get", + "kind": "replay", + "id": mcp_result["replay_id"], + }, + ) + assert read["result_session_id"] == mcp_result["session_id"] + refused = await mcp.call_tool( + "kitaru_workflow_start", + { + "request": { + "operation": "experiment_run", + "experiment_id": experiment["id"], + "cohort_version_id": cohort_version["id"], + "agent_version_id": str(version_id), + } + }, + ) + assert refused.is_error + print("MCP experiment, pagination and read-only mode passed", flush=True) + return { + "experiment_run_id": run_id, + **mcp_result, + "paginated_nodes": len(mcp_nodes), + } + + +async def check(output: Path) -> None: + """Exercise all headless paths and clean up owned resources even on failure.""" + node = shutil.which("node") + assert node, "Node 22.22 or 26 is required" + db_name = f"kitaru_memory_{uuid.uuid4().hex[:10]}" + directory = Path(tempfile.mkdtemp(prefix="kitaru-memory-proof-")) + print(f"Proof logs: {directory}", flush=True) + skills = directory / "skills" / "triage" + skills.mkdir(parents=True) + (skills / "SKILL.md").write_text( + "---\nname: triage\ndescription: HISTORICAL_SKILL.\n---\n" + "Use historical knowledge.\n" + ) + (directory / "production.json").write_text( + json.dumps({"preference": "historical-blue"}) + ) + stop = asyncio.Event() + server = None + worker_task = None + await ensure_postgres() + await create_database(db_name) + try: + port = get_free_port() + url = f"http://127.0.0.1:{port}" + server = start_server( + db_name, port, directory / "server.log", auth_scheme="local" + ) + await wait_for_health(url, server, directory / "server.log") + key = await bootstrap_api_key(url) + assert key + os.environ["KITARU_API_URL"] = url + os.environ["KITARU_API_KEY"] = key + os.environ.pop("KITARU_API_TOKEN", None) + worker = Worker( + WorkerConfig( + name=f"memory-proof-{uuid.uuid4().hex[:8]}", + concurrency=1, + poll_interval=0.1, + heartbeat_interval=0.5, + blob_cache_root=directory / "worker-blobs", + payload_cache_root=directory / "worker-payloads", + ) + ) + worker_task = asyncio.create_task(worker.run(stop)) + async with KitaruAPIClient(base_url=url, api_key=key) as client: + unrelated = await client.blobs.upload( + b"unrelated private content", media_type="text/plain" + ) + agent = await client.agents.create(AgentCreateRequest(name="memory-proof")) + version = await client.agents.create_version( + agent.id, + AgentVersionCreateRequest( + run_spec=RunSpec( + command=f'"{node}" "{ARTIFACT}"', + timeout_seconds=90, + env={ + "CHECK_AGENT_ID": str(agent.id), + "CHECK_DIRECTORY": str(directory), + "CHECK_UNRELATED_BLOB": str(unrelated.id), + }, + ) + ), + ) + job = await client.session_runs.create( + SessionRunCreateRequest( + agent_version_id=version.id, + inputs=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Update the preference.", + }, + { + "type": "file", + "data": "https://files.invalid/historical.pdf", + "mimeType": "application/pdf", + }, + ], + } + ], + ) + ) + job = await await_job(client, job.id, "baseline", 120) + if job.status != JobStatus.COMPLETED: + tasks = await client.jobs.list_tasks(job.id) + raise AssertionError( + [(task.kind, task.status, task.error) for task in tasks.items] + ) + sessions = [session async for session in client.sessions.iter()] + assert len(sessions) == 1 + baseline = await client.sessions.get(sessions[0].id) + assert baseline.inputs["mastra_memory_replay"]["complete"] is True + assert len(json.dumps(baseline.inputs).encode()) > 32768 + connection = await asyncpg.connect( + host=DB_HOST, + port=DB_PORT, + user=DB_USER, + password=DB_PWD, + database=db_name, + ) + try: + offloaded = await connection.fetchval( + "SELECT inputs_blob_id FROM session WHERE id = $1", baseline.id + ) + assert offloaded, ( + "Large session inputs did not use server payload offload" + ) + finally: + await connection.close() + (directory / "production.json").write_text( + json.dumps({"preference": "production-today"}) + ) + evaluator_blob = await client.blobs.upload( + (Path(__file__).parent / "evaluators.py").read_bytes(), + media_type="text/x-python", + filename="evaluator.py", + ) + evaluator = await client.evaluators.create( + EvaluatorCreateRequest(name="memory-completed") + ) + await client.evaluators.create_version( + evaluator.id, + EvaluatorVersionCreateRequest( + source=ScriptPluginSource( + blob_id=evaluator_blob.id, entrypoint="evaluate_outcome" + ) + ), + ) + config = EvaluatorConfig(evaluator="memory-completed", version=1) + replay = await client.replays.create( + ReplayCreateRequest( + baseline_session_id=baseline.id, + override=ReplayOverride(**OVERRIDE), + tool_policy=ToolPolicy.model_validate(POLICY), + evaluators=[config], + ) + ) + sdk = await finish_replay(client, replay.id) + print("SDK replay passed", flush=True) + created = await run_cli( + "replay", + "create", + str(baseline.id), + "--evaluator", + "memory-completed@1", + "--override", + json.dumps(OVERRIDE), + "--tool-policy", + json.dumps(POLICY), + ) + cli = await finish_replay(client, uuid.UUID(created["item"]["id"])) + cli_read = await run_cli("replay", "get", cli["replay_id"]) + assert cli_read["item"]["status"] == "completed" + assert cli_read["item"]["result_session_id"] == cli["session_id"] + cli_nodes = await run_cli( + "session", + "nodes", + cli["session_id"], + "--include-payloads", + "--size", + "100", + ) + assert {item["id"] for item in cli_nodes["items"]} == set(cli["node_ids"]) + assert all( + item["inputs"] and item["attributes"]["request_complete"] + for item in cli_nodes["items"] + if item["node_type"] == "llm_call" + ) + assert any(item["name"] == "memory_mutation" for item in cli_nodes["items"]) + print("CLI replay passed", flush=True) + mcp_result = await check_mcp( + client, url, directory, agent.id, version.id, evaluator.id, baseline.id + ) + incomplete = copy.deepcopy(baseline.inputs) + incomplete["mastra_memory_replay"]["complete"] = False + broken = await client.sessions.create( + SessionCreateRequest( + agent_id=agent.id, + agent_version_id=version.id, + origin=SessionOrigin.RECORDED, + status=SessionStatus.COMPLETED, + inputs=incomplete, + outputs={"text": "incomplete fixture"}, + ) + ) + failed = await client.replays.create( + ReplayCreateRequest( + baseline_session_id=broken.id, + evaluators=[config], + override=ReplayOverride(**OVERRIDE), + ) + ) + assert failed.job_id + failed_job = await await_job( + client, failed.job_id, "incomplete replay", 120 + ) + assert failed_job.status == JobStatus.FAILED, failed_job + failed_tasks = await client.jobs.list_tasks(failed_job.id) + assert any( + "incomplete" in (task.error or "").lower() + for task in failed_tasks.items + ) + failed = await client.replays.get(failed.id) + assert failed.status == "failed" + assert failed.result_session_id is None + failure_cli = await run_cli("replay", "get", str(failed.id)) + assert failure_cli["item"]["status"] == "failed" + reports = [ + json.loads(path.read_text()) + for path in directory.glob("*.json") + if path.name != "production.json" + ] + good = [report for report in reports if report.get("result") == "passed"] + assert len(good) == 4, reports + rejected = [ + report for report in reports if report.get("result") == "failed" + ] + assert rejected and all(report["actor_calls"] == 0 for report in rejected) + assert all(report["unrelated_blob_status"] == 403 for report in good) + assert all( + not report["task_inputs_in_environment"] + for report in good + if report["replay_id"] + ) + server_log = (directory / "server.log").read_text() + for report in good: + expected = 1 if report["replay_id"] else 0 + report["task_spec_requests"] = server_log.count( + f'"GET /api/v1/tasks/{report["task_id"]}/spec HTTP/1.1" 200 OK' + ) + assert report["task_spec_requests"] == expected, report + connection = await asyncpg.connect( + host=DB_HOST, + port=DB_PORT, + user=DB_USER, + password=DB_PWD, + database=db_name, + ) + try: + offloaded_requests = await connection.fetchval( + "SELECT count(*) FROM session_node " + "WHERE session_id = ANY($1::uuid[]) " + "AND node_type = 'llm_call' AND inputs_blob_id IS NOT NULL", + [uuid.UUID(item["session_id"]) for item in (sdk, cli, mcp_result)], + ) + assert offloaded_requests == 6, offloaded_requests + finally: + await connection.close() + proof = { + "result": "passed", + "baseline_session_id": str(baseline.id), + "offloaded_input_blob_id": str(offloaded), + "offloaded_request_nodes": offloaded_requests, + "sdk": sdk, + "cli": cli, + "mcp": mcp_result, + "incomplete_replay_id": str(failed.id), + "read_only_mutation_denied": True, + "task_reports": reports, + "provider_calls": 0, + "logs": str(directory), + } + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(proof, indent=2) + "\n") + print( + f"PASS: authenticated memory replay proof saved to {output}", flush=True + ) + finally: + stop.set() + if worker_task is not None: + try: + await asyncio.wait_for( + asyncio.gather(worker_task, return_exceptions=True), 15 + ) + except TimeoutError: + worker_task.cancel() + await asyncio.gather(worker_task, return_exceptions=True) + if server is not None: + server.terminate() + try: + server.wait(timeout=10) + except subprocess.TimeoutExpired: + server.kill() + server.wait(timeout=10) + await drop_database(db_name) + for name in ("worker-blobs", "worker-payloads"): + shutil.rmtree(directory / name, ignore_errors=True) + print( + "Removed the owned server, worker, database and worker caches.", flush=True + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output", type=Path, default=Path("/tmp/kitaru-memory-proof.json") + ) + asyncio.run(check(parser.parse_args().output)) From a17ec54649ec8368723054a3b81214541efe3d3a Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 20:05:07 +0200 Subject: [PATCH 013/115] Document and package isolated memory replay --- changelog.d/mastra-memory-replay.added.md | 1 + docs/book/adapters/mastra.md | 149 ++++++++++++++++- packages/mastra/README.md | 144 ++++++++++++++++- scripts/fixtures/mastra-memory-smoke.mjs | 186 ++++++++++++++++++++++ scripts/fixtures/mastra-memory-types.ts | 51 ++++++ scripts/smoke-typescript-packages.mjs | 42 ++++- 6 files changed, 563 insertions(+), 10 deletions(-) create mode 100644 changelog.d/mastra-memory-replay.added.md create mode 100644 scripts/fixtures/mastra-memory-smoke.mjs create mode 100644 scripts/fixtures/mastra-memory-types.ts diff --git a/changelog.d/mastra-memory-replay.added.md b/changelog.d/mastra-memory-replay.added.md new file mode 100644 index 000000000..467f51dba --- /dev/null +++ b/changelog.d/mastra-memory-replay.added.md @@ -0,0 +1 @@ +- Add opt-in Mastra streamed replay with isolated native thread memory. `createMemoryReplayAgent()` records starting working and observational memory, effective actor requests, controlled file bytes, and ordered memory changes, then replays from that state without reading or writing production memory. Requires Mastra core 1.67.0 and memory 1.30.0; existing history-only recordings must be recorded again to use this path. diff --git a/docs/book/adapters/mastra.md b/docs/book/adapters/mastra.md index 611c88a3f..0e87415f5 100644 --- a/docs/book/adapters/mastra.md +++ b/docs/book/adapters/mastra.md @@ -5,7 +5,7 @@ icon: robot # Mastra -The Kitaru Mastra adapter wraps an existing Mastra `Agent` and records `generate()` calls and supported streams as Kitaru [sessions](../concepts/agents-and-sessions.md). Mastra still runs the agent and Kitaru returns the native Mastra result unchanged. +The Kitaru Mastra adapter wraps an existing Mastra `Agent` and records `generate()` calls and supported streams as Kitaru [sessions](../concepts/agents-and-sessions.md). Mastra still runs the agent and Kitaru returns the native Mastra result unchanged. For thread-scoped working and observational memory, use the opt-in [isolated memory replay factory](#isolated-memory-replay). {% hint style="warning" %} `@zenml-io/kitaru-mastra` supports Node `>=22.22.0 <23 || >=26 <27`. `Agent.generate()` supports `@mastra/core >=1.51.0 <1.68.0`; recorded `Agent.stream()` calls require a stable Mastra 1.67.x release. @@ -112,7 +112,7 @@ Each call creates isolated recording state and: Each LLM node records the requested Kitaru model, the model and provider reported by Mastra, token usage, finish information, and provider metadata. Kitaru stores cost only when you provide a `costCalculator`; it does not calculate model prices on the server. -Step nodes do not record model inputs because Mastra repeats the full prompt and message history in each provider request. Step outputs include the finish reason, text, tool calls, tool results, tripwire details, and warnings. Tool inputs are the arguments requested by the model, before a tool schema applies defaults or coercion. +Ordinary `KitaruAgent` step nodes do not record model inputs because Mastra repeats the full prompt and message history in each provider request. Step outputs include the finish reason, text, tool calls, tool results, tripwire details, and warnings. Tool inputs are the arguments requested by the model, before a tool schema applies defaults or coercion. Recording uses bounded JSON conversion. Tool strings are limited to 4096 characters, arrays and objects to 100 items, and nesting to 8 levels by default. Set larger limits on the wrapper when a tool needs its full arguments and result for history replay: @@ -163,7 +163,7 @@ History matching uses the tool name and original JSON arguments. The Mastra impo A completed history match replays its result, including `null`, without executing the live tool. A failed match throws `ToolPolicyError` with its stored error text and does not execute the live tool. A tool call whose stored arguments or result were explicitly marked incomplete is a history miss and follows `on_miss`; with `passthrough`, this executes the live tool. Older recordings without fidelity flags remain readable, but Kitaru cannot verify whether their tool results were truncated. Re-record them before relying on history replay. Imported trace payloads retain their original values, although the executing adapter must record complete arguments for the lookup to match. -Before a replay starts, the adapter inventories configured tools, function-valued tools resolved from the run's `requestContext`, and per-run `clientTools` and `toolsets`. It rejects tools without a local `execute` function, approval-gated runs, sandboxed tools, and tool keys that Mastra would rename before exposing them to the model. Tools added only during execution and tools executed by a provider remain outside this preflight check and are not supported replay targets. +Before a replay starts, `KitaruAgent` inventories configured tools, function-valued tools resolved from the run's `requestContext`, and per-run `clientTools` and `toolsets`. It rejects tools without a local `execute` function, approval-gated runs, sandboxed tools, and tool keys that Mastra would rename before exposing them to the model. Tools added only during execution and tools executed by a provider remain outside this preflight check and are not supported replay targets. A tool-policy failure aborts the replay and records the session as failed. Replay forces `toolCallConcurrency: 1` and aborts Mastra's generation loop as soon as a tool hook fails, so a later model step or sibling tool cannot continue after the policy failure. Kitaru does not recreate the original exception class or convert a matched failure into a native tool-error result. On Mastra 1.67, a failed streaming policy may settle the native stream with no text instead of rejecting it; inspect the recorded replay session for the failure. @@ -171,7 +171,7 @@ A tool-policy failure aborts the replay and records the session as failed. Repla Replay is execution, not a transaction. A passthrough tool can complete an external side effect before a later model or recording failure, and Kitaru cannot roll it back. Use application-level idempotency keys for side-effecting tools, or choose static or history policies when replay must suppress execution. {% endhint %} -## Memory behavior +## History-only memory with `KitaruAgent` A supplied message array and recalled thread history are different inputs. An array contains only the messages the caller supplied; Mastra can still recall additional history when the invocation selects a memory thread. @@ -181,6 +181,142 @@ Replay removes per-run `memory`, `threadId`, `resourceId`, and `savePerStep` val A missing, incomplete, or lossy snapshot produces an actionable unsupported-replay error before model execution. Record the invocation again with this adapter, or supply its complete recorded message array without live memory selectors. An explicit array without memory selectors continues to replay directly. Old recordings do not acquire missing history automatically. Their raw inputs do not identify whether memory was used, so removing memory settings from the replay entrypoint cannot establish that those inputs are complete. Record legacy memory-dependent invocations again before replaying them. Prompt and system-instruction overrides on conversation snapshots remain unsupported because replacing them can discard part of the recorded context; record a new invocation with the desired messages instead. +## Isolated memory replay + +Use `createMemoryReplayAgent()` when a consumed stream needs thread-scoped schema working memory, observational memory, or controlled input processors. This opt-in factory requires exactly `@mastra/core@1.67.0` and `@mastra/memory@1.30.0`. The existing `KitaruAgent` wrapper keeps its history-only memory behavior. + +```bash +pnpm add @zenml-io/kitaru-mastra @mastra/core@1.67.0 @mastra/memory@1.30.0 zod +``` + +The factory creates a fresh native agent for each invocation. For a baseline, it binds native memory to your source storage and records the starting state before recall. For a replay, it restores that state into a separate in-memory store. Native working-memory tools and observation/reflection jobs then evolve the isolated state as the model runs again. Replay never calls `sourceMemory()`. + +The following binding uses a process-local store. Supply your existing public memory storage domain and its complete configuration for a persistent application: + +```ts +import { InMemoryStore } from "@mastra/core/storage"; +import { Memory } from "@mastra/memory"; +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, +} from "@zenml-io/kitaru-mastra"; +import { z } from "zod"; + +const store = new InMemoryStore(); +const sourceMemory = new Memory({ + storage: store, + options: { + semanticRecall: false, + workingMemory: { + enabled: true, + scope: "thread", + schema: z.object({ preference: z.string() }), + }, + }, +}); +// Share this same instance with every writer, for the lifetime of the store. +const exclusiveAccess = createProcessLocalMemoryAccess(); +const recorded = createMemoryReplayAgent( + ({ memory }) => ({ + id: "support", + name: "Support", + memory, + instructions: () => "Remember the user's preferences.", + model: () => "openai/gpt-5-mini", + defaultOptions: () => ({ maxSteps: 3 }), + }), + { + agentId: process.env.KITARU_AGENT_ID!, + requestedModelId: "openai/gpt-5-mini", + allowedReplayModels: ["openai/gpt-5-mini"], + sourceMemory: () => ({ + domain: store.stores.memory!, + configuration: sourceMemory.getMergedThreadConfig(), + settled: () => sourceMemory.settled(), + exclusiveAccess, + }), + resolveModel: (id) => { + if (id !== "openai/gpt-5-mini") throw new Error(`Unknown model: ${id}`); + return "openai/gpt-5-mini"; + }, + }, +); +const output = await recorded.stream("My preference is green.", { + memory: { thread: "support-thread", resource: "customer-123" }, + context: [{ role: "system", content: "The customer is asking about preferences." }], +}); +await output.consumeStream(); +await store.close(); +``` + +Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KITARU_AGENT_ID`, and the model provider credential. Register the compiled command as the agent version's run specification to run it through a worker. The same command serves baseline and replay tasks; the worker supplies the recorded input and replay identity. + +All source-memory writers must participate in the same exclusive-access mechanism. The process-local helper is suitable only when every writer shares that instance in one process. Use a distributed implementation of `MastraExclusiveMemoryAccess` when other processes can write. Its `acquire()` method must hold access until the returned release function runs. `settled()` must join pending work on the source `Memory` instance; it is not a lock. Configure working and observational memory with `scope: "thread"`; resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` are outside this contract. + +Dynamic `instructions`, `model`, and `defaultOptions` resolve once during baseline setup. Replay uses their recorded values instead of calling those resolvers again. `resolveModel` must resolve the recorded actor, observer, and reflector model identifiers as well as any allowed actor override. A `system_prompt` override replaces only application instructions and retains recorded extra system context. Model and model-setting overrides affect the actor; observation and reflection retain their recorded configuration. Raw-input `prompt` overrides are rejected; record a new baseline to change invocation input. + +### Create and inspect a memory replay + +After the worker records a complete baseline, create a replay using an existing evaluator: + +```bash +kitaru replay create \ + --evaluator your-evaluator@1 \ + --override '{"system_prompt":"Use the recorded preferences when answering."}' \ + --tool-policy '{"default":{"type":"history","scope":"baseline","on_miss":"fail"},"tools":{}}' \ + --output json +kitaru job watch +kitaru replay get --output json +kitaru session get --output json +kitaru session nodes --include-payloads --output json +``` + +Read `result_session_id` from the replay, check the session's final status, and inspect its model-request and memory-mutation nodes. `session nodes` returns one page; follow `page.next_cursor` with `--cursor` while `page.has_more` is true. + +The Python SDK uses the same replay request. Given an authenticated `client`, a baseline UUID, and an existing evaluator: + +```python +from kitaru.api_models.v1.plugin import EvaluatorConfig +from kitaru.api_models.v1.replay import ReplayCreateRequest +from kitaru.api_models.v1.replay_config import ReplayOverride, ToolPolicy + +replay = await client.replays.create( + ReplayCreateRequest( + baseline_session_id=baseline_id, + override=ReplayOverride(system_prompt="Use the recorded preferences."), + tool_policy=ToolPolicy.model_validate({ + "default": {"type": "history", "scope": "baseline", "on_miss": "fail"}, + "tools": {}, + }), + evaluators=[EvaluatorConfig(evaluator="your-evaluator", version=1)], + ) +) +``` + +Use `client.replays.get(replay.id)` to follow completion and obtain the result session. Fetch its inputs with `client.sessions.get()` and iterate nodes with `client.sessions.iter_nodes()` and `SessionNodeListParams(include_payloads=True)`. + +The native MCP server starts these replays through an experiment. Use `kitaru_cohorts_manage` to create a cohort and a version containing the baseline session, `kitaru_experiments_manage` to configure the same override, policy, and evaluator, then `kitaru_workflow_start` with `operation: "experiment_run"`, the experiment ID, cohort-version ID, and agent-version ID. No separate MCP replay-creation tool is required. + +Inspect the run with `kitaru_activity_read`: get `kind: "experiment_run"`, list `kind: "replay"` filtered by `experiment_run_id`, then get its result session. To read evidence, use `operation: "list_children"`, `kind: "session_nodes"`, `parent_id: ""`, and `include_payloads: true`; follow the returned cursor until all pages have been read. A read-only MCP connection can inspect these results but cannot start experiments. + +### Files, skills, and processors + +Pass a static `inputProcessors` array in the factory configuration. File processors must use the factory's supplied `resolveFile`; declare every allowed URL in the adapter's `files` list and provide a baseline `resolveFile` that returns `{ bytes: Uint8Array, mediaType: string }`. Kitaru records those bytes and serves them from the recorded input during replay. An undeclared URL fails instead of fetching live content. + +For skills, set `skillsDirectory` to the directory containing your skill folders and use the factory's supplied `workspace`. Kitaru reads skill files into an immutable native workspace and records their paths, sizes, and hashes. Deploy the same skill artifact with the replay command. Changed files, missing files, and symlinks are rejected before execution. + +The factory must use the supplied memory and workspace instances. Processors and tools are application code: their dependencies must use these supplied bindings for replay isolation. Kitaru does not sandbox arbitrary callbacks or prevent code from opening another database connection or making a network request. Workflows, subagents, provider-executed tools, approval/resume modes, dynamic tool inventories, `prepareStep`, output processors, and secondary structured-output models remain unsupported. + +### Tool policies and evidence + +Native memory tools execute against the isolated replay store, including under `history` with `on_miss: "fail"`. External tools, including tools added by a processor, follow the replay tool policy. A tool named `updateWorkingMemory` does not acquire the native-memory exemption by name. Use history with a failing miss when external tools must not execute. + +Session inputs contain a version-2 `mastra_memory_replay` envelope with the raw invocation, initial thread/resource/messages, observational state and buffers, effective configuration, approved request context, and controlled file bytes. Dates, URLs, and binary values retain their types. Set `captureRequestContext` to select the context the run needs; never include credentials. Unsupported values, redaction, or exceeding the 1,048,576-character serialized payload bound make the input incomplete. Older history-only snapshots cannot recover this state and must be recorded again with the factory. + +Unlike ordinary wrapper recording, this path records the effective actor prompt, tools, tool choice, and supported settings for each provider attempt, including failed retries. Request attributes include attempt identity, memory revision, source provenance, and evidence completeness. `memory_mutation` span nodes record ordered native storage changes and link them to the active actor attempt when one exists. This evidence describes the request sent at the adapter's model boundary, not a provider's internal processing. + +Consume the stream through completion so Kitaru can join native background memory work, flush evidence, release source access, and close the isolated replay store. Inspect the replay session's status and evidence completeness as well as native output: recording problems do not replace application output, and Mastra can settle a stream after a policy failure. Missing or incomplete initial state fails replay before model execution. A failed or incomplete recording is not proof that all evidence was saved. + ## Structured output Schema-only structured output is supported by both `generate()` and `stream()` and remains available on the returned Mastra result: @@ -224,12 +360,13 @@ The adapter supports: - `Agent.generate()` calls on Mastra 1.51 through 1.67. - Ordinary consumed `Agent.stream()` calls and replay on stable Mastra 1.67.x, with schema-only structured output. +- Opt-in isolated native memory replay through `createMemoryReplayAgent()` on exact Mastra core 1.67.0 and memory 1.30.0. - Local function tools, including function-valued tools resolved from the run's `requestContext`. - Per-run model, system-instruction, model-setting, and input overrides. - Passthrough, static, and same-adapter history tool policies. - Schema-only structured output, plus per-run secondary structuring models with strict validation for `generate()`. -Streaming does not support approval or resume modes, background or `untilIdle` execution, or secondary structured-output models. The adapter does not support workflows, subagents, MCP tools, provider-native tool replay, dynamic instructions, or LLM tool policy. `prepareStep` and input processors are rejected during replay because they can replace the model, prompt, or tools after policy preflight. +Streaming does not support approval or resume modes, background or `untilIdle` execution, or secondary structured-output models. The existing `KitaruAgent` wrapper does not support workflows, subagents, MCP tools, provider-native tool replay, dynamic instructions, or LLM tool policy. Its replay path rejects `prepareStep` and input processors because they can replace the model, prompt, or tools after policy preflight. The opt-in memory factory supports the narrower dynamic-configuration and processor contract described in [Isolated memory replay](#isolated-memory-replay). ## Import existing Mastra traces @@ -252,7 +389,7 @@ These commands use the server selected by `kitaru login`. Pass `--server URL` to ### Import for inspection or replay -Select an existing agent version that represents the exported run. For replay, its registered Node command must use the context-capable `KitaruAgent` described in [Memory behavior](#memory-behavior), with the same callable tool names and compatible schemas. An importer preserves the trace; it does not supply runnable agent code. +Select an existing agent version that represents the exported run. For replay, its registered Node command must use the context-capable `KitaruAgent` described in [History-only memory with `KitaruAgent`](#history-only-memory-with-kitaruagent), with the same callable tool names and compatible schemas. An importer preserves the trace; it does not supply runnable agent code. For inspection and evaluation, import the file without replay parameters: diff --git a/packages/mastra/README.md b/packages/mastra/README.md index a5e209d5f..49bf5acb3 100644 --- a/packages/mastra/README.md +++ b/packages/mastra/README.md @@ -2,6 +2,8 @@ Experimental recording and replay support for Mastra. `generate()` supports `@mastra/core >=1.51.0 <1.68.0`; recorded and replayed `stream()` calls require a stable `@mastra/core 1.67.x` release. +For native working and observational memory, use the opt-in [isolated memory replay factory](#isolated-memory-replay) on exact Mastra core 1.67.0 and memory 1.30.0. + This adapter depends on the framework-neutral `@zenml-io/kitaru` package, whose repository directory is `packages/core/`. The packages are versioned and released together. ```bash @@ -67,7 +69,7 @@ Failed sessions store a bounded failure category rather than the raw provider or Mastra 1.67 continues model execution in the background when the application stops reading or cancels its reader. Kitaru records the eventual finish callback and completed result in that case. Kitaru does not drain the returned reader itself. After queued steps settle, the finish callback chooses the terminal status once. An error or abort observed before that decision records failure; a later abort cannot reverse completion because the API does not reopen terminal sessions. -## Recording +## Recording with `KitaruAgent` Each call creates isolated run state and: @@ -111,7 +113,7 @@ const recorded = new KitaruAgent(existingAgent, { Each LLM node carries a `cost` attribute recording where the number came from: `disabled` with no calculator, `estimated` for a calculated value, and `unavailable` when the calculator throws or returns nothing. A throwing calculator never fails the run. -## Replay +## Replay with `KitaruAgent` When `KITARU_REPLAY_ID` is set, `generate()` or `stream()` fetches the replay and applies its model, system-instruction, model-parameter, and tool-policy overrides through public per-run options and tool hooks. The stream entrypoint still returns Mastra's native result; callers must consume it through completion so Kitaru can finalize the replay session. @@ -143,6 +145,142 @@ Supported tool policies are passthrough, static, and history, including `fail`, History matching is guaranteed only for traces recorded and replayed through this Mastra adapter. Another framework may validate, default, or serialize the same logical tool input differently, so cross-framework history replay is not a compatibility promise. +## Isolated memory replay + +Use `createMemoryReplayAgent()` when a consumed stream needs thread-scoped schema working memory, observational memory, or controlled input processors. This opt-in factory requires exactly `@mastra/core@1.67.0` and `@mastra/memory@1.30.0`. The existing `KitaruAgent` wrapper keeps its history-only memory behavior. + +```bash +pnpm add @zenml-io/kitaru-mastra @mastra/core@1.67.0 @mastra/memory@1.30.0 zod +``` + +The factory creates a fresh native agent for each invocation. For a baseline, it binds native memory to your source storage and records the starting state before recall. For a replay, it restores that state into a separate in-memory store. Native working-memory tools and observation/reflection jobs then evolve the isolated state as the model runs again. Replay never calls `sourceMemory()`. + +The following binding uses a process-local store. Supply your existing public memory storage domain and its complete configuration for a persistent application: + +```ts +import { InMemoryStore } from "@mastra/core/storage"; +import { Memory } from "@mastra/memory"; +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, +} from "@zenml-io/kitaru-mastra"; +import { z } from "zod"; + +const store = new InMemoryStore(); +const sourceMemory = new Memory({ + storage: store, + options: { + semanticRecall: false, + workingMemory: { + enabled: true, + scope: "thread", + schema: z.object({ preference: z.string() }), + }, + }, +}); +// Share this same instance with every writer, for the lifetime of the store. +const exclusiveAccess = createProcessLocalMemoryAccess(); +const recorded = createMemoryReplayAgent( + ({ memory }) => ({ + id: "support", + name: "Support", + memory, + instructions: () => "Remember the user's preferences.", + model: () => "openai/gpt-5-mini", + defaultOptions: () => ({ maxSteps: 3 }), + }), + { + agentId: process.env.KITARU_AGENT_ID!, + requestedModelId: "openai/gpt-5-mini", + allowedReplayModels: ["openai/gpt-5-mini"], + sourceMemory: () => ({ + domain: store.stores.memory!, + configuration: sourceMemory.getMergedThreadConfig(), + settled: () => sourceMemory.settled(), + exclusiveAccess, + }), + resolveModel: (id) => { + if (id !== "openai/gpt-5-mini") throw new Error(`Unknown model: ${id}`); + return "openai/gpt-5-mini"; + }, + }, +); +const output = await recorded.stream("My preference is green.", { + memory: { thread: "support-thread", resource: "customer-123" }, + context: [{ role: "system", content: "The customer is asking about preferences." }], +}); +await output.consumeStream(); +await store.close(); +``` + +Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KITARU_AGENT_ID`, and the model provider credential. Register the compiled command as the agent version's run specification to run it through a worker. The same command serves baseline and replay tasks; the worker supplies the recorded input and replay identity. + +All source-memory writers must participate in the same exclusive-access mechanism. The process-local helper is suitable only when every writer shares that instance in one process. Use a distributed implementation of `MastraExclusiveMemoryAccess` when other processes can write. Its `acquire()` method must hold access until the returned release function runs. `settled()` must join pending work on the source `Memory` instance; it is not a lock. Configure working and observational memory with `scope: "thread"`; resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` are outside this contract. + +Dynamic `instructions`, `model`, and `defaultOptions` resolve once during baseline setup. Replay uses their recorded values instead of calling those resolvers again. `resolveModel` must resolve the recorded actor, observer, and reflector model identifiers as well as any allowed actor override. A `system_prompt` override replaces only application instructions and retains recorded extra system context. Model and model-setting overrides affect the actor; observation and reflection retain their recorded configuration. Raw-input `prompt` overrides are rejected; record a new baseline to change invocation input. + +### Create and inspect a memory replay + +After the worker records a complete baseline, create a replay using an existing evaluator: + +```bash +kitaru replay create \ + --evaluator your-evaluator@1 \ + --override '{"system_prompt":"Use the recorded preferences when answering."}' \ + --tool-policy '{"default":{"type":"history","scope":"baseline","on_miss":"fail"},"tools":{}}' \ + --output json +kitaru job watch +kitaru replay get --output json +kitaru session get --output json +kitaru session nodes --include-payloads --output json +``` + +Read `result_session_id` from the replay, check the session's final status, and inspect its model-request and memory-mutation nodes. `session nodes` returns one page; follow `page.next_cursor` with `--cursor` while `page.has_more` is true. + +The Python SDK uses the same replay request. Given an authenticated `client`, a baseline UUID, and an existing evaluator: + +```python +from kitaru.api_models.v1.plugin import EvaluatorConfig +from kitaru.api_models.v1.replay import ReplayCreateRequest +from kitaru.api_models.v1.replay_config import ReplayOverride, ToolPolicy + +replay = await client.replays.create( + ReplayCreateRequest( + baseline_session_id=baseline_id, + override=ReplayOverride(system_prompt="Use the recorded preferences."), + tool_policy=ToolPolicy.model_validate({ + "default": {"type": "history", "scope": "baseline", "on_miss": "fail"}, + "tools": {}, + }), + evaluators=[EvaluatorConfig(evaluator="your-evaluator", version=1)], + ) +) +``` + +Use `client.replays.get(replay.id)` to follow completion and obtain the result session. Fetch its inputs with `client.sessions.get()` and iterate nodes with `client.sessions.iter_nodes()` and `SessionNodeListParams(include_payloads=True)`. + +The native MCP server starts these replays through an experiment. Use `kitaru_cohorts_manage` to create a cohort and a version containing the baseline session, `kitaru_experiments_manage` to configure the same override, policy, and evaluator, then `kitaru_workflow_start` with `operation: "experiment_run"`, the experiment ID, cohort-version ID, and agent-version ID. No separate MCP replay-creation tool is required. + +Inspect the run with `kitaru_activity_read`: get `kind: "experiment_run"`, list `kind: "replay"` filtered by `experiment_run_id`, then get its result session. To read evidence, use `operation: "list_children"`, `kind: "session_nodes"`, `parent_id: ""`, and `include_payloads: true`; follow the returned cursor until all pages have been read. A read-only MCP connection can inspect these results but cannot start experiments. + +### Files, skills, and processors + +Pass a static `inputProcessors` array in the factory configuration. File processors must use the factory's supplied `resolveFile`; declare every allowed URL in the adapter's `files` list and provide a baseline `resolveFile` that returns `{ bytes: Uint8Array, mediaType: string }`. Kitaru records those bytes and serves them from the recorded input during replay. An undeclared URL fails instead of fetching live content. + +For skills, set `skillsDirectory` to the directory containing your skill folders and use the factory's supplied `workspace`. Kitaru reads skill files into an immutable native workspace and records their paths, sizes, and hashes. Deploy the same skill artifact with the replay command. Changed files, missing files, and symlinks are rejected before execution. + +The factory must use the supplied memory and workspace instances. Processors and tools are application code: their dependencies must use these supplied bindings for replay isolation. Kitaru does not sandbox arbitrary callbacks or prevent code from opening another database connection or making a network request. Workflows, subagents, provider-executed tools, approval/resume modes, dynamic tool inventories, `prepareStep`, output processors, and secondary structured-output models remain unsupported. + +### Tool policies and evidence + +Native memory tools execute against the isolated replay store, including under `history` with `on_miss: "fail"`. External tools, including tools added by a processor, follow the replay tool policy. A tool named `updateWorkingMemory` does not acquire the native-memory exemption by name. Use history with a failing miss when external tools must not execute. + +Session inputs contain a version-2 `mastra_memory_replay` envelope with the raw invocation, initial thread/resource/messages, observational state and buffers, effective configuration, approved request context, and controlled file bytes. Dates, URLs, and binary values retain their types. Set `captureRequestContext` to select the context the run needs; never include credentials. Unsupported values, redaction, or exceeding the 1,048,576-character serialized payload bound make the input incomplete. Older history-only snapshots cannot recover this state and must be recorded again with the factory. + +Unlike ordinary wrapper recording, this path records the effective actor prompt, tools, tool choice, and supported settings for each provider attempt, including failed retries. Request attributes include attempt identity, memory revision, source provenance, and evidence completeness. `memory_mutation` span nodes record ordered native storage changes and link them to the active actor attempt when one exists. This evidence describes the request sent at the adapter's model boundary, not a provider's internal processing. + +Consume the stream through completion so Kitaru can join native background memory work, flush evidence, release source access, and close the isolated replay store. Inspect the replay session's status and evidence completeness as well as native output: recording problems do not replace application output, and Mastra can settle a stream after a policy failure. Missing or incomplete initial state fails replay before model execution. A failed or incomplete recording is not proof that all evidence was saved. + ## Callback composition Per-run Mastra hooks replace configured hooks. During replay, Kitaru evaluates its policy first. Passthrough calls then invoke an explicitly supplied configured hook followed by the caller's per-run hook. Kitaru-mocked calls do not invoke user tool hooks. Step recording completes before configured and caller `onStepFinish` callbacks. @@ -155,7 +293,7 @@ Mastra merges per-run model settings with configured defaults. Kitaru can replac Recorded payloads preserve JSON values, convert dates to ISO strings, bigints to decimal strings, and errors to `{name, message}`. Functions, symbols, circular references, and non-finite numbers are replaced with a marker instead of failing the run, because a recording problem must not break the agent. Replay tool inputs go through the same bounded converter that records them, so a history cache key computed during replay matches the key the server computed from the recorded call. Serialization never changes the Mastra result or tool output returned to the application. -## Current scope +## Existing wrapper scope This experimental release supports `Agent.generate()` with Mastra `>=1.51.0 <1.68.0` and consumed `Agent.stream()` calls, including replay, on stable Mastra 1.67.x. Streaming supports local function tools and schema-only structured output. It rejects user `prepareStep` and input processors, approval and resume modes, background or `untilIdle` execution, and secondary structured-output models before native execution. Both replay entrypoints reject `prepareStep` and input processors because they can replace the model, prompt, or tools after preflight. Workflows, subagents, MCP tools, provider-native tool replay, dynamic instructions, and LLM tool policy are intentionally not implemented. diff --git a/scripts/fixtures/mastra-memory-smoke.mjs b/scripts/fixtures/mastra-memory-smoke.mjs new file mode 100644 index 000000000..77c7612e2 --- /dev/null +++ b/scripts/fixtures/mastra-memory-smoke.mjs @@ -0,0 +1,186 @@ +import assert from "node:assert/strict"; +import { InMemoryStore } from "@mastra/core/storage"; +import { Memory } from "@mastra/memory"; +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, +} from "@zenml-io/kitaru-mastra"; +import { z } from "zod"; + +const calls = []; +const requests = []; +const replayId = "018f0000-0000-7000-8000-000000000102"; +const sessionId = "018f0000-0000-7000-8000-000000000101"; +const originalFetch = globalThis.fetch; +// Only the Kitaru transport is replaced. Memory, tools and streaming are native. +globalThis.fetch = async (input, init = {}) => { + const url = new URL(String(input)); + const method = init.method ?? "GET"; + const body = init.body ? JSON.parse(String(init.body)) : undefined; + calls.push({ body, method, path: url.pathname }); + if (method === "POST" && url.pathname === "/api/v1/sessions") + return Response.json( + { id: sessionId, origin: "recorded", status: "in_progress" }, + { status: 201 }, + ); + if (method === "POST" && url.pathname.endsWith("/nodes")) + return Response.json([]); + if (method === "PATCH" && url.pathname.endsWith(sessionId)) + return Response.json({ id: sessionId, origin: "recorded", status: body.status }); + if (method === "GET" && url.pathname.endsWith(replayId)) + return Response.json({ + id: replayId, + baseline_session_id: sessionId, + job_id: "018f0000-0000-7000-8000-000000000104", + override: { system_prompt: "Replay instructions" }, + status: "pending", + tool_policy: { + default: { type: "history", scope: "baseline", on_miss: "fail" }, + tools: {}, + }, + }); + throw new Error(`Unexpected memory smoke request: ${method} ${url.pathname}`); +}; + +let replaying = false; +let modelCalls = 0; +const model = { + modelId: "memory-smoke-model", + provider: "package-smoke", + specificationVersion: "v2", + supportedUrls: {}, + doGenerate: async () => { + throw new Error("Expected native streaming"); + }, + doStream: async (request) => { + requests.push(request); + const toolStep = ++modelCalls % 2 === 1; + const parts = toolStep + ? [ + { + type: "tool-call", + toolCallId: `memory-${modelCalls}`, + toolName: "updateWorkingMemory", + input: JSON.stringify({ + memory: { + preference: replaying ? "replay-green" : "baseline-red", + }, + }), + }, + ] + : [ + { type: "text-start", id: "text" }, + { type: "text-delta", id: "text", delta: "done" }, + { type: "text-end", id: "text" }, + ]; + return { + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: "stream-start", warnings: [] }); + for (const part of parts) controller.enqueue(part); + controller.enqueue({ + type: "finish", + finishReason: toolStep ? "tool-calls" : "stop", + usage: { inputTokens: 5, outputTokens: 2, totalTokens: 7 }, + }); + controller.close(); + }, + }), + }; + }, +}; +const store = new InMemoryStore(); +const sourceMemory = new Memory({ + storage: store, + options: { + semanticRecall: false, + workingMemory: { + enabled: true, + scope: "thread", + schema: z.object({ preference: z.string() }), + }, + }, +}); +const selector = { + threadId: "memory-smoke-thread", + resourceId: "memory-smoke-resource", +}; +const exclusiveAccess = createProcessLocalMemoryAccess(); +try { + await sourceMemory.createThread({ ...selector, title: "Smoke" }); + await sourceMemory.updateWorkingMemory({ + ...selector, + workingMemory: '{"preference":"historical-blue"}', + }); + const agent = createMemoryReplayAgent( + ({ memory }) => ({ + id: "memory-smoke", + name: "Memory smoke", + memory, + model, + instructions: "Baseline instructions", + defaultOptions: { maxSteps: 3 }, + }), + { + agentId: sessionId, + apiUrl: "https://api.example", + requestedModelId: "package-smoke/memory-smoke-model", + sourceMemory() { + assert.equal(replaying, false, "Replay consulted production memory"); + return { + domain: store.stores.memory, + configuration: sourceMemory.getMergedThreadConfig(), + settled: () => sourceMemory.settled(), + exclusiveAccess, + }; + }, + resolveModel: () => model, + }, + ); + await ( + await agent.stream("Change the preference", { + memory: { thread: selector.threadId, resource: selector.resourceId }, + }) + ).consumeStream(); + const input = calls.find( + (call) => call.path === "/api/v1/sessions" && call.method === "POST", + ).body.inputs; + assert.equal(input.mastra_memory_replay.complete, true); + await sourceMemory.updateWorkingMemory({ + ...selector, + workingMemory: '{"preference":"production-today"}', + }); + replaying = true; + process.env.KITARU_REPLAY_ID = replayId; + process.env.KITARU_TASK_INPUTS = JSON.stringify(input); + await (await agent.stream("Ignored caller input")).consumeStream(); + assert.equal(modelCalls, 4); + assert.match(JSON.stringify(requests[2].prompt), /historical-blue/); + assert.doesNotMatch(JSON.stringify(requests[2].prompt), /production-today/); + assert.match(JSON.stringify(requests[2].prompt), /Replay instructions/); + assert.match( + await sourceMemory.getWorkingMemory(selector), + /production-today/, + ); + assert.equal( + calls.filter((call) => call.method === "PATCH").at(-1).body.status, + "completed", + ); + const nodes = calls + .filter((call) => call.path.endsWith("/nodes")) + .flatMap((call) => call.body.nodes); + assert.ok( + nodes.some( + (node) => + node.name === "memory_mutation" && + JSON.stringify(node.inputs).includes("replay-green"), + ), + ); + assert.equal(nodes.filter((node) => node.node_type === "llm_call").length, 4); + console.log("Packed Mastra memory recording and isolated replay passed"); +} finally { + delete process.env.KITARU_REPLAY_ID; + delete process.env.KITARU_TASK_INPUTS; + globalThis.fetch = originalFetch; + await store.close(); +} diff --git a/scripts/fixtures/mastra-memory-types.ts b/scripts/fixtures/mastra-memory-types.ts new file mode 100644 index 000000000..a7d16118f --- /dev/null +++ b/scripts/fixtures/mastra-memory-types.ts @@ -0,0 +1,51 @@ +import type { Agent } from "@mastra/core/agent"; +import { InMemoryStore } from "@mastra/core/storage"; +import { Memory } from "@mastra/memory"; +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, + type MemoryReplayAgentOptions, +} from "@zenml-io/kitaru-mastra"; +import { z } from "zod"; + +const store = new InMemoryStore(); +const sourceMemory = new Memory({ + storage: store, + options: { + semanticRecall: false, + workingMemory: { + enabled: true, + scope: "thread", + schema: z.object({ preference: z.string() }), + }, + }, +}); +const exclusiveAccess = createProcessLocalMemoryAccess(); +const options: MemoryReplayAgentOptions = { + agentId: "018f0000-0000-7000-8000-000000000103", + requestedModelId: "openai/gpt-5-mini", + allowedReplayModels: ["openai/gpt-5-mini"], + sourceMemory: () => ({ + domain: store.stores.memory!, + configuration: sourceMemory.getMergedThreadConfig(), + settled: () => sourceMemory.settled(), + exclusiveAccess, + }), + resolveModel: (id) => { + if (id !== "openai/gpt-5-mini") throw new Error(`Unknown model: ${id}`); + return "openai/gpt-5-mini"; + }, +}; +const recorded = createMemoryReplayAgent( + ({ memory }) => ({ + id: "support", + name: "Support", + memory, + instructions: () => "Remember the user's preferences.", + model: () => "openai/gpt-5-mini", + defaultOptions: () => ({ maxSteps: 3 }), + }), + options, +); +const nativeStream: Agent["stream"] = recorded.stream; +void nativeStream; diff --git a/scripts/smoke-typescript-packages.mjs b/scripts/smoke-typescript-packages.mjs index f6080f408..fc56d789c 100644 --- a/scripts/smoke-typescript-packages.mjs +++ b/scripts/smoke-typescript-packages.mjs @@ -1,4 +1,10 @@ -import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { + copyFileSync, + mkdirSync, + mkdtempSync, + rmSync, + writeFileSync, +} from "node:fs"; import { tmpdir } from "node:os"; import { join, resolve } from "node:path"; import { spawnSync } from "node:child_process"; @@ -441,11 +447,25 @@ function smokeConsumer({ artifactRoot, mastraVersion, npmCache }) { npmCache, ...tarballs, `@mastra/core@${mastraVersion}`, + ...(mastraVersion === upperMastraVersion + ? ["@mastra/memory@1.30.0"] + : []), "ai@7.0.65", "zod@3.25.76", ], consumerRoot, ); + if (mastraVersion === lowerMastraVersion) { + run( + process.execPath, + [ + "--input-type=module", + "--eval", + "try { import.meta.resolve('@mastra/memory'); throw new Error('Optional memory dependency unexpectedly installed'); } catch (error) { if (error.code !== 'ERR_MODULE_NOT_FOUND') throw error; }", + ], + consumerRoot, + ); + } run(process.execPath, ["index.mjs"], consumerRoot); run( join(repositoryRoot, "node_modules", ".bin", "tsc"), @@ -461,6 +481,26 @@ function smokeConsumer({ artifactRoot, mastraVersion, npmCache }) { ); } run(process.execPath, ["stream.mjs", mastraVersion], consumerRoot); + if (mastraVersion === upperMastraVersion) { + copyFileSync( + join(repositoryRoot, "scripts/fixtures/mastra-memory-smoke.mjs"), + join(consumerRoot, "memory.mjs"), + ); + copyFileSync( + join(repositoryRoot, "scripts/fixtures/mastra-memory-types.ts"), + join(consumerRoot, "memory.ts"), + ); + writeFileSync( + join(consumerRoot, "tsconfig.memory.json"), + JSON.stringify({ extends: "./tsconfig.stream.json", include: ["memory.ts"] }), + ); + run( + join(repositoryRoot, "node_modules", ".bin", "tsc"), + ["-p", "tsconfig.memory.json"], + consumerRoot, + ); + run(process.execPath, ["memory.mjs"], consumerRoot); + } } const outputDirectory = parseOutputDirectory(process.argv.slice(2)); From b1b86d5cc423382dc8f93ffbda83831f08a7c592 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 20:12:34 +0200 Subject: [PATCH 014/115] Stop replay after native memory storage failure --- packages/mastra/src/stateful-agent.ts | 10 ++ .../mastra/test/stateful-overrides.test.ts | 94 +++++++++++++++++++ 2 files changed, 104 insertions(+) diff --git a/packages/mastra/src/stateful-agent.ts b/packages/mastra/src/stateful-agent.ts index 5b1de65ec..734fdee5a 100644 --- a/packages/mastra/src/stateful-agent.ts +++ b/packages/mastra/src/stateful-agent.ts @@ -227,6 +227,14 @@ export function createMemoryReplayAgent( if (!state) throw new Error("Memory recorder has not initialized."); return state; }; + const onIncomplete = (reason: string): void => { + // Mastra converts tool storage errors into results and otherwise continues. + // Recording-only incompleteness must preserve ordinary native execution. + if (reason !== "Native memory storage mutation failed.") return; + const error = new Error(reason); + state?.storeFailure(error); + abort.abort(error); + }; const writeNode = async (node: SessionNodeCreateRequest) => { const active = getState(); await active.enqueueStep(async () => { @@ -268,6 +276,7 @@ export function createMemoryReplayAgent( ), resolveModel: options.resolveModel, recordMutation, + onIncomplete, getRequestId: () => requestCapture?.currentRequestId, }); } else { @@ -279,6 +288,7 @@ export function createMemoryReplayAgent( domain: source.domain, exclusiveAccess: source.exclusiveAccess, recordMutation, + onIncomplete, getRequestId: () => requestCapture?.currentRequestId, }); const { Memory } = await import("@mastra/memory"); diff --git a/packages/mastra/test/stateful-overrides.test.ts b/packages/mastra/test/stateful-overrides.test.ts index 923e693e0..ef1c15590 100644 --- a/packages/mastra/test/stateful-overrides.test.ts +++ b/packages/mastra/test/stateful-overrides.test.ts @@ -619,3 +619,97 @@ it("releases the source lease on setup failure and cancellation", async () => { await release(); await runtime.store.close(); }); + +it("stops later actor and tool work after native memory storage fails", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const api = installTestApi(); + const lease = createProcessLocalMemoryAccess(); + const { createTool } = await import("@mastra/core/tools"); + const { z } = await import("zod/v4"); + const external = vi.fn(async () => "side effect"); + let actorCalls = 0; + const patchThread = runtime.domain.patchThread.bind(runtime.domain); + vi.spyOn(runtime.domain, "patchThread").mockImplementation( + async (...args) => { + if (actorCalls > 0) throw new Error("Native memory write failed"); + return patchThread(...args); + }, + ); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => + ++actorCalls === 1 + ? streamParts( + [ + { + type: "tool-call", + toolName: "updateWorkingMemory", + toolCallId: "memory-failure", + input: JSON.stringify({ memory: { preference: "green" } }), + }, + { + type: "tool-call", + toolName: "external", + toolCallId: "later-tool", + input: "{}", + }, + ], + "tool-calls", + ) + : textStream("incorrect continuation"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "write-failure", + name: "Write failure", + instructions: "Update memory", + model, + memory, + defaultOptions: { maxSteps: 3, toolCallConcurrency: 1 }, + tools: { + external: createTool({ + id: "external", + description: "A side effect", + inputSchema: z.object({}), + execute: external, + }), + }, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: lease, + }), + resolveModel: () => model, + }, + ); + try { + const output = await adapter.stream("Remember green", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await output.consumeStream(); + await vi.waitFor(() => + expect( + api.calls.filter((call) => call.method === "PATCH").at(-1)?.body + ?.status, + ).toBe("failed"), + ); + expect(actorCalls).toBe(1); + expect(external).not.toHaveBeenCalled(); + const release = await lease.acquire({ + threadId: THREAD, + resourceId: RESOURCE, + }); + await release(); + } finally { + await runtime.memory.settled(); + await runtime.store.close(); + } +}); From 2a61a2924b061f47dad82c6fb1d50dc5cc1dc741 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 20:12:35 +0200 Subject: [PATCH 015/115] Format replay documentation examples --- docs/book/adapters/mastra.md | 10 ++++++---- packages/mastra/README.md | 10 ++++++---- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/docs/book/adapters/mastra.md b/docs/book/adapters/mastra.md index 0e87415f5..94d84a8cb 100644 --- a/docs/book/adapters/mastra.md +++ b/docs/book/adapters/mastra.md @@ -284,10 +284,12 @@ replay = await client.replays.create( ReplayCreateRequest( baseline_session_id=baseline_id, override=ReplayOverride(system_prompt="Use the recorded preferences."), - tool_policy=ToolPolicy.model_validate({ - "default": {"type": "history", "scope": "baseline", "on_miss": "fail"}, - "tools": {}, - }), + tool_policy=ToolPolicy.model_validate( + { + "default": {"type": "history", "scope": "baseline", "on_miss": "fail"}, + "tools": {}, + } + ), evaluators=[EvaluatorConfig(evaluator="your-evaluator", version=1)], ) ) diff --git a/packages/mastra/README.md b/packages/mastra/README.md index 49bf5acb3..2a9b4e49c 100644 --- a/packages/mastra/README.md +++ b/packages/mastra/README.md @@ -248,10 +248,12 @@ replay = await client.replays.create( ReplayCreateRequest( baseline_session_id=baseline_id, override=ReplayOverride(system_prompt="Use the recorded preferences."), - tool_policy=ToolPolicy.model_validate({ - "default": {"type": "history", "scope": "baseline", "on_miss": "fail"}, - "tools": {}, - }), + tool_policy=ToolPolicy.model_validate( + { + "default": {"type": "history", "scope": "baseline", "on_miss": "fail"}, + "tools": {}, + } + ), evaluators=[EvaluatorConfig(evaluator="your-evaluator", version=1)], ) ) From 898bd5214aec9aeb800bacca31e49d86f965ac0a Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 20:50:52 +0200 Subject: [PATCH 016/115] Guard memory replay context and cleanup --- packages/mastra/src/memory-snapshot.ts | 61 ++- packages/mastra/src/stateful-agent.ts | 93 +++- .../mastra/test/memory-replay-safety.test.ts | 424 ++++++++++++++++++ 3 files changed, 561 insertions(+), 17 deletions(-) create mode 100644 packages/mastra/test/memory-replay-safety.test.ts diff --git a/packages/mastra/src/memory-snapshot.ts b/packages/mastra/src/memory-snapshot.ts index b7f6d5b70..e63f47c4a 100644 --- a/packages/mastra/src/memory-snapshot.ts +++ b/packages/mastra/src/memory-snapshot.ts @@ -1,5 +1,10 @@ import { createHash } from "node:crypto"; import type { MastraDBMessage, StorageThreadType } from "@mastra/core/memory"; +import { + MASTRA_AUTH_TOKEN_KEY, + MASTRA_RESOURCE_ID_KEY, + MASTRA_THREAD_ID_KEY, +} from "@mastra/core/request-context"; import type { ObservationalMemoryRecord, StorageResourceType, @@ -431,6 +436,18 @@ function validateConfiguration( configuration: unknown, ): asserts configuration is Record { requireValue(isRecord(configuration), "Malformed resolved configuration."); + function containsTransport(value: unknown): boolean { + if (Array.isArray(value)) return value.some(containsTransport); + if (!isRecord(value)) return false; + return Object.entries(value).some( + ([key, item]) => + /^(headers|abortsignal)$/i.test(key) || containsTransport(item), + ); + } + requireValue( + !containsTransport(configuration), + "Replay configuration contains transport metadata.", + ); const memory = configuration.memoryConfig ?? configuration.memory; if (!isRecord(memory)) return; requireValue( @@ -447,6 +464,38 @@ function validateConfiguration( } } +/** Require native context selectors to match the leased and captured memory. */ +export function validateMemoryReplayContext( + selector: Pick, + requestContext: Record, +): void { + requireValue( + !Object.hasOwn(requestContext, MASTRA_AUTH_TOKEN_KEY), + "Native authentication tokens are not replayable request context.", + ); + validateMemoryReplaySelectors(selector, requestContext); +} + +/** Validate selectors before a capture callback can omit middleware overrides. */ +export function validateMemoryReplaySelectors( + selector: Pick, + requestContext: Record, +): void { + for (const [key, expected] of [ + [MASTRA_THREAD_ID_KEY, selector.threadId], + [MASTRA_RESOURCE_ID_KEY, selector.resourceId], + ] as const) { + const value = requestContext[key]; + requireValue( + value === undefined || + value === null || + value === "" || + value === expected, + "Request-context memory selectors differ from the captured selectors.", + ); + } +} + /** Build safe diagnostic evidence even when complete replay prerequisites are unavailable. */ export function createMemoryReplayEnvelope( input: MastraMemoryReplayInput, @@ -464,7 +513,6 @@ export function createMemoryReplayEnvelope( }); try { validateMemorySnapshot(input.initialSnapshot); - validateConfiguration(input.configuration); const envelope: MastraMemoryReplayEnvelope = { version: 2, complete: true, @@ -540,6 +588,17 @@ export function decodeMemoryReplayEnvelope( validateConfiguration(configuration); const requestContext = decodeMemoryValue(value.requestContext as JsonValue); requireValue(isRecord(requestContext), "Malformed recorded request context."); + validateMemoryReplayContext(initialSnapshot, requestContext); + if (isRecord(configuration.runOptions)) { + const memory = configuration.runOptions.memory; + requireValue(isRecord(memory), "Missing invocation memory selectors."); + const threadId = isRecord(memory.thread) ? memory.thread.id : memory.thread; + requireValue( + threadId === initialSnapshot.threadId && + memory.resource === initialSnapshot.resourceId, + "Invocation memory selectors differ from the captured selectors.", + ); + } const urls = new Set(); const files = value.files.map((file) => { requireValue( diff --git a/packages/mastra/src/stateful-agent.ts b/packages/mastra/src/stateful-agent.ts index 734fdee5a..3afb96ee8 100644 --- a/packages/mastra/src/stateful-agent.ts +++ b/packages/mastra/src/stateful-agent.ts @@ -38,6 +38,8 @@ import { type MastraMemorySnapshot, MEMORY_REPLAY_KEY, restoreMemoryReplayEnvelope, + validateMemoryReplayContext, + validateMemoryReplaySelectors, } from "./memory-snapshot.js"; import { assertStableToolName } from "./replay-guards.js"; import { @@ -179,6 +181,26 @@ function assertSupportedConfiguration( ); } +class MemoryReplayRequestContext extends RequestContext { + constructor( + private readonly selector: Pick< + MastraMemorySnapshot, + "threadId" | "resourceId" + >, + ) { + super(); + } + + override set(key: string, value: unknown): void { + validateMemoryReplayContext(this.selector, { [key]: value }); + super.set(key, value); + } + + override setRaw(key: string, value: unknown): void { + this.set(key, value); + } +} + /** Construct each streamed invocation with historical configuration and isolated replay memory. */ export function createMemoryReplayAgent( factory: MemoryReplayAgentFactory, @@ -220,6 +242,25 @@ export function createMemoryReplayAgent( throw new Error( "Memory replay supports system_prompt overrides; replacing raw invocation input requires a new recording.", ); + const selector = historical?.initialSnapshot ?? getSelector(callerOptions); + const liveContext = callerOptions.requestContext ?? new RequestContext(); + if (!historical) + validateMemoryReplaySelectors( + selector, + Object.fromEntries(liveContext.entries()), + ); + const recordedContext = + historical?.requestContext ?? + options.captureRequestContext?.(liveContext) ?? + Object.fromEntries(liveContext.entries()); + const safeContext = requireRecord( + decodeMemoryValue(encodeMemoryValue(recordedContext)), + "request context", + ); + validateMemoryReplayContext(selector, safeContext); + const requestContext = new MemoryReplayRequestContext(selector); + for (const [key, value] of Object.entries(safeContext)) + requestContext.set(key, value); const abort = new AbortController(); let state: AdapterRunState | undefined; let requestCapture: ReturnType | undefined; @@ -281,7 +322,6 @@ export function createMemoryReplayAgent( }); } else { const source = await options.sourceMemory(); - const selector = getSelector(callerOptions); const binding = createMemoryCaptureBinding({ invocationId, ...selector, @@ -353,19 +393,6 @@ export function createMemoryReplayAgent( throw new Error( "Agent factory must use its supplied pinned workspace.", ); - const liveContext = callerOptions.requestContext ?? new RequestContext(); - const recordedContext = - historical?.requestContext ?? - options.captureRequestContext?.(liveContext) ?? - Object.fromEntries(liveContext.entries()); - // Validate before any dynamic resolver can observe an unrecordable value. - const safeContext = requireRecord( - decodeMemoryValue(encodeMemoryValue(recordedContext)), - "request context", - ); - const requestContext = new RequestContext(); - for (const [key, value] of Object.entries(safeContext)) - requestContext.set(key, value); const dynamic = { requestContext, mastra: options.mastra }; const instructions = historical ? historical.configuration.instructions @@ -394,6 +421,14 @@ export function createMemoryReplayAgent( ? await config.defaultOptions(dynamic) : (config.defaultOptions ?? {}); const defaults = requireRecord(resolvedDefaults, "default options"); + // Dynamic resolvers receive the mutable context that native Mastra uses. + const effectiveContext = requireRecord( + decodeMemoryValue( + encodeMemoryValue(Object.fromEntries(requestContext.entries())), + ), + "request context", + ); + validateMemoryReplayContext(selector, effectiveContext); const { deepMerge } = await import("@mastra/core/utils"); const callerData = { ...callerOptions }; delete callerData.requestContext; @@ -407,6 +442,14 @@ export function createMemoryReplayAgent( ) : deepMerge(defaults, callerData); assertSupportedConfiguration(config, effective); + const effectiveSelector = getSelector(effective); + if ( + effectiveSelector.threadId !== selector.threadId || + effectiveSelector.resourceId !== selector.resourceId + ) + throw new Error( + "Invocation memory selectors differ from the captured selectors.", + ); if (record(effective.memory) && effective.memory.options !== undefined) throw new Error( "Per-call memory.options are unsupported. Set the complete memory configuration in sourceMemory instead.", @@ -436,7 +479,7 @@ export function createMemoryReplayAgent( rawInput: invocationInput, initialSnapshot: runtime.initialSnapshot as MastraMemorySnapshot, configuration, - requestContext: safeContext, + requestContext: effectiveContext, files: files.files, }); if (!envelope.complete && historical) @@ -572,7 +615,25 @@ export function createMemoryReplayAgent( }, }); } catch (error) { - await runtime.finish(); + try { + await runtime.finish(); + } catch (cleanupError) { + if (options.onRecordingError) { + void Promise.resolve() + .then(() => + options.onRecordingError?.({ + error: cleanupError, + sessionId: state?.sessionId, + stage: "complete", + }), + ) + .catch(() => undefined); + } else { + console.warn( + "Kitaru memory cleanup failed after the invocation failed", + ); + } + } throw error; } } diff --git a/packages/mastra/test/memory-replay-safety.test.ts b/packages/mastra/test/memory-replay-safety.test.ts new file mode 100644 index 000000000..c82301ea4 --- /dev/null +++ b/packages/mastra/test/memory-replay-safety.test.ts @@ -0,0 +1,424 @@ +import { + MASTRA_AUTH_TOKEN_KEY, + MASTRA_RESOURCE_ID_KEY, + MASTRA_THREAD_ID_KEY, + RequestContext, +} from "@mastra/core/request-context"; +import { InMemoryStore } from "@mastra/core/storage"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { afterEach, expect, it, vi } from "vitest"; +import { createProcessLocalMemoryAccess } from "../src/memory-binding.js"; +import { + createMemoryReplayEnvelope, + decodeMemoryReplayEnvelope, + type MastraMemoryReplayInput, +} from "../src/memory-snapshot.js"; +import { + createMemoryReplayAgent, + type MemoryReplayAgentFactory, + type MemoryReplayAgentOptions, +} from "../src/stateful-agent.js"; +import { textStream } from "./helpers/memory-agent.js"; +import { AGENT_ID, installTestApi, REPLAY_ID } from "./helpers.js"; + +const stores: InMemoryStore[] = []; +afterEach(async () => { + for (const store of stores.splice(0)) await store.close(); + vi.restoreAllMocks(); + vi.unstubAllGlobals(); + vi.unstubAllEnvs(); +}); + +function input(): MastraMemoryReplayInput { + return { + invocationId: "safety-test", + rawInput: "hello", + initialSnapshot: { + threadId: "thread", + resourceId: "resource", + thread: null, + resource: null, + messages: [], + records: [], + }, + configuration: { + runOptions: { memory: { thread: "thread", resource: "resource" } }, + }, + requestContext: {}, + files: [], + }; +} + +function fixture( + factory?: MemoryReplayAgentFactory, + overrides: Partial = {}, +) { + const store = new InMemoryStore(); + stores.push(store); + const domain = store.stores.memory; + if (!domain) throw new Error("Missing native memory domain"); + const modelCall = vi.fn(async () => textStream("done")); + const model = new MastraLanguageModelV2Mock({ + provider: "fixture", + modelId: "actor", + doStream: modelCall, + }); + const lease = createProcessLocalMemoryAccess(); + const acquire = vi.spyOn(lease, "acquire"); + const adapter = createMemoryReplayAgent( + factory ?? + (({ memory }) => ({ + id: "safety", + name: "Safety", + memory, + model, + instructions: "Answer", + })), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + apiKey: "fixture", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: async () => {}, + domain, + configuration: { semanticRecall: false }, + exclusiveAccess: lease, + }), + resolveModel: () => model, + ...overrides, + }, + ); + return { adapter, modelCall, acquire, lease }; +} + +it.each(["defaultOptions", "runOptions", "memoryConfig"])( + "excludes transport credentials from %s during capture and decode", + (key) => { + const original = input(); + const unsafe = { + providerOptions: { + openai: { websocket: { headers: { "x-custom-access": "CREDENTIAL" } } }, + }, + }; + const clean = createMemoryReplayEnvelope(original); + const envelope = createMemoryReplayEnvelope({ + ...original, + configuration: { ...original.configuration, [key]: unsafe }, + }); + expect(envelope.complete).toBe(false); + expect(JSON.stringify(envelope)).not.toContain("CREDENTIAL"); + expect(() => + decodeMemoryReplayEnvelope({ + ...clean, + configuration: { ...original.configuration, [key]: unsafe }, + }), + ).toThrow(/transport/i); + }, +); + +it("rejects native auth tokens during capture and decode", () => { + const original = input(); + const requestContext = { [MASTRA_AUTH_TOKEN_KEY]: "CREDENTIAL" }; + const clean = createMemoryReplayEnvelope(original); + const unsafe = createMemoryReplayEnvelope({ ...original, requestContext }); + expect(unsafe.complete).toBe(false); + expect(JSON.stringify(unsafe)).not.toContain("CREDENTIAL"); + expect(() => + decodeMemoryReplayEnvelope({ ...clean, requestContext }), + ).toThrow(/auth/i); +}); + +it("rejects default auth-token context capture before recording or execution", async () => { + const api = installTestApi(); + const { adapter, modelCall } = fixture(); + const requestContext = new RequestContext(); + requestContext.set(MASTRA_AUTH_TOKEN_KEY, "CREDENTIAL"); + await expect( + adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + requestContext, + }), + ).rejects.toThrow(/auth/i); + expect(modelCall).not.toHaveBeenCalled(); + expect(api.calls).toEqual([]); +}); + +it.each([MASTRA_THREAD_ID_KEY, MASTRA_RESOURCE_ID_KEY])( + "rejects a mismatched %s before acquiring the source lease", + async (key) => { + const api = installTestApi(); + const { adapter, modelCall, acquire } = fixture(); + const requestContext = new RequestContext(); + requestContext.set(key, "other"); + await expect( + adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + requestContext, + }), + ).rejects.toThrow(/selector/i); + expect(acquire).not.toHaveBeenCalled(); + expect(modelCall).not.toHaveBeenCalled(); + expect(api.calls).toEqual([]); + }, +); + +it("does not let selective context capture bypass middleware memory selectors", async () => { + const api = installTestApi(); + const { adapter, acquire } = fixture(undefined, { + captureRequestContext: () => ({ locale: "en" }), + }); + const requestContext = new RequestContext(); + requestContext.set(MASTRA_THREAD_ID_KEY, "authorized-thread"); + await expect( + adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + requestContext, + }), + ).rejects.toThrow(/selector/i); + expect(acquire).not.toHaveBeenCalled(); + expect(api.calls).toEqual([]); +}); + +it("allows selective capture to exclude live authentication tokens", async () => { + const api = installTestApi(); + const { adapter } = fixture(undefined, { + captureRequestContext: () => ({ locale: "en" }), + }); + const requestContext = new RequestContext(); + requestContext.set(MASTRA_AUTH_TOKEN_KEY, "CREDENTIAL"); + const result = await adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + requestContext, + }); + await result.consumeStream(); + expect(JSON.stringify(api.calls)).not.toContain("CREDENTIAL"); + expect( + api.calls.filter((call) => call.method === "PATCH").at(-1)?.body?.status, + ).toBe("completed"); +}); + +it("records and replays matching middleware and invocation memory selectors", async () => { + const api = installTestApi(); + const { adapter, modelCall, acquire } = fixture(); + const requestContext = new RequestContext(); + requestContext.set(MASTRA_THREAD_ID_KEY, "thread"); + requestContext.set(MASTRA_RESOURCE_ID_KEY, "resource"); + const baseline = await adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + requestContext, + }); + await baseline.consumeStream(); + const recorded = api.calls.find((call) => call.path === "/api/v1/sessions") + ?.body?.inputs; + expect(recorded).toHaveProperty("mastra_memory_replay.complete", true); + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(recorded)); + const replay = await adapter.stream("ignored", { + memory: { thread: "live-thread", resource: "live-resource" }, + }); + await replay.consumeStream(); + expect(modelCall).toHaveBeenCalledTimes(2); + expect(acquire).toHaveBeenCalledTimes(1); + expect( + api.calls + .filter((call) => call.method === "PATCH") + .map((call) => call.body?.status), + ).toEqual(["completed", "completed"]); +}); + +it.each([MASTRA_THREAD_ID_KEY, MASTRA_RESOURCE_ID_KEY, MASTRA_AUTH_TOKEN_KEY])( + "rejects late processor writes to %s", + async (key) => { + vi.spyOn(console, "error").mockImplementation(() => {}); + const api = installTestApi(); + const modelCall = vi.fn(async () => textStream("done")); + const { adapter } = fixture(({ memory }) => ({ + id: "processor-mutation", + name: "Processor mutation", + memory, + model: new MastraLanguageModelV2Mock({ doStream: modelCall }), + instructions: "Answer", + inputProcessors: [ + { + id: "context-mutation", + processInput({ requestContext, messages }) { + if (!requestContext) throw new Error("Missing request context"); + requestContext.setRaw(key, "FORBIDDEN"); + return messages; + }, + }, + ], + })); + await expect( + adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + }), + ).rejects.toThrow(/processor/i); + expect(modelCall).not.toHaveBeenCalled(); + expect(JSON.stringify(api.calls)).not.toContain("FORBIDDEN"); + expect( + api.calls.filter((call) => call.method === "PATCH").at(-1)?.body?.status, + ).toBe("failed"); + }, +); + +it.each([MASTRA_THREAD_ID_KEY, MASTRA_RESOURCE_ID_KEY, MASTRA_AUTH_TOKEN_KEY])( + "rejects dynamic resolver mutations of %s before native execution", + async (key) => { + const api = installTestApi(); + const { adapter, modelCall } = fixture(({ memory }) => ({ + id: "mutation", + name: "Mutation", + memory, + model: new MastraLanguageModelV2Mock({ + doStream: async () => textStream("done"), + }), + instructions: ({ requestContext }) => { + requestContext.set(key, "other"); + return "Answer"; + }, + })); + await expect( + adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + }), + ).rejects.toThrow(/selector|auth/i); + expect(modelCall).not.toHaveBeenCalled(); + expect(api.calls).toEqual([]); + }, +); + +it("never uploads provider transport credentials in session inputs", async () => { + const api = installTestApi(); + const { adapter } = fixture(); + const result = await adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + providerOptions: { + openai: { websocket: { headers: { "x-custom-access": "CREDENTIAL" } } }, + }, + }); + await result.consumeStream(); + const session = api.calls.find((call) => call.path === "/api/v1/sessions"); + expect(session?.body?.inputs).toHaveProperty( + "mastra_memory_replay.complete", + false, + ); + expect(JSON.stringify(api.calls)).not.toContain("CREDENTIAL"); +}); + +it("rejects late auth-token additions by default-option resolvers", async () => { + const api = installTestApi(); + const { adapter } = fixture(({ memory }) => ({ + id: "defaults", + name: "Defaults", + memory, + instructions: "Answer", + model: new MastraLanguageModelV2Mock({}), + defaultOptions: ({ requestContext }) => { + requestContext.setRaw(MASTRA_AUTH_TOKEN_KEY, "CREDENTIAL"); + return {}; + }, + })); + await expect( + adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + }), + ).rejects.toThrow(/auth/i); + expect(api.calls).toEqual([]); +}); + +it("marks nested abort signals as unsupported transport configuration", () => { + const original = input(); + const configuration = { nested: [{ abortSignal: "transport-state" }] }; + expect( + createMemoryReplayEnvelope({ ...original, configuration }).complete, + ).toBe(false); + expect(() => + decodeMemoryReplayEnvelope({ + ...createMemoryReplayEnvelope(original), + configuration, + }), + ).toThrow(/transport/i); +}); + +it("rejects replay envelopes with selectors inconsistent with the snapshot", () => { + const envelope = createMemoryReplayEnvelope(input()); + expect(() => + decodeMemoryReplayEnvelope({ + ...envelope, + requestContext: { [MASTRA_THREAD_ID_KEY]: "other" }, + }), + ).toThrow(/selector/i); + expect(() => + decodeMemoryReplayEnvelope({ + ...envelope, + configuration: { + runOptions: { memory: { thread: "other", resource: "resource" } }, + }, + }), + ).toThrow(/selector/i); +}); + +it("preserves setup errors when settling memory also rejects", async () => { + const original = new Error("Original setup error"); + const cleanup = new Error("Memory cleanup error"); + const onRecordingError = vi.fn(); + const { adapter, lease } = fixture( + ({ memory }) => { + vi.spyOn(memory, "settled").mockRejectedValue(cleanup); + throw original; + }, + { onRecordingError }, + ); + await expect( + adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + }), + ).rejects.toBe(original); + await vi.waitFor(() => + expect(onRecordingError).toHaveBeenCalledWith( + expect.objectContaining({ error: cleanup, stage: "complete" }), + ), + ); + const release = await lease.acquire({ + threadId: "thread", + resourceId: "resource", + }); + await release(); +}); + +it("preserves native stream errors when cleanup and diagnostics reject", async () => { + installTestApi(); + const original = new Error("Native stream failed"); + const cleanup = new Error("Cleanup failed"); + const onRecordingError = vi.fn(async () => { + throw new Error("Diagnostic failed"); + }); + vi.spyOn(Agent.prototype, "stream").mockRejectedValue(original); + const { adapter } = fixture( + ({ memory }) => { + vi.spyOn(memory, "settled").mockRejectedValue(cleanup); + return { + id: "failure", + name: "Failure", + memory, + instructions: "Answer", + model: new MastraLanguageModelV2Mock({}), + }; + }, + { onRecordingError }, + ); + await expect( + adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + }), + ).rejects.toBe(original); + await vi.waitFor(() => + expect(onRecordingError).toHaveBeenCalledWith( + expect.objectContaining({ error: cleanup, stage: "complete" }), + ), + ); +}); + +import { Agent } from "@mastra/core/agent"; From 23038f1133121cce1a6362dcf1a10a86bfa03542 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 23 Sep 2026 20:51:07 +0200 Subject: [PATCH 017/115] Isolate the optional Mastra memory entrypoint --- changelog.d/mastra-memory-replay.added.md | 2 +- devtools/check_mastra_memory_replay.mjs | 2 +- docs/book/adapters/mastra.md | 6 +- packages/mastra/README.md | 6 +- packages/mastra/package.json | 4 ++ packages/mastra/src/index.ts | 34 ---------- packages/mastra/src/memory.ts | 34 ++++++++++ packages/mastra/test/processor-replay.test.ts | 2 +- .../mastra/test/stateful-overrides.test.ts | 2 +- scripts/fixtures/mastra-memory-smoke.mjs | 2 +- scripts/fixtures/mastra-memory-types.ts | 2 +- scripts/smoke-typescript-packages.mjs | 62 +++++++++++++++++++ 12 files changed, 114 insertions(+), 44 deletions(-) create mode 100644 packages/mastra/src/memory.ts diff --git a/changelog.d/mastra-memory-replay.added.md b/changelog.d/mastra-memory-replay.added.md index 467f51dba..167874569 100644 --- a/changelog.d/mastra-memory-replay.added.md +++ b/changelog.d/mastra-memory-replay.added.md @@ -1 +1 @@ -- Add opt-in Mastra streamed replay with isolated native thread memory. `createMemoryReplayAgent()` records starting working and observational memory, effective actor requests, controlled file bytes, and ordered memory changes, then replays from that state without reading or writing production memory. Requires Mastra core 1.67.0 and memory 1.30.0; existing history-only recordings must be recorded again to use this path. +- Add opt-in Mastra streamed replay with isolated native thread memory. `createMemoryReplayAgent()` from `@zenml-io/kitaru-mastra/memory` records starting working and observational memory, effective actor requests, controlled file bytes, and ordered memory changes, then replays from that state without reading or writing production memory. Requires Mastra core 1.67.0 and memory 1.30.0; existing history-only recordings must be recorded again to use this path. diff --git a/devtools/check_mastra_memory_replay.mjs b/devtools/check_mastra_memory_replay.mjs index a35ec87cd..2bba5be47 100644 --- a/devtools/check_mastra_memory_replay.mjs +++ b/devtools/check_mastra_memory_replay.mjs @@ -7,7 +7,7 @@ import { pathToFileURL } from "node:url"; import { createMemoryReplayAgent, createProcessLocalMemoryAccess, -} from "../packages/mastra/dist/index.js"; +} from "../packages/mastra/dist/memory.js"; const require = createRequire( new URL("../packages/mastra/package.json", import.meta.url), diff --git a/docs/book/adapters/mastra.md b/docs/book/adapters/mastra.md index 94d84a8cb..e9a326425 100644 --- a/docs/book/adapters/mastra.md +++ b/docs/book/adapters/mastra.md @@ -183,7 +183,7 @@ A missing, incomplete, or lossy snapshot produces an actionable unsupported-repl ## Isolated memory replay -Use `createMemoryReplayAgent()` when a consumed stream needs thread-scoped schema working memory, observational memory, or controlled input processors. This opt-in factory requires exactly `@mastra/core@1.67.0` and `@mastra/memory@1.30.0`. The existing `KitaruAgent` wrapper keeps its history-only memory behavior. +Import `createMemoryReplayAgent()` from `@zenml-io/kitaru-mastra/memory` when a consumed stream needs thread-scoped schema working memory, observational memory, or controlled input processors. This opt-in factory requires exactly `@mastra/core@1.67.0` and `@mastra/memory@1.30.0`. The existing `KitaruAgent` wrapper stays at the package root, keeps its history-only memory behavior, and does not require `@mastra/memory`. ```bash pnpm add @zenml-io/kitaru-mastra @mastra/core@1.67.0 @mastra/memory@1.30.0 zod @@ -199,7 +199,7 @@ import { Memory } from "@mastra/memory"; import { createMemoryReplayAgent, createProcessLocalMemoryAccess, -} from "@zenml-io/kitaru-mastra"; +} from "@zenml-io/kitaru-mastra/memory"; import { z } from "zod"; const store = new InMemoryStore(); @@ -253,6 +253,8 @@ Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KIT All source-memory writers must participate in the same exclusive-access mechanism. The process-local helper is suitable only when every writer shares that instance in one process. Use a distributed implementation of `MastraExclusiveMemoryAccess` when other processes can write. Its `acquire()` method must hold access until the returned release function runs. `settled()` must join pending work on the source `Memory` instance; it is not a lock. Configure working and observational memory with `scope: "thread"`; resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` are outside this contract. +Keep `memory.thread` and `memory.resource` equal to any reserved Mastra thread/resource IDs in `requestContext`; a mismatch is rejected before native execution. Use `captureRequestContext` to select safe values: native authentication tokens are rejected, and transport headers in recorded configuration make the replay envelope incomplete. Use `resolveModel` to reconstruct credential-bearing model instances for replay. + Dynamic `instructions`, `model`, and `defaultOptions` resolve once during baseline setup. Replay uses their recorded values instead of calling those resolvers again. `resolveModel` must resolve the recorded actor, observer, and reflector model identifiers as well as any allowed actor override. A `system_prompt` override replaces only application instructions and retains recorded extra system context. Model and model-setting overrides affect the actor; observation and reflection retain their recorded configuration. Raw-input `prompt` overrides are rejected; record a new baseline to change invocation input. ### Create and inspect a memory replay diff --git a/packages/mastra/README.md b/packages/mastra/README.md index 2a9b4e49c..2d57a6c0f 100644 --- a/packages/mastra/README.md +++ b/packages/mastra/README.md @@ -147,7 +147,7 @@ History matching is guaranteed only for traces recorded and replayed through thi ## Isolated memory replay -Use `createMemoryReplayAgent()` when a consumed stream needs thread-scoped schema working memory, observational memory, or controlled input processors. This opt-in factory requires exactly `@mastra/core@1.67.0` and `@mastra/memory@1.30.0`. The existing `KitaruAgent` wrapper keeps its history-only memory behavior. +Import `createMemoryReplayAgent()` from `@zenml-io/kitaru-mastra/memory` when a consumed stream needs thread-scoped schema working memory, observational memory, or controlled input processors. This opt-in factory requires exactly `@mastra/core@1.67.0` and `@mastra/memory@1.30.0`. The existing `KitaruAgent` wrapper stays at the package root, keeps its history-only memory behavior, and does not require `@mastra/memory`. ```bash pnpm add @zenml-io/kitaru-mastra @mastra/core@1.67.0 @mastra/memory@1.30.0 zod @@ -163,7 +163,7 @@ import { Memory } from "@mastra/memory"; import { createMemoryReplayAgent, createProcessLocalMemoryAccess, -} from "@zenml-io/kitaru-mastra"; +} from "@zenml-io/kitaru-mastra/memory"; import { z } from "zod"; const store = new InMemoryStore(); @@ -217,6 +217,8 @@ Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KIT All source-memory writers must participate in the same exclusive-access mechanism. The process-local helper is suitable only when every writer shares that instance in one process. Use a distributed implementation of `MastraExclusiveMemoryAccess` when other processes can write. Its `acquire()` method must hold access until the returned release function runs. `settled()` must join pending work on the source `Memory` instance; it is not a lock. Configure working and observational memory with `scope: "thread"`; resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` are outside this contract. +Keep `memory.thread` and `memory.resource` equal to any reserved Mastra thread/resource IDs in `requestContext`; a mismatch is rejected before native execution. Use `captureRequestContext` to select safe values: native authentication tokens are rejected, and transport headers in recorded configuration make the replay envelope incomplete. Use `resolveModel` to reconstruct credential-bearing model instances for replay. + Dynamic `instructions`, `model`, and `defaultOptions` resolve once during baseline setup. Replay uses their recorded values instead of calling those resolvers again. `resolveModel` must resolve the recorded actor, observer, and reflector model identifiers as well as any allowed actor override. A `system_prompt` override replaces only application instructions and retains recorded extra system context. Model and model-setting overrides affect the actor; observation and reflection retain their recorded configuration. Raw-input `prompt` overrides are rejected; record a new baseline to change invocation input. ### Create and inspect a memory replay diff --git a/packages/mastra/package.json b/packages/mastra/package.json index d1a080801..cc79da4c7 100644 --- a/packages/mastra/package.json +++ b/packages/mastra/package.json @@ -24,6 +24,10 @@ ".": { "types": "./dist/index.d.ts", "import": "./dist/index.js" + }, + "./memory": { + "types": "./dist/memory.d.ts", + "import": "./dist/memory.js" } }, "files": [ diff --git a/packages/mastra/src/index.ts b/packages/mastra/src/index.ts index 13d831dcf..07ffc713c 100644 --- a/packages/mastra/src/index.ts +++ b/packages/mastra/src/index.ts @@ -1,43 +1,9 @@ export { KitaruAgent } from "./agent.js"; -export type { - MastraExclusiveMemoryAccess, - MastraMemoryCaptureBinding, - MastraMemoryCaptureOptions, - MastraMemoryMutation, - MastraMemorySelector, -} from "./memory-binding.js"; -export { - createMemoryCaptureBinding, - createProcessLocalMemoryAccess, -} from "./memory-binding.js"; -export type { - MastraFileManifestEntry, - MastraMemoryReplayEnvelope, - MastraMemoryReplayInput, - MastraMemorySnapshot, - MastraRecordedFile, -} from "./memory-snapshot.js"; -export { - createMemoryReplayEnvelope, - decodeMemoryReplayEnvelope, - decodeMemoryValue, - encodeMemoryValue, - MEMORY_REPLAY_KEY, - restoreMemoryReplayEnvelope, - validateMemorySnapshot, -} from "./memory-snapshot.js"; export type { MastraEvaluatorOptions, RunnableMastraScorer, } from "./scorers.js"; export { createMastraEvaluator } from "./scorers.js"; -export type { - MemoryReplayAgentBindings, - MemoryReplayAgentFactory, - MemoryReplayAgentOptions, -} from "./stateful-agent.js"; - -export { createMemoryReplayAgent } from "./stateful-agent.js"; export type { ConfiguredAfterToolCall, ConfiguredBeforeToolCall, diff --git a/packages/mastra/src/memory.ts b/packages/mastra/src/memory.ts new file mode 100644 index 000000000..29ad78cfa --- /dev/null +++ b/packages/mastra/src/memory.ts @@ -0,0 +1,34 @@ +export type { + MastraExclusiveMemoryAccess, + MastraMemoryCaptureBinding, + MastraMemoryCaptureOptions, + MastraMemoryMutation, + MastraMemorySelector, +} from "./memory-binding.js"; +export { + createMemoryCaptureBinding, + createProcessLocalMemoryAccess, +} from "./memory-binding.js"; +export type { + MastraFileManifestEntry, + MastraMemoryReplayEnvelope, + MastraMemoryReplayInput, + MastraMemorySnapshot, + MastraRecordedFile, +} from "./memory-snapshot.js"; +export { + createMemoryReplayEnvelope, + decodeMemoryReplayEnvelope, + decodeMemoryValue, + encodeMemoryValue, + MEMORY_REPLAY_KEY, + restoreMemoryReplayEnvelope, + validateMemorySnapshot, +} from "./memory-snapshot.js"; +export type { + MemoryReplayAgentBindings, + MemoryReplayAgentFactory, + MemoryReplayAgentOptions, +} from "./stateful-agent.js"; + +export { createMemoryReplayAgent } from "./stateful-agent.js"; diff --git a/packages/mastra/test/processor-replay.test.ts b/packages/mastra/test/processor-replay.test.ts index 2b903ef65..a36f17d62 100644 --- a/packages/mastra/test/processor-replay.test.ts +++ b/packages/mastra/test/processor-replay.test.ts @@ -10,7 +10,7 @@ import { createMemoryReplayAgent, createProcessLocalMemoryAccess, MEMORY_REPLAY_KEY, -} from "../src/index.js"; +} from "../src/memory.js"; import { createMemoryRuntime, FILE_URL, diff --git a/packages/mastra/test/stateful-overrides.test.ts b/packages/mastra/test/stateful-overrides.test.ts index ef1c15590..8fed59648 100644 --- a/packages/mastra/test/stateful-overrides.test.ts +++ b/packages/mastra/test/stateful-overrides.test.ts @@ -103,7 +103,7 @@ import { createMemoryReplayAgent, createProcessLocalMemoryAccess, MEMORY_REPLAY_KEY, -} from "../src/index.js"; +} from "../src/memory.js"; import { AGENT_ID, installTestApi, diff --git a/scripts/fixtures/mastra-memory-smoke.mjs b/scripts/fixtures/mastra-memory-smoke.mjs index 77c7612e2..124042d9c 100644 --- a/scripts/fixtures/mastra-memory-smoke.mjs +++ b/scripts/fixtures/mastra-memory-smoke.mjs @@ -4,7 +4,7 @@ import { Memory } from "@mastra/memory"; import { createMemoryReplayAgent, createProcessLocalMemoryAccess, -} from "@zenml-io/kitaru-mastra"; +} from "@zenml-io/kitaru-mastra/memory"; import { z } from "zod"; const calls = []; diff --git a/scripts/fixtures/mastra-memory-types.ts b/scripts/fixtures/mastra-memory-types.ts index a7d16118f..35bb7653e 100644 --- a/scripts/fixtures/mastra-memory-types.ts +++ b/scripts/fixtures/mastra-memory-types.ts @@ -5,7 +5,7 @@ import { createMemoryReplayAgent, createProcessLocalMemoryAccess, type MemoryReplayAgentOptions, -} from "@zenml-io/kitaru-mastra"; +} from "@zenml-io/kitaru-mastra/memory"; import { z } from "zod"; const store = new InMemoryStore(); diff --git a/scripts/smoke-typescript-packages.mjs b/scripts/smoke-typescript-packages.mjs index fc56d789c..9129d2164 100644 --- a/scripts/smoke-typescript-packages.mjs +++ b/scripts/smoke-typescript-packages.mjs @@ -9,6 +9,7 @@ import { tmpdir } from "node:os"; import { join, resolve } from "node:path"; import { spawnSync } from "node:child_process"; import { fileURLToPath } from "node:url"; +import ts from "typescript"; import { loadTypescriptPackageMetadata } from "./typescript-packages.mjs"; const repositoryRoot = resolve(fileURLToPath(new URL("..", import.meta.url))); @@ -48,6 +49,66 @@ function assertPackageContents(tarball) { } } +function assertLegacyDeclarations(consumerRoot) { + const upstream = join(consumerRoot, "upstream-types.ts"); + const consumer = join(consumerRoot, "legacy-types.ts"); + // Load the same upstream declaration graph without the adapter. Mastra's + // published declarations can have errors of their own; only exact matches + // from that independent baseline may be excluded from the consumer result. + writeFileSync( + upstream, + [ + "@mastra/core/agent", + "@mastra/core/llm", + "@mastra/core/request-context", + "@mastra/core/stream", + "@mastra/core/tools", + "@mastra/core/evals", + "@zenml-io/kitaru", + "@zenml-io/kitaru/adapter", + ].map((name) => `import ${JSON.stringify(name)};`).join("\n"), + ); + writeFileSync( + consumer, + `import type { Agent } from "@mastra/core/agent"; +import { KitaruAgent, type KitaruAgentOptions } from "@zenml-io/kitaru-mastra"; +declare const agent: Agent; +declare const options: KitaruAgentOptions; +new KitaruAgent(agent, options); +`, + ); + const options = { + lib: ["lib.es2022.d.ts", "lib.dom.d.ts"], + module: ts.ModuleKind.NodeNext, + moduleResolution: ts.ModuleResolutionKind.NodeNext, + noEmit: true, + skipLibCheck: false, + strict: true, + types: [], + }; + const baseline = ts.createProgram([upstream], options); + const candidate = ts.createProgram([consumer], options); + const signature = (diagnostic) => JSON.stringify({ + file: diagnostic.file?.fileName, + start: diagnostic.start, + code: diagnostic.code, + message: ts.flattenDiagnosticMessageText(diagnostic.messageText, "\n"), + }); + const baselineDiagnostics = ts.getPreEmitDiagnostics(baseline); + const known = new Set(baselineDiagnostics.map(signature)); + const unexpected = ts.getPreEmitDiagnostics(candidate).filter((diagnostic) => + diagnostic.file?.fileName.includes("/node_modules/@zenml-io/") || + !known.has(signature(diagnostic)), + ); + const reachedMemory = candidate.getSourceFiles().some((file) => + /\/kitaru-mastra\/dist\/(?:memory|stateful)[^/]*\.d\.ts$/.test(file.fileName), + ); + if (reachedMemory || unexpected.length) { + throw new Error(`Legacy declaration isolation failed:\n${unexpected.map(signature).join("\n")}\nMemory declarations reached: ${reachedMemory}`); + } + console.log(`Legacy declarations: no adapter errors; ${baselineDiagnostics.length} upstream baseline diagnostics checked separately (skipLibCheck=false).`); +} + function writeConsumerFiles(consumerRoot) { writeFileSync( join(consumerRoot, "package.json"), @@ -465,6 +526,7 @@ function smokeConsumer({ artifactRoot, mastraVersion, npmCache }) { ], consumerRoot, ); + assertLegacyDeclarations(consumerRoot); } run(process.execPath, ["index.mjs"], consumerRoot); run( From e15a11dcdc75971dc2adf69180bb84ba470cf8dd Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 24 Sep 2026 07:50:18 +0200 Subject: [PATCH 018/115] Harden Mastra memory replay --- changelog.d/mastra-memory-replay.added.md | 2 +- devtools/check_mastra_memory_replay.mjs | 70 +- devtools/check_mastra_memory_replay.py | 69 +- docs/book/adapters/mastra.md | 51 +- openapi/openapi.json | 4 + packages/core/src/adapter/index.ts | 4 + packages/core/src/adapter/recorded-json.ts | 125 ++++ packages/core/src/adapter/replay.ts | 14 +- packages/core/src/adapter/run-recorder.ts | 37 +- packages/core/src/generated/openapi.ts | 5 + packages/core/test/adapter/replay.test.ts | 87 +++ .../core/test/adapter/run-recorder.test.ts | 57 ++ packages/mastra/README.md | 51 +- packages/mastra/src/memory-binding.ts | 402 +++++++++-- packages/mastra/src/memory-replay.ts | 57 +- packages/mastra/src/memory-snapshot.ts | 218 +++--- packages/mastra/src/memory.ts | 2 + packages/mastra/src/om-result-tape.ts | 211 ++++++ packages/mastra/src/request-capture.ts | 36 +- packages/mastra/src/stateful-agent.ts | 465 +++++++++++-- packages/mastra/src/stateful-files.ts | 294 +++++++- packages/mastra/src/stateful-tools.ts | 5 +- packages/mastra/src/stateful-workspace.ts | 24 +- packages/mastra/src/step-recorder.ts | 18 +- packages/mastra/src/stream-recording.ts | 241 +++++-- packages/mastra/src/tool-policies.ts | 30 +- .../mastra/test/helpers/file-memory-access.ts | 213 ++++++ .../test/helpers/file-memory-worker.mjs | 49 ++ packages/mastra/test/implicit-om.test.ts | 123 ++++ packages/mastra/test/memory-binding.test.ts | 241 ++++++- .../mastra/test/memory-lease-process.test.ts | 189 ++++++ .../mastra/test/memory-replay-safety.test.ts | 291 +++++--- packages/mastra/test/memory-snapshot.test.ts | 125 +++- packages/mastra/test/om-result-tape.test.ts | 176 +++++ packages/mastra/test/processor-replay.test.ts | 188 +++++- packages/mastra/test/request-capture.test.ts | 34 + packages/mastra/test/stateful-files.test.ts | 182 ++++- .../mastra/test/stateful-overrides.test.ts | 639 ++++++++++++++++-- packages/mastra/test/stateful-tools.test.ts | 58 ++ .../mastra/test/stateful-workspace.test.ts | 11 + packages/mastra/test/stream-lifecycle.test.ts | 31 + src/kitaru/api_models/v1/session.py | 3 + src/kitaru/cli/app.py | 9 +- src/kitaru/client/exceptions.py | 13 + src/kitaru/mcp/errors.py | 9 +- .../db/repositories/session_repository.py | 22 + .../server/adapters/rest/mapping/sessions.py | 1 + .../interfaces/session_repository.py | 16 +- .../server/application/models/session.py | 1 + .../services/experiment_service.py | 6 +- .../application/services/replay_pipeline.py | 74 +- .../application/services/replay_service.py | 6 +- .../application/services/session_service.py | 30 +- src/kitaru/server/domain/session.py | 170 +++++ tests/cli/test_replays.py | 130 ++++ tests/conftest.py | 56 ++ tests/mcp/test_review_workflows.py | 265 ++++++++ tests/server/test_replay_pipeline.py | 174 ++++- tests/server/test_session_service.py | 167 +++++ tests/server/test_sessions_api.py | 36 + tests/server/test_sessions_api_pg.py | 82 +++ 61 files changed, 5874 insertions(+), 525 deletions(-) create mode 100644 packages/mastra/src/om-result-tape.ts create mode 100644 packages/mastra/test/helpers/file-memory-access.ts create mode 100644 packages/mastra/test/helpers/file-memory-worker.mjs create mode 100644 packages/mastra/test/implicit-om.test.ts create mode 100644 packages/mastra/test/memory-lease-process.test.ts create mode 100644 packages/mastra/test/om-result-tape.test.ts create mode 100644 packages/mastra/test/stateful-tools.test.ts diff --git a/changelog.d/mastra-memory-replay.added.md b/changelog.d/mastra-memory-replay.added.md index 167874569..2b6b2a1cc 100644 --- a/changelog.d/mastra-memory-replay.added.md +++ b/changelog.d/mastra-memory-replay.added.md @@ -1 +1 @@ -- Add opt-in Mastra streamed replay with isolated native thread memory. `createMemoryReplayAgent()` from `@zenml-io/kitaru-mastra/memory` records starting working and observational memory, effective actor requests, controlled file bytes, and ordered memory changes, then replays from that state without reading or writing production memory. Requires Mastra core 1.67.0 and memory 1.30.0; existing history-only recordings must be recorded again to use this path. +- Add opt-in Mastra streamed replay with isolated native thread memory. `createMemoryReplayAgent()` from `@zenml-io/kitaru-mastra/memory` records starting working and observational memory, ordered observation and reflection outputs, effective actor requests, controlled file bytes, and memory changes. Replay reuses those recorded outputs without reading or writing production memory. Requires Mastra core 1.67.0 and memory 1.30.0; existing history-only recordings must be recorded again to use this path. diff --git a/devtools/check_mastra_memory_replay.mjs b/devtools/check_mastra_memory_replay.mjs index 2bba5be47..9d8ca5886 100644 --- a/devtools/check_mastra_memory_replay.mjs +++ b/devtools/check_mastra_memory_replay.mjs @@ -23,7 +23,7 @@ const directory = process.env.CHECK_DIRECTORY; const replayId = process.env.KITARU_REPLAY_ID; const thread = "historical-thread"; const resource = "historical-resource"; -const fileUrl = "https://files.invalid/historical.pdf"; +const fileUrl = "https://files.invalid/historical.pdf?token=historical-secret"; const bytes = new Uint8Array(40000).fill(65); const report = { task_id: process.env.KITARU_TASK_ID, @@ -109,7 +109,24 @@ const actor = (id) => ({ return text("done"); }, }); -const models = { actor: actor("actor"), replacement: actor("replacement") }; +const observationModel = (id) => ({ + specificationVersion: "v2", + supportedUrls: {}, + provider: "fixture", + modelId: id, + doGenerate: async () => { + throw new Error("Unexpected OM generate call"); + }, + doStream: async () => { + throw new Error("Unexpected OM model call"); + }, +}); +const models = { + actor: actor("actor"), + replacement: actor("replacement"), + observer: observationModel("observer"), + reflector: observationModel("reflector"), +}; const store = new InMemoryStore(); const memory = new Memory({ storage: store, @@ -121,6 +138,11 @@ const memory = new Memory({ scope: "thread", schema: z.object({ preference: z.string() }), }, + observationalMemory: { + scope: "thread", + observation: { model: "fixture/observer", messageTokens: 100000 }, + reflection: { model: "fixture/reflector", observationTokens: 100000 }, + }, }, }); const domain = store.stores.memory; @@ -138,6 +160,44 @@ await memory.updateWorkingMemory({ resourceId: resource, workingMemory: JSON.stringify(currentProduction), }); +const nestedHotels = Array.from({ length: 1500 }, (_, hotel) => ({ + id: hotel, + details: Object.fromEntries( + Array.from({ length: 10 }, (_, field) => [`field${field}`, field]), + ), +})); +await domain.saveMessages({ + messages: Array.from({ length: 830 }, (_, index) => ({ + id: `historical-${index}`, + role: "user", + content: { + format: 2, + parts: [ + { + type: "text", + text: index === 0 ? "HISTORICAL_MESSAGE: " + "x".repeat(1_100_000) : `history ${index}`, + }, + ], + ...(index === 0 ? { metadata: { hotels: nestedHotels } } : {}), + }, + createdAt: new Date(Date.UTC(2026, 0, 1, 0, 0, index)), + threadId: thread, + resourceId: resource, + })), +}); +const omRecord = await domain.initializeObservationalMemory({ + threadId: thread, + resourceId: resource, + scope: "thread", + config: { fixture: true }, +}); +await domain.updateActiveObservations({ + id: omRecord.id, + observations: "HISTORICAL_OBSERVATION: likes blue.", + tokenCount: 10, + lastObservedAt: new Date(Date.UTC(2026, 0, 1, 0, 0, 0)), + observedMessageIds: ["historical-0"], +}); const sourceDomain = new Proxy(domain, { get(target, key) { const value = Reflect.get(target, key, target); @@ -171,7 +231,11 @@ const agent = createMemoryReplayAgent( id: "historical-file", async processInput({ messages }) { report.processor_calls++; - const file = await resolveFile(fileUrl); + const filePart = messages + .flatMap((message) => message.content.parts) + .find((part) => part.type === "file"); + assert(filePart, "Expected a file in native processor input"); + const file = await resolveFile(String(filePart.data)); assert.deepEqual(file.bytes, bytes); return messages.map((message) => ({ ...message, diff --git a/devtools/check_mastra_memory_replay.py b/devtools/check_mastra_memory_replay.py index 2f8d0f618..e0e1184dd 100644 --- a/devtools/check_mastra_memory_replay.py +++ b/devtools/check_mastra_memory_replay.py @@ -56,6 +56,7 @@ from kitaru.api_models.v1.session_node import SessionNodeListParams from kitaru.api_models.v1.session_run import SessionRunCreateRequest from kitaru.client.api_client import KitaruAPIClient +from kitaru.client.exceptions import APIError from kitaru.worker import Worker, WorkerConfig ARTIFACT = Path(__file__).with_suffix(".mjs").resolve() @@ -446,7 +447,7 @@ async def check(output: Path) -> None: }, { "type": "file", - "data": "https://files.invalid/historical.pdf", + "data": "https://files.invalid/historical.pdf?token=historical-secret", "mimeType": "application/pdf", }, ], @@ -462,9 +463,30 @@ async def check(output: Path) -> None: ) sessions = [session async for session in client.sessions.iter()] assert len(sessions) == 1 - baseline = await client.sessions.get(sessions[0].id) - assert baseline.inputs["mastra_memory_replay"]["complete"] is True - assert len(json.dumps(baseline.inputs).encode()) > 32768 + deadline = asyncio.get_running_loop().time() + 30 + while True: + baseline = await client.sessions.get(sessions[0].id) + if baseline.metadata.get("mastra_replay_state") != "pending": + break + assert asyncio.get_running_loop().time() < deadline, ( + "Baseline memory evidence did not finalize" + ) + await asyncio.sleep(0.2) + assert baseline.metadata["mastra_replay_state"] == "eligible" + envelope = baseline.inputs["mastra_memory_replay"] + assert envelope["complete"] is True + assert len(envelope["initialSnapshot"]["messages"]) == 830 + assert envelope["initialSnapshot"]["records"] + first_message = next( + message + for message in envelope["initialSnapshot"]["messages"] + if message["id"] == "historical-0" + ) + hotels = first_message["content"]["metadata"]["hotels"] + assert len(hotels) == 1500 + assert all(len(hotel["details"]) == 10 for hotel in hotels) + assert len(json.dumps(baseline.inputs).encode()) > 1_048_576 + assert "historical-secret" not in json.dumps(baseline.inputs) connection = await asyncpg.connect( host=DB_HOST, port=DB_PORT, @@ -551,34 +573,26 @@ async def check(output: Path) -> None: SessionCreateRequest( agent_id=agent.id, agent_version_id=version.id, + framework="mastra", origin=SessionOrigin.RECORDED, status=SessionStatus.COMPLETED, inputs=incomplete, outputs={"text": "incomplete fixture"}, ) ) - failed = await client.replays.create( - ReplayCreateRequest( - baseline_session_id=broken.id, - evaluators=[config], - override=ReplayOverride(**OVERRIDE), + try: + await client.replays.create( + ReplayCreateRequest( + baseline_session_id=broken.id, + evaluators=[config], + override=ReplayOverride(**OVERRIDE), + ) ) - ) - assert failed.job_id - failed_job = await await_job( - client, failed.job_id, "incomplete replay", 120 - ) - assert failed_job.status == JobStatus.FAILED, failed_job - failed_tasks = await client.jobs.list_tasks(failed_job.id) - assert any( - "incomplete" in (task.error or "").lower() - for task in failed_tasks.items - ) - failed = await client.replays.get(failed.id) - assert failed.status == "failed" - assert failed.result_session_id is None - failure_cli = await run_cli("replay", "get", str(failed.id)) - assert failure_cli["item"]["status"] == "failed" + except APIError as error: + assert error.status_code == 409 + assert "mastra_replay_incomplete" in error.detail + else: + raise AssertionError("Incomplete replay was scheduled") reports = [ json.loads(path.read_text()) for path in directory.glob("*.json") @@ -586,10 +600,9 @@ async def check(output: Path) -> None: ] good = [report for report in reports if report.get("result") == "passed"] assert len(good) == 4, reports - rejected = [ + assert not [ report for report in reports if report.get("result") == "failed" ] - assert rejected and all(report["actor_calls"] == 0 for report in rejected) assert all(report["unrelated_blob_status"] == 403 for report in good) assert all( not report["task_inputs_in_environment"] @@ -628,7 +641,7 @@ async def check(output: Path) -> None: "sdk": sdk, "cli": cli, "mcp": mcp_result, - "incomplete_replay_id": str(failed.id), + "incomplete_baseline_id": str(broken.id), "read_only_mutation_denied": True, "task_reports": reports, "provider_calls": 0, diff --git a/docs/book/adapters/mastra.md b/docs/book/adapters/mastra.md index e9a326425..020283fdb 100644 --- a/docs/book/adapters/mastra.md +++ b/docs/book/adapters/mastra.md @@ -189,7 +189,9 @@ Import `createMemoryReplayAgent()` from `@zenml-io/kitaru-mastra/memory` when a pnpm add @zenml-io/kitaru-mastra @mastra/core@1.67.0 @mastra/memory@1.30.0 zod ``` -The factory creates a fresh native agent for each invocation. For a baseline, it binds native memory to your source storage and records the starting state before recall. For a replay, it restores that state into a separate in-memory store. Native working-memory tools and observation/reflection jobs then evolve the isolated state as the model runs again. Replay never calls `sourceMemory()`. +The factory creates a fresh native agent for each invocation. A baseline uses your source storage and records the starting state before recall, then records the observer and reflector model outputs produced during that invocation. Replay restores the starting state into a separate in-memory store, runs the actor again, and supplies those recorded outputs to native observational memory (OM). Working-memory tools and native OM storage updates therefore run against isolated state. Replay never calls `sourceMemory()` or makes fresh observer/reflector provider calls. + +Recorded OM outputs match by call order, phase (observer or reflector), and model method. A missing, extra, reordered, or unused call fails replay with `KITARU_REPLAY_DIVERGED:mastra_om_call_order`. Changed OM input alone does not fail replay: Kitaru reuses the recorded output and records an `om_input_mismatch` span. This lets you compare actor instruction/model changes with recorded OM results, but does not measure how a fresh observer or reflector would respond to the changed conversation. The following binding uses a process-local store. Supply your existing public memory storage domain and its complete configuration for a persistent application: @@ -246,20 +248,43 @@ const output = await recorded.stream("My preference is green.", { context: [{ role: "system", content: "The customer is asking about preferences." }], }); await output.consumeStream(); -await store.close(); +// Keep the application and source store alive while recording finalizes. +// Inspect session eligibility before shutting down or starting a replay. ``` Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KITARU_AGENT_ID`, and the model provider credential. Register the compiled command as the agent version's run specification to run it through a worker. The same command serves baseline and replay tasks; the worker supplies the recorded input and replay identity. -All source-memory writers must participate in the same exclusive-access mechanism. The process-local helper is suitable only when every writer shares that instance in one process. Use a distributed implementation of `MastraExclusiveMemoryAccess` when other processes can write. Its `acquire()` method must hold access until the returned release function runs. `settled()` must join pending work on the source `Memory` instance; it is not a lock. Configure working and observational memory with `scope: "thread"`; resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` are outside this contract. +### Source ownership and supported configuration + +All writers to a source thread must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. Overlapping turns can continue natively, but their recordings become ineligible for replay. `settled()` joins pending work on the source `Memory` instance; it does not provide exclusive access. + +A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must preserve ownership through the final eligible-session update, mark an overlapping or unowned write unsafe through `markUnsafeWrite()`, and retain unsafe state across process loss. Call `resetAfterQuiescence()` only after all possible stale writers have stopped. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. + +Schema working memory requires explicit `scope: "thread"`. An observational-memory configuration object may omit `scope`, using Mastra's implicit thread scope, or set it to `"thread"`. Supply explicit observer/reflector model identities, either shared through `observationalMemory.model` or in the phase configuration. Resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` remain unsupported. + +Keep `memory.thread` and `memory.resource` consistent with reserved Mastra thread/resource IDs in `requestContext`. A mismatch cannot produce an eligible recording. Baseline callbacks receive the original live context. `captureRequestContext` selects only approved, replay-relevant JSON values; it does not remove values from the live context. Replay receives that recorded projection. A nonempty context without an explicit projection makes the recording ineligible. For example, return `{ locale: context.get("locale") }` when locale is the only value replay needs, or `{}` when none are needed. + +Never include authentication tokens, credentials, or signed URLs in the projection. Credential-like context keys are rejected, and transport headers in recorded configuration make the envelope incomplete. These checks cannot identify every secret hidden in an arbitrary string; choose recorded fields explicitly. Use `resolveModel` to reconstruct model instances from locally configured credentials. + +Dynamic `instructions`, `model`, and `defaultOptions` resolve during baseline setup; replay uses their recorded values. `resolveModel` must resolve the recorded actor, observer, and reflector model identifiers and any allowed actor override. OM identifiers must resolve to native stream-capable model objects, whose provider methods Kitaru intercepts to reuse recorded results during replay. A `system_prompt` override replaces only application instructions and retains recorded extra system context. Model and model-setting overrides affect the actor; observation and reflection retain their recorded configuration and outputs. Raw-input `prompt` overrides are rejected; record a new baseline to change invocation input. -Keep `memory.thread` and `memory.resource` equal to any reserved Mastra thread/resource IDs in `requestContext`; a mismatch is rejected before native execution. Use `captureRequestContext` to select safe values: native authentication tokens are rejected, and transport headers in recorded configuration make the replay envelope incomplete. Use `resolveModel` to reconstruct credential-bearing model instances for replay. +### Recording readiness -Dynamic `instructions`, `model`, and `defaultOptions` resolve once during baseline setup. Replay uses their recorded values instead of calling those resolvers again. `resolveModel` must resolve the recorded actor, observer, and reflector model identifiers as well as any allowed actor override. A `system_prompt` override replaces only application instructions and retains recorded extra system context. Model and model-setting overrides affect the actor; observation and reflection retain their recorded configuration. Raw-input `prompt` overrides are rejected; record a new baseline to change invocation input. +Native output and replay readiness are separate. Consume the baseline stream normally. Once the actor finishes, Kitaru finalizes recording in the background: it joins native memory work, records OM results and memory evidence, verifies source ownership, and persists the final input. Keep the application and source storage alive until that finalization finishes. `consumeStream()` alone is not a recording-completion barrier for a baseline. + +Inspect the baseline session's metadata and status: + +- `mastra_replay_state: "pending"`: recording has not finalized; do not replay yet. +- `mastra_replay_state: "eligible"`: the complete version-3 input and evidence were persisted for replay. +- `mastra_replay_state: "ineligible"`: recording could not establish a complete, isolated baseline. Inspect `mastra_replay_reason` and `mastra_native_state` to distinguish recording failure from native execution failure, then record a new baseline after resolving the cause. + +Recording-only problems do not replace the baseline's native answer. Diagnostics use `KITARU_RECORDING_INCOMPLETE:`, such as `memory_mutation_failed`, `context_mutated_after_capture`, or `memory_evidence_incomplete`. A Kitaru outage can prevent even these diagnostics from being persisted; missing status updates are not evidence of successful recording. A pending baseline older than 30 minutes is reported as `mastra_replay_abandoned` when replay is requested; this does not cancel a native turn or release a source lease. + +The server rejects pending, ineligible, or incomplete memory baselines before scheduling a replay or experiment run. CLI and MCP conflict responses expose a safe `mastra_replay_*` reason. Do not retry a pending session by supplying provisional inputs yourself. ### Create and inspect a memory replay -After the worker records a complete baseline, create a replay using an existing evaluator: +The SDK, CLI, and native MCP server support this workflow without a frontend. First inspect the baseline with `kitaru session get --output json` and wait for `metadata.mastra_replay_state` to become `eligible`. Then create a replay using an existing evaluator: ```bash kitaru replay create \ @@ -273,7 +298,7 @@ kitaru session get --output json kitaru session nodes --include-payloads --output json ``` -Read `result_session_id` from the replay, check the session's final status, and inspect its model-request and memory-mutation nodes. `session nodes` returns one page; follow `page.next_cursor` with `--cursor` while `page.has_more` is true. +Read `result_session_id` from the replay, check the session's final status, and inspect its model-request and memory-mutation nodes. `session nodes` returns one page; pass a non-null `page.next_cursor` back with `--cursor` until it is null. The Python SDK uses the same replay request. Given an authenticated `client`, a baseline UUID, and an existing evaluator: @@ -305,9 +330,11 @@ Inspect the run with `kitaru_activity_read`: get `kind: "experiment_run"`, list ### Files, skills, and processors -Pass a static `inputProcessors` array in the factory configuration. File processors must use the factory's supplied `resolveFile`; declare every allowed URL in the adapter's `files` list and provide a baseline `resolveFile` that returns `{ bytes: Uint8Array, mediaType: string }`. Kitaru records those bytes and serves them from the recorded input during replay. An undeclared URL fails instead of fetching live content. +Pass a static `inputProcessors` array in the factory configuration. File processors must use the factory's supplied `resolveFile`; declare every allowed URL in the adapter's `files` list and provide a baseline resolver returning `{ bytes: Uint8Array, mediaType: string }`. + +Kitaru replaces declared file URLs, including signed URLs, with `kitaru-file://sha256/...` content references in recorded input and records the bytes and media type. Replay resolves these references from recorded content, verifies the hash, and does not fetch the original URL. A processor must pass the file reference from its current input to the injected resolver; do not close over the original signed URL. Undeclared URLs and missing or altered recorded content fail instead of falling back to a network request. Capture accepts at most 64 distinct file URLs and 16 MiB of file bytes in total, subject to the replay input limit below. -For skills, set `skillsDirectory` to the directory containing your skill folders and use the factory's supplied `workspace`. Kitaru reads skill files into an immutable native workspace and records their paths, sizes, and hashes. Deploy the same skill artifact with the replay command. Changed files, missing files, and symlinks are rejected before execution. +For skills, set `skillsDirectory` to the directory containing your skill folders and use the factory's supplied `workspace`. Kitaru reads skill files into an immutable native workspace and records their paths, sizes, and hashes. Deploy the same skill artifact with the replay command. The skills tree is limited to 1 MiB of file content and 10,000 files/directories. Changed files, missing files, and symlinks are rejected before replay execution. The factory must use the supplied memory and workspace instances. Processors and tools are application code: their dependencies must use these supplied bindings for replay isolation. Kitaru does not sandbox arbitrary callbacks or prevent code from opening another database connection or making a network request. Workflows, subagents, provider-executed tools, approval/resume modes, dynamic tool inventories, `prepareStep`, output processors, and secondary structured-output models remain unsupported. @@ -315,11 +342,13 @@ The factory must use the supplied memory and workspace instances. Processors and Native memory tools execute against the isolated replay store, including under `history` with `on_miss: "fail"`. External tools, including tools added by a processor, follow the replay tool policy. A tool named `updateWorkingMemory` does not acquire the native-memory exemption by name. Use history with a failing miss when external tools must not execute. -Session inputs contain a version-2 `mastra_memory_replay` envelope with the raw invocation, initial thread/resource/messages, observational state and buffers, effective configuration, approved request context, and controlled file bytes. Dates, URLs, and binary values retain their types. Set `captureRequestContext` to select the context the run needs; never include credentials. Unsupported values, redaction, or exceeding the 1,048,576-character serialized payload bound make the input incomplete. Older history-only snapshots cannot recover this state and must be recorded again with the factory. +Eligible session inputs contain a version-3 `mastra_memory_replay` envelope with the invocation, initial thread/resource/messages, observational state and buffers, effective configuration, approved request context, controlled file bytes, and ordered OM results (`omTape`). Supported dates and binary values retain their types; declared file URLs become content references. Older history-only snapshots cannot recover this state. OM recordings without recorded results must be recorded again with the factory. + +Inputs remain bounded and self-contained. Unsupported values, redaction, or exceeding the 16 MiB serialized UTF-8 JSON limit make a recording ineligible. The replay input also has a 200,000-item budget and maximum depth of 64. Each binary value is limited to 8 MiB before base64 encoding; encoded files and OM outputs consume the shared JSON budget. These larger bounds apply to the memory replay input, not to every ordinary recorded node. The adapter reads the full initial thread history; `lastMessages` does not make recording unbounded or restrict that snapshot to the actor's recall window. A pre-turn capture that has not finished after five seconds becomes ineligible so a stalled storage read does not indefinitely delay the native answer. Large documents or long threads may therefore require a smaller baseline. Unlike ordinary wrapper recording, this path records the effective actor prompt, tools, tool choice, and supported settings for each provider attempt, including failed retries. Request attributes include attempt identity, memory revision, source provenance, and evidence completeness. `memory_mutation` span nodes record ordered native storage changes and link them to the active actor attempt when one exists. This evidence describes the request sent at the adapter's model boundary, not a provider's internal processing. -Consume the stream through completion so Kitaru can join native background memory work, flush evidence, release source access, and close the isolated replay store. Inspect the replay session's status and evidence completeness as well as native output: recording problems do not replace application output, and Mastra can settle a stream after a policy failure. Missing or incomplete initial state fails replay before model execution. A failed or incomplete recording is not proof that all evidence was saved. +Consume replay streams through completion and inspect the replay session's final status and evidence completeness. Replay finalization waits for isolated memory work and closes its store. Mastra can settle a native stream after a policy failure, so native output alone does not establish replay success. Missing or incomplete starting state fails replay before model execution; a later OM call mismatch can fail after actor execution has begun. A failed or incomplete recording is not proof that all evidence was saved. ## Structured output diff --git a/openapi/openapi.json b/openapi/openapi.json index 43cfdad91..98da269a0 100644 --- a/openapi/openapi.json +++ b/openapi/openapi.json @@ -9486,6 +9486,10 @@ "description": "New error.", "title": "Error" }, + "inputs": { + "description": "Final replay inputs for a pending Mastra recording.", + "title": "Inputs" + }, "metadata": { "anyOf": [ { diff --git a/packages/core/src/adapter/index.ts b/packages/core/src/adapter/index.ts index d6d1a814d..1f28cbff2 100644 --- a/packages/core/src/adapter/index.ts +++ b/packages/core/src/adapter/index.ts @@ -9,9 +9,12 @@ export { boundedRecorderConversion, boundedRecorderJson, boundRecordedSize, + MAX_MASTRA_REPLAY_ITEMS, + MAX_MASTRA_REPLAY_JSON_BYTES, MAX_RECORDED_PAYLOAD_CHARS, MAX_RECORDED_STRING_CHARS, normalizeRecordingLimits, + projectMastraReplayInput, projectRecordedInput, projectRecordedMetadata, recordedPayloadConversion, @@ -19,6 +22,7 @@ export { recordedToolPayloadConversion, recordedToolPayloadJson, runResultSummary, + strictMastraReplayValue, strictRecordedJson, } from "./recorded-json.js"; export type { ReplayContext } from "./replay.js"; diff --git a/packages/core/src/adapter/recorded-json.ts b/packages/core/src/adapter/recorded-json.ts index 7b859ac27..e65fa2908 100644 --- a/packages/core/src/adapter/recorded-json.ts +++ b/packages/core/src/adapter/recorded-json.ts @@ -8,6 +8,9 @@ const MAX_RECORDED_JSON_CHARS = 65_536; export const MAX_RECORDED_PAYLOAD_CHARS = 1_048_576; const MAX_RECORDED_PAYLOAD_ITEMS = 10_000; const MAX_RECORDED_PAYLOAD_DEPTH = 64; +/** Mastra's captured memory can exceed generic recorder limits. */ +export const MAX_MASTRA_REPLAY_JSON_BYTES = 16 * 1_048_576; +export const MAX_MASTRA_REPLAY_ITEMS = 200_000; // Leave room for the session node and step envelope in toRecorderJson's // 10,000-item ceiling after the tool value is embedded in it. const MAX_BOUNDED_TOOL_ITEMS = 9_000; @@ -29,6 +32,17 @@ const SECRET_KEYS: ReadonlySet = new Set([ "secret", "token", ]); +const MASTRA_REPLAY_SENSITIVE_KEYS: ReadonlySet = new Set([ + ...SECRET_KEYS, + "headers", + "abortsignal", +]); + +function isSecretUrlQueryKey(key: string): boolean { + return /(?:^|[-_])(?:api[-_]?key|authorization|cookie|password|secret|token|signature|credential|sig)$/i.test( + key, + ); +} /** * Keys whose value is a blob, a transport envelope, or a framework context @@ -112,6 +126,7 @@ interface CloneOptions { maxStringChars: number; path: string; rejectLongStrings: boolean; + rejectUrlCredentials?: boolean; sensitiveKeyMode: SensitiveKeyMode; sensitiveKeys: ReadonlySet; } @@ -128,6 +143,23 @@ function spendBudget(options: CloneOptions, characters: number): void { } function boundedString(value: string, options: CloneOptions): JsonValue { + if (options.rejectUrlCredentials) { + for (const match of value.matchAll(/https?:\/\/[^\s"'<>]+/gi)) { + let url: URL; + try { + url = new URL(match[0]); + } catch { + continue; + } + if ( + url.username || + url.password || + [...url.searchParams.keys()].some(isSecretUrlQueryKey) + ) { + throw new TypeError(`${options.path} contains URL credentials`); + } + } + } if (value.length <= options.maxStringChars) { spendBudget(options, value.length); return value; @@ -506,6 +538,99 @@ export function projectRecordedInput( return converted; } +/** Strict, independently bounded JSON for the Mastra historical read-set. */ +export function strictMastraReplayValue( + value: unknown, + path = "Mastra memory replay", +): JsonValue { + const options: CloneOptions = { + budget: { + chars: MAX_MASTRA_REPLAY_JSON_BYTES * 2, + items: MAX_MASTRA_REPLAY_ITEMS, + }, + lossy: false, + maxDepth: MAX_RECORDED_PAYLOAD_DEPTH, + maxItems: MAX_MASTRA_REPLAY_ITEMS, + maxStringChars: MAX_MASTRA_REPLAY_JSON_BYTES, + path, + rejectLongStrings: true, + rejectUrlCredentials: true, + sensitiveKeyMode: "reject", + sensitiveKeys: MASTRA_REPLAY_SENSITIVE_KEYS, + }; + let converted: JsonValue; + try { + converted = convert(value, options); + } catch (error) { + if ( + error instanceof TypeError && + /recorded item count|maximum array length|maximum object size/.test( + error.message, + ) + ) + throw new TypeError( + `${path} exceeds maximum item count ${MAX_MASTRA_REPLAY_ITEMS}`, + ); + if (error instanceof TypeError && /recorded JSON size/.test(error.message)) + throw new TypeError( + `${path} exceeds maximum JSON bytes ${MAX_MASTRA_REPLAY_JSON_BYTES}`, + ); + throw error; + } + if (options.lossy) throw new TypeError(`${path} contains unsupported values`); + if ( + Buffer.byteLength(JSON.stringify(converted), "utf8") > + MAX_MASTRA_REPLAY_JSON_BYTES + ) + throw new TypeError( + `${path} exceeds maximum JSON bytes ${MAX_MASTRA_REPLAY_JSON_BYTES}`, + ); + return converted; +} + +/** Require a real version-3 envelope before using the larger Mastra bound. */ +export function projectMastraReplayInput(value: unknown): JsonValue { + if ( + typeof value !== "object" || + value === null || + !isPlainObject(value) || + !Object.hasOwn(value, "mastra_memory_replay") || + Object.keys(value).some( + (key) => + !["mastra_memory_replay", "system_prompt", "prompt"].includes(key), + ) + ) + throw new TypeError("Mastra replay input requires a version-3 envelope"); + const envelope = (value as Record).mastra_memory_replay; + if ( + typeof envelope !== "object" || + envelope === null || + !isPlainObject(envelope) + ) + throw new TypeError( + "Mastra replay input requires a complete version-3 envelope", + ); + const fields = envelope as Record; + if ( + fields.version !== 3 || + fields.complete !== true || + !Array.isArray(fields.reasons) || + fields.reasons.length !== 0 || + typeof fields.invocationId !== "string" || + fields.invocationId.length === 0 || + !Object.hasOwn(fields, "rawInput") || + !Object.hasOwn(fields, "initialSnapshot") || + !Object.hasOwn(fields, "configuration") || + !Object.hasOwn(fields, "requestContext") || + !Array.isArray(fields.files) || + !Array.isArray(fields.omTape) + ) + throw new TypeError( + "Mastra replay input requires a complete version-3 envelope", + ); + return strictMastraReplayValue(value, "Mastra replay input"); +} + /** * Convert a value with narrow bounds, rejecting anything that exceeds them. */ diff --git a/packages/core/src/adapter/replay.ts b/packages/core/src/adapter/replay.ts index 1bcf9b03c..ddaec57de 100644 --- a/packages/core/src/adapter/replay.ts +++ b/packages/core/src/adapter/replay.ts @@ -2,6 +2,7 @@ import type { KitaruEnvironmentVariables } from "../environment.js"; import { toRecorderJson } from "../json.js"; import type { JsonValue, ReplayOverride, ReplaySpec } from "../types.js"; import { isRecord, isUuid } from "../validation.js"; +import { projectMastraReplayInput } from "./recorded-json.js"; import type { AdapterClient } from "./run-state.js"; function parseUuidEnvironment( @@ -220,6 +221,7 @@ export async function resolveReplayContext(options: { callerInput: unknown; client: AdapterClient; environment?: KitaruEnvironmentVariables; + recordedInputProjector?: (input: unknown) => Promise | unknown; requestedModelId: string; }): Promise { const environment = options.environment ?? process.env; @@ -241,6 +243,10 @@ export async function resolveReplayContext(options: { ) : undefined; const effective = resolveEffectiveInputs(workerInput, override); + const recordedInput = + !spec && options.recordedInputProjector + ? await options.recordedInputProjector(effective.recorded) + : effective.recorded; const replacementModelId = modelReplacement( override, options.requestedModelId, @@ -252,7 +258,13 @@ export async function resolveReplayContext(options: { } return { - effectiveInput: toRecorderJson(effective.recorded), + effectiveInput: + spec && + isRecord(recordedInput) && + isRecord(recordedInput.mastra_memory_replay) && + recordedInput.mastra_memory_replay.version === 3 + ? projectMastraReplayInput(recordedInput) + : toRecorderJson(recordedInput), effectiveRuntimeInput: effective.runtime, override, replayId, diff --git a/packages/core/src/adapter/run-recorder.ts b/packages/core/src/adapter/run-recorder.ts index 75a680803..0b9f83642 100644 --- a/packages/core/src/adapter/run-recorder.ts +++ b/packages/core/src/adapter/run-recorder.ts @@ -26,6 +26,7 @@ function rootNode( options: { endedAt?: string; error?: string; + inputs?: JsonValue; output?: JsonValue; startedAt: string; status: "completed" | "failed" | "in_progress"; @@ -36,7 +37,8 @@ function rootNode( ended_at: options.endedAt, error: options.error, external_id: ROOT_NODE_EXTERNAL_ID, - inputs: state.effectiveInput, + inputs: + options.inputs === undefined ? state.effectiveInput : options.inputs, name: "run", node_type: "span", outputs: options.output ?? null, @@ -62,6 +64,7 @@ export interface RunRecorderOptions { effectiveInput: JsonValue; effectiveModelSettings?: Record; framework: string; + metadata?: Record; name?: string; replayId?: string; requestedModelId: string; @@ -97,6 +100,7 @@ export class RunRecorder { agent_version_id: options.agentVersionId, framework: options.framework, inputs: options.effectiveInput, + ...(options.metadata ? { metadata: options.metadata } : {}), name: options.name, origin: options.replayId ? "replay" : "recorded", outputs: null, @@ -134,7 +138,13 @@ export class RunRecorder { } } - async complete(result: unknown): Promise { + async complete( + result: unknown, + options: { + inputs?: JsonValue; + metadata?: Record; + } = {}, + ): Promise { await this.state.awaitSteps(); // The run has finished by the time its result is recorded, so a result too // large or too circular to record is bounded instead of turning a @@ -145,6 +155,7 @@ export class RunRecorder { nodes: [ rootNode(this.state, { endedAt, + inputs: options.inputs, output: serializedOutput, startedAt: this.#startedAt, status: "completed", @@ -153,28 +164,39 @@ export class RunRecorder { }); await this.#client.updateSession(this.state.sessionId, { ended_at: endedAt, + ...("inputs" in options ? { inputs: options.inputs } : {}), + ...(options.metadata ? { metadata: options.metadata } : {}), outputs: serializedOutput, status: "completed", }); } - async fail(error: unknown): Promise { + async fail( + error: unknown, + metadata?: Record, + ): Promise { this.state.storeFailure(error); // Let queued step writes land before the failed ledger and the closing // node, so a late step cannot arrive after the session is marked failed. await bestEffort(() => this.state.awaitSteps()); await bestEffort(() => flushFailedPolicyOutcomes(this.state)); - await this.#closeFailed(error); + await this.#closeFailed(error, metadata); } - async failRecording(error: unknown): Promise { + async failRecording( + error: unknown, + metadata?: Record, + ): Promise { // A telemetry failure must not enter application state. Tool hooks use // state.failure to stop execution after policy or runtime failures. await bestEffort(() => this.state.awaitSteps()); - await this.#closeFailed(error); + await this.#closeFailed(error, metadata); } - async #closeFailed(error: unknown): Promise { + async #closeFailed( + error: unknown, + metadata?: Record, + ): Promise { const endedAt = new Date().toISOString(); await bestEffort(() => this.#client.upsertSessionNodes(this.state.sessionId, { @@ -192,6 +214,7 @@ export class RunRecorder { this.#client.updateSession(this.state.sessionId, { ended_at: endedAt, error: errorText(error), + ...(metadata ? { metadata } : {}), status: "failed", }), ); diff --git a/packages/core/src/generated/openapi.ts b/packages/core/src/generated/openapi.ts index 5d820f604..09ea70183 100644 --- a/packages/core/src/generated/openapi.ts +++ b/packages/core/src/generated/openapi.ts @@ -9258,6 +9258,11 @@ export interface components { * @description New error. */ error?: string | null; + /** + * Inputs + * @description Final replay inputs for a pending Mastra recording. + */ + inputs?: unknown; /** * Metadata * @description New metadata. diff --git a/packages/core/test/adapter/replay.test.ts b/packages/core/test/adapter/replay.test.ts index 5d2f42d4b..3ee929a13 100644 --- a/packages/core/test/adapter/replay.test.ts +++ b/packages/core/test/adapter/replay.test.ts @@ -11,6 +11,93 @@ import { import { fakeClient, REPLAY_ID, replay, TASK_ID } from "./helpers.js"; describe("adapter replay preparation", () => { + it("projects signed baseline input without changing native runtime input", async () => { + const source = { file: "https://files.invalid/a?token=secret" }; + const context = await resolveReplayContext({ + callerInput: "caller", + client: fakeClient({ taskInput: source }), + environment: { KITARU_TASK_ID: TASK_ID }, + recordedInputProjector: (input) => ({ + ...(input as typeof source), + file: "kitaru-file://sha256/recorded", + }), + requestedModelId: "requested", + }); + expect(context.effectiveInput).toEqual({ + file: "kitaru-file://sha256/recorded", + }); + expect(context.effectiveRuntimeInput).toEqual(source); + }); + + it("accepts a large version-3 Mastra task input while preserving generic limits", async () => { + const envelope = { + version: 3, + complete: true, + reasons: [], + invocationId: "invocation", + rawInput: null, + initialSnapshot: { values: Array.from({ length: 15_000 }, (_, i) => i) }, + configuration: {}, + requestContext: {}, + files: [], + omTape: [], + }; + const context = await resolveReplayContext({ + callerInput: "caller", + client: fakeClient({ + replay: replay(), + taskInput: { mastra_memory_replay: envelope }, + }), + environment: { KITARU_REPLAY_ID: REPLAY_ID, KITARU_TASK_ID: TASK_ID }, + requestedModelId: "requested", + }); + expect(context.effectiveInput).toEqual({ mastra_memory_replay: envelope }); + + await expect( + resolveReplayContext({ + callerInput: { values: Array.from({ length: 15_000 }, (_, i) => i) }, + client: fakeClient(), + environment: {}, + requestedModelId: "requested", + }), + ).rejects.toThrow(/maximum item count 10000/); + }); + + it("rejects secret-bearing and malformed large Mastra task inputs", async () => { + const envelope = { + version: 3, + complete: true, + reasons: [], + invocationId: "invocation", + rawInput: null, + initialSnapshot: { values: Array.from({ length: 15_000 }, (_, i) => i) }, + configuration: {}, + requestContext: {}, + files: [], + omTape: [], + }; + for (const changed of [ + { ...envelope, requestContext: { token: "secret" } }, + { ...envelope, requestContext: { headers: { "x-auth": "secret" } } }, + { + ...envelope, + rawInput: "https://files.invalid/a?X-Amz-Signature=secret", + }, + { ...envelope, omTape: null }, + ]) { + await expect( + resolveReplayContext({ + callerInput: "caller", + client: fakeClient({ + replay: replay(), + taskInput: { mastra_memory_replay: changed }, + }), + environment: { KITARU_REPLAY_ID: REPLAY_ID, KITARU_TASK_ID: TASK_ID }, + requestedModelId: "requested", + }), + ).rejects.toThrow(); + } + }); it("keeps caller input and legacy override outside replay", async () => { const client = fakeClient(); const context = await resolveReplayContext({ diff --git a/packages/core/test/adapter/run-recorder.test.ts b/packages/core/test/adapter/run-recorder.test.ts index 5a3807ff0..d4c9632eb 100644 --- a/packages/core/test/adapter/run-recorder.test.ts +++ b/packages/core/test/adapter/run-recorder.test.ts @@ -16,6 +16,50 @@ async function recorder(client: AdapterClient): Promise { } describe("normalized run lifecycle", () => { + it("creates a pending recording with metadata atomically", async () => { + const client = fakeClient(); + await RunRecorder.create({ + adapterVersion: "test-adapter", + agentId: "018f0000-0000-7000-8000-000000000100", + client, + effectiveInput: { prompt: "hello" }, + framework: "mastra", + metadata: { mastra_replay_state: "pending" }, + requestedModelId: "requested-model", + }); + expect(client.created[0]?.metadata).toEqual({ + mastra_replay_state: "pending", + }); + }); + + it("publishes final inputs and eligibility with completion", async () => { + const client = fakeClient(); + const run = await recorder(client); + await run.initialize(); + const finalInputs = { + mastra_memory_replay: { version: 3, complete: true }, + }; + await run.complete( + { text: "done" }, + { + inputs: finalInputs, + metadata: { + mastra_replay_state: "eligible", + mastra_native_state: "completed", + }, + }, + ); + expect(client.updates.at(-1)).toMatchObject({ + inputs: finalInputs, + metadata: { + mastra_replay_state: "eligible", + mastra_native_state: "completed", + }, + status: "completed", + }); + expect(client.nodes.at(-1)?.nodes[0]?.inputs).toEqual(finalInputs); + }); + it("creates, records, and completes one run", async () => { const client = fakeClient(); const run = await recorder(client); @@ -318,6 +362,19 @@ describe("normalized run lifecycle", () => { }); }); + it("persists an ineligible recording reason in its terminal update", async () => { + const client = fakeClient(); + const run = await recorder(client); + await run.failRecording(new Error("capture failed"), { + mastra_replay_state: "ineligible", + mastra_replay_reason: "capture_incomplete", + }); + expect(client.updates.at(-1)?.metadata).toEqual({ + mastra_replay_state: "ineligible", + mastra_replay_reason: "capture_incomplete", + }); + }); + it("creates no nodes when session creation fails", async () => { const client = fakeClient(); client.createSession = async () => { diff --git a/packages/mastra/README.md b/packages/mastra/README.md index 2d57a6c0f..7b3f691a9 100644 --- a/packages/mastra/README.md +++ b/packages/mastra/README.md @@ -153,7 +153,9 @@ Import `createMemoryReplayAgent()` from `@zenml-io/kitaru-mastra/memory` when a pnpm add @zenml-io/kitaru-mastra @mastra/core@1.67.0 @mastra/memory@1.30.0 zod ``` -The factory creates a fresh native agent for each invocation. For a baseline, it binds native memory to your source storage and records the starting state before recall. For a replay, it restores that state into a separate in-memory store. Native working-memory tools and observation/reflection jobs then evolve the isolated state as the model runs again. Replay never calls `sourceMemory()`. +The factory creates a fresh native agent for each invocation. A baseline uses your source storage and records the starting state before recall, then records the observer and reflector model outputs produced during that invocation. Replay restores the starting state into a separate in-memory store, runs the actor again, and supplies those recorded outputs to native observational memory (OM). Working-memory tools and native OM storage updates therefore run against isolated state. Replay never calls `sourceMemory()` or makes fresh observer/reflector provider calls. + +Recorded OM outputs match by call order, phase (observer or reflector), and model method. A missing, extra, reordered, or unused call fails replay with `KITARU_REPLAY_DIVERGED:mastra_om_call_order`. Changed OM input alone does not fail replay: Kitaru reuses the recorded output and records an `om_input_mismatch` span. This lets you compare actor instruction/model changes with recorded OM results, but does not measure how a fresh observer or reflector would respond to the changed conversation. The following binding uses a process-local store. Supply your existing public memory storage domain and its complete configuration for a persistent application: @@ -210,20 +212,43 @@ const output = await recorded.stream("My preference is green.", { context: [{ role: "system", content: "The customer is asking about preferences." }], }); await output.consumeStream(); -await store.close(); +// Keep the application and source store alive while recording finalizes. +// Inspect session eligibility before shutting down or starting a replay. ``` Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KITARU_AGENT_ID`, and the model provider credential. Register the compiled command as the agent version's run specification to run it through a worker. The same command serves baseline and replay tasks; the worker supplies the recorded input and replay identity. -All source-memory writers must participate in the same exclusive-access mechanism. The process-local helper is suitable only when every writer shares that instance in one process. Use a distributed implementation of `MastraExclusiveMemoryAccess` when other processes can write. Its `acquire()` method must hold access until the returned release function runs. `settled()` must join pending work on the source `Memory` instance; it is not a lock. Configure working and observational memory with `scope: "thread"`; resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` are outside this contract. +### Source ownership and supported configuration + +All writers to a source thread must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. Overlapping turns can continue natively, but their recordings become ineligible for replay. `settled()` joins pending work on the source `Memory` instance; it does not provide exclusive access. + +A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must preserve ownership through the final eligible-session update, mark an overlapping or unowned write unsafe through `markUnsafeWrite()`, and retain unsafe state across process loss. Call `resetAfterQuiescence()` only after all possible stale writers have stopped. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. + +Schema working memory requires explicit `scope: "thread"`. An observational-memory configuration object may omit `scope`, using Mastra's implicit thread scope, or set it to `"thread"`. Supply explicit observer/reflector model identities, either shared through `observationalMemory.model` or in the phase configuration. Resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` remain unsupported. + +Keep `memory.thread` and `memory.resource` consistent with reserved Mastra thread/resource IDs in `requestContext`. A mismatch cannot produce an eligible recording. Baseline callbacks receive the original live context. `captureRequestContext` selects only approved, replay-relevant JSON values; it does not remove values from the live context. Replay receives that recorded projection. A nonempty context without an explicit projection makes the recording ineligible. For example, return `{ locale: context.get("locale") }` when locale is the only value replay needs, or `{}` when none are needed. + +Never include authentication tokens, credentials, or signed URLs in the projection. Credential-like context keys are rejected, and transport headers in recorded configuration make the envelope incomplete. These checks cannot identify every secret hidden in an arbitrary string; choose recorded fields explicitly. Use `resolveModel` to reconstruct model instances from locally configured credentials. + +Dynamic `instructions`, `model`, and `defaultOptions` resolve during baseline setup; replay uses their recorded values. `resolveModel` must resolve the recorded actor, observer, and reflector model identifiers and any allowed actor override. OM identifiers must resolve to native stream-capable model objects, whose provider methods Kitaru intercepts to reuse recorded results during replay. A `system_prompt` override replaces only application instructions and retains recorded extra system context. Model and model-setting overrides affect the actor; observation and reflection retain their recorded configuration and outputs. Raw-input `prompt` overrides are rejected; record a new baseline to change invocation input. -Keep `memory.thread` and `memory.resource` equal to any reserved Mastra thread/resource IDs in `requestContext`; a mismatch is rejected before native execution. Use `captureRequestContext` to select safe values: native authentication tokens are rejected, and transport headers in recorded configuration make the replay envelope incomplete. Use `resolveModel` to reconstruct credential-bearing model instances for replay. +### Recording readiness -Dynamic `instructions`, `model`, and `defaultOptions` resolve once during baseline setup. Replay uses their recorded values instead of calling those resolvers again. `resolveModel` must resolve the recorded actor, observer, and reflector model identifiers as well as any allowed actor override. A `system_prompt` override replaces only application instructions and retains recorded extra system context. Model and model-setting overrides affect the actor; observation and reflection retain their recorded configuration. Raw-input `prompt` overrides are rejected; record a new baseline to change invocation input. +Native output and replay readiness are separate. Consume the baseline stream normally. Once the actor finishes, Kitaru finalizes recording in the background: it joins native memory work, records OM results and memory evidence, verifies source ownership, and persists the final input. Keep the application and source storage alive until that finalization finishes. `consumeStream()` alone is not a recording-completion barrier for a baseline. + +Inspect the baseline session's metadata and status: + +- `mastra_replay_state: "pending"`: recording has not finalized; do not replay yet. +- `mastra_replay_state: "eligible"`: the complete version-3 input and evidence were persisted for replay. +- `mastra_replay_state: "ineligible"`: recording could not establish a complete, isolated baseline. Inspect `mastra_replay_reason` and `mastra_native_state` to distinguish recording failure from native execution failure, then record a new baseline after resolving the cause. + +Recording-only problems do not replace the baseline's native answer. Diagnostics use `KITARU_RECORDING_INCOMPLETE:`, such as `memory_mutation_failed`, `context_mutated_after_capture`, or `memory_evidence_incomplete`. A Kitaru outage can prevent even these diagnostics from being persisted; missing status updates are not evidence of successful recording. A pending baseline older than 30 minutes is reported as `mastra_replay_abandoned` when replay is requested; this does not cancel a native turn or release a source lease. + +The server rejects pending, ineligible, or incomplete memory baselines before scheduling a replay or experiment run. CLI and MCP conflict responses expose a safe `mastra_replay_*` reason. Do not retry a pending session by supplying provisional inputs yourself. ### Create and inspect a memory replay -After the worker records a complete baseline, create a replay using an existing evaluator: +The SDK, CLI, and native MCP server support this workflow without a frontend. First inspect the baseline with `kitaru session get --output json` and wait for `metadata.mastra_replay_state` to become `eligible`. Then create a replay using an existing evaluator: ```bash kitaru replay create \ @@ -237,7 +262,7 @@ kitaru session get --output json kitaru session nodes --include-payloads --output json ``` -Read `result_session_id` from the replay, check the session's final status, and inspect its model-request and memory-mutation nodes. `session nodes` returns one page; follow `page.next_cursor` with `--cursor` while `page.has_more` is true. +Read `result_session_id` from the replay, check the session's final status, and inspect its model-request and memory-mutation nodes. `session nodes` returns one page; pass a non-null `page.next_cursor` back with `--cursor` until it is null. The Python SDK uses the same replay request. Given an authenticated `client`, a baseline UUID, and an existing evaluator: @@ -269,9 +294,11 @@ Inspect the run with `kitaru_activity_read`: get `kind: "experiment_run"`, list ### Files, skills, and processors -Pass a static `inputProcessors` array in the factory configuration. File processors must use the factory's supplied `resolveFile`; declare every allowed URL in the adapter's `files` list and provide a baseline `resolveFile` that returns `{ bytes: Uint8Array, mediaType: string }`. Kitaru records those bytes and serves them from the recorded input during replay. An undeclared URL fails instead of fetching live content. +Pass a static `inputProcessors` array in the factory configuration. File processors must use the factory's supplied `resolveFile`; declare every allowed URL in the adapter's `files` list and provide a baseline resolver returning `{ bytes: Uint8Array, mediaType: string }`. + +Kitaru replaces declared file URLs, including signed URLs, with `kitaru-file://sha256/...` content references in recorded input and records the bytes and media type. Replay resolves these references from recorded content, verifies the hash, and does not fetch the original URL. A processor must pass the file reference from its current input to the injected resolver; do not close over the original signed URL. Undeclared URLs and missing or altered recorded content fail instead of falling back to a network request. Capture accepts at most 64 distinct file URLs and 16 MiB of file bytes in total, subject to the replay input limit below. -For skills, set `skillsDirectory` to the directory containing your skill folders and use the factory's supplied `workspace`. Kitaru reads skill files into an immutable native workspace and records their paths, sizes, and hashes. Deploy the same skill artifact with the replay command. Changed files, missing files, and symlinks are rejected before execution. +For skills, set `skillsDirectory` to the directory containing your skill folders and use the factory's supplied `workspace`. Kitaru reads skill files into an immutable native workspace and records their paths, sizes, and hashes. Deploy the same skill artifact with the replay command. The skills tree is limited to 1 MiB of file content and 10,000 files/directories. Changed files, missing files, and symlinks are rejected before replay execution. The factory must use the supplied memory and workspace instances. Processors and tools are application code: their dependencies must use these supplied bindings for replay isolation. Kitaru does not sandbox arbitrary callbacks or prevent code from opening another database connection or making a network request. Workflows, subagents, provider-executed tools, approval/resume modes, dynamic tool inventories, `prepareStep`, output processors, and secondary structured-output models remain unsupported. @@ -279,11 +306,13 @@ The factory must use the supplied memory and workspace instances. Processors and Native memory tools execute against the isolated replay store, including under `history` with `on_miss: "fail"`. External tools, including tools added by a processor, follow the replay tool policy. A tool named `updateWorkingMemory` does not acquire the native-memory exemption by name. Use history with a failing miss when external tools must not execute. -Session inputs contain a version-2 `mastra_memory_replay` envelope with the raw invocation, initial thread/resource/messages, observational state and buffers, effective configuration, approved request context, and controlled file bytes. Dates, URLs, and binary values retain their types. Set `captureRequestContext` to select the context the run needs; never include credentials. Unsupported values, redaction, or exceeding the 1,048,576-character serialized payload bound make the input incomplete. Older history-only snapshots cannot recover this state and must be recorded again with the factory. +Eligible session inputs contain a version-3 `mastra_memory_replay` envelope with the invocation, initial thread/resource/messages, observational state and buffers, effective configuration, approved request context, controlled file bytes, and ordered OM results (`omTape`). Supported dates and binary values retain their types; declared file URLs become content references. Older history-only snapshots cannot recover this state. OM recordings without recorded results must be recorded again with the factory. + +Inputs remain bounded and self-contained. Unsupported values, redaction, or exceeding the 16 MiB serialized UTF-8 JSON limit make a recording ineligible. The replay input also has a 200,000-item budget and maximum depth of 64. Each binary value is limited to 8 MiB before base64 encoding; encoded files and OM outputs consume the shared JSON budget. These larger bounds apply to the memory replay input, not to every ordinary recorded node. The adapter reads the full initial thread history; `lastMessages` does not make recording unbounded or restrict that snapshot to the actor's recall window. A pre-turn capture that has not finished after five seconds becomes ineligible so a stalled storage read does not indefinitely delay the native answer. Large documents or long threads may therefore require a smaller baseline. Unlike ordinary wrapper recording, this path records the effective actor prompt, tools, tool choice, and supported settings for each provider attempt, including failed retries. Request attributes include attempt identity, memory revision, source provenance, and evidence completeness. `memory_mutation` span nodes record ordered native storage changes and link them to the active actor attempt when one exists. This evidence describes the request sent at the adapter's model boundary, not a provider's internal processing. -Consume the stream through completion so Kitaru can join native background memory work, flush evidence, release source access, and close the isolated replay store. Inspect the replay session's status and evidence completeness as well as native output: recording problems do not replace application output, and Mastra can settle a stream after a policy failure. Missing or incomplete initial state fails replay before model execution. A failed or incomplete recording is not proof that all evidence was saved. +Consume replay streams through completion and inspect the replay session's final status and evidence completeness. Replay finalization waits for isolated memory work and closes its store. Mastra can settle a native stream after a policy failure, so native output alone does not establish replay success. Missing or incomplete starting state fails replay before model execution; a later OM call mismatch can fail after actor execution has begun. A failed or incomplete recording is not proof that all evidence was saved. ## Callback composition diff --git a/packages/mastra/src/memory-binding.ts b/packages/mastra/src/memory-binding.ts index 5780e89bf..998a5321e 100644 --- a/packages/mastra/src/memory-binding.ts +++ b/packages/mastra/src/memory-binding.ts @@ -13,12 +13,48 @@ export interface MastraMemorySelector { resourceId: string; } -/** All writers must participate in this application's lease, including other processes. */ +export interface MastraMemoryLeaseOptions { + /** Bound acquisition so a live answer does not wait for a delayed reflection. */ + waitMs?: number; + /** Cancel a waiting acquisition without releasing another writer's lease. */ + signal?: AbortSignal; + /** An advisory notification; verifyEligibility is the authoritative check. */ + onConflict?: () => void; +} + +/** A callable release keeps existing direct lease users source compatible. */ +export interface MastraMemoryLease { + (): Promise; + /** Read the shared coordination state, including conflict and lease loss. */ + verifyEligibility(): Promise; +} + +/** + * Coordinate every writer of a source thread across all application processes. + * + * The implementation must atomically poison eligibility for the thread when a + * competing native turn proceeds without ownership or an owner loses its lease. + * Poison must survive process loss and prevent a later acquisition from becoming + * eligible until all possible stale writers have quiesced. A timeout or failed + * coordination call must fail closed for replay eligibility, while native Mastra + * storage writes still run. Keep ownership through the final eligible-session + * update, then release only after no delayed source write remains possible. + * + * The process-local helper below is valid only when every writer shares one + * instance in one process. A production multi-server application must provide + * its own implementation backed by shared atomic storage. + */ export interface MastraExclusiveMemoryAccess { acquire( selector: MastraMemorySelector, - onConflict?: () => void, - ): Promise<() => Promise>; + options?: MastraMemoryLeaseOptions, + ): Promise; + /** Persist an unsafe-write marker before an unowned native write proceeds. + * An unknown selector poisons every thread until global quiescence is proven. + */ + markUnsafeWrite(selector?: MastraMemorySelector): Promise; + /** Clear a persistent loss marker only after all writers are proven stopped. */ + resetAfterQuiescence(selector?: MastraMemorySelector): Promise; } /** @@ -26,19 +62,99 @@ export interface MastraExclusiveMemoryAccess { * and shares this instance; distributed writers require a distributed lease. */ export function createProcessLocalMemoryAccess(): MastraExclusiveMemoryAccess { - const activeThreads = new Map void) | undefined>(); + type Turn = { onConflict?: () => void; invalidated: boolean }; + type ThreadState = { + turns: Set; + poisoned: boolean; + persistentLoss: boolean; + }; + const threads = new Map(); + let unknownWriterPoisoned = false; return { - async acquire({ threadId }, onConflict) { - if (activeThreads.has(threadId)) { - activeThreads.get(threadId)?.(); - throw new Error("Exclusive source-thread ownership is unavailable."); + async acquire({ threadId }, options = {}) { + if (options.signal?.aborted) + throw new Error("Exclusive source-thread ownership was cancelled."); + const state = threads.get(threadId) ?? { + turns: new Set(), + poisoned: false, + persistentLoss: false, + }; + threads.set(threadId, state); + if (state.turns.size > 0) { + state.poisoned = true; + for (const turn of state.turns) { + turn.invalidated = true; + turn.onConflict?.(); + } } - activeThreads.set(threadId, onConflict); + const owner = { + onConflict: options.onConflict, + invalidated: + state.poisoned || state.persistentLoss || unknownWriterPoisoned, + }; + state.turns.add(owner); let released = false; - return async () => { - if (!released) activeThreads.delete(threadId); + const release = async () => { + if (!released) state.turns.delete(owner); released = true; + if (state.turns.size === 0 && !state.persistentLoss) + threads.delete(threadId); }; + return Object.assign(release, { + async verifyEligibility() { + return ( + !released && + !owner.invalidated && + !state.poisoned && + !state.persistentLoss && + !unknownWriterPoisoned + ); + }, + }); + }, + async markUnsafeWrite(selector) { + if (!selector) { + unknownWriterPoisoned = true; + for (const state of threads.values()) { + state.poisoned = true; + state.persistentLoss = true; + for (const turn of state.turns) { + turn.invalidated = true; + turn.onConflict?.(); + } + } + return; + } + const { threadId } = selector; + const state = threads.get(threadId) ?? { + turns: new Set(), + poisoned: false, + persistentLoss: false, + }; + threads.set(threadId, state); + state.poisoned = true; + // The unsafe writer may outlive every current lease. Do not let the + // final release erase the conflict before that writer quiesces. + state.persistentLoss = true; + for (const turn of state.turns) { + turn.invalidated = true; + turn.onConflict?.(); + } + }, + async resetAfterQuiescence(selector) { + if (!selector) { + if ([...threads.values()].some((state) => state.turns.size > 0)) + throw new Error("Source-thread writers are still active."); + threads.clear(); + unknownWriterPoisoned = false; + return; + } + const { threadId } = selector; + const state = threads.get(threadId); + if (!state) return; + if (state.turns.size > 0) + throw new Error("Source-thread writers are still active."); + threads.delete(threadId); }, }; } @@ -61,6 +177,12 @@ export interface MastraMemoryCaptureOptions extends MastraMemorySelector { recordMutation: (event: MastraMemoryMutation) => Promise; getRequestId?: () => string | undefined; onIncomplete?: (reason: string) => void; + /** Replace captured file URLs in evidence without changing native writes. */ + sanitizeEvidence?: (value: T) => T; + leaseWaitMs?: number; + leaseSignal?: AbortSignal; + /** Bound pre-turn storage reads so capture cannot stall a native answer. */ + captureWaitMs?: number; } export interface MastraMemoryCaptureBinding { @@ -73,9 +195,29 @@ export interface MastraMemoryCaptureBinding { }): Promise; markIncomplete(reason: string): void; drain(): Promise; + /** Check shared ownership immediately before persisting eligible inputs. */ + verifyEligibility(): Promise; release(): Promise; } +async function boundedCoordination( + operation: Promise, + waitMs: number, + timeoutMessage = "Source-thread coordination timed out.", +): Promise { + let timer: ReturnType | undefined; + try { + return await Promise.race([ + operation, + new Promise((_resolve, reject) => { + timer = setTimeout(() => reject(new Error(timeoutMessage)), waitMs); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } +} + // The pinned public MemoryStorage mutation inventory. Delegation binds `this` to // the original domain, so a native method's own helper calls record only once. const MUTATIONS = new Set([ @@ -110,6 +252,55 @@ const MUTATIONS = new Set([ "updateObservationalMemoryConfig", ]); +function record(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function modelIdentity(model: unknown): string { + if (typeof model === "string" && model) return model; + if ( + record(model) && + typeof model.modelId === "string" && + typeof model.provider === "string" + ) + return `${model.provider}/${model.modelId}`; + if (record(model) && typeof model.id === "string") return model.id; + throw new Error("Observational-memory model has no stable identity."); +} + +/** Keep provider clients out of evidence without changing native storage calls. */ +function mutationEvidenceValue(method: PropertyKey, value: unknown): unknown { + if (method !== "initializeObservationalMemory") return value; + function project(input: unknown): unknown { + if (!record(input) || !record(input.config)) return input; + const config = { ...input.config }; + if (config.model !== undefined) config.model = modelIdentity(config.model); + for (const name of ["observation", "reflection"]) { + const phase = config[name]; + if (!record(phase)) continue; + const projected = { ...phase }; + if (projected.model !== undefined) + projected.model = modelIdentity(projected.model); + if (Array.isArray(projected.extractors)) + projected.extractors = projected.extractors.map((extractor) => { + if ( + record(extractor) && + extractor.internal === true && + typeof extractor.slug === "string" && + ["current-task", "suggested-response", "thread-title"].includes( + extractor.slug, + ) + ) + return { mastraBuiltinExtractor: extractor.slug }; + throw new Error("Unsupported observational-memory extractor."); + }); + config[name] = projected; + } + return { ...input, config }; + } + return Array.isArray(value) ? value.map(project) : project(value); +} + /** Capture one native invocation without changing the shared source domain or Agent. */ export function createMemoryCaptureBinding( options: MastraMemoryCaptureOptions, @@ -119,11 +310,12 @@ export function createMemoryCaptureBinding( let capturing = false; let readingSnapshot = false; let released = false; - let releaseLease: (() => Promise) | undefined; + let lease: MastraMemoryLease | undefined; let mutations = Promise.resolve(); let evidence = Promise.resolve(); const reasons: string[] = []; const methods = new Map(); + const waitMs = options.leaseWaitMs ?? 100; function markIncomplete(reason: string): void { if (reasons.includes(reason)) return; @@ -150,7 +342,10 @@ export function createMemoryCaptureBinding( let complete = true; let requestId: string | undefined; try { - encodedArguments = encodeMemoryValue(args); + const evidence = mutationEvidenceValue(property, args); + encodedArguments = encodeMemoryValue( + options.sanitizeEvidence?.(evidence) ?? evidence, + ); requestId = options.getRequestId?.(); } catch { complete = false; @@ -172,6 +367,46 @@ export function createMemoryCaptureBinding( const result = mutations.then(async () => { let output: unknown; try { + if (released || !lease) { + markIncomplete( + "Memory mutation occurred without source-thread ownership.", + ); + try { + await boundedCoordination( + options.exclusiveAccess.markUnsafeWrite(options), + waitMs, + ); + } catch { + markIncomplete("Unsafe memory write could not be fenced."); + } + } else { + try { + if ( + !(await boundedCoordination( + lease.verifyEligibility(), + waitMs, + )) + ) { + markIncomplete("Exclusive source-thread ownership was lost."); + await boundedCoordination( + options.exclusiveAccess.markUnsafeWrite(options), + waitMs, + ); + } + } catch { + markIncomplete( + "Exclusive source-thread ownership could not be verified.", + ); + try { + await boundedCoordination( + options.exclusiveAccess.markUnsafeWrite(options), + waitMs, + ); + } catch { + markIncomplete("Unsafe memory write could not be fenced."); + } + } + } output = await Reflect.apply(value, target, args); } catch (error) { markIncomplete("Native memory storage mutation failed."); @@ -182,7 +417,10 @@ export function createMemoryCaptureBinding( revision += 1; let encodedResult: JsonValue = null; try { - encodedResult = encodeMemoryValue(output); + const evidence = mutationEvidenceValue(property, output); + encodedResult = encodeMemoryValue( + options.sanitizeEvidence?.(evidence) ?? evidence, + ); } catch { complete = false; markIncomplete( @@ -249,6 +487,21 @@ export function createMemoryCaptureBinding( } } + async function verifyEligibility(): Promise { + if (released || !lease) { + markIncomplete("Exclusive source-thread ownership is unavailable."); + return; + } + try { + if (!(await boundedCoordination(lease.verifyEligibility(), waitMs))) + markIncomplete("Exclusive source-thread ownership was lost."); + } catch { + markIncomplete( + "Exclusive source-thread ownership could not be verified.", + ); + } + } + return { domain, get revision() { @@ -269,61 +522,112 @@ export function createMemoryCaptureBinding( capturing = true; try { try { - releaseLease = await options.exclusiveAccess.acquire(options, () => - markIncomplete( - "Exclusive source-thread ownership was invalidated by an overlapping invocation.", - ), + const timeout = AbortSignal.timeout(waitMs); + const signal = options.leaseSignal + ? AbortSignal.any([timeout, options.leaseSignal]) + : timeout; + const attempted = options.exclusiveAccess.acquire(options, { + waitMs, + signal, + onConflict: () => + markIncomplete( + "Exclusive source-thread ownership was invalidated by an overlapping invocation.", + ), + }); + let accepted = false; + // A backend that ignores cancellation must not retain ownership if + // its acquire resolves after the caller has resumed natively. + void attempted.then( + async (lateLease) => { + if (!accepted && signal.aborted) { + try { + await lateLease(); + } catch { + markIncomplete("Late source-thread lease release failed."); + } + } + }, + () => undefined, ); + lease = await Promise.race([ + attempted, + new Promise((_resolve, reject) => { + if (signal.aborted) reject(signal.reason); + else + signal.addEventListener("abort", () => reject(signal.reason), { + once: true, + }); + }), + ]); + accepted = true; + await verifyEligibility(); + if (reasons.length) return undefined; } catch { markIncomplete("Exclusive source-thread ownership is unavailable."); return undefined; } - await memory.settled(); - await mutations; - readingSnapshot = true; - const thread = await options.domain.getThreadById({ - threadId: options.threadId, - }); - const resource = await options.domain.getResourceById({ - resourceId: options.resourceId, - }); - const { messages } = await options.domain.listMessages({ - threadId: options.threadId, - perPage: false, - }); - const records = await options.domain.getObservationalMemoryHistory( - options.threadId, - options.resourceId, + const capture = (async () => { + await memory.settled(); + await mutations; + readingSnapshot = true; + try { + const thread = await options.domain.getThreadById({ + threadId: options.threadId, + }); + const resource = await options.domain.getResourceById({ + resourceId: options.resourceId, + }); + const { messages } = await options.domain.listMessages({ + threadId: options.threadId, + perPage: false, + }); + const records = await options.domain.getObservationalMemoryHistory( + options.threadId, + options.resourceId, + ); + const snapshot = { + threadId: options.threadId, + resourceId: options.resourceId, + thread, + resource, + messages, + records, + }; + // No storage-owned objects or Dates escape the explicit codec. + const copy = decodeMemoryValue(encodeMemoryValue(snapshot)); + validateMemorySnapshot(copy); + return copy; + } finally { + readingSnapshot = false; + } + })(); + const copy = await boundedCoordination( + capture, + options.captureWaitMs ?? 5_000, + "Initial memory capture timed out.", ); - const snapshot = { - threadId: options.threadId, - resourceId: options.resourceId, - thread, - resource, - messages, - records, - }; - // Copy through the explicit codec: no storage-owned objects or Dates escape. - const copy = decodeMemoryValue(encodeMemoryValue(snapshot)); - validateMemorySnapshot(copy); + await verifyEligibility(); return reasons.length === 0 ? copy : undefined; - } catch { + } catch (error) { markIncomplete( - "Initial memory capture failed: unsupported, altered, or Unjoined observational-memory state.", + error instanceof Error && + error.message === "Initial memory capture timed out." + ? error.message + : "Initial memory capture failed: unsupported, altered, or Unjoined observational-memory state.", ); return undefined; } finally { capturing = false; - readingSnapshot = false; } }, drain, + verifyEligibility, async release() { if (released) return; await drain(); released = true; try { - await releaseLease?.(); + await lease?.(); } catch { markIncomplete("Exclusive source-thread lease release failed."); } diff --git a/packages/mastra/src/memory-replay.ts b/packages/mastra/src/memory-replay.ts index 1cebcc662..efa77aebf 100644 --- a/packages/mastra/src/memory-replay.ts +++ b/packages/mastra/src/memory-replay.ts @@ -12,6 +12,7 @@ import { type MastraMemorySnapshot, validateMemorySnapshot, } from "./memory-snapshot.js"; +import type { createOMResultTape } from "./om-result-tape.js"; function record(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); @@ -68,8 +69,12 @@ function checkConfiguration(config: Record): void { const feature = config[key]; if (feature === undefined || feature === false) continue; if (!record(feature)) - unsupported(`${key} requires explicit thread-scoped configuration.`); - if (feature.enabled !== false && feature.scope !== "thread") + unsupported(`${key} requires thread-scoped configuration.`); + if ( + feature.enabled !== false && + feature.scope !== undefined && + feature.scope !== "thread" + ) unsupported(`${key} must use thread scope.`); } } @@ -133,6 +138,42 @@ export async function restoreMemoryConfiguration( return copy as MemoryConfigInternal; } +/** Bind each native OM phase to its own result tape without changing actor calls. */ +export async function bindOMResultModels( + config: MemoryConfigInternal, + resolveModel: (id: string) => Promise | MastraModelConfig, + tape: ReturnType, +): Promise { + const om = config.observationalMemory; + if (!om || (record(om) && om.enabled === false)) return config; + const source = om === true ? {} : (om as Record); + const top = source.model; + const bound: Record = { ...source, scope: "thread" }; + delete bound.model; + for (const [name, phase] of [ + ["observation", "observer"], + ["reflection", "reflector"], + ] as const) { + const settings = record(source[name]) ? { ...source[name] } : {}; + const other = source[name === "observation" ? "reflection" : "observation"]; + const modelIdentity = + settings.model ?? top ?? (record(other) ? other.model : undefined); + if (modelIdentity === undefined) + unsupported( + "OM requires an explicit observer and reflector model identity.", + ); + const model = + typeof modelIdentity === "string" + ? await resolveModel(modelIdentity) + : modelIdentity; + if (!record(model) || typeof model.doStream !== "function") + unsupported("OM model resolution did not return a stream-capable model."); + settings.model = tape.instrument(model, phase); + bound[name] = settings; + } + return { ...config, observationalMemory: bound } as MemoryConfigInternal; +} + export interface IsolatedMemoryReplayOptions { invocationId: string; initialSnapshot: MastraMemorySnapshot; @@ -141,6 +182,7 @@ export interface IsolatedMemoryReplayOptions { recordMutation: MastraMemoryCaptureOptions["recordMutation"]; getRequestId?: MastraMemoryCaptureOptions["getRequestId"]; onIncomplete?: MastraMemoryCaptureOptions["onIncomplete"]; + omTape?: ReturnType; } /** Restore historical state into a fresh store; no production store is accepted. */ @@ -152,10 +194,16 @@ export async function createIsolatedMemoryReplay( const snapshot = decodeMemoryValue( encodeMemoryValue(options.initialSnapshot), ) as MastraMemorySnapshot; - const configuration = await restoreMemoryConfiguration( + let configuration = await restoreMemoryConfiguration( options.configuration, options.resolveModel, ); + if (options.omTape) + configuration = await bindOMResultModels( + configuration, + options.resolveModel, + options.omTape, + ); const { Memory } = await import("@mastra/memory"); const { InMemoryStore, MastraCompositeStore } = await import( "@mastra/core/storage" @@ -214,6 +262,9 @@ export async function createIsolatedMemoryReplay( })(); return finished; }, + release(): Promise { + return finished ?? Promise.resolve(); + }, }; } catch (error) { await store.close(); diff --git a/packages/mastra/src/memory-snapshot.ts b/packages/mastra/src/memory-snapshot.ts index e63f47c4a..a250c7efb 100644 --- a/packages/mastra/src/memory-snapshot.ts +++ b/packages/mastra/src/memory-snapshot.ts @@ -9,11 +9,13 @@ import type { ObservationalMemoryRecord, StorageResourceType, } from "@mastra/core/storage"; -import { type JsonValue, toRecorderJson } from "@zenml-io/kitaru"; +import type { JsonValue } from "@zenml-io/kitaru"; import { - MAX_RECORDED_PAYLOAD_CHARS, - recordedToolPayloadConversion, + MAX_MASTRA_REPLAY_ITEMS, + MAX_MASTRA_REPLAY_JSON_BYTES, + strictMastraReplayValue, } from "@zenml-io/kitaru/adapter"; +import { fileReference } from "./stateful-files.js"; export const MEMORY_REPLAY_KEY = "mastra_memory_replay"; const CODEC_KEY = "$mastra"; @@ -50,11 +52,12 @@ export interface MastraMemoryReplayInput { configuration: Record; requestContext: Record; files: MastraRecordedFile[]; + omTape?: JsonValue[]; } export interface MastraMemoryReplayEnvelope { [key: string]: JsonValue; - version: 2; + version: 3; complete: boolean; reasons: string[]; invocationId: string; @@ -63,6 +66,7 @@ export interface MastraMemoryReplayEnvelope { configuration: JsonValue; requestContext: JsonValue; files: MastraFileManifestEntry[]; + omTape: JsonValue[]; } class MemoryReplayError extends Error {} @@ -79,16 +83,6 @@ function requireValue(condition: unknown, reason: string): asserts condition { if (!condition) throw unsupported(reason); } -function checkBudget(value: JsonValue): void { - try { - toRecorderJson(value); - } catch { - throw unsupported( - "Memory value exceeds the replay JSON depth/item limits.", - ); - } -} - function hash(bytes: Uint8Array): string { return createHash("sha256").update(bytes).digest("hex"); } @@ -99,8 +93,8 @@ function binary(bytes: Uint8Array): { sha256: string; } { requireValue( - bytes.byteLength <= MAX_RECORDED_PAYLOAD_CHARS, - "Binary content exceeds the replay payload limit.", + bytes.byteLength <= 8 * 1_048_576, + "Binary content exceeds maximum file bytes 8388608.", ); return { base64: Buffer.from(bytes).toString("base64"), @@ -115,6 +109,7 @@ function readBinary(value: Record): Uint8Array { typeof value.length === "number" && Number.isSafeInteger(value.length) && value.length >= 0 && + value.length <= 8 * 1_048_576 && typeof value.sha256 === "string", "Malformed binary content.", ); @@ -130,15 +125,27 @@ function readBinary(value: Record): Uint8Array { function validateUrl(value: string): URL { const url = new URL(value); + if (url.protocol === "kitaru-file:") { + requireValue( + /^kitaru-file:\/\/sha256\/[a-f0-9]{64}$/.test(value), + "Malformed captured file reference.", + ); + return url; + } requireValue( - !url.username && !url.password, + (url.protocol === "https:" || url.protocol === "http:") && + !url.username && + !url.password, "URL credentials are not replayable.", ); - const query = recordedToolPayloadConversion( - Object.fromEntries(url.searchParams), - "Mastra file URL", + requireValue( + ![...url.searchParams.keys()].some((key) => + /(?:^|[-_])(?:api[-_]?key|authorization|cookie|password|secret|token|signature|credential|sig)$/i.test( + key, + ), + ), + "URL query credentials are not replayable.", ); - requireValue(!query.lossy, "URL query credentials are not replayable."); return url; } @@ -148,15 +155,15 @@ export function encodeMemoryValue(value: unknown): JsonValue { const active = new Set(); function visit(current: unknown, depth: number): JsonValue { requireValue( - ++items <= 10_000 && depth < 64, - "Memory value exceeds the replay depth/item limits.", + ++items <= MAX_MASTRA_REPLAY_ITEMS && depth < 64, + `Memory value exceeds maximum items ${MAX_MASTRA_REPLAY_ITEMS} or depth 64.`, ); if (current === undefined) return { [CODEC_KEY]: "undefined" }; if (current === null || typeof current === "boolean") return current; if (typeof current === "string") { requireValue( - current.length <= MAX_RECORDED_PAYLOAD_CHARS, - "Memory value exceeds the replay payload limit.", + current.length <= MAX_MASTRA_REPLAY_JSON_BYTES, + `Memory value exceeds maximum string length ${MAX_MASTRA_REPLAY_JSON_BYTES}.`, ); return current; } @@ -209,29 +216,12 @@ export function encodeMemoryValue(value: unknown): JsonValue { } } const encoded = visit(value, 0); - const converted = recordedToolPayloadConversion( - encoded, - "Mastra memory replay", - ); - requireValue( - !converted.lossy, - "Memory content was altered by credential protection or replay payload bounds.", - ); - checkBudget(converted.value); - return converted.value; + return strictMastraReplayValue(encoded); } /** Decode an already bounded value, rejecting ambiguous or damaged codec records. */ export function decodeMemoryValue(value: JsonValue): unknown { - const converted = recordedToolPayloadConversion( - value, - "Mastra memory replay", - ); - requireValue( - !converted.lossy, - "Memory content was altered by credential protection or replay payload bounds.", - ); - checkBudget(converted.value); + const converted = strictMastraReplayValue(value); function visit(current: JsonValue): unknown { if (Array.isArray(current)) return current.map(visit); if (!isRecord(current)) return current; @@ -264,7 +254,7 @@ export function decodeMemoryValue(value: JsonValue): unknown { ]), ); } - return visit(converted.value); + return visit(converted); } /** Validate the complete native state before an isolated store receives any writes. */ @@ -448,20 +438,57 @@ function validateConfiguration( !containsTransport(configuration), "Replay configuration contains transport metadata.", ); - const memory = configuration.memoryConfig ?? configuration.memory; - if (!isRecord(memory)) return; - requireValue( - memory.semanticRecall === undefined || memory.semanticRecall === false, - "Semantic recall is outside isolated memory replay scope.", - ); - for (const key of ["workingMemory", "observationalMemory"]) { - const feature = memory[key]; - if (isRecord(feature) && feature.enabled !== false) - requireValue( - feature.scope === "thread", - "Only explicitly thread-scoped memory is replayable.", - ); + for (const memory of [configuration.memoryConfig, configuration.memory]) { + if (!isRecord(memory)) continue; + requireValue( + memory.semanticRecall === undefined || memory.semanticRecall === false, + "Semantic recall is outside isolated memory replay scope.", + ); + for (const key of ["workingMemory", "observationalMemory"]) { + const feature = memory[key]; + if (isRecord(feature) && feature.enabled !== false) + requireValue( + feature.scope === "thread" || + (key === "observationalMemory" && feature.scope === undefined), + "Only thread-scoped memory is replayable.", + ); + } + } +} + +function normalizeReplayConfiguration( + configuration: Record, +): Record { + const normalized = { ...configuration }; + for (const key of ["memoryConfig", "memory"]) { + const memory = normalized[key]; + if ( + !isRecord(memory) || + !isRecord(memory.observationalMemory) || + memory.observationalMemory.enabled === false || + memory.observationalMemory.scope !== undefined + ) + continue; + normalized[key] = { + ...memory, + observationalMemory: { ...memory.observationalMemory, scope: "thread" }, + }; } + return normalized; +} + +function usesObservationalMemory( + configuration: Record, +): boolean { + return [configuration.memoryConfig, configuration.memory].some((memory) => { + if (!isRecord(memory)) return false; + const feature = memory.observationalMemory; + return ( + feature !== undefined && + feature !== false && + (!isRecord(feature) || feature.enabled !== false) + ); + }); } /** Require native context selectors to match the leased and captured memory. */ @@ -501,7 +528,7 @@ export function createMemoryReplayEnvelope( input: MastraMemoryReplayInput, ): MastraMemoryReplayEnvelope { const incomplete = (reason: string): MastraMemoryReplayEnvelope => ({ - version: 2, + version: 3, complete: false, reasons: [reason], invocationId: "", @@ -510,60 +537,73 @@ export function createMemoryReplayEnvelope( configuration: null, requestContext: null, files: [], + omTape: [], }); try { validateMemorySnapshot(input.initialSnapshot); const envelope: MastraMemoryReplayEnvelope = { - version: 2, + version: 3, complete: true, reasons: [], invocationId: input.invocationId, rawInput: encodeMemoryValue(input.rawInput), initialSnapshot: encodeMemoryValue(input.initialSnapshot), - configuration: encodeMemoryValue(input.configuration), + configuration: encodeMemoryValue( + normalizeReplayConfiguration(input.configuration), + ), requestContext: encodeMemoryValue(input.requestContext), files: input.files.map((file) => ({ url: file.url, mediaType: file.mediaType, ...binary(file.bytes), })), + omTape: input.omTape === undefined ? [] : input.omTape, }; // The combined envelope, including encoded bytes and metadata, shares one budget. - const converted = recordedToolPayloadConversion( + const converted = strictMastraReplayValue( envelope, "Mastra memory replay envelope", ); - requireValue( - !converted.lossy, - "Envelope exceeds replay limits or contains credentials.", - ); - decodeMemoryReplayEnvelope(converted.value); - return converted.value as unknown as MastraMemoryReplayEnvelope; + decodeConvertedMemoryReplayEnvelope(converted); + return converted as MastraMemoryReplayEnvelope; } catch (error) { return incomplete( error instanceof MemoryReplayError ? error.message - : "Memory replay prerequisites could not be captured safely.", + : error instanceof Error && /exceeds maximum/.test(error.message) + ? error.message + : "Memory replay prerequisites could not be captured safely.", ); } } +/** Produce the immutable final envelope after recorded OM work has settled. */ +export function finalizeMemoryReplayEnvelope( + envelope: MastraMemoryReplayEnvelope, + omTape: JsonValue[], +): MastraMemoryReplayEnvelope { + const final = strictMastraReplayValue( + { ...envelope, omTape }, + "Mastra memory replay envelope", + ) as MastraMemoryReplayEnvelope; + decodeConvertedMemoryReplayEnvelope(final); + return final; +} + export function decodeMemoryReplayEnvelope( input: unknown, ): MastraMemoryReplayInput { - const converted = recordedToolPayloadConversion( - input, - "Mastra memory replay envelope", - ); - requireValue( - !converted.lossy, - "Envelope exceeds replay limits or contains credentials.", - ); - checkBudget(converted.value); - const value = converted.value; + const value = strictMastraReplayValue(input, "Mastra memory replay envelope"); + return decodeConvertedMemoryReplayEnvelope(value); +} + +/** Validate a value already copied through the strict replay codec. */ +function decodeConvertedMemoryReplayEnvelope( + value: JsonValue, +): MastraMemoryReplayInput { requireValue( isRecord(value) && - value.version === 2 && + (value.version === 2 || value.version === 3) && value.complete === true && Array.isArray(value.reasons) && value.reasons.length === 0 && @@ -572,6 +612,10 @@ export function decodeMemoryReplayEnvelope( Array.isArray(value.files), "Missing, incomplete, or unknown version of memory replay envelope.", ); + requireValue( + value.version === 2 || Array.isArray(value.omTape), + "Malformed recorded observational-memory tape.", + ); for (const key of [ "rawInput", "initialSnapshot", @@ -586,6 +630,10 @@ export function decodeMemoryReplayEnvelope( validateMemorySnapshot(initialSnapshot); const configuration = decodeMemoryValue(value.configuration as JsonValue); validateConfiguration(configuration); + requireValue( + value.version !== 2 || !usesObservationalMemory(configuration), + "mastra_om_tape_missing", + ); const requestContext = decodeMemoryValue(value.requestContext as JsonValue); requireValue(isRecord(requestContext), "Malformed recorded request context."); validateMemoryReplayContext(initialSnapshot, requestContext); @@ -610,11 +658,18 @@ export function decodeMemoryReplayEnvelope( "Malformed or duplicate recorded file.", ); validateUrl(file.url); + const bytes = readBinary(file); + if (value.version === 3) + requireValue( + /^kitaru-file:\/\/sha256\/[a-f0-9]{64}$/.test(file.url) && + file.url === fileReference({ mediaType: file.mediaType, bytes }), + "Recorded file must use its captured content reference.", + ); urls.add(file.url); return { url: file.url, mediaType: file.mediaType, - bytes: readBinary(file), + bytes, }; }); return { @@ -624,6 +679,7 @@ export function decodeMemoryReplayEnvelope( configuration, requestContext, files, + ...(value.version === 3 ? { omTape: value.omTape as JsonValue[] } : {}), }; } diff --git a/packages/mastra/src/memory.ts b/packages/mastra/src/memory.ts index 29ad78cfa..205a2a245 100644 --- a/packages/mastra/src/memory.ts +++ b/packages/mastra/src/memory.ts @@ -2,6 +2,8 @@ export type { MastraExclusiveMemoryAccess, MastraMemoryCaptureBinding, MastraMemoryCaptureOptions, + MastraMemoryLease, + MastraMemoryLeaseOptions, MastraMemoryMutation, MastraMemorySelector, } from "./memory-binding.js"; diff --git a/packages/mastra/src/om-result-tape.ts b/packages/mastra/src/om-result-tape.ts new file mode 100644 index 000000000..d311f1da2 --- /dev/null +++ b/packages/mastra/src/om-result-tape.ts @@ -0,0 +1,211 @@ +import { createHash } from "node:crypto"; +import type { JsonValue } from "@zenml-io/kitaru"; +import { + MAX_MASTRA_REPLAY_ITEMS, + MAX_MASTRA_REPLAY_JSON_BYTES, +} from "@zenml-io/kitaru/adapter"; +import { decodeMemoryValue, encodeMemoryValue } from "./memory-snapshot.js"; + +export type OMPhase = "observer" | "reflector"; +type OMMethod = "doGenerate" | "doStream"; + +export interface OMResultEntry { + phase: OMPhase; + ordinal: number; + method: OMMethod; + inputFingerprint: string; + output: JsonValue; +} + +export class MastraOMDivergenceError extends Error { + readonly code = "mastra_om_diverged"; + + constructor(reason: string) { + super(`Recorded Mastra observational memory diverged: ${reason}`); + } +} + +function fingerprint(input: unknown): string { + try { + return createHash("sha256").update(JSON.stringify(input)).digest("hex"); + } catch { + return "unavailable"; + } +} + +function countJsonItems(value: JsonValue, remaining: number): number { + let count = 0; + function visit(current: JsonValue): void { + if (++count > remaining) throw new Error("OM stream item limit exceeded"); + if (Array.isArray(current)) { + for (const item of current) visit(item); + } else if (current !== null && typeof current === "object") { + for (const item of Object.values(current)) visit(item); + } + } + visit(value); + return count; +} + +interface ModelLike { + doGenerate?: (input: unknown) => Promise; + doStream?: (input: unknown) => Promise; +} + +/** Intercept only OM model calls; the actor model remains untouched. */ +export function createOMResultTape( + recorded: readonly OMResultEntry[] | undefined, + onIncomplete: (reason: string) => void, + onInputMismatch?: (entry: OMResultEntry) => void, +) { + const entries: Array = recorded + ? [...recorded] + : []; + const pending = new Set>(); + let next = 0; + let incomplete = false; + + function failCapture(): void { + incomplete = true; + try { + onIncomplete("Observational-memory result could not be recorded safely."); + } catch { + // Diagnostic sinks never alter the provider's stream. + } + } + + function take( + phase: OMPhase, + method: OMMethod, + input: unknown, + ): OMResultEntry { + const entry = entries[next++]; + if ( + !entry || + entry.phase !== phase || + entry.method !== method || + entry.ordinal !== next - 1 + ) + throw new MastraOMDivergenceError( + "missing, extra, or reordered model call", + ); + if (entry.inputFingerprint !== fingerprint(input)) onInputMismatch?.(entry); + return entry; + } + + function instrument(model: T, phase: OMPhase): T { + const proxy = new Proxy(model, { + get(target, key) { + const value = Reflect.get(target, key, target); + if (key !== "doGenerate" && key !== "doStream") + return typeof value === "function" ? value.bind(target) : value; + const method = key as OMMethod; + if (recorded) { + return async (input: unknown) => { + const entry = take(phase, method, input); + const output = decodeMemoryValue(entry.output); + if (method === "doGenerate") return output; + if (!Array.isArray(output)) + throw new MastraOMDivergenceError("invalid recorded stream"); + return { + stream: new ReadableStream({ + start(controller) { + for (const chunk of output) controller.enqueue(chunk); + controller.close(); + }, + }), + }; + }; + } + return async (input: unknown) => { + const ordinal = next++; + const inputFingerprint = fingerprint(input); + const result = await Reflect.apply( + value as (input: unknown) => Promise, + target, + [input], + ); + if (method === "doGenerate") { + try { + entries[ordinal] = { + phase, + ordinal, + method, + inputFingerprint, + output: encodeMemoryValue(result), + }; + } catch { + failCapture(); + } + return result; + } + const stream = (result as { stream?: ReadableStream }) + ?.stream; + if (!(stream instanceof ReadableStream)) { + failCapture(); + return result; + } + const [native, capture] = stream.tee(); + const work = (async () => { + const reader = capture.getReader(); + try { + const chunks: JsonValue[] = []; + let capturedBytes = 2; // JSON array brackets. + let capturedItems = 1; // JSON array itself. + while (true) { + const item = await reader.read(); + if (item.done) break; + const encoded = encodeMemoryValue(item.value); + capturedItems += countJsonItems( + encoded, + MAX_MASTRA_REPLAY_ITEMS - capturedItems, + ); + capturedBytes += + Buffer.byteLength(JSON.stringify(encoded), "utf8") + + (chunks.length > 0 ? 1 : 0); + if (capturedBytes > MAX_MASTRA_REPLAY_JSON_BYTES) + throw new Error("OM stream byte limit exceeded"); + chunks.push(encoded); + } + entries[ordinal] = { + phase, + ordinal, + method, + inputFingerprint, + output: chunks, + }; + } catch { + failCapture(); + // Tee cancellation may wait for the native branch to finish. + void reader.cancel().catch(() => {}); + } finally { + reader.releaseLock(); + } + })(); + pending.add(work); + void work.finally(() => pending.delete(work)); + return { ...(result as object), stream: native }; + }; + }, + }); + return proxy; + } + + async function finish(): Promise { + while (pending.size > 0) await Promise.all([...pending]); + if (recorded) { + if (next !== entries.length) + throw new MastraOMDivergenceError("recorded model call was not made"); + } else if ( + incomplete || + Array.from({ length: next }, (_, ordinal) => entries[ordinal]).some( + (entry) => !entry, + ) + ) { + throw new Error("Observational-memory result tape is incomplete."); + } + return entries as OMResultEntry[]; + } + + return { instrument, finish }; +} diff --git a/packages/mastra/src/request-capture.ts b/packages/mastra/src/request-capture.ts index 3eef999e7..94649c82a 100644 --- a/packages/mastra/src/request-capture.ts +++ b/packages/mastra/src/request-capture.ts @@ -33,6 +33,7 @@ export interface RequestCaptureOptions { error: unknown, ) => void | Promise; onCaptureError?: (error: unknown) => void; + sanitizeEvidence?: (value: T) => T; } export interface RequestStepContext { @@ -102,7 +103,9 @@ export function createRequestCapture(options: RequestCaptureOptions) { reasons: string[], ): JsonValue { try { - const encoded = encodeMemoryValue(value); + const encoded = encodeMemoryValue( + options.sanitizeEvidence ? options.sanitizeEvidence(value) : value, + ); // Provider options can contain custom transport headers whose keys are // not recognizable credential names. Do not persist that transport bag. function containsTransport(current: JsonValue): boolean { @@ -141,17 +144,38 @@ export function createRequestCapture(options: RequestCaptureOptions) { try { const list = step.messageList; const sources = list?.makeMessageSourceChecker(); + const tagged = Object.entries( + list?.getPersisted.taggedSystemMessages ?? {}, + ).flatMap(([tag, messages]) => + messages.map((message) => ({ role: message.role, tag })), + ); + const allSystems = list?.getSystemMessages() ?? []; + const systems = [ + ...allSystems + .slice(0, allSystems.length - tagged.length) + .map((message) => ({ role: message.role, tag: undefined })), + ...tagged, + ]; provenance = convert( { - applicationInstructions: step.applicationInstructions, - extraContext: step.extraContext, - systemMessages: list?.serializeForSpan().systemMessages ?? [], + version: 2, + ...(step.applicationInstructions === undefined + ? {} + : { applicationInstructionsRef: "effective_request.prompt" }), + ...(step.extraContext === undefined + ? {} + : { extraContextRef: "mastra_memory_replay.configuration" }), + systemMessages: systems.map((message, index) => ({ + index, + tag: message.tag, + role: message.role, + })), messages: - list?.get.all.db().map((message) => ({ + list?.get.all.db().map((message, index) => ({ + index, id: message.id, source: sources?.getSource(message) ?? null, role: message.role, - content: message.content, })) ?? [], }, "Prompt provenance", diff --git a/packages/mastra/src/stateful-agent.ts b/packages/mastra/src/stateful-agent.ts index 3afb96ee8..a62758884 100644 --- a/packages/mastra/src/stateful-agent.ts +++ b/packages/mastra/src/stateful-agent.ts @@ -7,12 +7,9 @@ import type { InputProcessor } from "@mastra/core/processors"; import { RequestContext } from "@mastra/core/request-context"; import type { MemoryStorage } from "@mastra/core/storage"; import type { Memory } from "@mastra/memory"; +import { KitaruClient, type SessionNodeCreateRequest } from "@zenml-io/kitaru"; import { - type JsonValue, - KitaruClient, - type SessionNodeCreateRequest, -} from "@zenml-io/kitaru"; -import { + type AdapterClient, type AdapterRunState, normalizeRecordingLimits, parseModelSettings, @@ -27,6 +24,7 @@ import { } from "./memory-binding.js"; import { assertMemoryReplayVersions, + bindOMResultModels, createIsolatedMemoryReplay, getMemoryModelId, serializeMemoryConfiguration, @@ -35,19 +33,26 @@ import { createMemoryReplayEnvelope, decodeMemoryValue, encodeMemoryValue, + finalizeMemoryReplayEnvelope, type MastraMemorySnapshot, MEMORY_REPLAY_KEY, restoreMemoryReplayEnvelope, validateMemoryReplayContext, validateMemoryReplaySelectors, } from "./memory-snapshot.js"; +import { createOMResultTape, type OMResultEntry } from "./om-result-tape.js"; import { assertStableToolName } from "./replay-guards.js"; import { createRequestCapture, type RequestEvidence, requestEvidenceAttributes, } from "./request-capture.js"; -import { createCapturedFiles, restoreCapturedFiles } from "./stateful-files.js"; +import { + createCapturedFiles, + createRecordedEvidenceSanitizer, + restoreCapturedFiles, + type UnsafeEvidenceReason, +} from "./stateful-files.js"; import { bindMemoryToolIdentity, createStatefulToolProcessors, @@ -66,6 +71,33 @@ interface MastraMemorySource { exclusiveAccess: MastraExclusiveMemoryAccess; } +function recordingClient( + client: KitaruClient, + sanitize: (value: T) => T, + unsafeReason: () => string | undefined, +): AdapterClient { + return { + createSession: (request) => client.createSession(sanitize(request)), + getReplay: client.getReplay.bind(client), + getTaskSpec: client.getTaskSpec.bind(client), + lookupToolResult: client.lookupToolResult.bind(client), + upsertSessionNodes: (sessionId, request) => + client.upsertSessionNodes(sessionId, sanitize(request)), + updateSession: (sessionId, request) => { + const safe = sanitize(request); + const reason = unsafeReason(); + if (reason && safe.metadata?.mastra_replay_state === "eligible") { + safe.metadata = { + ...safe.metadata, + mastra_replay_state: "ineligible", + mastra_replay_reason: reason, + }; + } + return client.updateSession(sessionId, safe); + }, + }; +} + export interface MemoryReplayAgentOptions extends KitaruAgentOptions { /** Registry context passed to baseline dynamic configuration resolvers. */ mastra?: Mastra; @@ -99,11 +131,14 @@ function requireRecord(value: unknown, label: string): Record { throw new Error(`Unsupported Mastra memory replay: missing ${label}.`); return value; } +class MemoryReplayContextError extends Error {} + function getSelector(options: RuntimeStreamOptions) { - const memory = requireRecord( - options.memory, - "memory thread/resource selectors", - ); + const memory = record(options.memory) ? options.memory : undefined; + if (!memory) + throw new MemoryReplayContextError( + "Memory replay requires explicit memory.thread and memory.resource strings.", + ); const threadId = typeof memory.thread === "string" ? memory.thread @@ -111,7 +146,7 @@ function getSelector(options: RuntimeStreamOptions) { ? memory.thread.id : undefined; if (typeof threadId !== "string" || typeof memory.resource !== "string") - throw new Error( + throw new MemoryReplayContextError( "Memory replay requires explicit memory.thread and memory.resource strings.", ); return { threadId, resourceId: memory.resource }; @@ -215,28 +250,145 @@ export function createMemoryReplayAgent( apiUrl: options.apiUrl, timeoutMs: options.timeoutMs, }); - async function stream( + async function runNativeBaseline( + rawInput: unknown, + callerOptions: RuntimeStreamOptions, + ): Promise { + const { MastraCompositeStore } = await import("@mastra/core/storage"); + const { Memory } = await import("@mastra/memory"); + const source = await options.sourceMemory(); + let unsafeSelector: { threadId: string; resourceId: string } | undefined; + try { + unsafeSelector = getSelector(callerOptions); + } catch { + // An implicit native selector can come from factory defaults. + } + try { + let timeout: ReturnType | undefined; + try { + await Promise.race([ + source.exclusiveAccess.markUnsafeWrite(unsafeSelector), + new Promise((_resolve, reject) => { + timeout = setTimeout( + () => reject(new Error("Source-thread unsafe marker timed out.")), + 100, + ); + }), + ]); + } finally { + if (timeout) clearTimeout(timeout); + } + } catch (error) { + reportLocalRecordingError(error); + } + const memory = new Memory({ + storage: new MastraCompositeStore({ + id: `kitaru-native-${globalThis.crypto.randomUUID()}`, + domains: { memory: source.domain }, + }), + options: source.configuration, + }); + const workspace = options.skillsDirectory + ? await loadSkillsWorkspace(options.skillsDirectory) + : undefined; + const config = await factory({ + memory, + resolveFile: + options.resolveFile ?? + (async () => { + throw new Error("Missing controlled file resolver."); + }), + workspace: workspace?.workspace, + }); + const native = new Agent({ ...config, memory }) as unknown as { + stream(input: unknown, options: RuntimeStreamOptions): Promise; + }; + return native.stream(rawInput, callerOptions); + } + + function reportLocalRecordingError(error: unknown): void { + if (options.onRecordingError) { + void Promise.resolve() + .then(() => options.onRecordingError?.({ error, stage: "complete" })) + .catch(() => undefined); + } else { + console.warn("Kitaru memory recording is unavailable for this turn"); + } + } + + async function reportSetupFailure( + error: unknown, + reasonCode: string, + ): Promise { + try { + const session = await client.createSession({ + agent_id: options.agentId, + agent_version_id: options.agentVersionId, + adapter_version: ( + createRequire(import.meta.url)("../package.json") as { + version: string; + } + ).version, + framework: "mastra", + inputs: { + [MEMORY_REPLAY_KEY]: { version: 2, complete: false }, + }, + metadata: { + mastra_replay_state: "ineligible", + mastra_replay_reason: reasonCode, + mastra_native_state: "started", + }, + name: options.sessionName, + origin: "recorded", + outputs: null, + started_at: new Date().toISOString(), + status: "in_progress", + }); + await client.updateSession(session.id, { + error: `KITARU_RECORDING_INCOMPLETE:${reasonCode}`, + ended_at: new Date().toISOString(), + status: "failed", + }); + } catch { + reportLocalRecordingError(error); + } + } + + async function recordedStream( rawInput: unknown, callerOptions: RuntimeStreamOptions = {}, + markRecordingStreamEntered: () => void, ): Promise { assertMemoryReplayVersions(); const { resolveModelConfig } = await import("@mastra/core/llm"); const { MastraCompositeStore } = await import("@mastra/core/storage"); const startedAt = new Date().toISOString(); const invocationId = globalThis.crypto.randomUUID(); + let baselineFiles: + | Awaited> + | undefined; const replay = await resolveReplayContext({ allowedReplayModels: options.allowedReplayModels, - callerInput: encodeMemoryValue(rawInput), + callerInput: rawInput, client, + recordedInputProjector: async (input) => { + baselineFiles = await createCapturedFiles( + options.files ?? [], + options.resolveFile ?? + (async () => { + throw new Error("Missing controlled file resolver."); + }), + ); + return encodeMemoryValue(baselineFiles.replaceDeclaredFileUrls(input)); + }, requestedModelId: options.requestedModelId, }); const historical = restoreMemoryReplayEnvelope(replay.effectiveInput); const invocationInput = - historical?.rawInput ?? - decodeMemoryValue(replay.effectiveRuntimeInput as JsonValue); + historical?.rawInput ?? replay.effectiveRuntimeInput; if (Boolean(replay.spec) !== Boolean(historical)) throw new Error( - "Memory replay requires a complete version-2 recorded invocation and an active Kitaru replay.", + "Memory replay requires a complete recorded invocation and an active Kitaru replay.", ); if (replay.override?.prompt != null) throw new Error( @@ -244,23 +396,36 @@ export function createMemoryReplayAgent( ); const selector = historical?.initialSnapshot ?? getSelector(callerOptions); const liveContext = callerOptions.requestContext ?? new RequestContext(); - if (!historical) - validateMemoryReplaySelectors( - selector, - Object.fromEntries(liveContext.entries()), - ); - const recordedContext = - historical?.requestContext ?? - options.captureRequestContext?.(liveContext) ?? - Object.fromEntries(liveContext.entries()); + if (!historical) { + try { + validateMemoryReplaySelectors( + selector, + Object.fromEntries(liveContext.entries()), + ); + } catch { + throw new MemoryReplayContextError( + "Request context contains incompatible memory selectors.", + ); + } + } + const recordedContext = historical?.requestContext ?? {}; const safeContext = requireRecord( decodeMemoryValue(encodeMemoryValue(recordedContext)), "request context", ); - validateMemoryReplayContext(selector, safeContext); - const requestContext = new MemoryReplayRequestContext(selector); - for (const [key, value] of Object.entries(safeContext)) - requestContext.set(key, value); + try { + validateMemoryReplayContext(selector, safeContext); + } catch { + throw new MemoryReplayContextError( + "Request context is incompatible with recorded memory selectors.", + ); + } + const requestContext = historical + ? new MemoryReplayRequestContext(selector) + : liveContext; + if (historical) + for (const [key, value] of Object.entries(safeContext)) + requestContext.set(key, value); const abort = new AbortController(); let state: AdapterRunState | undefined; let requestCapture: ReturnType | undefined; @@ -269,9 +434,9 @@ export function createMemoryReplayAgent( return state; }; const onIncomplete = (reason: string): void => { - // Mastra converts tool storage errors into results and otherwise continues. - // Recording-only incompleteness must preserve ordinary native execution. - if (reason !== "Native memory storage mutation failed.") return; + // Baseline diagnosis must not alter Mastra's native storage result. + if (!historical || reason !== "Native memory storage mutation failed.") + return; const error = new Error(reason); state?.storeFailure(error); abort.abort(error); @@ -301,12 +466,32 @@ export function createMemoryReplayAgent( evidence_complete: event.complete, }, }); + const omCaptureErrors: string[] = []; + const omMismatches: OMResultEntry[] = []; + const omTape = createOMResultTape( + historical?.omTape as OMResultEntry[] | undefined, + (reason) => omCaptureErrors.push(reason), + (entry) => omMismatches.push(entry), + ); let runtime: { memory: Memory; binding: MastraMemoryCaptureBinding; initialSnapshot: MastraMemorySnapshot | undefined; finish(): Promise; + release(): Promise; }; + let unsafeEvidenceReason: UnsafeEvidenceReason | undefined; + let markUnknownOnBinding: (() => void) | undefined; + const markUnknownCredentialUrl = (reason: UnsafeEvidenceReason) => { + if (!unsafeEvidenceReason || reason === "credential_url") + unsafeEvidenceReason = reason; + markUnknownOnBinding?.(); + }; + const sanitizer = historical + ? createRecordedEvidenceSanitizer(new Map(), markUnknownCredentialUrl) + : baselineFiles?.evidenceSanitizer(markUnknownCredentialUrl); + if (!sanitizer) + throw new Error("Controlled evidence sanitizer was not initialized."); if (historical) { runtime = await createIsolatedMemoryReplay({ invocationId, @@ -319,6 +504,7 @@ export function createMemoryReplayAgent( recordMutation, onIncomplete, getRequestId: () => requestCapture?.currentRequestId, + omTape, }); } else { const source = await options.sourceMemory(); @@ -327,6 +513,7 @@ export function createMemoryReplayAgent( ...selector, domain: source.domain, exclusiveAccess: source.exclusiveAccess, + sanitizeEvidence: sanitizer.replace, recordMutation, onIncomplete, getRequestId: () => requestCapture?.currentRequestId, @@ -337,7 +524,11 @@ export function createMemoryReplayAgent( id: `kitaru-baseline-${invocationId}`, domains: { memory: binding.domain }, }), - options: source.configuration, + options: await bindOMResultModels( + source.configuration, + options.resolveModel, + omTape, + ), }); const initialSnapshot = await binding.captureInitial(source); let finished: Promise | undefined; @@ -347,27 +538,32 @@ export function createMemoryReplayAgent( initialSnapshot, finish() { finished ??= (async () => { - try { - await memory.settled(); - await binding.drain(); - } finally { - await binding.release(); - } + await memory.settled(); + await binding.drain(); })(); return finished; }, + release: () => binding.release(), }; } + markUnknownOnBinding = () => + runtime.binding.markIncomplete( + "Recorded evidence contains an uncaptured credential URL or unsupported value.", + ); + if (unsafeEvidenceReason) markUnknownOnBinding(); try { const files = historical ? restoreCapturedFiles(historical.files) - : await createCapturedFiles( - options.files ?? [], - options.resolveFile ?? - (async () => { - throw new Error("Missing controlled file resolver."); - }), - ); + : baselineFiles; + if (!files) + throw new Error("Controlled file capture was not initialized."); + const evidenceClient = recordingClient(client, sanitizer.replace, () => + unsafeEvidenceReason === "credential_url" + ? "credential_url_uncaptured" + : unsafeEvidenceReason === "unsupported_value" + ? "recorded_evidence_unsupported" + : undefined, + ); const workspace = options.skillsDirectory ? await loadSkillsWorkspace( options.skillsDirectory, @@ -421,14 +617,39 @@ export function createMemoryReplayAgent( ? await config.defaultOptions(dynamic) : (config.defaultOptions ?? {}); const defaults = requireRecord(resolvedDefaults, "default options"); - // Dynamic resolvers receive the mutable context that native Mastra uses. - const effectiveContext = requireRecord( - decodeMemoryValue( - encodeMemoryValue(Object.fromEntries(requestContext.entries())), - ), - "request context", - ); - validateMemoryReplayContext(selector, effectiveContext); + // Baseline resolvers receive the original mutable context. Only the + // application's approved projection enters replay evidence. + const projectedContext = historical + ? Object.fromEntries(requestContext.entries()) + : (options.captureRequestContext?.(requestContext) ?? {}); + let effectiveContext: Record; + try { + effectiveContext = requireRecord( + decodeMemoryValue(encodeMemoryValue(projectedContext)), + "request context", + ); + validateMemoryReplayContext(selector, effectiveContext); + } catch { + throw new MemoryReplayContextError( + "Request context cannot be captured for memory replay.", + ); + } + if ( + Object.keys(effectiveContext).some((key) => + /auth|credential|jwt|key|password|secret|token/i.test(key), + ) + ) + throw new MemoryReplayContextError( + "Unsupported replay request context credential key.", + ); + if ( + !historical && + !options.captureRequestContext && + [...requestContext.entries()].length > 0 + ) + runtime.binding.markIncomplete( + "Request context was not captured safely.", + ); const { deepMerge } = await import("@mastra/core/utils"); const callerData = { ...callerOptions }; delete callerData.requestContext; @@ -447,7 +668,7 @@ export function createMemoryReplayAgent( effectiveSelector.threadId !== selector.threadId || effectiveSelector.resourceId !== selector.resourceId ) - throw new Error( + throw new MemoryReplayContextError( "Invocation memory selectors differ from the captured selectors.", ); if (record(effective.memory) && effective.memory.options !== undefined) @@ -474,9 +695,16 @@ export function createMemoryReplayAgent( serializeMemoryConfiguration(runtime.memory.getMergedThreadConfig()), ...(workspace ? { workspaceManifest: workspace.manifest } : {}), }; + let recordedRawInput = invocationInput; + if (!historical) { + if (!baselineFiles) + throw new Error("Controlled file capture was not initialized."); + recordedRawInput = + baselineFiles.replaceDeclaredFileUrls(invocationInput); + } const envelope = createMemoryReplayEnvelope({ invocationId, - rawInput: invocationInput, + rawInput: recordedRawInput, initialSnapshot: runtime.initialSnapshot as MastraMemorySnapshot, configuration, requestContext: effectiveContext, @@ -502,6 +730,7 @@ export function createMemoryReplayAgent( }); const capture = createRequestCapture({ invocationId, + sanitizeEvidence: sanitizer.replace, getMemoryRevision: () => runtime.binding.revision, onFailedAttempt: writeAttempt, onCaptureError: () => @@ -513,15 +742,40 @@ export function createMemoryReplayAgent( const policy = createStatefulToolProcessors({ tokens: owned.tokens, getState, + sanitizeEvidence: sanitizer.replace, abort(reason) { state?.storeFailure(reason); abort.abort(reason); }, adapter: options, }); + const contextAtCapture = new Map(requestContext.entries()); const requestProcessor: InputProcessor = { id: "kitaru-effective-request", async processInputStep(args) { + const currentContext = + args.requestContext instanceof RequestContext + ? args.requestContext + : requestContext; + const current = new Map( + [...currentContext.entries()].filter( + ([key]) => key !== "MastraMemory", + ), + ); + if ( + current.size !== contextAtCapture.size || + [...current].some( + ([key, value]) => !Object.is(contextAtCapture.get(key), value), + ) + ) { + runtime.binding.markIncomplete( + "Request context changed after replay capture.", + ); + if (historical) + throw new Error( + "Unsupported Mastra memory replay: request context changed after capture.", + ); + } capture.beginStep({ stepNumber: args.stepNumber, messageList: args.messageList, @@ -557,7 +811,12 @@ export function createMemoryReplayAgent( onError: callerOptions.onError, onAbort: callerOptions.onAbort, onStepFinish: callerOptions.onStepFinish, - requestContext, + // An empty adapter-created context changes Mastra's native stream lifecycle. + ...(historical || + callerOptions.requestContext || + [...requestContext.entries()].length > 0 + ? { requestContext } + : {}), abortSignal: callerOptions.abortSignal ? AbortSignal.any([callerOptions.abortSignal, abort.signal]) : abort.signal, @@ -565,6 +824,7 @@ export function createMemoryReplayAgent( const version = createRequire(import.meta.url)("../package.json") as { version: string; }; + markRecordingStreamEntered(); return await streamWithRecording({ adapterVersion: version.version, agent: agent as unknown as Parameters< @@ -572,16 +832,27 @@ export function createMemoryReplayAgent( >[0]["agent"], callerMessages: invocationInput, callerOptions: runtimeOptions, - client, + client: evidenceClient, options, replayInput: replay.effectiveInput, replay, + nativeFallback: async (error) => { + try { + await runtime.finish(); + await runtime.release(); + } catch (cleanupError) { + reportLocalRecordingError(cleanupError); + } + reportLocalRecordingError(error); + return runNativeBaseline(rawInput, callerOptions); + }, requestedModelId: replay.replacementModelId ?? String(configuration.modelId), sessionName: options.sessionName, startedAt, stateful: { input: { [MEMORY_REPLAY_KEY]: envelope }, + sanitizeEvidence: sanitizer.replace, initialize(value) { state = value; }, @@ -596,6 +867,21 @@ export function createMemoryReplayAgent( async finish() { await runtime.finish(); await capture.drain(); + if (!historical) await runtime.binding.verifyEligibility(); + const omResults = await omTape.finish(); + for (const reason of omCaptureErrors) + runtime.binding.markIncomplete(reason); + if (omMismatches.length) + await writeNode({ + external_id: `${invocationId}:om-input-mismatch`, + parent_external_id: ROOT_NODE_EXTERNAL_ID, + node_type: "span", + name: "om_input_mismatch", + status: "completed", + inputs: null, + outputs: null, + attributes: { count: omMismatches.length }, + }); for (const pending of capture.flushUnfinished()) await writeAttempt( pending, @@ -603,20 +889,42 @@ export function createMemoryReplayAgent( ); if (runtime.binding.incompleteReasons.length) { const message = runtime.binding.incompleteReasons.join(" "); - if ( - runtime.binding.incompleteReasons.includes( - "Native memory storage mutation failed.", - ) + const reasonCode = runtime.binding.incompleteReasons.includes( + "Native memory storage mutation failed.", ) - throw new Error(message); - throw new StatefulRecordingError(message); + ? "memory_mutation_failed" + : runtime.binding.incompleteReasons.includes( + "Request context changed after replay capture.", + ) + ? "context_mutated_after_capture" + : "memory_evidence_incomplete"; + throw new StatefulRecordingError(message, reasonCode); } + if (!envelope.complete) + throw new StatefulRecordingError( + envelope.reasons.join(" "), + "capture_prerequisite_failed", + ); + return { + [MEMORY_REPLAY_KEY]: finalizeMemoryReplayEnvelope( + envelope, + omResults.map((entry) => ({ + phase: entry.phase, + ordinal: entry.ordinal, + method: entry.method, + inputFingerprint: entry.inputFingerprint, + output: entry.output, + })), + ), + }; }, + release: () => runtime.release(), }, }); } catch (error) { try { await runtime.finish(); + await runtime.release(); } catch (cleanupError) { if (options.onRecordingError) { void Promise.resolve() @@ -637,5 +945,34 @@ export function createMemoryReplayAgent( throw error; } } + async function stream( + rawInput: unknown, + callerOptions: RuntimeStreamOptions = {}, + ): Promise { + let enteredRecordingStream = false; + try { + return await recordedStream(rawInput, callerOptions, () => { + enteredRecordingStream = true; + }); + } catch (error) { + if ( + enteredRecordingStream || + process.env.KITARU_REPLAY_ID || + process.env.KITARU_TASK_INPUTS || + process.env.KITARU_OVERRIDE + ) + throw error; + const reasonCode = + error instanceof Error && + /requires @mastra\/(?:core|memory)@/.test(error.message) + ? "version_mismatch" + : error instanceof MemoryReplayContextError + ? "context_unsupported" + : "capture_setup_failed"; + const nativeResult = await runNativeBaseline(rawInput, callerOptions); + void reportSetupFailure(error, reasonCode); + return nativeResult; + } + } return { stream: stream as Agent["stream"] }; } diff --git a/packages/mastra/src/stateful-files.ts b/packages/mastra/src/stateful-files.ts index 176cd5f1b..24c343648 100644 --- a/packages/mastra/src/stateful-files.ts +++ b/packages/mastra/src/stateful-files.ts @@ -1,3 +1,4 @@ +import { createHash } from "node:crypto"; import type { MastraRecordedFile } from "./memory-snapshot.js"; export interface ResolvedMemoryFile { @@ -7,47 +8,306 @@ export interface ResolvedMemoryFile { export type MemoryFileResolver = (url: string) => Promise; -/** Bind processors to the invocation's declared content, with no live fallback. */ +export interface RecordedEvidenceSanitizer { + replace(value: T): T; +} + +export type UnsafeEvidenceReason = "credential_url" | "unsupported_value"; + +const FILE_REFERENCE = /^kitaru-file:\/\/sha256\/[a-f0-9]{64}$/; +const MAX_FILE_BYTES = 8 * 1024 * 1024; +const MAX_TOTAL_FILE_BYTES = 16 * 1024 * 1024; +const MAX_RECORDED_FILES = 64; +const URL_IN_TEXT = /https?:\/\/[^\s"'<>]+/gi; +const CREDENTIAL_QUERY_KEY = + /^(?:token|access[_-]?token|id[_-]?token|refresh[_-]?token|api[_-]?key|signature|sig|secret|auth|authorization|credential|x-amz-[\w-]+|x-goog-[\w-]+)$/i; + +function hasUrlCredential(url: string): boolean { + try { + const parsed = new URL(url); + return ( + Boolean(parsed.username || parsed.password) || + [...parsed.searchParams.keys()].some((key) => + CREDENTIAL_QUERY_KEY.test(key), + ) + ); + } catch { + return /[?&](?:token|signature|sig|secret|auth|credential)=/i.test(url); + } +} + +export function fileReference(file: ResolvedMemoryFile): string { + const digest = createHash("sha256") + .update(file.mediaType) + .update("\0") + .update(file.bytes) + .digest("hex"); + return `kitaru-file://sha256/${digest}`; +} + +function copiedFile(file: MastraRecordedFile): MastraRecordedFile { + if ( + !FILE_REFERENCE.test(file.url) || + !(file.bytes instanceof Uint8Array) || + file.bytes.byteLength > MAX_FILE_BYTES || + typeof file.mediaType !== "string" || + !file.mediaType || + fileReference(file) !== file.url + ) + throw new Error("Unsupported Mastra memory replay: invalid recorded file."); + return { ...file, bytes: new Uint8Array(file.bytes) }; +} + +/** Replay resolves only captured content references, never network URLs. */ export function restoreCapturedFiles(recorded: readonly MastraRecordedFile[]) { - const files = recorded.map((file) => ({ - ...file, - bytes: new Uint8Array(file.bytes), - })); + if ( + recorded.length > MAX_RECORDED_FILES || + recorded.reduce( + (size, file) => + size + (file.bytes instanceof Uint8Array ? file.bytes.byteLength : 0), + 0, + ) > MAX_TOTAL_FILE_BYTES + ) + throw new Error( + "Unsupported Mastra memory replay: file capture limit exceeded.", + ); + const files = recorded.map(copiedFile); const lookup = new Map(files.map((file) => [file.url, file])); if (lookup.size !== files.length) - throw new Error("Duplicate recorded file URL"); + throw new Error( + "Unsupported Mastra memory replay: duplicate recorded file.", + ); return { files, - resolveFile: async (url: string): Promise => { - const file = lookup.get(url); + resolveFile: async (reference: string): Promise => { + const file = lookup.get(reference); if (!file) throw new Error( - "Unsupported Mastra memory replay: file URL was not recorded.", + "Unsupported Mastra memory replay: file reference was not recorded.", ); return { bytes: new Uint8Array(file.bytes), mediaType: file.mediaType }; }, }; } -/** Fetch each declared file before recording the immutable session input. */ +/** Sanitize every persisted evidence field, including URLs inside prompt text. */ +export function createRecordedEvidenceSanitizer( + references: ReadonlyMap, + onUnsafeEvidence: (reason: UnsafeEvidenceReason) => void, +): RecordedEvidenceSanitizer { + const declared = [...references].sort( + (left, right) => right[0].length - left[0].length, + ); + function replaceString(value: string): string { + let replaced = value; + for (const [url, reference] of declared) { + replaced = replaced.split(url).join(reference); + } + return replaced.replace(URL_IN_TEXT, (url) => { + if (!hasUrlCredential(url)) return url; + onUnsafeEvidence("credential_url"); + return "[redacted credential URL]"; + }); + } + function replace(value: T): T { + const active = new Set(); + function visit(current: unknown): unknown { + if (typeof current === "string") return replaceString(current); + if (current instanceof URL) return replaceString(current.href); + if (current instanceof Date || current instanceof Uint8Array) + return current; + if ( + current instanceof ArrayBuffer || + typeof current === "function" || + typeof current === "symbol" + ) { + onUnsafeEvidence("unsupported_value"); + return "[unrecordable evidence value]"; + } + if (current === null || typeof current !== "object") return current; + if (active.has(current)) { + onUnsafeEvidence("unsupported_value"); + return "[unrecordable circular evidence]"; + } + active.add(current); + try { + if (Array.isArray(current)) { + if ( + Object.values(Object.getOwnPropertyDescriptors(current)).some( + (descriptor) => descriptor.enumerable && !("value" in descriptor), + ) + ) { + onUnsafeEvidence("unsupported_value"); + return "[unrecordable accessor evidence]"; + } + return current.map(visit); + } + if ( + (Object.getPrototypeOf(current) !== Object.prototype && + Object.getPrototypeOf(current) !== null) || + Reflect.ownKeys(current).some((key) => typeof key !== "string") + ) { + onUnsafeEvidence("unsupported_value"); + return "[unrecordable evidence value]"; + } + const descriptors = Object.getOwnPropertyDescriptors(current); + if ( + Object.values(descriptors).some( + (descriptor) => descriptor.enumerable && !("value" in descriptor), + ) + ) { + onUnsafeEvidence("unsupported_value"); + return "[unrecordable accessor evidence]"; + } + return Object.fromEntries( + Object.entries(current).map(([key, item]) => [ + replaceString(key), + visit(item), + ]), + ); + } finally { + active.delete(current); + } + } + return visit(value) as T; + } + return { replace }; +} + +/** Capture declared URLs once and convert persisted input to secret-free references. */ export async function createCapturedFiles( urls: readonly string[], resolveFile: MemoryFileResolver, ) { - const files: MastraRecordedFile[] = []; - for (const url of new Set(urls)) { - const resolved = await resolveFile(url); + const originalToReference = new Map(); + const filesByReference = new Map(); + const uniqueUrls = new Set(urls); + if (uniqueUrls.size > MAX_RECORDED_FILES) + throw new Error( + "Unsupported Mastra memory replay: file count limit exceeded.", + ); + let totalBytes = 0; + for (const url of uniqueUrls) { + let resolved: ResolvedMemoryFile; + try { + resolved = await resolveFile(url); + } catch { + // Resolver errors can contain a signed URL. Keep them out of diagnostics. + throw new Error("Controlled file capture failed."); + } if ( !(resolved.bytes instanceof Uint8Array) || typeof resolved.mediaType !== "string" || !resolved.mediaType ) throw new TypeError("File resolver must return bytes and mediaType"); - files.push({ - url, + if (resolved.bytes.byteLength > MAX_FILE_BYTES) + throw new Error("Unsupported Mastra memory replay: file exceeds 8 MiB."); + totalBytes += resolved.bytes.byteLength; + if (totalBytes > MAX_TOTAL_FILE_BYTES) + throw new Error("Unsupported Mastra memory replay: files exceed 16 MiB."); + const file = { + url: fileReference(resolved), bytes: new Uint8Array(resolved.bytes), mediaType: resolved.mediaType, - }); + }; + originalToReference.set(url, file.url); + filesByReference.set(file.url, file); + } + const captured = restoreCapturedFiles([...filesByReference.values()]); + const declared = [...originalToReference].sort( + (left, right) => right[0].length - left[0].length, + ); + function referenceFor(url: string): string { + const reference = originalToReference.get(url); + if (!reference) + throw new Error("Unsupported Mastra memory replay: undeclared file URL."); + return reference; } - return restoreCapturedFiles(files); + function replaceDeclaredFileUrls(value: T): T { + const active = new Set(); + function visit(current: unknown, filePart: boolean): unknown { + if (typeof current === "string") { + const reference = originalToReference.get(current); + if (reference) return reference; + if (filePart && /^https?:\/\//i.test(current)) + throw new Error( + "Unsupported Mastra memory replay: undeclared file URL.", + ); + let replaced = current; + for (const [url, fileReference] of declared) + replaced = replaced.split(url).join(fileReference); + return replaced; + } + if (current instanceof URL) { + const reference = originalToReference.get(current.href); + if (reference) return new URL(reference); + if (filePart && /^https?:$/i.test(current.protocol)) + throw new Error( + "Unsupported Mastra memory replay: undeclared file URL.", + ); + return current; + } + if ( + current === null || + typeof current !== "object" || + current instanceof Date || + current instanceof Uint8Array + ) + return current; + if (active.has(current)) + throw new Error( + "Unsupported Mastra memory replay: circular file input.", + ); + active.add(current); + try { + if (Array.isArray(current)) { + if ( + Reflect.ownKeys(current).some((key) => { + const descriptor = Object.getOwnPropertyDescriptor(current, key); + return descriptor && !("value" in descriptor); + }) + ) + throw new Error( + "Unsupported Mastra memory replay: file input accessor.", + ); + return current.map((item) => visit(item, filePart)); + } + if ( + (Object.getPrototypeOf(current) !== Object.prototype && + Object.getPrototypeOf(current) !== null) || + Reflect.ownKeys(current).some((key) => typeof key !== "string") + ) + throw new Error("Unsupported Mastra memory replay: file input type."); + const descriptors = Object.getOwnPropertyDescriptors(current); + const part = + filePart || + descriptors.type?.value === "file" || + descriptors.type?.value === "image"; + const entries = Object.entries(descriptors); + return Object.fromEntries( + entries.map(([key, descriptor]) => { + if (!descriptor.enumerable || !("value" in descriptor)) + throw new Error( + "Unsupported Mastra memory replay: file input accessor.", + ); + return [key, visit(descriptor.value, part)]; + }), + ); + } finally { + active.delete(current); + } + } + return visit(value, false) as T; + } + return { + files: captured.files, + referenceFor, + evidenceSanitizer: ( + onUnsafeEvidence: (reason: UnsafeEvidenceReason) => void, + ) => createRecordedEvidenceSanitizer(originalToReference, onUnsafeEvidence), + replaceDeclaredFileUrls, + resolveFile: async (url: string): Promise => + captured.resolveFile(referenceFor(url)), + }; } diff --git a/packages/mastra/src/stateful-tools.ts b/packages/mastra/src/stateful-tools.ts index 7fe71b783..19a4a52c7 100644 --- a/packages/mastra/src/stateful-tools.ts +++ b/packages/mastra/src/stateful-tools.ts @@ -46,6 +46,7 @@ export function createStatefulToolProcessors(options: { getState(): AdapterRunState; abort(reason: unknown): void; adapter: KitaruAgentOptions; + sanitizeEvidence?: (value: T) => T; }) { const trusted = new WeakSet<(...args: never[]) => unknown>(); const wrappersByName = new Map< @@ -55,12 +56,9 @@ export function createStatefulToolProcessors(options: { (...args: unknown[]) => Promise > >(); - let inspected = false; const first: InputProcessor = { id: "kitaru-memory-tool-identity", processInputStep({ tools }) { - if (inspected) return; - inspected = true; for (const tool of Object.values(tools ?? {})) { if ( record(tool) && @@ -105,6 +103,7 @@ export function createStatefulToolProcessors(options: { configuredBeforeToolCall: options.adapter.configuredBeforeToolCall, configuredAfterToolCall: options.adapter.configuredAfterToolCall, limits: options.adapter.recordingLimits, + sanitizeEvidence: options.sanitizeEvidence, }); wrapper = async (input: unknown, context: unknown) => { const event = { toolName: name, input, context, metadata: {} }; diff --git a/packages/mastra/src/stateful-workspace.ts b/packages/mastra/src/stateful-workspace.ts index 7940941e6..e4ed2bc43 100644 --- a/packages/mastra/src/stateful-workspace.ts +++ b/packages/mastra/src/stateful-workspace.ts @@ -1,11 +1,13 @@ import { createHash } from "node:crypto"; -import { lstat, readdir, readFile } from "node:fs/promises"; +import { constants } from "node:fs"; +import { lstat, open, readdir } from "node:fs/promises"; import { join, posix } from "node:path"; import type { SkillSource, SkillSourceEntry } from "@mastra/core/workspace"; import { MAX_RECORDED_PAYLOAD_CHARS } from "@zenml-io/kitaru/adapter"; export interface SkillsManifest { files: { path: string; length: number; sha256: string }[]; + directories: string[]; } /** Read and pin an artifact's skills. Native tools subsequently read only these bytes. */ @@ -24,6 +26,7 @@ export async function loadSkillsWorkspace( throw new Error("Unsupported Mastra skills symlink"); if (info.isDirectory()) { const entries = await readdir(path, { withFileTypes: true }); + entries.sort((a, b) => (a.name < b.name ? -1 : a.name > b.name ? 1 : 0)); directories.set( relative, entries.map((entry) => ({ @@ -31,14 +34,24 @@ export async function loadSkillsWorkspace( type: entry.isDirectory() ? "directory" : "file", })), ); - for (const entry of entries.sort((a, b) => - a.name < b.name ? -1 : a.name > b.name ? 1 : 0, - )) + for (const entry of entries) await visit(posix.join(relative, entry.name)); } else if (info.isFile()) { if (size + info.size > MAX_RECORDED_PAYLOAD_CHARS) throw new Error("Skills content exceeds the supported replay limit"); - const content = await readFile(path); + // O_NOFOLLOW closes the gap between lstat and the read if a file is swapped. + const handle = await open( + path, + constants.O_RDONLY | constants.O_NOFOLLOW, + ); + let content: Buffer; + try { + if (!(await handle.stat()).isFile()) + throw new Error("Unsupported Mastra skills file type"); + content = await handle.readFile(); + } finally { + await handle.close(); + } size += content.length; files.set(relative, content); } else throw new Error("Unsupported Mastra skills file type"); @@ -55,6 +68,7 @@ export async function loadSkillsWorkspace( length: content.length, sha256: createHash("sha256").update(content).digest("hex"), })), + directories: [...directories.keys()].sort(), }; if ( expectedManifest && diff --git a/packages/mastra/src/step-recorder.ts b/packages/mastra/src/step-recorder.ts index 0b20fe14f..cf216f0ac 100644 --- a/packages/mastra/src/step-recorder.ts +++ b/packages/mastra/src/step-recorder.ts @@ -123,11 +123,15 @@ function usageTokens(usage: unknown): SessionNodeCreateRequest["tokens"] { function stepOutputs( step: RecordedStep, tools: readonly NormalizedToolCall[], + sanitizeEvidence?: (value: T) => T, ): JsonValue { return boundedRecorderJson( { finish_reason: step.finishReason ?? null, - text: boundedRecorderJson(step.text, "model step text"), + text: boundedRecorderJson( + sanitizeEvidence?.(step.text) ?? step.text, + "model step text", + ), tool_calls: tools.map((tool) => ({ args: tool.inputs, toolCallId: tool.callId, @@ -147,7 +151,10 @@ function stepOutputs( ], ), tripwire: projectRecordedMetadata(step.tripwire), - warnings: boundedRecorderJson(step.warnings, "model warnings"), + warnings: boundedRecorderJson( + sanitizeEvidence?.(step.warnings) ?? step.warnings, + "model warnings", + ), }, "model step output", ); @@ -178,6 +185,7 @@ export async function recordStep( costCalculator?: KitaruCostCalculator, limits?: RecordingLimits, requestEvidence?: RequestEvidence, + sanitizeEvidence?: (value: T) => T, ): Promise { const calls = step.toolCalls.flatMap((item) => { const call = toolCallPayload(item); @@ -192,13 +200,13 @@ export async function recordStep( const tools: NormalizedToolCall[] = calls.map((call) => { const result = results.get(call.toolCallId); const inputs = boundedRecorderConversion( - call.args, + sanitizeEvidence?.(call.args) ?? call.args, `tool '${call.toolName}' input`, limits, ); const recordedResult = result ? boundedRecorderConversion( - result.result, + sanitizeEvidence?.(result.result) ?? result.result, `tool '${call.toolName}' output`, limits, ) @@ -257,7 +265,7 @@ export async function recordStep( model: servedModelId, modelSettings: requestEvidence?.modelSettings, startedAt: requestEvidence?.startedAt, - outputs: stepOutputs(step, tools), + outputs: stepOutputs(step, tools, sanitizeEvidence), provider, tokens, tools, diff --git a/packages/mastra/src/stream-recording.ts b/packages/mastra/src/stream-recording.ts index 0b3e9291f..c6e829f43 100644 --- a/packages/mastra/src/stream-recording.ts +++ b/packages/mastra/src/stream-recording.ts @@ -1,6 +1,7 @@ import { createRequire } from "node:module"; -import type { JsonValue, KitaruClient } from "@zenml-io/kitaru"; +import type { JsonValue } from "@zenml-io/kitaru"; import { + type AdapterClient, type AdapterRunState, parseModelSettings, type ReplayContext, @@ -18,6 +19,7 @@ import { restoreConversationContext, unsupportedContext, } from "./conversation-context.js"; +import { MastraOMDivergenceError } from "./om-result-tape.js"; import { assertReplayToolCoverage, stripLiveMemoryOptions, @@ -39,13 +41,22 @@ type StreamAgent = { listConfiguredInputProcessors?: (requestContext?: unknown) => unknown; }; -export class StatefulRecordingError extends Error {} +export class StatefulRecordingError extends Error { + constructor( + message: string, + readonly reasonCode = "capture_incomplete", + ) { + super(message); + } +} export interface StatefulStreamRecording { input: JsonValue; + sanitizeEvidence?: (value: T) => T; initialize(state: AdapterRunState): void; takeRequest(): RequestEvidence | undefined; - finish(): Promise; + finish(): Promise; + release(): Promise; } interface StreamRecordingOptions { @@ -54,13 +65,15 @@ interface StreamRecordingOptions { agent: StreamAgent; callerMessages: unknown; callerOptions: RuntimeStreamOptions; - client: KitaruClient; + client: AdapterClient; options: KitaruAgentOptions; replayInput: JsonValue; requestedModelId: string; sessionName?: string; startedAt: string; replay: ReplayContext; + nativeFallback?: (error: unknown) => Promise; + markNativeStart?: () => void; } const ERROR_STEP_GRACE_MS = 250; @@ -178,6 +191,15 @@ function getTripwireReason(value: unknown): string | undefined { } function getSafeStreamError(error: unknown): Error { + if ( + error instanceof MastraOMDivergenceError || + (error instanceof Error && + (error.message.includes( + "Recorded Mastra observational memory diverged:", + ) || + error.message === "KITARU_REPLAY_DIVERGED:mastra_om_call_order")) + ) + return new Error("KITARU_REPLAY_DIVERGED:mastra_om_call_order"); const name = error instanceof Error && error.name.length <= MAX_STREAM_ERROR_NAME_LENGTH && @@ -202,7 +224,9 @@ class StreamLifecycle { #recordingError?: { error: unknown; stage: StreamRecordingErrorStage; + reasonCode: string; }; + #nativeState: "pending" | "completed" | "failed" = "pending"; #stepTail: Promise = Promise.resolve(); constructor( @@ -213,7 +237,6 @@ class StreamLifecycle { async record(step: RecordedStep): Promise { if (this.#recordingError !== undefined) return; - let writeFailure: { error: unknown } | undefined; const write = this.#stepTail.then(() => recordStep( this.recorder.state, @@ -221,16 +244,13 @@ class StreamLifecycle { this.options.costCalculator, this.options.recordingLimits, this.stateful?.takeRequest(), + this.stateful?.sanitizeEvidence, ), ); this.#stepTail = write.catch((error: unknown) => { - writeFailure = { error }; - this.requestRecordingFailure("step", error); + this.requestRecordingFailure("step", error, "recording_step_failed"); }); await this.#stepTail; - if (writeFailure !== undefined) { - await this.cleanup(writeFailure.error, "recording"); - } } async complete(result: unknown): Promise { @@ -238,6 +258,18 @@ class StreamLifecycle { if (this.recorder.state.failure !== undefined) { this.requestFailure(this.recorder.state.failure); } + if (this.stateful && !this.recorder.state.spec) { + // Mastra can start observational work after the actor finishes. Keep its + // recorder and source lease alive without delaying the native stream. + void this.finalize(true, result).catch((error: unknown) => { + this.requestRecordingFailure( + "complete", + error, + "recording_finalization_failed", + ); + }); + return; + } await this.finalize(true, result); } @@ -251,6 +283,14 @@ class StreamLifecycle { await this.finalize(false); } + markNativeCompleted(): void { + if (this.#nativeState === "pending") this.#nativeState = "completed"; + } + + markNativeFailed(): void { + this.#nativeState = "failed"; + } + deferFailure(error: unknown): void { if (this.#completionStarted || this.#deferredFailure !== undefined) return; // Mastra normally follows onError with a failed step. Error-only paths, @@ -271,9 +311,10 @@ class StreamLifecycle { private requestRecordingFailure( stage: StreamRecordingErrorStage, error: unknown, + reasonCode: string, ): void { if (this.#recordingError !== undefined) return; - this.#recordingError = { error, stage }; + this.#recordingError = { error, stage, reasonCode }; this.notify(stage, error); } @@ -286,48 +327,108 @@ class StreamLifecycle { private async finalize(complete: boolean, result?: unknown): Promise { this.#finalizerPromise ??= (async () => { - await this.#stepTail; try { - await this.stateful?.finish(); - } catch (error) { - if (error instanceof StatefulRecordingError) - this.requestRecordingFailure("complete", error); - else this.requestFailure(error); - } - if (this.#recordingError !== undefined) { - await this.cleanup(this.#recordingError.error, "recording"); - return; - } - if (!this.#failureRequested && complete) { - // The API cannot reopen a terminal session. Choose completion once all - // queued steps settle; later aborts cannot reverse this terminal write. - this.#completionStarted = true; - let completionFailure: { error: unknown } | undefined; + await this.#stepTail; + let finalInput: JsonValue | undefined; try { - await this.recorder.complete(result); + finalInput = await this.stateful?.finish(); } catch (error) { - completionFailure = { error }; - this.requestRecordingFailure("complete", error); + if ( + this.recorder.state.spec && + error instanceof MastraOMDivergenceError + ) + this.requestFailure(getSafeStreamError(error)); + else + this.requestRecordingFailure( + "complete", + error, + error instanceof StatefulRecordingError + ? error.reasonCode + : "recording_finalization_failed", + ); } - if (completionFailure !== undefined) { - await this.cleanup(completionFailure.error, "recording"); + if (this.#failureRequested) { + await this.cleanup(this.#failureReason, "run"); + return; } - } - if (this.#failureRequested) { - await this.cleanup(this.#failureReason, "run"); + if (this.#recordingError !== undefined) { + await this.cleanup(this.#recordingError.error, "recording"); + return; + } + if (!this.#failureRequested && complete) { + this.#completionStarted = true; + try { + await this.recorder.complete( + result, + this.stateful && !this.recorder.state.spec + ? { + inputs: finalInput, + metadata: { + mastra_replay_state: "eligible", + mastra_native_state: "completed", + }, + } + : undefined, + ); + } catch (error) { + this.requestRecordingFailure( + "complete", + error, + "recording_completion_failed", + ); + await this.cleanup(error, "recording"); + } + } + if (this.#failureRequested) + await this.cleanup(this.#failureReason, "run"); + } finally { + await this.stateful + ?.release() + .catch((error: unknown) => this.notify("complete", error)); } })(); await this.#finalizerPromise; } private cleanup(error: unknown, kind: "recording" | "run"): Promise { - this.#cleanupPromise ??= ( - kind === "recording" - ? this.recorder.failRecording( - new Error("Kitaru stream recording failed"), - ) - : this.recorder.fail(error) - ).catch(() => undefined); + this.#cleanupPromise ??= (async () => { + const reasonCode = this.#recordingError?.reasonCode; + const safeError = getSafeStreamError(error); + const omDiverged = + safeError.message === "KITARU_REPLAY_DIVERGED:mastra_om_call_order"; + const metadata: Record | undefined = + this.stateful && !this.recorder.state.spec + ? { + mastra_replay_state: "ineligible", + mastra_replay_reason: + reasonCode ?? + (kind === "run" ? "native_run_failed" : "recording_failed"), + mastra_native_state: this.#nativeState, + } + : this.stateful && this.recorder.state.spec + ? { + mastra_replay_state: omDiverged ? "diverged" : "failed", + mastra_replay_reason: omDiverged + ? "mastra_om_call_order" + : "replay_failed", + } + : undefined; + if (kind === "recording") { + await this.recorder.failRecording( + new Error(`KITARU_RECORDING_INCOMPLETE:${reasonCode ?? "unknown"}`), + metadata, + ); + } else { + await this.recorder.fail( + reasonCode + ? new Error( + `${safeError.message}; KITARU_RECORDING_INCOMPLETE:${reasonCode}`, + ) + : error, + metadata, + ); + } + })().catch(() => undefined); return this.#cleanupPromise; } @@ -351,7 +452,7 @@ class StreamLifecycle { } } -export async function streamWithRecording({ +async function recordedStreamWithRecording({ stateful, adapterVersion, agent, @@ -364,6 +465,7 @@ export async function streamWithRecording({ sessionName, startedAt, replay, + markNativeStart, }: StreamRecordingOptions): Promise { assertStreamSupported(agent); const resolvedDefaults = @@ -489,6 +591,14 @@ export async function streamWithRecording({ effectiveInput: recordedInput, effectiveModelSettings: serializedSettings(effective.modelSettings), framework: "mastra", + ...(stateful && !replay.spec + ? { + metadata: { + mastra_replay_state: "pending", + mastra_native_state: "pending", + }, + } + : {}), name: sessionName, replayId: replay.replayId, requestedModelId, @@ -499,7 +609,18 @@ export async function streamWithRecording({ try { await recorder.initialize(); } catch (error) { - await recorder.fail(error).catch(() => undefined); + await recorder + .failRecording( + new Error("KITARU_RECORDING_INCOMPLETE:capture_setup_failed"), + stateful && !replay.spec + ? { + mastra_replay_state: "ineligible", + mastra_replay_reason: "capture_setup_failed", + mastra_native_state: "pending", + } + : undefined, + ) + .catch(() => undefined); throw error; } stateful?.initialize(recorder.state); @@ -536,7 +657,10 @@ export async function streamWithRecording({ let modelError: unknown; effective.onStepFinish = async (step) => { const active = await initialize(); - if (step.finishReason === "error") active.cancelDeferredFailure(); + if (step.finishReason === "error") { + active.cancelDeferredFailure(); + active.markNativeFailed(); + } const pendingModelError = step.finishReason === "error" ? modelError : undefined; const recordedStep = @@ -569,18 +693,21 @@ export async function streamWithRecording({ try { await callerFinish?.(event); } catch (error) { + active.markNativeFailed(); await active.fail(getSafeStreamError(error)); throw error; } + active.markNativeCompleted(); const tripwire = getTripwireReason(event); if (tripwire) { await active.fail(new Error(tripwire)); return; } - const summary = runResultSummary(event, { + const rawSummary = runResultSummary(event, { structuredOutputField: effective.structuredOutput === undefined ? undefined : "object", }); + const summary = stateful?.sanitizeEvidence?.(rawSummary) ?? rawSummary; if ( effective.structuredOutput !== undefined && isRecord(summary) && @@ -601,6 +728,7 @@ export async function streamWithRecording({ modelError ??= event.error; const active = lifecycle ?? (await initializePromise?.catch(() => undefined)); + active?.markNativeFailed(); active?.deferFailure( active.recorder.state.failure ?? getSafeStreamError(modelError), ); @@ -608,6 +736,7 @@ export async function streamWithRecording({ }; effective.onAbort = async (event) => { const active = await initialize(); + active.markNativeFailed(); await active.fail( active.recorder.state.failure ?? new Error("Mastra stream aborted"), ); @@ -631,10 +760,12 @@ export async function streamWithRecording({ }; try { + markNativeStart?.(); return await agent.stream(effectiveMessages, effective); } catch (error) { const active = lifecycle ?? (await initializePromise?.catch(() => undefined)); + active?.markNativeFailed(); const replayFailure = replay.spec ? active?.recorder.state.failure : undefined; @@ -642,3 +773,21 @@ export async function streamWithRecording({ throw replayFailure ?? error; } } + +export async function streamWithRecording( + options: StreamRecordingOptions, +): Promise { + let nativeStarted = false; + try { + return await recordedStreamWithRecording({ + ...options, + markNativeStart: () => { + nativeStarted = true; + }, + }); + } catch (error) { + if (!nativeStarted && !options.replay.spec && options.nativeFallback) + return options.nativeFallback(error); + throw error; + } +} diff --git a/packages/mastra/src/tool-policies.ts b/packages/mastra/src/tool-policies.ts index 7db6e7f59..da555ffa7 100644 --- a/packages/mastra/src/tool-policies.ts +++ b/packages/mastra/src/tool-policies.ts @@ -27,6 +27,7 @@ interface ToolHookOptions { configuredAfterToolCall?: ConfiguredAfterToolCall; configuredBeforeToolCall?: ConfiguredBeforeToolCall; limits?: RecordingLimits; + sanitizeEvidence?: (value: T) => T; state: AdapterRunState; } @@ -65,16 +66,27 @@ async function invokePassthroughBeforeHooks( configuredHook?: ConfiguredBeforeToolCall, callerHook?: ToolHooks["beforeToolCall"], limits?: RecordingLimits, + sanitizeEvidence?: (value: T) => T, ): Promise> { try { const configuredResult = await configuredHook?.(hookContext); if (isSkippedResult(configuredResult)) { - completeToolCall(state, callId, configuredResult.output, limits); + completeToolCall( + state, + callId, + sanitizeEvidence?.(configuredResult.output) ?? configuredResult.output, + limits, + ); return configuredResult; } const callerResult = await callerHook?.(hookContext); if (isSkippedResult(callerResult)) { - completeToolCall(state, callId, callerResult.output, limits); + completeToolCall( + state, + callId, + sanitizeEvidence?.(callerResult.output) ?? callerResult.output, + limits, + ); return callerResult; } return undefined; @@ -91,6 +103,7 @@ export function createToolHooks(options: ToolHookOptions): ToolHooks { configuredAfterToolCall, configuredBeforeToolCall, limits, + sanitizeEvidence, state, } = options; @@ -104,7 +117,7 @@ export function createToolHooks(options: ToolHookOptions): ToolHooks { } const callId = toolCallId(hookContext.context); const converted = boundedRecorderConversion( - hookContext.input, + sanitizeEvidence?.(hookContext.input) ?? hookContext.input, `tool '${hookContext.toolName}' input`, limits, ); @@ -113,7 +126,8 @@ export function createToolHooks(options: ToolHookOptions): ToolHooks { callId, inputs: converted.value, inputsLossy: converted.lossy, - originalInputs: hookContext.input, + originalInputs: + sanitizeEvidence?.(hookContext.input) ?? hookContext.input, toolName: hookContext.toolName, }); if (decision.type !== "execute") { @@ -137,6 +151,7 @@ export function createToolHooks(options: ToolHookOptions): ToolHooks { configuredBeforeToolCall, callerHooks?.beforeToolCall, limits, + sanitizeEvidence, ); } catch (error) { abortReplay?.(error); @@ -152,7 +167,12 @@ export function createToolHooks(options: ToolHookOptions): ToolHooks { if (hookContext.error !== undefined) { failToolCall(state, callId, hookContext.error); } else { - completeToolCall(state, callId, hookContext.output, limits); + completeToolCall( + state, + callId, + sanitizeEvidence?.(hookContext.output) ?? hookContext.output, + limits, + ); } try { diff --git a/packages/mastra/test/helpers/file-memory-access.ts b/packages/mastra/test/helpers/file-memory-access.ts new file mode 100644 index 000000000..f482ce3dd --- /dev/null +++ b/packages/mastra/test/helpers/file-memory-access.ts @@ -0,0 +1,213 @@ +import { createHash, randomUUID } from "node:crypto"; +import { + mkdir, + open, + readdir, + readFile, + rm, + writeFile, +} from "node:fs/promises"; +import { join } from "node:path"; +import type { + MastraExclusiveMemoryAccess, + MastraMemoryLease, + MastraMemoryLeaseOptions, + MastraMemorySelector, +} from "../../src/memory-binding.js"; + +const pause = (ms: number) => + new Promise((resolve) => setTimeout(resolve, ms)); + +function isExists(error: unknown): boolean { + return ( + typeof error === "object" && + error !== null && + "code" in error && + error.code === "EEXIST" + ); +} + +/** Test-only shared lease: atomic directory creation and persistent poison. */ +export function createFileMemoryAccess( + root: string, +): MastraExclusiveMemoryAccess & { + simulateLeaseLoss(selector: MastraMemorySelector): Promise; +} { + function directory(selector: MastraMemorySelector): string { + return join( + root, + createHash("sha256").update(selector.threadId).digest("hex"), + ); + } + const globalPoison = join(root, "unknown-writer-poison"); + + async function withMutex( + selector: MastraMemorySelector, + run: (dir: string) => Promise, + ): Promise { + const dir = directory(selector); + await mkdir(dir, { recursive: true }); + const mutex = join(dir, "mutex"); + const deadline = Date.now() + 1000; + while (true) { + try { + await mkdir(mutex); + break; + } catch (error) { + if (!isExists(error) || Date.now() >= deadline) throw error; + await pause(2); + } + } + try { + return await run(dir); + } finally { + await rm(mutex, { recursive: true, force: true }); + } + } + + async function exists(path: string): Promise { + try { + await readFile(path); + return true; + } catch (error) { + if ( + typeof error === "object" && + error !== null && + "code" in error && + error.code === "ENOENT" + ) + return false; + throw error; + } + } + + async function poison(dir: string, persistent: boolean): Promise { + const marker = join(dir, persistent ? "persistent-loss" : "poison"); + try { + const file = await open(marker, "wx"); + await file.close(); + } catch (error) { + if (!isExists(error)) throw error; + } + if (persistent) await poison(dir, false); + } + + async function activeTurns(dir: string): Promise { + const names = await readdir(dir); + return names.filter((name) => name.startsWith("turn-")); + } + + async function makeLease( + selector: MastraMemorySelector, + token: string, + owns: boolean, + ): Promise { + let released = false; + const release = async () => { + if (released) return; + await withMutex(selector, async (dir) => { + if (owns) { + const current = await readFile(join(dir, "owner"), "utf8").catch( + () => undefined, + ); + if (current === token) await rm(join(dir, "owner"), { force: true }); + } else { + await rm(join(dir, `turn-${token}`), { force: true }); + } + if ( + !(await exists(join(dir, "owner"))) && + (await activeTurns(dir)).length === 0 && + !(await exists(join(dir, "persistent-loss"))) + ) + await rm(join(dir, "poison"), { force: true }); + }); + released = true; + }; + return Object.assign(release, { + async verifyEligibility() { + if (released || !owns) return false; + return withMutex(selector, async (dir) => { + if ( + (await exists(globalPoison)) || + (await exists(join(dir, "poison"))) + ) + return false; + return (await readFile(join(dir, "owner"), "utf8")) === token; + }); + }, + }); + } + + return { + async acquire(selector, options: MastraMemoryLeaseOptions = {}) { + const waitMs = options.waitMs ?? 100; + const deadline = Date.now() + waitMs; + const token = randomUUID(); + while (true) { + if (options.signal?.aborted) + throw new Error("Source-thread lease wait cancelled."); + const status = await withMutex(selector, async (dir) => { + if ( + (await exists(globalPoison)) || + (await exists(join(dir, "poison"))) + ) { + await writeFile(join(dir, `turn-${token}`), ""); + return "denied"; + } + if (await exists(join(dir, "owner"))) return "busy"; + await writeFile(join(dir, "owner"), token, { flag: "wx" }); + return "owned"; + }); + if (status === "owned") return makeLease(selector, token, true); + if (status === "denied") return makeLease(selector, token, false); + if (Date.now() >= deadline) { + await withMutex(selector, async (dir) => { + await poison(dir, false); + await writeFile(join(dir, `turn-${token}`), ""); + }); + return makeLease(selector, token, false); + } + await pause(Math.min(5, deadline - Date.now())); + } + }, + async markUnsafeWrite(selector) { + if (!selector) { + await writeFile(globalPoison, "unsafe write"); + return; + } + await withMutex(selector, async (dir) => { + await poison(dir, true); + }); + }, + async resetAfterQuiescence(selector) { + if (!selector) { + for (const name of await readdir(root)) { + if (name === "unknown-writer-poison") continue; + const dir = join(root, name); + if ( + (await exists(join(dir, "owner"))) || + (await activeTurns(dir)).length > 0 + ) + throw new Error("Source-thread writers are still active."); + } + await rm(globalPoison, { force: true }); + return; + } + await withMutex(selector, async (dir) => { + if ( + (await exists(join(dir, "owner"))) || + (await activeTurns(dir)).length > 0 + ) + throw new Error("Source-thread writers are still active."); + await rm(join(dir, "persistent-loss"), { force: true }); + await rm(join(dir, "poison"), { force: true }); + }); + }, + async simulateLeaseLoss(selector) { + await withMutex(selector, async (dir) => { + await poison(dir, true); + await rm(join(dir, "owner"), { force: true }); + }); + }, + }; +} diff --git a/packages/mastra/test/helpers/file-memory-worker.mjs b/packages/mastra/test/helpers/file-memory-worker.mjs new file mode 100644 index 000000000..9adccf477 --- /dev/null +++ b/packages/mastra/test/helpers/file-memory-worker.mjs @@ -0,0 +1,49 @@ +import { createFileMemoryAccess } from "./file-memory-access.ts"; + +const access = createFileMemoryAccess(process.env.KITARU_LEASE_TEST_ROOT); +const leases = new Map(); +const waits = new Map(); + +process.on("message", async (message) => { + const { id, action, selector } = message; + try { + let result; + if (action === "acquire") { + const abort = new AbortController(); + waits.set(id, abort); + try { + const lease = await access.acquire(selector, { + waitMs: message.waitMs, + signal: abort.signal, + }); + leases.set(message.name, lease); + result = await lease.verifyEligibility(); + } finally { + waits.delete(id); + } + } else if (action === "cancel") { + waits.get(message.waitId)?.abort(); + result = true; + } else if (action === "verify") { + result = await leases.get(message.name).verifyEligibility(); + } else if (action === "release") { + await leases.get(message.name)(); + leases.delete(message.name); + result = true; + } else if (action === "lose") { + await access.simulateLeaseLoss(selector); + result = true; + } else if (action === "nativeWrite") { + await access.markUnsafeWrite(selector); + result = true; + } else if (action === "reset") { + await access.resetAfterQuiescence(selector); + result = true; + } else { + throw new Error(`Unknown worker action: ${action}`); + } + process.send?.({ id, ok: true, result }); + } catch (error) { + process.send?.({ id, ok: false, error: String(error) }); + } +}); diff --git a/packages/mastra/test/implicit-om.test.ts b/packages/mastra/test/implicit-om.test.ts new file mode 100644 index 000000000..4a80d2fcc --- /dev/null +++ b/packages/mastra/test/implicit-om.test.ts @@ -0,0 +1,123 @@ +import { InMemoryStore } from "@mastra/core/storage"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { Memory } from "@mastra/memory"; +import { expect, it, vi } from "vitest"; +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, +} from "../src/memory.js"; +import { textStream } from "./helpers/memory-agent.js"; +import { AGENT_ID, installTestApi, REPLAY_ID } from "./helpers.js"; + +it("records a first turn with implicit thread-scoped OM and string models", async () => { + const store = new InMemoryStore(); + const domain = store.stores.memory; + if (!domain) throw new Error("Missing memory store"); + const source = new Memory({ + storage: store, + options: { + lastMessages: 20, + observationalMemory: { + observation: { + model: "fixture/observer", + messageTokens: 100, + bufferTokens: 20, + bufferActivation: 1, + }, + reflection: { model: "fixture/reflector", observationTokens: 100 }, + }, + }, + }); + const actor = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("first answer"), + }); + const observe = vi.fn(async () => + textStream( + "First turn.Continue.", + ), + ); + const observer = new MastraLanguageModelV2Mock({ + modelId: "observer", + provider: "fixture", + doStream: observe, + }); + const reflect = vi.fn(async () => + textStream("First turn."), + ); + const reflector = new MastraLanguageModelV2Mock({ + modelId: "reflector", + provider: "fixture", + doStream: reflect, + }); + const api = installTestApi(); + const agent = createMemoryReplayAgent( + ({ memory }) => ({ + id: "first-turn", + name: "First turn", + instructions: "Answer", + model: actor, + memory, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => source.settled(), + domain, + configuration: source.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: (id) => + id === "fixture/observer" + ? observer + : id === "fixture/reflector" + ? reflector + : actor, + }, + ); + try { + const result = await agent.stream( + "This is the first message. ".repeat(80), + { + memory: { thread: "new-thread", resource: "owner" }, + }, + ); + await result.consumeStream(); + expect(await result.text).toBe("first answer"); + await vi.waitFor(() => + expect( + api.calls.some( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", + ), + ).toBe(true), + ); + const input = api.calls.find( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", + )?.body?.inputs as Record; + expect(input.mastra_memory_replay?.version).toBe(3); + expect(input.mastra_memory_replay?.omTape.length).toBeGreaterThan(0); + const observerCalls = observe.mock.calls.length; + const reflectorCalls = reflect.mock.calls.length; + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(input)); + const replay = await agent.stream("ignored"); + await replay.consumeStream(); + expect(await replay.text).toBe("first answer"); + expect(observe).toHaveBeenCalledTimes(observerCalls); + expect(reflect).toHaveBeenCalledTimes(reflectorCalls); + } finally { + vi.unstubAllEnvs(); + vi.unstubAllGlobals(); + await source.settled(); + await store.close(); + } +}); diff --git a/packages/mastra/test/memory-binding.test.ts b/packages/mastra/test/memory-binding.test.ts index 856707e36..be0db6987 100644 --- a/packages/mastra/test/memory-binding.test.ts +++ b/packages/mastra/test/memory-binding.test.ts @@ -1,7 +1,9 @@ +import { Extractor } from "@mastra/memory"; import { expect, it, vi } from "vitest"; import { createMemoryCaptureBinding, createProcessLocalMemoryAccess, + type MastraMemoryLease, } from "../src/memory-binding.js"; import { createMemoryRuntime, @@ -71,6 +73,49 @@ it("rejects shared-thread overlap while allowing independent source threads", as await two.binding.release(); }); +it("shows why separate process-local helpers cannot qualify for multi-server replay", async () => { + const first = createProcessLocalMemoryAccess(); + const second = createProcessLocalMemoryAccess(); + const selector = { threadId: THREAD, resourceId: RESOURCE }; + const firstLease = await first.acquire(selector); + const secondLease = await second.acquire(selector); + expect(await firstLease.verifyEligibility()).toBe(true); + expect(await secondLease.verifyEligibility()).toBe(true); + await firstLease(); + await secondLease(); +}); + +it("poisons all threads when a native fallback cannot identify its selector", async () => { + const access = createProcessLocalMemoryAccess(); + const first = await access.acquire({ + threadId: "first", + resourceId: RESOURCE, + }); + const second = await access.acquire({ + threadId: "second", + resourceId: RESOURCE, + }); + await access.markUnsafeWrite(); + expect(await first.verifyEligibility()).toBe(false); + expect(await second.verifyEligibility()).toBe(false); + await expect(access.resetAfterQuiescence()).rejects.toThrow(/active/); + await first(); + await second(); + const later = await access.acquire({ + threadId: "third", + resourceId: RESOURCE, + }); + expect(await later.verifyEligibility()).toBe(false); + await later(); + await access.resetAfterQuiescence(); + const recovered = await access.acquire({ + threadId: "third", + resourceId: RESOURCE, + }); + expect(await recovered.verifyEligibility()).toBe(true); + await recovered(); +}); + it("preserves native mutation results despite evidence persistence failure", async () => { const { runtime, binding, recordMutation } = await fixture(); await binding.captureInitial(runtime.memory); @@ -103,6 +148,128 @@ it("invalidates the first recording when a conflicting invocation cannot get its expect(one.binding.incompleteReasons.join()).toMatch(/overlapping/); await one.binding.release(); await two.binding.release(); + const next = await fixture("next", access); + expect(await next.binding.captureInitial(next.runtime.memory)).toBeDefined(); + await next.binding.release(); +}); + +it("preserves a native write after lost ownership and blocks eligibility until quiescence reset", async () => { + const access = createProcessLocalMemoryAccess(); + const first = await fixture("first", access); + const second = await fixture("second", access); + await first.binding.captureInitial(first.runtime.memory); + expect( + await second.binding.captureInitial(second.runtime.memory), + ).toBeUndefined(); + const native = await first.binding.domain.updateThread({ + id: THREAD, + title: "native write still succeeds", + }); + expect(native.title).toBe("native write still succeeds"); + await first.binding.verifyEligibility(); + expect(first.binding.incompleteReasons.join()).toMatch(/ownership/); + await first.binding.release(); + await second.binding.release(); + const next = await fixture("after", access); + expect( + await next.binding.captureInitial(next.runtime.memory), + ).toBeUndefined(); + await next.binding.release(); + await access.resetAfterQuiescence({ threadId: THREAD, resourceId: RESOURCE }); + const recovered = await fixture("recovered", access); + expect( + await recovered.binding.captureInitial(recovered.runtime.memory), + ).toBeDefined(); + await recovered.binding.release(); +}); + +it("keeps an untracked late write unsafe until an explicit quiescence reset", async () => { + const access = createProcessLocalMemoryAccess(); + const first = await fixture("first", access); + await first.binding.captureInitial(first.runtime.memory); + await first.binding.release(); + const native = await first.binding.domain.updateThread({ + id: THREAD, + title: "late native write", + }); + expect(native.title).toBe("late native write"); + const next = await fixture("next", access); + expect( + await next.binding.captureInitial(next.runtime.memory), + ).toBeUndefined(); + await next.binding.release(); + await access.resetAfterQuiescence({ threadId: THREAD, resourceId: RESOURCE }); + const recovered = await fixture("recovered", access); + expect( + await recovered.binding.captureInitial(recovered.runtime.memory), + ).toBeDefined(); + await recovered.binding.release(); +}); + +it("bounds a hung acquisition and releases a lease returned after cancellation", async () => { + const release = Object.assign( + vi.fn(async () => {}), + { + verifyEligibility: vi.fn(async () => true), + }, + ) satisfies MastraMemoryLease; + let resolveAcquire: ((lease: MastraMemoryLease) => void) | undefined; + const acquire = vi.fn( + () => + new Promise((resolve) => { + resolveAcquire = resolve; + }), + ); + const runtime = createMemoryRuntime(); + await seedMemory(runtime); + const binding = createMemoryCaptureBinding({ + invocationId: "timeout", + domain: runtime.domain, + threadId: THREAD, + resourceId: RESOURCE, + exclusiveAccess: { + acquire, + markUnsafeWrite: async () => {}, + resetAfterQuiescence: async () => {}, + }, + recordMutation: async () => {}, + leaseWaitMs: 20, + }); + const start = Date.now(); + expect(await binding.captureInitial(runtime.memory)).toBeUndefined(); + expect(Date.now() - start).toBeLessThan(200); + required(resolveAcquire)(release); + await vi.waitFor(() => expect(release).toHaveBeenCalledOnce()); + await binding.release(); +}); + +it("bounds a stalled pre-turn source read without changing native storage", async () => { + const runtime = createMemoryRuntime(); + await seedMemory(runtime); + const native = runtime.domain.updateThread.bind(runtime.domain); + const binding = createMemoryCaptureBinding({ + invocationId: "stalled-read", + domain: runtime.domain, + threadId: THREAD, + resourceId: RESOURCE, + exclusiveAccess: createProcessLocalMemoryAccess(), + recordMutation: async () => {}, + captureWaitMs: 20, + }); + const read = vi.spyOn(runtime.domain, "listMessages").mockImplementation( + async () => + new Promise(() => { + /* A storage read that never settles. */ + }), + ); + const start = Date.now(); + expect(await binding.captureInitial(runtime.memory)).toBeUndefined(); + expect(Date.now() - start).toBeLessThan(200); + expect(binding.incompleteReasons.join()).toMatch(/capture timed out/i); + await binding.release(); + read.mockRestore(); + const updated = await native({ id: THREAD, title: "native still works" }); + expect(updated.title).toBe("native still works"); }); it("does not return a coherent initial snapshot after overlap during capture", async () => { @@ -135,7 +302,7 @@ it("serializes overlapping native mutations and preserves original storage error }); const first = binding.domain.updateThread({ id: THREAD, title: "first" }); const second = binding.domain.updateThread({ id: THREAD, title: "second" }); - await Promise.resolve(); + await vi.waitFor(() => expect(calls).toEqual(["first"])); expect(calls).toEqual(["first"]); required(unblock)(); await Promise.all([first, second]); @@ -200,6 +367,78 @@ it("records OM flags, buffers, config, activation and working memory with stable await binding.release(); }); +it("records first-turn OM initialization by model and built-in extractor identity", async () => { + const { runtime, binding, recordMutation } = await fixture(); + await binding.captureInitial(runtime.memory); + const observer = runtime.observer.model; + const reflector = runtime.reflector.model; + const original = { + threadId: "first-turn-thread", + resourceId: RESOURCE, + scope: "thread" as const, + config: { + scope: "thread", + observation: { + model: observer, + extractors: [ + new Extractor( + { name: "current-task", instructions: "Continue." }, + true, + ), + ], + }, + reflection: { model: reflector }, + }, + }; + const native = await binding.domain.initializeObservationalMemory(original); + expect((native.config.observation as Record).model).toBe( + observer, + ); + await binding.drain(); + expect(binding.incompleteReasons).toEqual([]); + const event = recordMutation.mock.calls[0]?.[0] as { + arguments: unknown[]; + result: Record; + }; + expect(event.arguments).toMatchObject([ + { + config: { + observation: { + model: "fixture/observer", + extractors: [{ mastraBuiltinExtractor: "current-task" }], + }, + reflection: { model: "fixture/reflector" }, + }, + }, + ]); + expect(event.result).toMatchObject({ + config: { observation: { model: "fixture/observer" } }, + }); + await binding.release(); +}); + +it("keeps custom OM extractor execution native while refusing incomplete replay evidence", async () => { + const { runtime, binding } = await fixture(); + await binding.captureInitial(runtime.memory); + const native = await binding.domain.initializeObservationalMemory({ + threadId: "custom-extractor-thread", + resourceId: RESOURCE, + scope: "thread", + config: { + observation: { + model: runtime.observer.model, + extractors: [ + new Extractor({ name: "custom", instructions: "Extract a value." }), + ], + }, + reflection: { model: runtime.reflector.model }, + }, + }); + expect(native.threadId).toBe("custom-extractor-thread"); + expect(binding.incompleteReasons.join()).toMatch(/arguments|result/); + await binding.release(); +}); + it("marks credential-altered mutation evidence incomplete but keeps native arguments and results", async () => { const { runtime, binding, recordMutation } = await fixture(); await binding.captureInitial(runtime.memory); diff --git a/packages/mastra/test/memory-lease-process.test.ts b/packages/mastra/test/memory-lease-process.test.ts new file mode 100644 index 000000000..edc404564 --- /dev/null +++ b/packages/mastra/test/memory-lease-process.test.ts @@ -0,0 +1,189 @@ +import { type ChildProcess, fork } from "node:child_process"; +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { fileURLToPath } from "node:url"; +import { afterEach, expect, it } from "vitest"; + +const workerPath = fileURLToPath( + new URL("./helpers/file-memory-worker.mjs", import.meta.url), +); +const thread = { threadId: "shared-thread", resourceId: "resource" }; +const other = { threadId: "other-thread", resourceId: "resource" }; +let root: string | undefined; +const children: ChildProcess[] = []; +let nextId = 0; + +type Reply = { id: number; ok: boolean; result?: unknown; error?: string }; + +async function worker(): Promise<{ + child: ChildProcess; + request(message: Record): Promise; + start(message: Record): { + id: number; + result: Promise; + }; +}> { + root ??= await mkdtemp(join(tmpdir(), "kitaru-lease-")); + const child = fork(workerPath, [], { + env: { ...process.env, KITARU_LEASE_TEST_ROOT: root }, + execArgv: ["--experimental-strip-types"], + stdio: ["ignore", "pipe", "pipe", "ipc"], + }); + children.push(child); + const pending = new Map< + number, + { resolve(value: unknown): void; reject(error: Error): void } + >(); + child.on("message", (value: Reply) => { + const waiting = pending.get(value.id); + if (!waiting) return; + pending.delete(value.id); + if (value.ok) waiting.resolve(value.result); + else waiting.reject(new Error(value.error)); + }); + child.on("exit", (code) => { + for (const waiting of pending.values()) + waiting.reject(new Error(`Lease worker exited with ${code}`)); + pending.clear(); + }); + function start(message: Record) { + const id = ++nextId; + const result = new Promise((resolve, reject) => { + pending.set(id, { resolve, reject }); + child.send({ id, ...message }); + }); + return { id, result }; + } + return { child, start, request: (message) => start(message).result }; +} + +afterEach(async () => { + for (const child of children.splice(0)) child.kill(); + if (root) await rm(root, { recursive: true, force: true }); + root = undefined; +}); + +it("invalidates both processes on bounded same-thread contention, then recovers after both finish", async () => { + const first = await worker(); + const second = await worker(); + expect( + await first.request({ action: "acquire", name: "first", selector: thread }), + ).toBe(true); + expect( + await second.request({ + action: "acquire", + name: "second", + selector: thread, + waitMs: 40, + }), + ).toBe(false); + expect( + await first.request({ action: "verify", name: "first", selector: thread }), + ).toBe(false); + await Promise.all([ + first.request({ action: "release", name: "first", selector: thread }), + second.request({ action: "release", name: "second", selector: thread }), + ]); + expect( + await second.request({ action: "acquire", name: "next", selector: thread }), + ).toBe(true); + await second.request({ action: "release", name: "next", selector: thread }); +}); + +it("poisons a lost holder before a stale native write and blocks a successor", async () => { + const first = await worker(); + const second = await worker(); + expect( + await first.request({ action: "acquire", name: "first", selector: thread }), + ).toBe(true); + await second.request({ action: "lose", selector: thread }); + expect( + await first.request({ action: "verify", name: "first", selector: thread }), + ).toBe(false); + await first.request({ action: "nativeWrite", selector: thread }); + expect( + await second.request({ action: "acquire", name: "next", selector: thread }), + ).toBe(false); + await Promise.all([ + first.request({ action: "release", name: "first", selector: thread }), + second.request({ action: "release", name: "next", selector: thread }), + ]); + expect( + await second.request({ + action: "acquire", + name: "still-unsafe", + selector: thread, + }), + ).toBe(false); + await second.request({ + action: "release", + name: "still-unsafe", + selector: thread, + }); + await second.request({ action: "reset", selector: thread }); + expect( + await second.request({ + action: "acquire", + name: "recovered", + selector: thread, + }), + ).toBe(true); + await second.request({ + action: "release", + name: "recovered", + selector: thread, + }); +}); + +it("allows another thread and cancels a waiter without releasing the holder", async () => { + const first = await worker(); + const second = await worker(); + await first.request({ action: "acquire", name: "first", selector: thread }); + expect( + await second.request({ action: "acquire", name: "other", selector: other }), + ).toBe(true); + const waiting = second.start({ + action: "acquire", + name: "cancelled", + selector: thread, + waitMs: 500, + }); + await second.request({ action: "cancel", waitId: waiting.id }); + await expect(waiting.result).rejects.toThrow(/cancelled/); + expect( + await first.request({ action: "verify", name: "first", selector: thread }), + ).toBe(true); + await Promise.all([ + first.request({ action: "release", name: "first", selector: thread }), + second.request({ action: "release", name: "other", selector: other }), + ]); +}); + +it("poisons every process when a fallback writer has no selector", async () => { + const first = await worker(); + const second = await worker(); + expect( + await first.request({ action: "acquire", name: "first", selector: thread }), + ).toBe(true); + await second.request({ action: "nativeWrite" }); + expect(await first.request({ action: "verify", name: "first" })).toBe(false); + await first.request({ action: "release", name: "first", selector: thread }); + expect( + await second.request({ action: "acquire", name: "other", selector: other }), + ).toBe(false); + await second.request({ action: "release", name: "other", selector: other }); + await first.request({ action: "reset" }); + expect( + await first.request({ + action: "acquire", + name: "recovered", + selector: other, + }), + ).toBe(true); + await first.request({ + action: "release", + name: "recovered", + selector: other, + }); +}); diff --git a/packages/mastra/test/memory-replay-safety.test.ts b/packages/mastra/test/memory-replay-safety.test.ts index c82301ea4..44a76f08f 100644 --- a/packages/mastra/test/memory-replay-safety.test.ts +++ b/packages/mastra/test/memory-replay-safety.test.ts @@ -1,3 +1,5 @@ +import { createHash } from "node:crypto"; +import { Agent } from "@mastra/core/agent"; import { MASTRA_AUTH_TOKEN_KEY, MASTRA_RESOURCE_ID_KEY, @@ -49,6 +51,49 @@ function input(): MastraMemoryReplayInput { }; } +it("stores captured file references and refuses signed source URLs", () => { + const bytes = new Uint8Array([1, 2, 3]); + const digest = createHash("sha256") + .update("image/png") + .update("\0") + .update(bytes) + .digest("hex"); + const ref = `kitaru-file://sha256/${digest}`; + const safe = input(); + safe.rawInput = { file: new URL(ref) }; + safe.files = [{ url: ref, mediaType: "image/png", bytes }]; + const envelope = createMemoryReplayEnvelope(safe); + expect(envelope.complete, envelope.reasons.join("; ")).toBe(true); + expect(decodeMemoryReplayEnvelope(envelope).files[0]?.bytes).toEqual(bytes); + + const plain = createMemoryReplayEnvelope({ + ...safe, + files: [ + { url: "https://files.invalid/image.png", mediaType: "image/png", bytes }, + ], + }); + expect(plain.complete).toBe(false); + + const signed = createMemoryReplayEnvelope({ + ...safe, + rawInput: { + file: "https://files.invalid/image.png?X-Amz-Signature=SECRET", + }, + }); + expect(signed.complete).toBe(false); + expect(JSON.stringify(signed)).not.toContain("SECRET"); + + const captured = envelope.files[0]; + if (!captured) throw new Error("Missing captured file"); + const altered = { + ...envelope, + files: [{ ...captured, url: `kitaru-file://sha256/${"0".repeat(64)}` }], + }; + expect(() => decodeMemoryReplayEnvelope(altered)).toThrow( + /captured content reference/, + ); +}); + function fixture( factory?: MemoryReplayAgentFactory, overrides: Partial = {}, @@ -113,7 +158,7 @@ it.each(["defaultOptions", "runOptions", "memoryConfig"])( ...clean, configuration: { ...original.configuration, [key]: unsafe }, }), - ).toThrow(/transport/i); + ).toThrow(/transport|sensitive key/i); }, ); @@ -129,55 +174,80 @@ it("rejects native auth tokens during capture and decode", () => { ).toThrow(/auth/i); }); -it("rejects default auth-token context capture before recording or execution", async () => { +it("keeps native auth context while marking uncaptured context ineligible", async () => { const api = installTestApi(); const { adapter, modelCall } = fixture(); const requestContext = new RequestContext(); requestContext.set(MASTRA_AUTH_TOKEN_KEY, "CREDENTIAL"); - await expect( - adapter.stream("hello", { - memory: { thread: "thread", resource: "resource" }, - requestContext, - }), - ).rejects.toThrow(/auth/i); - expect(modelCall).not.toHaveBeenCalled(); - expect(api.calls).toEqual([]); + const result = await adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + requestContext, + }); + await result.consumeStream(); + expect(await result.text).toBe("done"); + expect(modelCall).toHaveBeenCalledTimes(1); + await vi.waitFor(() => + expect( + api.calls.find( + (call) => call.method === "PATCH" && call.body?.status === "failed", + )?.body?.metadata, + ).toMatchObject({ mastra_replay_state: "ineligible" }), + ); + expect(JSON.stringify(api.calls)).not.toContain("CREDENTIAL"); }); it.each([MASTRA_THREAD_ID_KEY, MASTRA_RESOURCE_ID_KEY])( - "rejects a mismatched %s before acquiring the source lease", + "keeps a mismatched %s native and rejects its recording before leasing", async (key) => { const api = installTestApi(); const { adapter, modelCall, acquire } = fixture(); const requestContext = new RequestContext(); requestContext.set(key, "other"); - await expect( - adapter.stream("hello", { - memory: { thread: "thread", resource: "resource" }, - requestContext, - }), - ).rejects.toThrow(/selector/i); + const result = await adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + requestContext, + }); + await result.consumeStream(); + expect(await result.text).toBe("done"); expect(acquire).not.toHaveBeenCalled(); - expect(modelCall).not.toHaveBeenCalled(); - expect(api.calls).toEqual([]); + expect(modelCall).toHaveBeenCalledTimes(1); + await vi.waitFor(() => + expect( + api.calls.find( + (call) => call.method === "POST" && call.path === "/api/v1/sessions", + )?.body?.metadata, + ).toMatchObject({ + mastra_replay_state: "ineligible", + mastra_replay_reason: "context_unsupported", + }), + ); }, ); -it("does not let selective context capture bypass middleware memory selectors", async () => { +it("does not make mismatched middleware selectors replayable through selective capture", async () => { const api = installTestApi(); const { adapter, acquire } = fixture(undefined, { captureRequestContext: () => ({ locale: "en" }), }); const requestContext = new RequestContext(); requestContext.set(MASTRA_THREAD_ID_KEY, "authorized-thread"); - await expect( - adapter.stream("hello", { - memory: { thread: "thread", resource: "resource" }, - requestContext, - }), - ).rejects.toThrow(/selector/i); + const result = await adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + requestContext, + }); + await result.consumeStream(); + expect(await result.text).toBe("done"); expect(acquire).not.toHaveBeenCalled(); - expect(api.calls).toEqual([]); + await vi.waitFor(() => + expect( + api.calls.find( + (call) => call.method === "POST" && call.path === "/api/v1/sessions", + )?.body?.metadata, + ).toMatchObject({ + mastra_replay_state: "ineligible", + mastra_replay_reason: "context_unsupported", + }), + ); }); it("allows selective capture to exclude live authentication tokens", async () => { @@ -193,14 +263,23 @@ it("allows selective capture to exclude live authentication tokens", async () => }); await result.consumeStream(); expect(JSON.stringify(api.calls)).not.toContain("CREDENTIAL"); - expect( - api.calls.filter((call) => call.method === "PATCH").at(-1)?.body?.status, - ).toBe("completed"); + await vi.waitFor(() => + expect( + api.calls.find( + (call) => call.method === "PATCH" && call.body?.status === "completed", + )?.body?.metadata, + ).toMatchObject({ mastra_native_state: "completed" }), + ); }); it("records and replays matching middleware and invocation memory selectors", async () => { const api = installTestApi(); - const { adapter, modelCall, acquire } = fixture(); + const { adapter, modelCall, acquire } = fixture(undefined, { + captureRequestContext: (context) => ({ + [MASTRA_THREAD_ID_KEY]: context.get(MASTRA_THREAD_ID_KEY), + [MASTRA_RESOURCE_ID_KEY]: context.get(MASTRA_RESOURCE_ID_KEY), + }), + }); const requestContext = new RequestContext(); requestContext.set(MASTRA_THREAD_ID_KEY, "thread"); requestContext.set(MASTRA_RESOURCE_ID_KEY, "resource"); @@ -220,54 +299,20 @@ it("records and replays matching middleware and invocation memory selectors", as await replay.consumeStream(); expect(modelCall).toHaveBeenCalledTimes(2); expect(acquire).toHaveBeenCalledTimes(1); - expect( - api.calls - .filter((call) => call.method === "PATCH") - .map((call) => call.body?.status), - ).toEqual(["completed", "completed"]); + await vi.waitFor(() => + expect( + api.calls.filter( + (call) => call.method === "PATCH" && call.body?.status === "completed", + ), + ).toHaveLength(2), + ); }); it.each([MASTRA_THREAD_ID_KEY, MASTRA_RESOURCE_ID_KEY, MASTRA_AUTH_TOKEN_KEY])( - "rejects late processor writes to %s", + "preserves native resolver mutations of %s while making recording ineligible", async (key) => { - vi.spyOn(console, "error").mockImplementation(() => {}); const api = installTestApi(); - const modelCall = vi.fn(async () => textStream("done")); const { adapter } = fixture(({ memory }) => ({ - id: "processor-mutation", - name: "Processor mutation", - memory, - model: new MastraLanguageModelV2Mock({ doStream: modelCall }), - instructions: "Answer", - inputProcessors: [ - { - id: "context-mutation", - processInput({ requestContext, messages }) { - if (!requestContext) throw new Error("Missing request context"); - requestContext.setRaw(key, "FORBIDDEN"); - return messages; - }, - }, - ], - })); - await expect( - adapter.stream("hello", { - memory: { thread: "thread", resource: "resource" }, - }), - ).rejects.toThrow(/processor/i); - expect(modelCall).not.toHaveBeenCalled(); - expect(JSON.stringify(api.calls)).not.toContain("FORBIDDEN"); - expect( - api.calls.filter((call) => call.method === "PATCH").at(-1)?.body?.status, - ).toBe("failed"); - }, -); - -it.each([MASTRA_THREAD_ID_KEY, MASTRA_RESOURCE_ID_KEY, MASTRA_AUTH_TOKEN_KEY])( - "rejects dynamic resolver mutations of %s before native execution", - async (key) => { - const api = installTestApi(); - const { adapter, modelCall } = fixture(({ memory }) => ({ id: "mutation", name: "Mutation", memory, @@ -279,13 +324,18 @@ it.each([MASTRA_THREAD_ID_KEY, MASTRA_RESOURCE_ID_KEY, MASTRA_AUTH_TOKEN_KEY])( return "Answer"; }, })); - await expect( - adapter.stream("hello", { - memory: { thread: "thread", resource: "resource" }, - }), - ).rejects.toThrow(/selector|auth/i); - expect(modelCall).not.toHaveBeenCalled(); - expect(api.calls).toEqual([]); + const result = await adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + }); + await result.consumeStream(); + expect(await result.text).toBe("done"); + await vi.waitFor(() => + expect( + api.calls.find( + (call) => call.method === "PATCH" && call.body?.status === "failed", + )?.body?.metadata, + ).toMatchObject({ mastra_replay_state: "ineligible" }), + ); }, ); @@ -307,25 +357,34 @@ it("never uploads provider transport credentials in session inputs", async () => expect(JSON.stringify(api.calls)).not.toContain("CREDENTIAL"); }); -it("rejects late auth-token additions by default-option resolvers", async () => { +it("preserves native default-option auth tokens while marking recording ineligible", async () => { const api = installTestApi(); const { adapter } = fixture(({ memory }) => ({ id: "defaults", name: "Defaults", memory, instructions: "Answer", - model: new MastraLanguageModelV2Mock({}), + model: new MastraLanguageModelV2Mock({ + doStream: async () => textStream("done"), + }), defaultOptions: ({ requestContext }) => { requestContext.setRaw(MASTRA_AUTH_TOKEN_KEY, "CREDENTIAL"); return {}; }, })); - await expect( - adapter.stream("hello", { - memory: { thread: "thread", resource: "resource" }, - }), - ).rejects.toThrow(/auth/i); - expect(api.calls).toEqual([]); + const result = await adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + }); + await result.consumeStream(); + expect(await result.text).toBe("done"); + await vi.waitFor(() => + expect( + api.calls.find( + (call) => call.method === "PATCH" && call.body?.status === "failed", + )?.body?.metadata, + ).toMatchObject({ mastra_replay_state: "ineligible" }), + ); + expect(JSON.stringify(api.calls)).not.toContain("CREDENTIAL"); }); it("marks nested abort signals as unsupported transport configuration", () => { @@ -339,7 +398,7 @@ it("marks nested abort signals as unsupported transport configuration", () => { ...createMemoryReplayEnvelope(original), configuration, }), - ).toThrow(/transport/i); + ).toThrow(/transport|sensitive key/i); }); it("rejects replay envelopes with selectors inconsistent with the snapshot", () => { @@ -421,4 +480,56 @@ it("preserves native stream errors when cleanup and diagnostics reject", async ( ); }); -import { Agent } from "@mastra/core/agent"; +it.each([MASTRA_THREAD_ID_KEY, MASTRA_RESOURCE_ID_KEY, MASTRA_AUTH_TOKEN_KEY])( + "keeps two native streams usable after a processor changes implicit context %s", + async (key) => { + const api = installTestApi(); + const modelCall = vi.fn(async () => textStream("done")); + const processorCall = vi.fn(); + const { adapter } = fixture(({ memory }) => ({ + id: "processor-mutation", + name: "Processor mutation", + memory, + model: new MastraLanguageModelV2Mock({ doStream: modelCall }), + instructions: "Answer", + inputProcessors: [ + { + id: "context-mutation", + processInput({ requestContext, messages }) { + processorCall(); + if (!requestContext) throw new Error("Missing request context"); + requestContext.setRaw(key, "FORBIDDEN"); + return messages; + }, + }, + ], + })); + for (let index = 1; index <= 2; index++) { + const result = await adapter.stream("hello", { + memory: { thread: "thread", resource: "resource" }, + }); + await result.consumeStream(); + expect(await result.text).toBe("done"); + await vi.waitFor(() => + expect( + api.calls.filter( + (call) => call.method === "PATCH" && call.body?.status === "failed", + ), + ).toHaveLength(index), + ); + } + expect(processorCall).toHaveBeenCalledTimes(2); + expect(modelCall).toHaveBeenCalledTimes(2); + expect( + api.calls + .filter( + (call) => call.method === "PATCH" && call.body?.status === "failed", + ) + .every( + (call) => + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "ineligible", + ), + ).toBe(true); + }, +); diff --git a/packages/mastra/test/memory-snapshot.test.ts b/packages/mastra/test/memory-snapshot.test.ts index 2bb069e91..46df0e3ad 100644 --- a/packages/mastra/test/memory-snapshot.test.ts +++ b/packages/mastra/test/memory-snapshot.test.ts @@ -1,3 +1,4 @@ +import { createHash } from "node:crypto"; import { expect, it } from "vitest"; import { createContextInput, @@ -8,15 +9,17 @@ import { decodeMemoryReplayEnvelope, decodeMemoryValue, encodeMemoryValue, + finalizeMemoryReplayEnvelope, } from "../src/memory-snapshot.js"; import { createMemoryRuntime, FILE_BYTES, - FILE_URL, seedMemory, snapshotMemory, } from "./helpers/memory-agent.js"; +const FILE_REF = `kitaru-file://sha256/${createHash("sha256").update("application/pdf\0").update(FILE_BYTES).digest("hex")}`; + function required(value: T | undefined | null): T { if (value === undefined || value === null) throw new Error("Missing fixture value"); @@ -29,7 +32,7 @@ async function fixture() { return { invocationId: "invocation-1", rawInput: [ - { role: "user", content: [{ type: "file", data: new URL(FILE_URL) }] }, + { role: "user", content: [{ type: "file", data: new URL(FILE_REF) }] }, ], initialSnapshot: { ...(await snapshotMemory(runtime, true)), @@ -46,7 +49,8 @@ async function fixture() { }, }, requestContext: { locale: "en" }, - files: [{ url: FILE_URL, mediaType: "application/pdf", bytes: FILE_BYTES }], + files: [{ url: FILE_REF, mediaType: "application/pdf", bytes: FILE_BYTES }], + omTape: [], }; } @@ -72,7 +76,7 @@ it.each(["version", "hash", "missing", "inflight"])( "rejects invalid replay prerequisites: %s", async (kind) => { const envelope = createMemoryReplayEnvelope(await fixture()); - if (kind === "version") envelope.version = 3 as 2; + if (kind === "version") envelope.version = 4 as 3; if (kind === "hash") required(envelope.files[0]).sha256 = "0".repeat(64); if (kind === "missing") envelope.initialSnapshot = {}; if (kind === "inflight") { @@ -89,9 +93,9 @@ it.each(["version", "hash", "missing", "inflight"])( it.each([ { apiKey: "private-value" }, { callback: () => "live" }, - { value: "a".repeat(1_048_576) }, + { value: "a".repeat(16_777_216) }, { value: new Map([["key", "value"]]) }, - { value: Array.from({ length: 10_001 }, () => 1) }, + { value: Array.from({ length: 200_001 }, () => 1) }, ])( "marks altered or oversized state incomplete without exposing credentials", async (configuration) => { @@ -165,23 +169,116 @@ it("round-trips an empty initial conversation", async () => { it("counts aggregate envelope items, depth, and binary expansion against the shared budget", async () => { const input = await fixture(); - expect( - createMemoryReplayEnvelope({ - ...input, - configuration: { values: Array.from({ length: 6000 }, () => 1) }, - requestContext: { values: Array.from({ length: 6000 }, () => 1) }, - }).complete, - ).toBe(false); + const tooMany = createMemoryReplayEnvelope({ + ...input, + configuration: { values: Array.from({ length: 110_000 }, () => 1) }, + requestContext: { values: Array.from({ length: 110_000 }, () => 1) }, + }); + expect(tooMany.complete).toBe(false); + expect(tooMany.reasons[0]).toMatch(/maximum item count 200000/); const deep = Array.from({ length: 65 }).reduce( (value) => ({ value }), null, ); expect(() => encodeMemoryValue(deep)).toThrow(/depth/); const withFile = await fixture(); - required(withFile.files[0]).bytes = new Uint8Array(800_000); + required(withFile.files[0]).bytes = new Uint8Array(13_000_000); expect(createMemoryReplayEnvelope(withFile).complete).toBe(false); }); +it.each([50, 830])( + "round-trips %i-message history with 15k nested fields and a file", + async (count) => { + const input = await fixture(); + const original = required(input.initialSnapshot.messages[0]); + input.initialSnapshot.messages = Array.from( + { length: count }, + (_, index) => ({ + ...original, + id: `message-${index}`, + content: { + ...original.content, + parts: [ + { + type: "text", + text: index === 0 ? "x".repeat(1_100_000) : `message ${index}`, + }, + ], + metadata: + index === 0 + ? { + hotels: Array.from({ length: 1_500 }, (_, hotel) => ({ + id: hotel, + details: Object.fromEntries( + Array.from({ length: 10 }, (_, field) => [ + `field${field}`, + field, + ]), + ), + })), + } + : undefined, + }, + }), + ); + const envelope = createMemoryReplayEnvelope(input); + expect(envelope.complete).toBe(true); + expect(JSON.stringify(envelope).length).toBeGreaterThan(1_048_576); + expect( + decodeMemoryReplayEnvelope(JSON.parse(JSON.stringify(envelope))), + ).toEqual(input); + }, +); + +it("finalizes a separate version-3 envelope with an ordered OM tape", async () => { + const provisional = createMemoryReplayEnvelope(await fixture()); + const final = finalizeMemoryReplayEnvelope(provisional, [ + { phase: "observation", ordinal: 0, output: "remember" }, + ]); + expect(final).not.toBe(provisional); + expect(provisional.omTape).toEqual([]); + expect(decodeMemoryReplayEnvelope(final).omTape).toEqual([ + { phase: "observation", ordinal: 0, output: "remember" }, + ]); + expect(() => + finalizeMemoryReplayEnvelope(provisional, [{ token: "secret" }]), + ).toThrow(); +}); + +it("normalizes implicit thread OM and rejects old OM envelopes without a tape", async () => { + const input = await fixture(); + delete (input.configuration.memory.observationalMemory as { scope?: string }) + .scope; + const current = createMemoryReplayEnvelope(input); + expect(current.complete).toBe(true); + expect( + decodeMemoryReplayEnvelope(current).configuration.memory, + ).toMatchObject({ + observationalMemory: { scope: "thread" }, + }); + const old = { ...current, version: 2 }; + delete (old as { omTape?: unknown }).omTape; + expect(() => decodeMemoryReplayEnvelope(old)).toThrow( + /mastra_om_tape_missing/, + ); + const splitConfig = { + ...old, + configuration: encodeMemoryValue({ + memoryConfig: {}, + memory: { observationalMemory: { scope: "thread" } }, + }), + }; + expect(() => decodeMemoryReplayEnvelope(splitConfig)).toThrow( + /mastra_om_tape_missing/, + ); + const workingOnly = await fixture(); + delete (workingOnly.configuration.memory as { observationalMemory?: unknown }) + .observationalMemory; + const oldWorking = { ...createMemoryReplayEnvelope(workingOnly), version: 2 }; + delete (oldWorking as { omTape?: unknown }).omTape; + expect(decodeMemoryReplayEnvelope(oldWorking).omTape).toBeUndefined(); +}); + it.each(["resourceScope", "semanticRecall", "buffer", "ids", "date", "url"])( "rejects malformed or out-of-scope state: %s", async (kind) => { diff --git a/packages/mastra/test/om-result-tape.test.ts b/packages/mastra/test/om-result-tape.test.ts new file mode 100644 index 000000000..01a695db4 --- /dev/null +++ b/packages/mastra/test/om-result-tape.test.ts @@ -0,0 +1,176 @@ +import { expect, it, vi } from "vitest"; +import { decodeMemoryValue } from "../src/memory-snapshot.js"; +import { + createOMResultTape, + MastraOMDivergenceError, +} from "../src/om-result-tape.js"; + +function model() { + return { + specificationVersion: "v2", + modelId: "fixture", + provider: "fixture", + doStream: vi.fn(async (_input: unknown) => ({ + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: "text-delta", textDelta: "remembered" }); + controller.enqueue({ type: "finish", finishReason: "stop" }); + controller.close(); + }, + }), + })), + }; +} + +async function collect(stream: ReadableStream): Promise { + const chunks: unknown[] = []; + const reader = stream.getReader(); + while (true) { + const item = await reader.read(); + if (item.done) break; + chunks.push(item.value); + } + reader.releaseLock(); + return chunks; +} + +it("reuses ordered recorded OM output without calling the live model", async () => { + const capture = createOMResultTape(undefined, () => { + throw new Error("unexpected incomplete result"); + }); + const native = model(); + const baseline = capture.instrument(native, "observer"); + const output = await baseline.doStream({ prompt: "before" }); + expect(await collect(output.stream)).toHaveLength(2); + const entries = await capture.finish(); + expect(entries).toMatchObject([{ phase: "observer", ordinal: 0 }]); + const live = model(); + const mismatch = vi.fn(); + const replay = createOMResultTape(entries, () => {}, mismatch); + const recorded = replay.instrument(live, "observer"); + const replayOutput = await recorded.doStream({ prompt: "after" }); + expect(await collect(replayOutput.stream)).toEqual( + await collect((await native.doStream({ prompt: "before" })).stream), + ); + expect(live.doStream).not.toHaveBeenCalled(); + expect(mismatch).toHaveBeenCalledTimes(1); + await replay.finish(); +}); + +it("rejects an extra or missing OM call", async () => { + const capture = createOMResultTape(undefined, () => {}); + const output = await capture.instrument(model(), "reflector").doStream({}); + await collect(output.stream); + const entries = await capture.finish(); + const extra = createOMResultTape(entries, () => {}); + const instrumented = extra.instrument(model(), "reflector"); + await collect((await instrumented.doStream({})).stream); + await expect(instrumented.doStream({})).rejects.toBeInstanceOf( + MastraOMDivergenceError, + ); + const missing = createOMResultTape(entries, () => {}); + await expect(missing.finish()).rejects.toBeInstanceOf( + MastraOMDivergenceError, + ); +}); + +it("captures non-JSON OM stream chunks without encoding codec tags twice", async () => { + const chunks = [ + { + date: new Date("2026-09-24T00:00:00.000Z"), + absent: undefined, + bytes: new Uint8Array([1, 2, 3]), + }, + ]; + const tape = createOMResultTape(undefined, () => { + throw new Error("unexpected incomplete result"); + }); + const wrapped = tape.instrument( + { + doStream: async (_input: unknown) => ({ + stream: new ReadableStream({ + start(controller) { + for (const chunk of chunks) controller.enqueue(chunk); + controller.close(); + }, + }), + }), + }, + "observer", + ); + const output = await wrapped.doStream({}); + expect(await collect(output.stream)).toEqual(chunks); + const entries = await tape.finish(); + expect(decodeMemoryValue(entries[0]?.output ?? null)).toEqual(chunks); +}); + +it("stops oversized OM capture while preserving every native stream chunk", async () => { + const total = 10; + let produced = 0; + const onIncomplete = vi.fn(); + const tape = createOMResultTape(undefined, onIncomplete); + const wrapped = tape.instrument( + { + doStream: async (_input: unknown) => ({ + stream: new ReadableStream( + { + pull(controller) { + if (produced === total) { + controller.close(); + } else { + controller.enqueue( + Array.from({ length: 60_000 }, (_, index) => index), + ); + produced += 1; + } + }, + }, + { highWaterMark: 0 }, + ), + }), + }, + "observer", + ); + const output = await wrapped.doStream({}); + await expect(tape.finish()).rejects.toThrow(/incomplete/); + expect(onIncomplete).toHaveBeenCalledTimes(1); + expect(produced).toBeLessThan(total); + const native = await collect(output.stream); + expect(native).toHaveLength(total); + expect((native[9] as number[])[59_999]).toBe(59_999); +}); + +it("stops OM capture at the aggregate byte limit without truncating native output", async () => { + const chunk = "x".repeat(1_048_576); + const total = 20; + let produced = 0; + const onIncomplete = vi.fn(); + const tape = createOMResultTape(undefined, onIncomplete); + const wrapped = tape.instrument( + { + doStream: async (_input: unknown) => ({ + stream: new ReadableStream( + { + pull(controller) { + if (produced === total) { + controller.close(); + } else { + controller.enqueue(chunk); + produced += 1; + } + }, + }, + { highWaterMark: 0 }, + ), + }), + }, + "observer", + ); + const output = await wrapped.doStream({}); + await expect(tape.finish()).rejects.toThrow(/incomplete/); + expect(onIncomplete).toHaveBeenCalledTimes(1); + expect(produced).toBeLessThan(total); + const native = await collect(output.stream); + expect(native).toHaveLength(total); + expect(native[total - 1]).toBe(chunk); +}); diff --git a/packages/mastra/test/processor-replay.test.ts b/packages/mastra/test/processor-replay.test.ts index a36f17d62..7c7c17253 100644 --- a/packages/mastra/test/processor-replay.test.ts +++ b/packages/mastra/test/processor-replay.test.ts @@ -34,6 +34,7 @@ afterEach(() => { }); it("runs the native file processor with historical bytes, skills and complete large request evidence", async () => { + const signedFileUrl = `${FILE_URL}?token=HISTORICAL_SECRET`; const directory = await mkdtemp(join(tmpdir(), "kitaru-stateful-files-")); await mkdir(join(directory, "triage")); await writeFile( @@ -91,7 +92,11 @@ it("runs the native file processor with historical bytes, skills and complete la id: "file-content", async processInput({ messages }) { processFile(); - const content = await resolveFile(FILE_URL); + const part = messages + .flatMap((message) => message.content.parts) + .find((item) => item.type === "file"); + if (part?.type !== "file") throw new Error("Missing file"); + const content = await resolveFile(String(part.data)); return messages.map((message) => ({ ...message, content: { @@ -126,7 +131,7 @@ it("runs the native file processor with historical bytes, skills and complete la : id.includes("reflector") ? runtime.reflector.model : model, - files: [FILE_URL], + files: [signedFileUrl], resolveFile: fetchFile, skillsDirectory: directory, }, @@ -140,7 +145,7 @@ it("runs the native file processor with historical bytes, skills and complete la { type: "text", text: "Please read" }, { type: "file", - data: new URL(FILE_URL), + data: new URL(signedFileUrl), mimeType: "application/pdf", }, ], @@ -152,8 +157,21 @@ it("runs the native file processor with historical bytes, skills and complete la }, ); await baseline.consumeStream(); + await vi.waitFor(() => + expect( + api.calls.some( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", + ), + ).toBe(true), + ); const input = api.calls.find( - (call) => call.path === "/api/v1/sessions" && call.method === "POST", + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", )?.body?.inputs; expect( (input as Record)[MEMORY_REPLAY_KEY] @@ -187,12 +205,174 @@ it("runs the native file processor with historical bytes, skills and complete la ), ).toBe(true); expect(JSON.stringify(modelNodes[1]?.inputs).length).toBeGreaterThan(40000); + expect(JSON.stringify(api.calls)).not.toContain("HISTORICAL_SECRET"); } finally { await runtime.store.close(); await rm(directory, { recursive: true, force: true }); } }); +it("keeps signed URLs native while every recorded node and output uses file references", async () => { + const signed = "https://files.invalid/report.pdf?token=NATIVE_SECRET"; + const runtime = createMemoryRuntime({ messageTokens: 100000 }); + await seedMemory(runtime); + const api = installTestApi(); + const toolInput = vi.fn(); + const reported = vi.fn(); + const modelRequests: unknown[] = []; + let calls = 0; + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async (args) => { + modelRequests.push(args); + return ++calls === 1 + ? streamParts( + [ + { + type: "tool-call", + toolCallId: "signed-file", + toolName: "readFile", + input: JSON.stringify({ url: signed }), + }, + ], + "tool-calls", + ) + : textStream(`Opened ${signed}`); + }, + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "signed-evidence", + name: "Signed evidence", + instructions: "Read the URL", + memory, + model, + defaultOptions: { maxSteps: 3 }, + tools: { + readFile: createTool({ + id: "readFile", + description: "Read a file", + inputSchema: z.object({ url: z.string() }), + execute: async ({ url }) => { + toolInput(url); + return { opened: url }; + }, + }), + }, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + onRecordingError: reported, + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: () => model, + files: [signed], + resolveFile: async () => ({ + bytes: new Uint8Array([1, 2, 3]), + mediaType: "application/pdf", + }), + }, + ); + try { + const output = await adapter.stream(`Please open ${signed}`, { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await output.consumeStream(); + expect(await output.text).toContain(signed); + expect(JSON.stringify(modelRequests)).toContain("NATIVE_SECRET"); + expect(toolInput).toHaveBeenCalledWith(signed); + await vi.waitFor(() => + expect( + api.calls.some( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", + ), + ).toBe(true), + ); + expect(reported).not.toHaveBeenCalled(); + expect(JSON.stringify(api.calls)).not.toContain("NATIVE_SECRET"); + const tool = api + .nodeBatches() + .flat() + .find((node) => node.node_type === "tool_call"); + expect(JSON.stringify(tool?.inputs)).toContain("kitaru-file://sha256/"); + expect(JSON.stringify(tool?.outputs)).toContain("kitaru-file://sha256/"); + const completion = api.calls.find( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", + ); + expect(JSON.stringify(completion?.body?.outputs)).toContain( + "kitaru-file://sha256/", + ); + } finally { + await runtime.store.close(); + } +}); + +it("marks a recording ineligible when model output contains an uncaptured signed URL", async () => { + const runtime = createMemoryRuntime({ messageTokens: 100000 }); + await seedMemory(runtime); + const api = installTestApi(); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => + textStream("https://unknown.invalid/a?token=UNDECLARED_SECRET"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "unknown-signed-evidence", + name: "Unknown signed evidence", + instructions: "Answer", + model, + memory, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: () => model, + }, + ); + try { + const output = await adapter.stream("Hello", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await output.consumeStream(); + expect(await output.text).toContain("UNDECLARED_SECRET"); + await vi.waitFor(() => + expect( + api.calls.some( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "ineligible", + ), + ).toBe(true), + ); + expect(JSON.stringify(api.calls)).not.toContain("UNDECLARED_SECRET"); + } finally { + await runtime.store.close(); + } +}); + it.each(["late", "memory-name-spoof", "copied-memory-id"])( "applies history failure to a %s processor tool before it can execute", async (kind) => { diff --git a/packages/mastra/test/request-capture.test.ts b/packages/mastra/test/request-capture.test.ts index 261e93ef1..9e785410d 100644 --- a/packages/mastra/test/request-capture.test.ts +++ b/packages/mastra/test/request-capture.test.ts @@ -173,6 +173,7 @@ it("uses public tags and message sources for provenance", () => { messageList.addSystem("remembered", "memory"); messageList.add([{ role: "user", content: "old" }], "memory"); messageList.add([{ role: "user", content: "extra" }], "context"); + const serialize = vi.spyOn(messageList, "serializeForSpan"); const capture = createRequestCapture({ invocationId: "sources", getMemoryRevision: () => 0, @@ -190,9 +191,42 @@ it("uses public tags and message sources for provenance", () => { systemMessages: [{ tag: "application" }, { tag: "memory" }], messages: [{ source: "memory" }, { source: "context" }], }); + expect(JSON.stringify(evidence.provenance)).not.toContain( + '"content":"old"', + ); + expect(JSON.stringify(evidence.provenance)).not.toContain( + '"content":"extra"', + ); + expect(serialize).not.toHaveBeenCalled(); }); }); +it("keeps a 500 KiB actor request complete without copying its text into provenance", async () => { + const capture = createRequestCapture({ + invocationId: "large-request", + getMemoryRevision: () => 0, + }); + const list = new MessageList(); + const content = "x".repeat(500 * 1024); + list.add([{ role: "user", content }], "memory"); + capture.beginStep({ + stepNumber: 0, + messageList: list, + extraContext: { context: content }, + }); + const model = capture.instrumentModel({ + specificationVersion: "v2", + modelId: "actor", + provider: "fixture", + doGenerate: async (_args: unknown) => "native", + }); + await model.doGenerate({ prompt: [{ role: "user", content }] }); + const evidence = required(capture.takeSuccessful()); + expect(evidence.complete).toBe(true); + expect(JSON.stringify(evidence.inputs)).toContain(content); + expect(JSON.stringify(evidence.provenance).length).toBeLessThan(1_000); +}); + it("contains capture and telemetry failures while preserving native errors", async () => { const error = new Error("provider failed"); const onCaptureError = vi.fn(() => { diff --git a/packages/mastra/test/stateful-files.test.ts b/packages/mastra/test/stateful-files.test.ts index aeaf8238b..13192c003 100644 --- a/packages/mastra/test/stateful-files.test.ts +++ b/packages/mastra/test/stateful-files.test.ts @@ -1,25 +1,187 @@ import { expect, it, vi } from "vitest"; +import { + decodeMemoryValue, + encodeMemoryValue, +} from "../src/memory-snapshot.js"; import { createCapturedFiles, + fileReference, restoreCapturedFiles, } from "../src/stateful-files.js"; -it("preloads once and supplies independent historical bytes without network fallback", async () => { +it("captures signed URLs under secret-free references and replays immutable bytes", async () => { const resolver = vi.fn(async () => ({ bytes: new Uint8Array([0, 255, 1]), mediaType: "application/pdf", })); - const captured = await createCapturedFiles( - ["https://files.invalid/a", "https://files.invalid/a"], - resolver, - ); + const signedUrl = "https://files.invalid/a?token=PRIVATE_TOKEN"; + const captured = await createCapturedFiles([signedUrl, signedUrl], resolver); + const reference = captured.referenceFor(signedUrl); + expect(reference).toMatch(/^kitaru-file:\/\/sha256\/[a-f0-9]{64}$/); + expect(JSON.stringify(captured.files)).not.toContain("PRIVATE_TOKEN"); + expect(JSON.stringify(captured.files)).not.toContain("files.invalid"); const replay = restoreCapturedFiles(captured.files); - (await captured.resolveFile("https://files.invalid/a")).bytes[0] = 9; - expect((await replay.resolveFile("https://files.invalid/a")).bytes).toEqual( + (await captured.resolveFile(signedUrl)).bytes[0] = 9; + expect((await replay.resolveFile(reference)).bytes).toEqual( new Uint8Array([0, 255, 1]), ); - await expect( - replay.resolveFile("https://files.invalid/missing"), - ).rejects.toThrow(/not recorded/); + await expect(replay.resolveFile(signedUrl)).rejects.toThrow(/not recorded/); expect(resolver).toHaveBeenCalledTimes(1); }); + +it("rewrites declared file and image URLs without changing baseline input", async () => { + const signedUrl = "https://files.invalid/a?X-Amz-Signature=SECRET"; + const captured = await createCapturedFiles([signedUrl], async () => ({ + bytes: new Uint8Array([1, 2]), + mediaType: "image/png", + })); + const baseline = { + messages: [ + { role: "user", content: [{ type: "image", image: new URL(signedUrl) }] }, + ], + context: { attachment: signedUrl }, + }; + const historical = captured.replaceDeclaredFileUrls(baseline); + expect(baseline.messages[0]?.content[0]?.image.href).toBe(signedUrl); + expect(historical.messages[0]?.content[0]?.image.href).toBe( + captured.referenceFor(signedUrl), + ); + expect(historical.context.attachment).toBe(captured.referenceFor(signedUrl)); + expect(JSON.stringify(historical)).not.toContain("SECRET"); + await expect( + Promise.resolve().then(() => + captured.replaceDeclaredFileUrls({ + type: "file", + data: "https://files.invalid/undeclared?token=SECRET", + }), + ), + ).rejects.toThrow(/undeclared file URL/); +}); + +it("scrubs declared URLs embedded in evidence text and flags unknown signed URLs", async () => { + const signed = "https://files.invalid/a?token=KNOWN_SECRET"; + const captured = await createCapturedFiles([signed], async () => ({ + bytes: new Uint8Array([1]), + mediaType: "image/png", + })); + const unknown = vi.fn(); + const sanitizer = captured.evidenceSanitizer(unknown); + const native = { + prompt: `Open ${signed} then https://other.invalid/b?token=UNKNOWN_SECRET`, + tool: { result: signed }, + }; + const persisted = sanitizer.replace(native); + expect(native.prompt).toContain("KNOWN_SECRET"); + expect(persisted.prompt).toContain(captured.referenceFor(signed)); + expect(persisted.tool.result).toBe(captured.referenceFor(signed)); + expect(JSON.stringify(persisted)).not.toMatch(/KNOWN_SECRET|UNKNOWN_SECRET/); + expect(unknown).toHaveBeenCalledTimes(1); +}); + +it("preserves Date and bytes through signed URL evidence projection", async () => { + const signed = "https://files.invalid/photo?token=SECRET"; + const captured = await createCapturedFiles([signed], async () => ({ + bytes: new Uint8Array([0, 255]), + mediaType: "image/png", + })); + const unknown = vi.fn(); + const date = new Date("2025-01-01T00:00:00.000Z"); + const bytes = new Uint8Array([1, 0, 255]); + const value = captured.evidenceSanitizer(unknown).replace({ + date, + bytes, + url: signed, + }); + expect(decodeMemoryValue(encodeMemoryValue(value))).toEqual({ + date, + bytes, + url: captured.referenceFor(signed), + }); + expect(unknown).not.toHaveBeenCalled(); + const unsupported = captured.evidenceSanitizer(unknown).replace({ + value: new Map([["key", signed]]), + }); + expect(unsupported.value).toBe("[unrecordable evidence value]"); + expect(unknown).toHaveBeenCalledOnce(); +}); + +it("rejects altered file bytes and hides resolver errors containing URL secrets", async () => { + const url = "https://files.invalid/a?token=SECRET"; + const captured = await createCapturedFiles([url], async () => ({ + bytes: new Uint8Array([1]), + mediaType: "image/png", + })); + const file = captured.files[0]; + if (!file) throw new Error("Missing captured file"); + file.bytes[0] = 2; + expect(() => restoreCapturedFiles(captured.files)).toThrow( + /invalid recorded file/, + ); + await expect( + createCapturedFiles([url], async () => { + throw new Error(`Failed to fetch ${url}`); + }), + ).rejects.toThrow(/^Controlled file capture failed\.$/); +}); + +it("rejects cyclic or accessor-backed replay input without reading the accessor", async () => { + const captured = await createCapturedFiles([], async () => { + throw new Error("unexpected resolver call"); + }); + const cycle: unknown[] = []; + cycle.push(cycle); + expect(() => captured.replaceDeclaredFileUrls(cycle)).toThrow(/circular/); + const getter = vi.fn(() => "https://files.invalid/a?token=SECRET"); + const part = Object.defineProperty({ type: "file" }, "data", { + enumerable: true, + get: getter, + }); + expect(() => captured.replaceDeclaredFileUrls(part)).toThrow(/accessor/); + expect(getter).not.toHaveBeenCalled(); +}); + +it("rejects a file over 8 MiB before copying it", async () => { + const bytes = new Uint8Array(8 * 1024 * 1024 + 1); + await expect( + createCapturedFiles(["https://files.invalid/large"], async () => ({ + bytes, + mediaType: "application/pdf", + })), + ).rejects.toThrow(/8 MiB/); + expect(() => + restoreCapturedFiles([ + { + url: fileReference({ bytes, mediaType: "application/pdf" }), + bytes, + mediaType: "application/pdf", + }, + ]), + ).toThrow(/invalid recorded file/); +}); + +it("rejects more than 16 MiB or 64 declared files before retaining them", async () => { + const bytes = new Uint8Array(6 * 1024 * 1024); + const resolveFile = vi.fn(async () => ({ bytes, mediaType: "image/png" })); + await expect( + createCapturedFiles( + [ + "https://files.invalid/1", + "https://files.invalid/2", + "https://files.invalid/3", + ], + resolveFile, + ), + ).rejects.toThrow(/16 MiB/); + expect(resolveFile).toHaveBeenCalledTimes(3); + resolveFile.mockClear(); + await expect( + createCapturedFiles( + Array.from( + { length: 65 }, + (_, index) => `https://files.invalid/${index}`, + ), + resolveFile, + ), + ).rejects.toThrow(/count limit/); + expect(resolveFile).not.toHaveBeenCalled(); +}); diff --git a/packages/mastra/test/stateful-overrides.test.ts b/packages/mastra/test/stateful-overrides.test.ts index 8fed59648..c83a21ee2 100644 --- a/packages/mastra/test/stateful-overrides.test.ts +++ b/packages/mastra/test/stateful-overrides.test.ts @@ -1,5 +1,9 @@ import { Agent } from "@mastra/core/agent"; import type { InputProcessor } from "@mastra/core/processors"; +import { + MASTRA_AUTH_TOKEN_KEY, + RequestContext, +} from "@mastra/core/request-context"; import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; import { afterEach, expect, it, vi } from "vitest"; import { @@ -99,6 +103,128 @@ it("retains invocation memory tool identity through public native conversion", a await runtime.store.close(); }); +it("executes a second-step memory tool during replay", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const { createTool } = await import("@mastra/core/tools"); + const { z } = await import("zod/v4"); + const access = createProcessLocalMemoryAccess(); + const api = installTestApi({ + replaySpec: { + id: REPLAY_ID, + baseline_session_id: ORIGINAL_SESSION_ID, + status: "pending", + override: null, + tool_policy: { + default: { type: "history", scope: "baseline", on_miss: "fail" }, + tools: { advance: { type: "passthrough" } }, + }, + }, + }); + const advance = vi.fn(async () => ({ next: true })); + let calls = 0; + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => { + const step = ++calls % 3; + if (step === 1) + return streamParts( + [ + { + type: "tool-call", + toolCallId: `advance-${calls}`, + toolName: "advance", + input: "{}", + }, + ], + "tool-calls", + ); + if (step === 2) + return streamParts( + [ + { + type: "tool-call", + toolCallId: `memory-${calls}`, + toolName: "updateWorkingMemory", + input: JSON.stringify({ memory: { preference: "green" } }), + }, + ], + "tool-calls", + ); + return textStream("done"); + }, + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "second-step-memory", + name: "Second-step memory", + instructions: "Update memory after advancing", + model, + memory, + defaultOptions: { maxSteps: 4 }, + tools: { + advance: createTool({ + id: "advance", + description: "Advance to the next step", + inputSchema: z.object({}), + execute: advance, + }), + }, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: access, + }), + resolveModel: (id) => + id.includes("observer") + ? runtime.observer.model + : id.includes("reflector") + ? runtime.reflector.model + : model, + }, + ); + try { + const baseline = await adapter.stream("Remember green", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await baseline.consumeStream(); + await vi.waitFor(() => + expect( + api.calls.some( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", + ), + ).toBe(true), + ); + const input = api.calls.find( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", + )?.body?.inputs; + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(input)); + const replay = await adapter.stream("ignored"); + await replay.consumeStream(); + expect(await replay.text).toBe("done"); + expect(advance).toHaveBeenCalledTimes(2); + expect(api.calls.some((call) => call.path.endsWith("/tool-lookup"))).toBe( + false, + ); + } finally { + await runtime.store.close(); + } +}); + import { createMemoryReplayAgent, createProcessLocalMemoryAccess, @@ -206,14 +332,40 @@ it("records and replays native evolving memory without re-resolving live configu }); await baseline.consumeStream(); expect(dynamicCalls).toEqual(["instructions", "model", "defaults"]); + await vi.waitFor(() => + expect( + baselineApi.calls.some( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", + ), + ).toBe(true), + ); const recorded = baselineApi.calls.find( - (call) => call.method === "POST" && call.path === "/api/v1/sessions", + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", )?.body?.inputs; expect(recorded).toHaveProperty(MEMORY_REPLAY_KEY); expect( - baselineApi.calls.filter((call) => call.method === "PATCH").at(-1)?.body - ?.status, - ).toBe("completed"); + baselineApi.calls.find( + (call) => call.method === "POST" && call.path === "/api/v1/sessions", + )?.body?.metadata, + ).toMatchObject({ + mastra_replay_state: "pending", + mastra_native_state: "pending", + }); + expect( + baselineApi.calls.filter((call) => call.method === "PATCH").at(-1)?.body, + ).toMatchObject({ + status: "completed", + metadata: { + mastra_replay_state: "eligible", + mastra_native_state: "completed", + }, + }); const baselineNodes = baselineApi.nodeBatches().flat(); expect(baselineNodes.some((node) => node.name === "memory_mutation")).toBe( true, @@ -225,9 +377,10 @@ it("records and replays native evolving memory without re-resolving live configu ).toBe(true); expect( baselineNodes.find((node) => node.node_type === "llm_call")?.attributes, - ).toHaveProperty("prompt_provenance.extraContext.context", [ - { role: "system", content: "Extra context. Original instructions" }, - ]); + ).toHaveProperty( + "prompt_provenance.extraContextRef", + "mastra_memory_replay.configuration", + ); await runtime.memory.updateWorkingMemory({ threadId: THREAD, resourceId: RESOURCE, @@ -269,7 +422,7 @@ it("records and replays native evolving memory without re-resolving live configu await runtime.store.close(); }); -it("waits for owned native observation before reporting completion and releases its lease", async () => { +it("finishes the native stream before delayed observation and finalizes replay later", async () => { let release!: () => void; let started!: () => void; const blocked = new Promise((resolve) => { @@ -326,17 +479,44 @@ it("waits for owned native observation before reporting completion and releases (call) => call.method === "PATCH" && call.body?.status === "completed", ), ).toBe(false); - release(); await consuming; + expect(await result.text).toBe("done"); expect( - api.calls.filter((call) => call.method === "PATCH").at(-1)?.body?.status, - ).toBe("completed"); + api.calls.some( + (call) => call.method === "PATCH" && call.body?.status === "completed", + ), + ).toBe(false); + release(); + await vi.waitFor(() => + expect( + api.calls.filter((call) => call.method === "PATCH").at(-1)?.body, + ).toMatchObject({ + status: "completed", + metadata: { mastra_replay_state: "eligible" }, + }), + ); const releaseLease = await lease.acquire({ threadId: THREAD, resourceId: RESOURCE, }); await releaseLease(); expect(runtime.observer.calls.length).toBeGreaterThan(0); + const observerCalls = runtime.observer.calls.length; + const reflectorCalls = runtime.reflector.calls.length; + const finalizedInputs = api.calls.find( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", + )?.body?.inputs; + expect(finalizedInputs).toHaveProperty(MEMORY_REPLAY_KEY); + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(finalizedInputs)); + const replay = await adapter.stream("ignored"); + await replay.consumeStream(); + expect(await replay.text).toBe("done"); + expect(runtime.observer.calls).toHaveLength(observerCalls); + expect(runtime.reflector.calls).toHaveLength(reflectorCalls); expect( api .nodeBatches() @@ -368,10 +548,11 @@ it("keeps baseline output when initial snapshot evidence fails and rejects its i configuration: runtime.memory.getMergedThreadConfig(), exclusiveAccess: createProcessLocalMemoryAccess(), })); + const modelCalls = vi.fn(async () => textStream("native output")); const model = new MastraLanguageModelV2Mock({ modelId: "actor", provider: "fixture", - doStream: async () => textStream("native output"), + doStream: modelCalls, }); const adapter = createMemoryReplayAgent( ({ memory }) => ({ @@ -398,14 +579,364 @@ it("keeps baseline output when initial snapshot evidence fails and rejects its i (call) => call.method === "POST" && call.path === "/api/v1/sessions", )?.body?.inputs as Record; expect(input[MEMORY_REPLAY_KEY]?.complete).toBe(false); + await vi.waitFor(() => + expect( + api.calls.find( + (call) => + call.method === "PATCH" && + call.body?.status === "failed" && + call.body?.metadata && + (call.body.metadata as Record) + .mastra_replay_state === "ineligible", + )?.body?.metadata, + ).toMatchObject({ + mastra_replay_state: "ineligible", + mastra_replay_reason: "memory_evidence_incomplete", + mastra_native_state: "completed", + }), + ); vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(input)); - await expect(adapter.stream("ignored")).rejects.toThrow(/incomplete/i); + await expect(adapter.stream("ignored")).rejects.toThrow( + /complete version-3/i, + ); expect(source).toHaveBeenCalledTimes(1); + expect(modelCalls).toHaveBeenCalledTimes(1); + await runtime.store.close(); +}); + +it("runs natively when request context cannot be captured safely", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const api = installTestApi(); + const access = createProcessLocalMemoryAccess(); + const unsafeWrite = vi.spyOn(access, "markUnsafeWrite"); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("native answer"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "unsupported-context", + name: "Unsupported context", + instructions: "Answer", + memory, + model, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: access, + }), + resolveModel: () => model, + captureRequestContext: () => ({ [MASTRA_AUTH_TOKEN_KEY]: "secret" }), + }, + ); + const output = await adapter.stream("Hello", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await output.consumeStream(); + expect(await output.text).toBe("native answer"); + await vi.waitFor(() => + expect( + api.calls.find( + (call) => call.method === "POST" && call.path === "/api/v1/sessions", + )?.body?.metadata, + ).toMatchObject({ + mastra_replay_state: "ineligible", + mastra_replay_reason: "context_unsupported", + }), + ); + expect(JSON.stringify(api.calls)).not.toContain("secret"); + expect(unsafeWrite).toHaveBeenCalledWith({ + threadId: THREAD, + resourceId: RESOURCE, + }); + const next = await access.acquire({ threadId: THREAD, resourceId: RESOURCE }); + expect(await next.verifyEligibility()).toBe(false); + await next(); + await runtime.memory.settled(); + await runtime.store.close(); +}); + +it("poisons all source threads when a native fallback has only an implicit selector", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + installTestApi(); + const access = createProcessLocalMemoryAccess(); + const unsafeWrite = vi.spyOn(access, "markUnsafeWrite"); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("native answer"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "implicit-native-fallback", + name: "Implicit native fallback", + instructions: "Answer", + memory, + model, + defaultOptions: { + memory: { thread: THREAD, resource: RESOURCE }, + }, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: access, + }), + resolveModel: () => model, + }, + ); + try { + const output = await adapter.stream("Hello"); + await output.consumeStream(); + expect(await output.text).toBe("native answer"); + expect(unsafeWrite).toHaveBeenCalledWith(undefined); + const next = await access.acquire({ + threadId: THREAD, + resourceId: RESOURCE, + }); + expect(await next.verifyEligibility()).toBe(false); + await next(); + } finally { + await runtime.store.close(); + } +}); + +it("keeps the native answer moving when unsafe-write coordination hangs", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + installTestApi(); + const reported = vi.fn(); + const access = { + ...createProcessLocalMemoryAccess(), + markUnsafeWrite: vi.fn(() => new Promise(() => undefined)), + }; + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("native answer"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "hung-unsafe-marker", + name: "Hung unsafe marker", + instructions: "Answer", + memory, + model, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + onRecordingError: reported, + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: access, + }), + resolveModel: () => model, + captureRequestContext: () => ({ [MASTRA_AUTH_TOKEN_KEY]: "secret" }), + }, + ); + try { + const output = await adapter.stream("Hello", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await output.consumeStream(); + expect(await output.text).toBe("native answer"); + expect(access.markUnsafeWrite).toHaveBeenCalledOnce(); + expect(reported).toHaveBeenCalledWith( + expect.objectContaining({ + error: expect.objectContaining({ + message: "Source-thread unsafe marker timed out.", + }), + }), + ); + } finally { + await runtime.store.close(); + } +}, 2000); + +it("classifies an unrelated resource setup error as capture setup failure", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const api = installTestApi(); + let factoryCalls = 0; + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("native answer"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => { + if (++factoryCalls === 1) + throw new Error("resource registry temporarily unavailable"); + return { + id: "resource-setup-failure", + name: "Resource setup failure", + instructions: "Answer", + model, + memory, + }; + }, + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: () => model, + }, + ); + try { + const output = await adapter.stream("Hello", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await output.consumeStream(); + expect(await output.text).toBe("native answer"); + await vi.waitFor(() => + expect( + api.calls.find( + (call) => call.method === "POST" && call.path === "/api/v1/sessions", + )?.body?.metadata, + ).toMatchObject({ mastra_replay_reason: "capture_setup_failed" }), + ); + } finally { + await runtime.store.close(); + } +}); + +it("passes the original request context to baseline resolvers while recording only approved values", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + const api = installTestApi(); + const seen: unknown[] = []; + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("native answer"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "native-context", + name: "Native context", + instructions: ({ requestContext }) => { + seen.push(requestContext.get("accessToken")); + return "Answer"; + }, + memory, + model, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: () => model, + captureRequestContext: (context) => ({ locale: context.get("locale") }), + }, + ); + const context = new RequestContext(); + context.set("accessToken", "private-token"); + context.set("locale", "nl"); + const output = await adapter.stream("Hello", { + memory: { thread: THREAD, resource: RESOURCE }, + requestContext: context, + }); + await output.consumeStream(); + expect(await output.text).toBe("native answer"); + expect(seen).toContain("private-token"); + expect(JSON.stringify(api.calls)).not.toContain("private-token"); + const recordedInputs = api.calls.find( + (call) => call.method === "POST" && call.path === "/api/v1/sessions", + )?.body?.inputs as Record | undefined; + expect(recordedInputs?.[MEMORY_REPLAY_KEY]?.requestContext).toMatchObject({ + locale: "nl", + }); + await runtime.memory.settled(); await runtime.store.close(); }); -it("rebuilds observation/reflection on the replay trajectory with an overridden actor model", async () => { +it("runs natively and reports locally when Kitaru session creation fails", async () => { + const runtime = createMemoryRuntime({ messageTokens: 10000 }); + await seedMemory(runtime); + installTestApi(); + const reported = vi.fn(); + const originalFetch = globalThis.fetch; + vi.stubGlobal( + "fetch", + vi.fn(async (input, init) => { + if ( + init?.method === "POST" && + new URL(String(input)).pathname === "/api/v1/sessions" + ) + throw new Error("Kitaru unavailable"); + return originalFetch(input, init); + }), + ); + const model = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("native answer"), + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "outage", + name: "Outage", + instructions: "Answer", + memory, + model, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + onRecordingError: reported, + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => runtime.memory.settled(), + domain: runtime.domain, + configuration: runtime.memory.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: () => model, + }, + ); + const output = await adapter.stream("Hello", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await output.consumeStream(); + expect(await output.text).toBe("native answer"); + await vi.waitFor(() => expect(reported).toHaveBeenCalledTimes(1)); + expect(reported.mock.calls[0]?.[0]).toMatchObject({ stage: "complete" }); + await runtime.memory.settled(); + await runtime.store.close(); +}); + +it("reports OM divergence when an override adds a new memory call", async () => { const runtime = createMemoryRuntime({ messageTokens: 600 }); await seedMemory(runtime); const api = installTestApi({ @@ -508,9 +1039,24 @@ it("rebuilds observation/reflection on the replay trajectory with an overridden memory: { thread: THREAD, resource: RESOURCE }, }); await baseline.consumeStream(); + await vi.waitFor(() => + expect( + api.calls.some( + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", + ), + ).toBe(true), + ); const input = api.calls.find( - (call) => call.method === "POST" && call.path === "/api/v1/sessions", + (call) => + call.method === "PATCH" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "eligible", )?.body?.inputs; + const observerCalls = runtime.observer.calls.length; + const reflectorCalls = runtime.reflector.calls.length; vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(input)); const nativeStream = Agent.prototype.stream; @@ -526,27 +1072,20 @@ it("rebuilds observation/reflection on the replay trajectory with an overridden expect(output).toBe(nativeResult); await output.consumeStream(); expect(await output.text).toBe("evolved"); - expect(requests).toHaveLength(3); - expect(runtime.observer.calls.length).toBeGreaterThan(0); - expect(runtime.reflector.calls.length).toBeGreaterThan(0); - expect(JSON.stringify(requests[2])).toContain("REFLECTED_REPLAY"); - const resultNodes = api.nodeBatches(api.sessionIds[1]).flat(); - expect( - resultNodes - .filter((node) => node.node_type === "llm_call") - .every((node) => node.model === "replacement"), - ).toBe(true); - expect( - resultNodes.some( - (node) => - node.name === "memory_mutation" && - (node.attributes as Record).memory_method === - "createReflectionGeneration", - ), - ).toBe(true); - expect( - api.calls.filter((call) => call.method === "PATCH").at(-1)?.body?.status, - ).toBe("completed"); + expect(runtime.observer.calls).toHaveLength(observerCalls); + expect(runtime.reflector.calls).toHaveLength(reflectorCalls); + await vi.waitFor(() => + expect( + api.calls.filter((call) => call.method === "PATCH").at(-1)?.body, + ).toMatchObject({ + status: "failed", + error: "KITARU_REPLAY_DIVERGED:mastra_om_call_order", + metadata: { + mastra_replay_state: "diverged", + mastra_replay_reason: "mastra_om_call_order", + }, + }), + ); expect( await runtime.memory.getWorkingMemory({ threadId: THREAD, @@ -620,7 +1159,7 @@ it("releases the source lease on setup failure and cancellation", async () => { await runtime.store.close(); }); -it("stops later actor and tool work after native memory storage fails", async () => { +it("preserves native continuation after a memory storage write fails", async () => { const runtime = createMemoryRuntime({ messageTokens: 10000 }); await seedMemory(runtime); const api = installTestApi(); @@ -658,7 +1197,7 @@ it("stops later actor and tool work after native memory storage fails", async () ], "tool-calls", ) - : textStream("incorrect continuation"), + : textStream("native continuation"), }); const adapter = createMemoryReplayAgent( ({ memory }) => ({ @@ -691,18 +1230,38 @@ it("stops later actor and tool work after native memory storage fails", async () }, ); try { + const onFinish = vi.fn(); const output = await adapter.stream("Remember green", { memory: { thread: THREAD, resource: RESOURCE }, + onFinish, }); await output.consumeStream(); + expect(await output.text).toBe("native continuation"); + expect(onFinish).toHaveBeenCalledTimes(1); + expect(onFinish.mock.calls[0]?.[0]?.steps?.at(-1)?.finishReason).toBe( + "stop", + ); await vi.waitFor(() => expect( api.calls.filter((call) => call.method === "PATCH").at(-1)?.body ?.status, ).toBe("failed"), ); - expect(actorCalls).toBe(1); - expect(external).not.toHaveBeenCalled(); + expect(actorCalls).toBe(2); + expect(external).toHaveBeenCalledTimes(1); + expect( + api.calls.find( + (call) => + call.method === "PATCH" && + call.body?.status === "failed" && + (call.body?.metadata as Record | undefined) + ?.mastra_replay_state === "ineligible", + )?.body?.metadata, + ).toMatchObject({ + mastra_replay_state: "ineligible", + mastra_replay_reason: "memory_mutation_failed", + mastra_native_state: "completed", + }); const release = await lease.acquire({ threadId: THREAD, resourceId: RESOURCE, diff --git a/packages/mastra/test/stateful-tools.test.ts b/packages/mastra/test/stateful-tools.test.ts new file mode 100644 index 000000000..9948de03d --- /dev/null +++ b/packages/mastra/test/stateful-tools.test.ts @@ -0,0 +1,58 @@ +import type { AdapterRunState } from "@zenml-io/kitaru/adapter"; +import { expect, it, vi } from "vitest"; +import { createStatefulToolProcessors } from "../src/stateful-tools.js"; + +it("trusts a fresh native memory executor on the second replay step", async () => { + const ledger = new Map(); + const lookup = vi.fn(async () => ({ match: null })); + const state = { + client: { lookupToolResult: lookup }, + failure: undefined, + replayId: "replay", + spec: { + tool_policy: { + default: { type: "history", scope: "baseline", on_miss: "fail" }, + tools: {}, + }, + }, + getToolCall: (id: string) => ledger.get(id), + setToolCall: (entry: { callId: string }) => ledger.set(entry.callId, entry), + getHistoryOccurrence: () => 0, + storeFailure: vi.fn(), + } as unknown as AdapterRunState; + const native = vi.fn(async () => ({ updated: true })); + const { first, last } = createStatefulToolProcessors({ + tokens: new Set(["memory-first", "memory-second"]), + getState: () => state, + abort: vi.fn(), + adapter: { agentId: "agent", requestedModelId: "model" }, + }); + const tool = (id: string, execute: (...args: unknown[]) => unknown) => ({ + id, + execute, + }); + await first.processInputStep?.({ + tools: { updateWorkingMemory: tool("memory-first", vi.fn()) }, + } as never); + await last.processInputStep?.({ + tools: { updateWorkingMemory: tool("memory-first", vi.fn()) }, + } as never); + const second = tool("memory-second", native); + await first.processInputStep?.({ + tools: { updateWorkingMemory: second }, + } as never); + const result = (await last.processInputStep?.({ + tools: { updateWorkingMemory: second }, + } as never)) as { + tools: { + updateWorkingMemory: { + execute(input: unknown, context: unknown): Promise; + }; + }; + }; + await expect( + result.tools.updateWorkingMemory.execute({}, { toolCallId: "second" }), + ).resolves.toEqual({ updated: true }); + expect(native).toHaveBeenCalledOnce(); + expect(lookup).not.toHaveBeenCalled(); +}); diff --git a/packages/mastra/test/stateful-workspace.test.ts b/packages/mastra/test/stateful-workspace.test.ts index 087c239bc..b58bd9a08 100644 --- a/packages/mastra/test/stateful-workspace.test.ts +++ b/packages/mastra/test/stateful-workspace.test.ts @@ -15,12 +15,23 @@ it("hashes skills by content and refuses changed content or symlinks", async () ); const baseline = await loadSkillsWorkspace(root); expect(baseline.manifest.files[0]?.path).toBe("triage/SKILL.md"); + expect(baseline.manifest.directories).toEqual([".", "triage"]); expect(await baseline.workspace.skills?.list()).toEqual( expect.arrayContaining([expect.objectContaining({ name: "triage" })]), ); expect( (await loadSkillsWorkspace(root, baseline.manifest)).manifest, ).toEqual(baseline.manifest); + await mkdir(join(root, "extra")); + await expect(loadSkillsWorkspace(root, baseline.manifest)).rejects.toThrow( + /changed/, + ); + await rm(join(root, "extra"), { recursive: true }); + await writeFile(join(root, "triage", "new.md"), "new instructions"); + await expect(loadSkillsWorkspace(root, baseline.manifest)).rejects.toThrow( + /changed/, + ); + await rm(join(root, "triage", "new.md")); await writeFile(path, "changed"); expect( JSON.stringify(await baseline.workspace.skills?.get("triage")), diff --git a/packages/mastra/test/stream-lifecycle.test.ts b/packages/mastra/test/stream-lifecycle.test.ts index 5e55fd74c..7c0425343 100644 --- a/packages/mastra/test/stream-lifecycle.test.ts +++ b/packages/mastra/test/stream-lifecycle.test.ts @@ -749,6 +749,37 @@ describe("stream recording lifecycle", () => { await vi.waitFor(() => expect(reported).toHaveBeenCalledTimes(1)); }); + it("keeps a native provider error ahead of a simultaneous recording failure", async () => { + const api = installTestApi(); + let nodeWrites = 0; + wrapFetch(({ method, path }) => { + if (method === "POST" && path.endsWith("/nodes")) { + nodeWrites += 1; + return nodeWrites === 2; + } + return false; + }); + const nativeError = new Error("native provider failed"); + const agent = Object.assign(new FakeAgent(), { + async stream(_messages: unknown, options: RuntimeStreamOptions = {}) { + await options.onStepFinish?.(textStep("before-provider-failure")); + throw nativeError; + }, + }); + const recorded = new KitaruAgent(agent, { + agentId: AGENT_ID, + apiUrl: "https://api.example", + requestedModelId: "provider-error-model", + }); + + await expect(recorded.stream("hello")).rejects.toBe(nativeError); + expect(api.calls.at(-1)?.body).toMatchObject({ + error: + "Mastra stream failed; KITARU_RECORDING_INCOMPLETE:recording_step_failed", + status: "failed", + }); + }); + it("reports a step write rejected after terminal completion", async () => { installTestApi(); enforceTerminalSessionTransitions(); diff --git a/src/kitaru/api_models/v1/session.py b/src/kitaru/api_models/v1/session.py index 93c040107..cc5b6d937 100644 --- a/src/kitaru/api_models/v1/session.py +++ b/src/kitaru/api_models/v1/session.py @@ -118,6 +118,9 @@ class SessionUpdateRequest(RequestModel): default=None, description="New session status." ) outputs: Any = Field(default=None, description="New session outputs.") + inputs: Any = Field( + default=None, description="Final replay inputs for a pending Mastra recording." + ) output_text_selector: str | None = Field( default=None, description="New output text selector." ) diff --git a/src/kitaru/cli/app.py b/src/kitaru/cli/app.py index fea3f2ad5..3ec407c06 100644 --- a/src/kitaru/cli/app.py +++ b/src/kitaru/cli/app.py @@ -99,7 +99,11 @@ from kitaru.client.control_plane import ControlPlaneLoginError from kitaru.client.credential_store import CredentialStore from kitaru.client.device_grant import DeviceLoginError -from kitaru.client.exceptions import APIError, InvalidServerResponseError +from kitaru.client.exceptions import ( + APIError, + InvalidServerResponseError, + parse_mastra_replay_reason, +) F = TypeVar("F", bound=Callable[..., Any]) _MACHINE_TRUE = {"1", "true", "yes", "on"} @@ -5282,6 +5286,9 @@ def _convert_error( if exception.status_code == 404: return CLIError("not_found", detail, details=details) if exception.status_code == 409: + reason = parse_mastra_replay_reason(detail) + if reason is not None: + details["reason"] = reason return CLIError("conflict", detail, details=details) if exception.status_code in {400, 413, 422}: return CLIError("invalid_arguments", detail, details=details) diff --git a/src/kitaru/client/exceptions.py b/src/kitaru/client/exceptions.py index 3c90251e8..8bd1bfca6 100644 --- a/src/kitaru/client/exceptions.py +++ b/src/kitaru/client/exceptions.py @@ -13,6 +13,8 @@ # permissions and limitations under the License. """Typed client exceptions.""" +import re + import httpx from kitaru.api_models.v1.agent import AgentResponse @@ -66,6 +68,17 @@ def __init__(self, status_code: int, detail: str) -> None: self.detail = detail +_MASTRA_REPLAY_REASON = re.compile( + r"Session [0-9a-f-]{36}: (mastra_replay_[a-z][a-z0-9_]{0,63})" +) + + +def parse_mastra_replay_reason(detail: str) -> str | None: + """Extract a safe replay refusal code from a session conflict detail.""" + match = _MASTRA_REPLAY_REASON.fullmatch(detail) + return match.group(1) if match is not None else None + + class AuthenticationError(APIError): """Authentication error.""" diff --git a/src/kitaru/mcp/errors.py b/src/kitaru/mcp/errors.py index ad58a08d3..d560706b1 100644 --- a/src/kitaru/mcp/errors.py +++ b/src/kitaru/mcp/errors.py @@ -13,7 +13,7 @@ from pydantic import ValidationError from kitaru.api_models.v1.base import JsonValue -from kitaru.client.exceptions import APIError +from kitaru.client.exceptions import APIError, parse_mastra_replay_reason from kitaru.mcp.connection import ConnectionConfigurationError from kitaru.mcp.models.common import ToolError, ToolResult from kitaru.mcp.redaction import redact, redact_data @@ -156,7 +156,12 @@ def map_exception(error: BaseException) -> MCPToolError: "rate_limited": "The Kitaru server rate limited the request.", "remote_failed": "The Kitaru server failed the request.", } - return MCPToolError(code, messages[code], retryable=retryable) + details = None + if status == 409: + reason = parse_mastra_replay_reason(error.detail) + if reason is not None: + details = {"reason": reason} + return MCPToolError(code, messages[code], retryable=retryable, details=details) if isinstance(error, httpx.TimeoutException): return MCPToolError("timeout", "The Kitaru request timed out.", retryable=True) if isinstance(error, httpx.TransportError): diff --git a/src/kitaru/server/adapters/db/repositories/session_repository.py b/src/kitaru/server/adapters/db/repositories/session_repository.py index e9293ae60..d5bd7a16e 100644 --- a/src/kitaru/server/adapters/db/repositories/session_repository.py +++ b/src/kitaru/server/adapters/db/repositories/session_repository.py @@ -38,6 +38,7 @@ from kitaru.server.adapters.db.orm.investigation_session import ( INVESTIGATION_SESSION_SESSION_ID_FOREIGN_KEY, ) +from kitaru.server.adapters.db.orm.orm_utils import split_payload from kitaru.server.adapters.db.orm.replay import ( REPLAY_BASELINE_SESSION_ID_FOREIGN_KEY, REPLAY_RESULT_SESSION_ID_FOREIGN_KEY, @@ -370,6 +371,27 @@ async def update(self, session: Session) -> Session: ) return row.to_domain(exclude={column.key for column in PAYLOAD_COLUMNS}) + async def finalize_replay_inputs(self, session: Session) -> Session: + """Persist final replay inputs with the terminal session transition. + + Args: + session: Validated pending Mastra session. + + Returns: + Stored session without payloads. + """ + row = await self._get_row(session.id, deferred_columns=PAYLOAD_COLUMNS) + row.apply_domain(session) + row.inputs, row.inputs_blob_id = split_payload(session.inputs) + await self._flush( + { + SESSION_IMPORTED_FROM_EXTERNAL_ID_AGENT_ID_UNIQUE_CONSTRAINT: lambda: ( + self._duplicate_external_id(session) + ) + } + ) + return row.to_domain(exclude={column.key for column in PAYLOAD_COLUMNS}) + async def delete(self, session_id: uuid.UUID) -> None: """Delete a session by id. diff --git a/src/kitaru/server/adapters/rest/mapping/sessions.py b/src/kitaru/server/adapters/rest/mapping/sessions.py index 473f1fbcd..cd1795793 100644 --- a/src/kitaru/server/adapters/rest/mapping/sessions.py +++ b/src/kitaru/server/adapters/rest/mapping/sessions.py @@ -150,6 +150,7 @@ def session_update_to_command(body: SessionUpdateRequest) -> SessionUpdate: values: dict[str, Any] = {} for field in ( "status", + "inputs", "outputs", "output_text_selector", "error", diff --git a/src/kitaru/server/application/interfaces/session_repository.py b/src/kitaru/server/application/interfaces/session_repository.py index 161fc7dc8..7506755f1 100644 --- a/src/kitaru/server/application/interfaces/session_repository.py +++ b/src/kitaru/server/application/interfaces/session_repository.py @@ -131,8 +131,9 @@ async def get_many( async def update(self, session: Session) -> Session: """Persist changes to an existing session. - The inputs are create-only and never written back. The outputs are - written back only when the session's outputs were changed. + The inputs are not written by ordinary updates; the guarded + ``finalize_replay_inputs`` operation is the sole exception. Outputs + are written only when changed. Args: session: Session with modified fields. @@ -148,6 +149,17 @@ async def update(self, session: Session) -> Session: """ ... + async def finalize_replay_inputs(self, session: Session) -> Session: + """Persist guarded replay inputs and terminal session fields together. + + Args: + session: Finalized pending Mastra session. + + Returns: + Stored session without payloads. + """ + ... + async def delete(self, session_id: uuid.UUID) -> None: """Delete a session by id. diff --git a/src/kitaru/server/application/models/session.py b/src/kitaru/server/application/models/session.py index 1a2563957..c891f230d 100644 --- a/src/kitaru/server/application/models/session.py +++ b/src/kitaru/server/application/models/session.py @@ -97,6 +97,7 @@ class SessionUpdate(FrozenModel): status: SessionStatus | None = None outputs: Any = None + inputs: Any = None output_text_selector: str | None = None error: str | None = None ended_at: datetime | None = None diff --git a/src/kitaru/server/application/services/experiment_service.py b/src/kitaru/server/application/services/experiment_service.py index c21b70fdf..28a5fc4ab 100644 --- a/src/kitaru/server/application/services/experiment_service.py +++ b/src/kitaru/server/application/services/experiment_service.py @@ -60,7 +60,10 @@ from kitaru.server.application.services.evaluator_resolution import ( validate_evaluators, ) -from kitaru.server.application.services.replay_pipeline import create_replay_pipelines +from kitaru.server.application.services.replay_pipeline import ( + create_replay_pipelines, + validate_replay_baselines, +) from kitaru.server.application.services.server_analytics import ServerAnalytics from kitaru.server.application.services.task_transitions import TaskTransitions from kitaru.server.domain.base import ValidationError @@ -474,6 +477,7 @@ async def start_run( assert run_spec is not None config.check_capabilities(run_spec.runtime_capabilities) sessions = await self._resolve_cohort_version_sessions(cohort_version.id) + await validate_replay_baselines(sessions, self._payload_store) number = await self._experiment_runs.get_max_number(experiment_id) + 1 run = ExperimentRun( diff --git a/src/kitaru/server/application/services/replay_pipeline.py b/src/kitaru/server/application/services/replay_pipeline.py index 4c74ad750..9724ea6c5 100644 --- a/src/kitaru/server/application/services/replay_pipeline.py +++ b/src/kitaru/server/application/services/replay_pipeline.py @@ -13,8 +13,11 @@ # permissions and limitations under the License. """Replay job and task composition, shared by standalone replays and run fan-out.""" +import re import uuid from collections.abc import Sequence +from datetime import UTC, datetime, timedelta +from typing import Any from kitaru.api_models.v1.job import JobKind, JobStatus from kitaru.api_models.v1.replay import BaselineEvaluationMode @@ -41,11 +44,76 @@ from kitaru.server.domain.job import Job from kitaru.server.domain.replay import Replay from kitaru.server.domain.replay_config import ReplayConfig -from kitaru.server.domain.session import Session +from kitaru.server.domain.session import ( + Session, + SessionReplayNotReady, + mastra_replay_uses_observational_memory, + mastra_replay_v3_complete, +) from kitaru.server.domain.task import AgentTask, EvaluationTask, Task from kitaru.server.utils import hash_params +def _record(value: Any) -> dict[str, Any] | None: + """Return a JSON object, if this value is one.""" + return value if isinstance(value, dict) else None + + +def _check_mastra_replay_ready(baseline: Session) -> None: + """Reject provisional or incomplete Mastra memory input before task creation.""" + if baseline.framework != "mastra": + return + state = baseline.metadata.get("mastra_replay_state") + if state == "pending": + started = baseline.started_at or baseline.created + stale = started is not None and datetime.now(UTC) - started > timedelta( + minutes=30 + ) + raise SessionReplayNotReady( + baseline.id, "mastra_replay_abandoned" if stale else "mastra_replay_pending" + ) + if state == "ineligible": + reason = baseline.metadata.get("mastra_replay_reason") + code = ( + reason + if isinstance(reason, str) and re.fullmatch(r"[a-z][a-z0-9_]{0,63}", reason) + else "capture_incomplete" + ) + raise SessionReplayNotReady(baseline.id, f"mastra_replay_{code}") + inputs = _record(baseline.inputs.value) if baseline.inputs is not None else None + envelope = _record(inputs.get("mastra_memory_replay")) if inputs else None + if envelope is None: + if state == "eligible": + raise SessionReplayNotReady(baseline.id, "mastra_replay_incomplete") + return # Existing history-only Mastra recordings have no memory envelope. + if baseline.status.value == "in_progress" or envelope.get("complete") is not True: + raise SessionReplayNotReady(baseline.id, "mastra_replay_incomplete") + if envelope.get("version") == 3 and not mastra_replay_v3_complete(envelope): + raise SessionReplayNotReady(baseline.id, "mastra_replay_incomplete") + if mastra_replay_uses_observational_memory(envelope) and not isinstance( + envelope.get("omTape"), list + ): + raise SessionReplayNotReady(baseline.id, "mastra_replay_tape_missing") + if state == "eligible" and envelope.get("version") != 3: + raise SessionReplayNotReady(baseline.id, "mastra_replay_incomplete") + + +async def validate_replay_baselines( + baselines: Sequence[Session], payload_store: PayloadStore +) -> None: + """Resolve and validate memory replay inputs before replay side effects. + + Args: + baselines: Sessions selected for replay. + payload_store: Store used to resolve offloaded session inputs. + """ + await payload_store.resolve( + [baseline.inputs for baseline in baselines if baseline.inputs is not None] + ) + for baseline in baselines: + _check_mastra_replay_ready(baseline) + + async def create_replay_pipelines( baselines: Sequence[Session], agent_version_id: uuid.UUID, @@ -96,9 +164,7 @@ async def create_replay_pipelines( if evaluate_baselines: for baseline in baselines: baseline.check_evaluate() - await payload_store.resolve( - [baseline.inputs for baseline in baselines if baseline.inputs is not None] - ) + await validate_replay_baselines(baselines, payload_store) jobs = [Job(owner_id=actor.account.id, kind=JobKind.REPLAY) for _ in baselines] replays = [ Replay( diff --git a/src/kitaru/server/application/services/replay_service.py b/src/kitaru/server/application/services/replay_service.py index 9620c9dc4..6a8ec99cf 100644 --- a/src/kitaru/server/application/services/replay_service.py +++ b/src/kitaru/server/application/services/replay_service.py @@ -53,7 +53,10 @@ resolve_runnable_agent_version, ) from kitaru.server.application.services.evaluator_resolution import validate_evaluators -from kitaru.server.application.services.replay_pipeline import create_replay_pipelines +from kitaru.server.application.services.replay_pipeline import ( + create_replay_pipelines, + validate_replay_baselines, +) from kitaru.server.application.services.server_analytics import ServerAnalytics from kitaru.server.domain.base import ValidationError from kitaru.server.domain.replay import ( @@ -180,6 +183,7 @@ async def create_replay( baseline = await self._sessions.get( command.baseline_session_id, include_payloads=True ) + await validate_replay_baselines([baseline], self._payload_store) agent_version_id = command.agent_version_id if agent_version_id is None: if baseline.agent_version_id is None: diff --git a/src/kitaru/server/application/services/session_service.py b/src/kitaru/server/application/services/session_service.py index fab19221e..774676a9c 100644 --- a/src/kitaru/server/application/services/session_service.py +++ b/src/kitaru/server/application/services/session_service.py @@ -404,12 +404,23 @@ async def update_session( await check_task_attempt(actor, self._tasks) session.check_update() fields = command.model_fields_set + target_status = session.status + if "status" in fields: + if command.status is None: + raise SessionStatusCannotBeCleared(session_id) + target_status = command.status + next_metadata = ( + (command.metadata if command.metadata is not None else {}) + if "metadata" in fields + else session.metadata + ) + session.check_replay_finalization( + status=target_status, + metadata=next_metadata, + inputs=command.inputs, + replacing_inputs="inputs" in fields, + ) if {"status", "outputs", "output_text_selector", "error", "ended_at"} & fields: - target_status = session.status - if "status" in fields: - if command.status is None: - raise SessionStatusCannotBeCleared(session_id) - target_status = command.status session.finish( status=target_status, output_text_selector=command.output_text_selector @@ -437,12 +448,15 @@ async def update_session( AnalyticsEvent.SESSION_COMPLETED, analytics_events.build_session_completed_properties(session), ) + if "inputs" in fields: + session.inputs = Payload.from_json(command.inputs) + await self._payload_store.offload([session.inputs], session.owner_id) if "name" in fields: session.update_name(command.name) if "metadata" in fields: - session.update_metadata( - command.metadata if command.metadata is not None else {} - ) + session.update_metadata(next_metadata) + if "inputs" in fields: + return await self._repository.finalize_replay_inputs(session) return await self._repository.update(session) async def delete_session(self, session_id: uuid.UUID, actor: AuthContext) -> None: diff --git a/src/kitaru/server/domain/session.py b/src/kitaru/server/domain/session.py index e3c5ee82d..d118bea57 100644 --- a/src/kitaru/server/domain/session.py +++ b/src/kitaru/server/domain/session.py @@ -13,6 +13,10 @@ # permissions and limitations under the License. """Session entity, rollups, and errors.""" +import base64 +import binascii +import hashlib +import re import uuid from collections.abc import Iterable from datetime import datetime @@ -201,6 +205,116 @@ def __init__(self, session_id: uuid.UUID) -> None: super().__init__(f"Session {session_id} does not accept updates") +class SessionReplayFinalizationInvalid(ValidationError): + """Raised when a Mastra replay input transition is incomplete or unauthorized.""" + + def __init__(self, session_id: uuid.UUID) -> None: + """Initialize the error. + + Args: + session_id: Id of the session. + """ + super().__init__(f"Session {session_id} has invalid Mastra replay finalization") + + +class SessionReplayNotReady(ConflictError): + """Raised when a Mastra baseline cannot be replayed yet.""" + + def __init__(self, session_id: uuid.UUID, reason: str) -> None: + """Initialize the error. + + Args: + session_id: Id of the baseline session. + reason: Stable replay eligibility reason code. + """ + super().__init__(f"Session {session_id}: {reason}") + self.reason = reason + + +def mastra_replay_uses_observational_memory(envelope: dict[str, Any]) -> bool: + """Return whether a recorded Mastra replay input enables observational memory.""" + config = envelope.get("configuration") + memory = config.get("memoryConfig") if isinstance(config, dict) else None + om = memory.get("observationalMemory") if isinstance(memory, dict) else None + return om is True or (isinstance(om, dict) and om.get("enabled") is not False) + + +def mastra_replay_v3_complete(envelope: dict[str, Any]) -> bool: + """Check the required shape of a finalized Mastra replay input.""" + snapshot = envelope.get("initialSnapshot") + config = envelope.get("configuration") + files = envelope.get("files") + return ( + envelope.get("version") == 3 + and envelope.get("complete") is True + and envelope.get("reasons") == [] + and isinstance(envelope.get("invocationId"), str) + and bool(envelope["invocationId"]) + and "rawInput" in envelope + and isinstance(snapshot, dict) + and isinstance(snapshot.get("threadId"), str) + and isinstance(snapshot.get("resourceId"), str) + and isinstance(snapshot.get("messages"), list) + and isinstance(snapshot.get("records"), list) + and isinstance(config, dict) + and isinstance(config.get("memoryConfig"), dict) + and isinstance(envelope.get("requestContext"), dict) + and isinstance(files, list) + and _mastra_replay_files_complete(files) + and isinstance(envelope.get("omTape"), list) + ) + + +def _mastra_replay_files_complete(files: list[Any]) -> bool: + """Validate bounded file references before publishing replay eligibility.""" + if len(files) > 64: + return False + seen: set[str] = set() + total_bytes = 0 + for file in files: + if not isinstance(file, dict): + return False + url = file.get("url") + media_type = file.get("mediaType") + encoded = file.get("base64") + length = file.get("length") + digest = file.get("sha256") + if ( + not isinstance(url, str) + or url in seen + or not isinstance(media_type, str) + or not media_type + or not isinstance(encoded, str) + or not isinstance(length, int) + or isinstance(length, bool) + or length < 0 + or length > 8 * 1_048_576 + or not isinstance(digest, str) + or re.fullmatch(r"[a-f0-9]{64}", digest) is None + ): + return False + try: + content = base64.b64decode(encoded, validate=True) + reference = ( + "kitaru-file://sha256/" + + hashlib.sha256(media_type.encode() + b"\0" + content).hexdigest() + ) + except (binascii.Error, ValueError, UnicodeEncodeError): + return False + if ( + base64.b64encode(content).decode("ascii") != encoded + or len(content) != length + or hashlib.sha256(content).hexdigest() != digest + or url != reference + ): + return False + total_bytes += length + if total_bytes > 16 * 1_048_576: + return False + seen.add(url) + return True + + class SessionRollups(FrozenModel): """Session rollup deltas.""" @@ -331,6 +445,62 @@ def check_update(self) -> None: if self.status != SessionStatus.IN_PROGRESS: raise SessionNotUpdatable(self.id) + def check_replay_finalization( + self, + status: SessionStatus, + metadata: dict[str, Any], + inputs: Any, + replacing_inputs: bool, + ) -> None: + """Require an atomic, one-time transition for final Mastra replay inputs. + + Args: + status: Session status after the update. + metadata: Session metadata after the update. + inputs: Replacement input value, when supplied. + replacing_inputs: Whether the request explicitly supplied inputs. + """ + is_mastra_recording = ( + self.framework == "mastra" and self.origin == SessionOrigin.RECORDED + ) + if not is_mastra_recording: + if replacing_inputs: + raise SessionReplayFinalizationInvalid(self.id) + return + prior = self.metadata.get("mastra_replay_state") + next_state = metadata.get("mastra_replay_state") + if prior != "pending": + if replacing_inputs: + raise SessionReplayFinalizationInvalid(self.id) + return + if status == SessionStatus.IN_PROGRESS: + if replacing_inputs or next_state != "pending": + raise SessionReplayFinalizationInvalid(self.id) + return + if next_state not in {"eligible", "ineligible"}: + raise SessionReplayFinalizationInvalid(self.id) + if replacing_inputs and ( + not isinstance(inputs, dict) + or not isinstance(inputs.get("mastra_memory_replay"), dict) + ): + raise SessionReplayFinalizationInvalid(self.id) + if next_state == "eligible": + envelope = ( + inputs.get("mastra_memory_replay") + if replacing_inputs and isinstance(inputs, dict) + else None + ) + if ( + status != SessionStatus.COMPLETED + or not isinstance(envelope, dict) + or not mastra_replay_v3_complete(envelope) + ): + raise SessionReplayFinalizationInvalid(self.id) + if mastra_replay_uses_observational_memory(envelope) and not isinstance( + envelope.get("omTape"), list + ): + raise SessionReplayFinalizationInvalid(self.id) + def check_evaluate(self) -> None: """Require the session to currently accept evaluations. diff --git a/tests/cli/test_replays.py b/tests/cli/test_replays.py index 6343530b7..7ed41dc22 100644 --- a/tests/cli/test_replays.py +++ b/tests/cli/test_replays.py @@ -21,9 +21,12 @@ ReplayResponse, ReplayStatus, ) +from kitaru.api_models.v1.session import SessionDetailResponse, SessionOrigin from kitaru.cli import app as app_module from kitaru.cli import replays +from kitaru.cli import sessions as cli_sessions from kitaru.cli.output import CLIError +from kitaru.client.exceptions import APIError class StubReplayClient: @@ -63,12 +66,24 @@ def __init__(self) -> None: updated=now, ) self.create_calls: list[ReplayCreateRequest] = [] + self.create_error: APIError | None = None + self.result_session: SessionDetailResponse | None = None self.create_idempotency_keys: list[str | None] = [] self.list_calls: list[ReplayListParams] = [] self.get_calls: list[uuid.UUID] = [] self.agents = self._Agents(self) self.evaluators = self._Evaluators(self) self.replays = self._Replays(self) + self.sessions = self._Sessions(self) + + class _Sessions: + def __init__(self, owner: "StubReplayClient") -> None: + self.owner = owner + + async def get(self, session_id: uuid.UUID) -> SessionDetailResponse: + assert self.owner.result_session is not None + assert session_id == self.owner.result_session.id + return self.owner.result_session class _Agents: def __init__(self, owner: "StubReplayClient") -> None: @@ -110,6 +125,8 @@ async def create( ) -> ReplayResponse: self.owner.create_calls.append(request) self.owner.create_idempotency_keys.append(idempotency_key) + if self.owner.create_error is not None: + raise self.owner.create_error return self.owner.replay async def list(self, params: ReplayListParams) -> Any: @@ -279,6 +296,119 @@ async def test_list_and_get_preserve_sdk_results() -> None: assert client.get_calls == [client.replay_id, client.replay_id] +async def test_get_completed_replay_links_to_result_session() -> None: + """A CLI replay read retains the session id needed for evidence inspection.""" + client = StubReplayClient() + result_session_id = uuid.uuid4() + client.replay.status = ReplayStatus.COMPLETED + client.replay.result_session_id = result_session_id + + result = await replays.get_replay(client, client.replay_id) + + assert result.item["status"] == "completed" + assert result.item["result_session_id"] == str(result_session_id) + + +@pytest.mark.parametrize( + ("state", "reason"), + [ + ("diverged", "mastra_om_call_order"), + ("failed", "replay_failed"), + ], +) +async def test_get_replay_result_session_exposes_failure_reason( + state: str, reason: str +) -> None: + """The linked result session retains a durable replay failure reason.""" + client = StubReplayClient() + now = datetime.now(UTC) + result_session_id = uuid.uuid4() + client.replay.status = ReplayStatus.FAILED + client.replay.result_session_id = result_session_id + client.result_session = SessionDetailResponse( + id=result_session_id, + owner_id=uuid.uuid4(), + agent_id=client.agent.id, + number=2, + origin=SessionOrigin.REPLAY, + status="failed", + inputs={}, + outputs={}, + metadata={"mastra_replay_state": state, "mastra_replay_reason": reason}, + llm_call_count=0, + tool_call_count=0, + created=now, + updated=now, + ) + + replay = await replays.get_replay(client, client.replay_id) + session = await cli_sessions.get_session(client, result_session_id) + + assert replay.item["result_session_id"] == str(result_session_id) + assert session.item["metadata"] == { + "mastra_replay_state": state, + "mastra_replay_reason": reason, + } + + +@pytest.mark.parametrize( + "reason", + ["mastra_replay_pending", "mastra_replay_abandoned", "mastra_replay_tape_missing"], +) +def test_public_create_reports_safe_mastra_ineligibility_reason( + argv_client: StubReplayClient, + capsys: pytest.CaptureFixture[str], + reason: str, +) -> None: + """A rejected replay exposes the server's safe reason in structured output.""" + client = argv_client + + client.create_error = APIError(409, f"Session {client.baseline_id}: {reason}") + assert ( + app_module.main( + [ + "replay", + "create", + str(client.baseline_id), + "--evaluator", + "quality@3", + "--output", + "json", + ] + ) + == 5 + ) + error = json.loads(capsys.readouterr().err)["error"] + assert error["kind"] == "conflict" + assert error["details"] == {"status_code": 409, "reason": reason} + + +def test_public_create_keeps_unrelated_conflict_generic( + argv_client: StubReplayClient, capsys: pytest.CaptureFixture[str] +) -> None: + """Only the exact safe Mastra reason grammar gets a reason field.""" + client = argv_client + + client.create_error = APIError(409, "unrelated conflict token=secret") + assert ( + app_module.main( + [ + "replay", + "create", + str(client.baseline_id), + "--evaluator", + "quality@3", + "--output", + "json", + ] + ) + == 5 + ) + error = json.loads(capsys.readouterr().err)["error"] + assert error["kind"] == "conflict" + assert error["details"] == {"status_code": 409} + + @pytest.fixture def argv_client(monkeypatch: pytest.MonkeyPatch) -> StubReplayClient: """Route public replay commands through one recording client.""" diff --git a/tests/conftest.py b/tests/conftest.py index 7c65ef3ae..8f7adf9fb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,6 +14,7 @@ """Shared test helpers and in-memory fakes.""" import asyncio +import base64 import hashlib import os import sys @@ -373,6 +374,43 @@ def imported_session( ) +@pytest.fixture +def complete_mastra_memory_replay_inputs() -> dict[str, Any]: + """Build a structurally complete v3 input for server finalization tests.""" + media_type = "text/plain" + content = b"recorded file" + reference = hashlib.sha256(media_type.encode() + b"\0" + content).hexdigest() + return { + "mastra_memory_replay": { + "version": 3, + "complete": True, + "reasons": [], + "invocationId": "test-invocation", + "rawInput": "hello", + "initialSnapshot": { + "threadId": "thread-1", + "resourceId": "resource-1", + "thread": None, + "resource": None, + "messages": [], + "records": [], + }, + "configuration": {"memoryConfig": {}}, + "requestContext": {}, + "files": [ + { + "url": f"kitaru-file://sha256/{reference}", + "mediaType": media_type, + "base64": base64.b64encode(content).decode("ascii"), + "length": len(content), + "sha256": hashlib.sha256(content).hexdigest(), + } + ], + "omTape": [], + } + } + + @pytest.fixture(autouse=True) def worker_api_env( request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch @@ -3106,6 +3144,24 @@ async def update(self, session: Session) -> Session: self._sessions[session.id] = updated return self._copy(updated, include_payloads=False) + async def finalize_replay_inputs(self, session: Session) -> Session: + """Persist final replay input and status in one fake repository update.""" + stored = self._sessions.get(session.id) + if stored is None: + raise SessionNotFound(session.id) + self._check_duplicate_external_id(session) + updated = session.model_copy( + update={ + "created": stored.created, + "updated": _renewed_timestamp(stored.updated), + "outputs": session.outputs + if "outputs" in session.model_fields_set + else stored.outputs, + } + ) + self._sessions[session.id] = updated + return self._copy(updated, include_payloads=False) + async def delete(self, session_id: uuid.UUID) -> None: """Delete a session by id. diff --git a/tests/mcp/test_review_workflows.py b/tests/mcp/test_review_workflows.py index cf827dc1a..fc9c81771 100644 --- a/tests/mcp/test_review_workflows.py +++ b/tests/mcp/test_review_workflows.py @@ -21,6 +21,10 @@ EvaluatorResponse, EvaluatorVersionResponse, ) +from kitaru.api_models.v1.experiment_run import ( + ExperimentRunProgress, + ExperimentRunResponse, +) from kitaru.api_models.v1.info import AuthScheme, ServerInfoResponse from kitaru.api_models.v1.insight import InsightResponse from kitaru.api_models.v1.investigation import ( @@ -29,6 +33,8 @@ ) from kitaru.api_models.v1.job import JobResponse from kitaru.api_models.v1.plugin import PackagePluginSource +from kitaru.api_models.v1.replay import BaselineEvaluationMode, ReplayResponse +from kitaru.api_models.v1.session import SessionDetailResponse, SessionOrigin from kitaru.api_models.v1.tag import ( TagCreateRequest, TagLinkCreateRequest, @@ -1098,6 +1104,265 @@ async def start_run( await handle_workflow_start(_get_state(client), request) +async def test_experiment_run_and_activity_reads_link_replay_to_sessions() -> None: + """The existing MCP workflow exposes eligibility and result-session links.""" + now = datetime.now(UTC) + experiment_id = uuid.uuid4() + cohort_version_id = uuid.uuid4() + agent_version_id = uuid.uuid4() + baseline_id = uuid.uuid4() + result_id = uuid.uuid4() + run = ExperimentRunResponse( + id=uuid.uuid4(), + owner_id=uuid.uuid4(), + experiment_id=experiment_id, + number=1, + status="completed", + cohort_version_id=cohort_version_id, + agent_version_id=agent_version_id, + evaluate_baselines=False, + baseline_evaluation_mode=BaselineEvaluationMode.NONE, + progress=ExperimentRunProgress( + pending=0, evaluating=0, completed=1, failed=0, canceled=0, total=1 + ), + created=now, + updated=now, + ) + replay = ReplayResponse( + id=uuid.uuid4(), + job_id=uuid.uuid4(), + experiment_run_id=run.id, + baseline_session_id=baseline_id, + result_session_id=result_id, + override=None, + tool_policy={"default": {"type": "passthrough"}, "tools": {}}, + evaluators=[], + evaluate_baselines=False, + baseline_evaluation_mode=BaselineEvaluationMode.NONE, + status="completed", + created=now, + updated=now, + ) + baseline = SessionDetailResponse( + id=baseline_id, + owner_id=uuid.uuid4(), + agent_id=uuid.uuid4(), + number=1, + origin="recorded", + status="completed", + inputs={}, + outputs={"text": "native answer"}, + metadata={ + "mastra_native_state": "completed", + "mastra_replay_state": "eligible", + }, + llm_call_count=1, + tool_call_count=0, + created=now, + updated=now, + ) + result_session = baseline.model_copy( + update={ + "id": result_id, + "origin": SessionOrigin.REPLAY, + "outputs": {"text": "replayed answer"}, + "metadata": {}, + } + ) + + async def start_run( + _experiment_id: uuid.UUID, _request: object, idempotency_key: str | None = None + ) -> ExperimentRunResponse: + del idempotency_key + assert _experiment_id == experiment_id + return run + + async def get_run(item_id: uuid.UUID) -> ExperimentRunResponse: + assert item_id == run.id + return run + + async def get_replay(item_id: uuid.UUID) -> ReplayResponse: + assert item_id == replay.id + return replay + + async def get_session(item_id: uuid.UUID) -> SessionDetailResponse: + return {baseline_id: baseline, result_id: result_session}[item_id] + + client = SimpleNamespace( + experiments=SimpleNamespace(start_run=start_run), + experiment_runs=SimpleNamespace(get=get_run), + replays=SimpleNamespace(get=get_replay), + sessions=SimpleNamespace(get=get_session), + ) + server, context = _get_context(client, CapabilityMode.STANDARD) + start = await server.call_tool( + "kitaru_workflow_start", + { + "request": { + "operation": "experiment_run", + "experiment_id": str(experiment_id), + "cohort_version_id": str(cohort_version_id), + "agent_version_id": str(agent_version_id), + } + }, + context, + ) + assert isinstance(start, CallToolResult) + assert start.is_error is False + assert start.structured_content is not None + assert start.structured_content["data"]["result"]["id"] == str(run.id) + + async def read(kind: str, item_id: uuid.UUID) -> dict[str, Any]: + response = await server.call_tool( + "kitaru_activity_read", + {"request": {"operation": "get", "kind": kind, "id": str(item_id)}}, + context, + ) + assert isinstance(response, CallToolResult) + assert response.is_error is False + assert response.structured_content is not None + return cast(dict[str, Any], response.structured_content["data"]) + + assert (await read("experiment_run", run.id))["progress"]["completed"] == 1 + replay_data = await read("replay", replay.id) + assert replay_data["result_session_id"] == str(result_id) + assert replay_data["job_id"] == str(replay.job_id) + assert (await read("session", baseline_id))["metadata"] == { + "mastra_native_state": "completed", + "mastra_replay_state": "eligible", + } + result_data = await read("session", result_id) + assert result_data["outputs"] == {"text": "replayed answer"} + + +@pytest.mark.parametrize( + ("state", "reason"), + [ + ("diverged", "mastra_om_call_order"), + ("failed", "replay_failed"), + ], +) +async def test_activity_result_session_exposes_durable_replay_failure_reason( + state: str, reason: str +) -> None: + """MCP activity reads retain safe result-session diagnostics.""" + now = datetime.now(UTC) + result_session = SessionDetailResponse( + id=uuid.uuid4(), + owner_id=uuid.uuid4(), + agent_id=uuid.uuid4(), + number=2, + origin=SessionOrigin.REPLAY, + status="failed", + inputs={}, + outputs={}, + metadata={"mastra_replay_state": state, "mastra_replay_reason": reason}, + llm_call_count=0, + tool_call_count=0, + created=now, + updated=now, + ) + + async def get_session(item_id: uuid.UUID) -> SessionDetailResponse: + assert item_id == result_session.id + return result_session + + client = SimpleNamespace(sessions=SimpleNamespace(get=get_session)) + server, context = _get_context(client, CapabilityMode.READ_ONLY) + result = await server.call_tool( + "kitaru_activity_read", + { + "request": { + "operation": "get", + "kind": "session", + "id": str(result_session.id), + } + }, + context, + ) + + assert isinstance(result, CallToolResult) + assert result.is_error is False + assert result.structured_content is not None + assert result.structured_content["data"]["metadata"] == { + "mastra_replay_state": state, + "mastra_replay_reason": reason, + } + + +@pytest.mark.parametrize( + "reason", + ["mastra_replay_pending", "mastra_replay_abandoned", "mastra_replay_tape_missing"], +) +async def test_experiment_run_start_reports_safe_mastra_replay_reason( + reason: str, +) -> None: + """MCP preserves a safe eligibility reason when the server rejects a run.""" + baseline_session_id = uuid.uuid4() + + async def start_run( + _experiment_id: uuid.UUID, _request: object, idempotency_key: str | None = None + ) -> object: + del idempotency_key + raise APIError(409, f"Session {baseline_session_id}: {reason}") + + client = SimpleNamespace(experiments=SimpleNamespace(start_run=start_run)) + server, context = _get_context(client, CapabilityMode.STANDARD) + result = await server.call_tool( + "kitaru_workflow_start", + { + "request": { + "operation": "experiment_run", + "experiment_id": str(uuid.uuid4()), + "cohort_version_id": str(uuid.uuid4()), + "agent_version_id": str(uuid.uuid4()), + } + }, + context, + ) + + assert isinstance(result, CallToolResult) + assert result.is_error is True + assert result.structured_content is not None + assert result.structured_content["error"]["code"] == "conflict" + assert result.structured_content["error"]["details"] == {"reason": reason} + assert json.loads(cast(TextContent, result.content[0]).text) == ( + result.structured_content + ) + + +async def test_experiment_run_start_does_not_echo_arbitrary_conflict_detail() -> None: + """An unrelated server conflict remains generic without leaking its detail.""" + + async def start_run( + _experiment_id: uuid.UUID, _request: object, idempotency_key: str | None = None + ) -> object: + del idempotency_key + raise APIError(409, "token=secret") + + client = SimpleNamespace(experiments=SimpleNamespace(start_run=start_run)) + server, context = _get_context(client, CapabilityMode.STANDARD) + result = await server.call_tool( + "kitaru_workflow_start", + { + "request": { + "operation": "experiment_run", + "experiment_id": str(uuid.uuid4()), + "cohort_version_id": str(uuid.uuid4()), + "agent_version_id": str(uuid.uuid4()), + } + }, + context, + ) + + assert isinstance(result, CallToolResult) + assert result.is_error is True + assert result.structured_content is not None + assert result.structured_content["error"]["code"] == "conflict" + assert result.structured_content["error"]["details"] is None + assert "secret" not in cast(TextContent, result.content[0]).text + + async def test_evaluator_resolution_uses_bounded_concurrency() -> None: active = 0 peak = 0 diff --git a/tests/server/test_replay_pipeline.py b/tests/server/test_replay_pipeline.py index 7ec098b3f..f0d4a0b01 100644 --- a/tests/server/test_replay_pipeline.py +++ b/tests/server/test_replay_pipeline.py @@ -15,6 +15,8 @@ import uuid from collections.abc import Sequence +from copy import deepcopy +from typing import Any import pytest @@ -65,7 +67,11 @@ from kitaru.server.domain.plugin import PluginKind, PluginVersion, ScriptPluginSource from kitaru.server.domain.replay import DuplicateReplayForBaseline from kitaru.server.domain.replay_config import ReplayOverride -from kitaru.server.domain.session import Session, SessionNotEvaluatable +from kitaru.server.domain.session import ( + Session, + SessionNotEvaluatable, + SessionReplayNotReady, +) from kitaru.server.domain.task import AgentTask, AgentTaskDetails, EvaluationTask from kitaru.server.filtering import FilterCondition @@ -121,6 +127,172 @@ async def _baseline_session( ) +async def test_pending_mastra_baseline_creates_no_replay( + services: ReplayServices, +) -> None: + """Refuse a pending memory recording before creating jobs or tasks.""" + version = await _agent_version_with_run_spec(services) + baseline = await create_session( + services.sessions, + ACTOR.account.id, + agent_id=version.agent_id, + agent_version_id=version.id, + origin=SessionOrigin.RECORDED, + status=SessionStatus.IN_PROGRESS, + framework="mastra", + inputs={"mastra_memory_replay": {"version": 3, "complete": False}}, + metadata={"mastra_replay_state": "pending"}, + ) + with pytest.raises(SessionReplayNotReady, match="mastra_replay_pending"): + await services.replay_service.create_replay( + ReplayCreate( + baseline_session_id=baseline.id, + evaluators=[], + baseline_evaluation_mode=BaselineEvaluationMode.NONE, + ), + actor=ACTOR, + ) + replays, _ = await services.replays.query(ReplayFilter()) + assert not replays + + +@pytest.mark.parametrize("invalid_field", ["raw_input", "file_hash"]) +async def test_malformed_eligible_mastra_baseline_creates_no_replay( + services: ReplayServices, + complete_mastra_memory_replay_inputs: dict[str, Any], + invalid_field: str, +) -> None: + """Refuse a stored eligible marker when replay input is malformed.""" + version = await _agent_version_with_run_spec(services) + invalid = deepcopy(complete_mastra_memory_replay_inputs) + if invalid_field == "raw_input": + del invalid["mastra_memory_replay"]["rawInput"] + else: + invalid["mastra_memory_replay"]["files"][0]["sha256"] = "0" * 64 + baseline = await create_session( + services.sessions, + ACTOR.account.id, + agent_id=version.agent_id, + agent_version_id=version.id, + origin=SessionOrigin.RECORDED, + status=SessionStatus.COMPLETED, + framework="mastra", + inputs=invalid, + metadata={"mastra_replay_state": "eligible"}, + ) + with pytest.raises(SessionReplayNotReady, match="mastra_replay_incomplete"): + await services.replay_service.create_replay( + ReplayCreate( + baseline_session_id=baseline.id, + evaluators=[], + baseline_evaluation_mode=BaselineEvaluationMode.NONE, + ), + actor=ACTOR, + ) + replays, _ = await services.replays.query(ReplayFilter()) + assert not replays + + +async def test_old_mastra_om_without_result_tape_is_not_replayable( + services: ReplayServices, +) -> None: + """A completed legacy OM envelope cannot silently rerun its observer.""" + version = await _agent_version_with_run_spec(services) + baseline = await create_session( + services.sessions, + ACTOR.account.id, + agent_id=version.agent_id, + agent_version_id=version.id, + origin=SessionOrigin.RECORDED, + status=SessionStatus.COMPLETED, + framework="mastra", + inputs={ + "mastra_memory_replay": { + "version": 2, + "complete": True, + "configuration": { + "memoryConfig": {"observationalMemory": {"scope": "thread"}} + }, + } + }, + ) + with pytest.raises(SessionReplayNotReady, match="mastra_replay_tape_missing"): + await services.replay_service.create_replay( + ReplayCreate( + baseline_session_id=baseline.id, + evaluators=[], + baseline_evaluation_mode=BaselineEvaluationMode.NONE, + ), + actor=ACTOR, + ) + + +async def test_old_mastra_working_memory_remains_replayable( + services: ReplayServices, +) -> None: + """The new tape guard does not reject completed working-memory-only v2 input.""" + version = await _agent_version_with_run_spec(services) + baseline = await create_session( + services.sessions, + ACTOR.account.id, + agent_id=version.agent_id, + agent_version_id=version.id, + origin=SessionOrigin.RECORDED, + status=SessionStatus.COMPLETED, + framework="mastra", + inputs={ + "mastra_memory_replay": { + "version": 2, + "complete": True, + "configuration": { + "memoryConfig": {"workingMemory": {"scope": "thread"}} + }, + } + }, + ) + bundle = await services.replay_service.create_replay( + ReplayCreate( + baseline_session_id=baseline.id, + evaluators=[], + baseline_evaluation_mode=BaselineEvaluationMode.NONE, + ), + actor=ACTOR, + ) + assert bundle.replay.baseline_session_id == baseline.id + + +async def test_experiment_rejects_pending_mastra_baseline_before_run( + services: ReplayServices, +) -> None: + """Cohort fan-out refuses pending memory input before it creates a run.""" + version = await _agent_version_with_run_spec(services) + experiment_id, _ = await _create_experiment_with_evaluator( + services, version.agent_id + ) + baseline = await create_session( + services.sessions, + ACTOR.account.id, + agent_id=version.agent_id, + agent_version_id=version.id, + origin=SessionOrigin.RECORDED, + framework="mastra", + inputs={"mastra_memory_replay": {"version": 3, "complete": False}}, + metadata={"mastra_replay_state": "pending"}, + ) + cohort = await _cohort_version(services, version.agent_id, [baseline.id]) + with pytest.raises(SessionReplayNotReady, match="mastra_replay_pending"): + await services.experiment_service.start_run( + experiment_id, + ExperimentRunCreate( + cohort_version_id=cohort.id, + agent_version_id=version.id, + baseline_evaluation_mode=BaselineEvaluationMode.NONE, + ), + actor=ACTOR, + ) + assert await services.experiment_runs.list_by_experiment(experiment_id) == [] + + async def _cohort_version( services: ReplayServices, agent_id: uuid.UUID, diff --git a/tests/server/test_session_service.py b/tests/server/test_session_service.py index 7013e7b4f..1b028ac19 100644 --- a/tests/server/test_session_service.py +++ b/tests/server/test_session_service.py @@ -14,6 +14,7 @@ """Tests for session use cases.""" import uuid +from copy import deepcopy from datetime import UTC, datetime, timedelta from typing import Any @@ -72,6 +73,7 @@ SessionInUse, SessionNotFound, SessionNotUpdatable, + SessionReplayFinalizationInvalid, SessionStatusCannotBeCleared, ) from kitaru.server.domain.task import ( @@ -452,6 +454,171 @@ async def test_update_session_clears_outputs_with_explicit_null( assert updated.status == SessionStatus.COMPLETED +async def test_finalize_pending_mastra_replay_inputs_atomically( + service: SessionService, + complete_mastra_memory_replay_inputs: dict[str, Any], +) -> None: + """Publish final replay input and eligibility in one guarded transition.""" + created = await service.create_session( + SessionCreate( + agent_id=uuid.uuid4(), + origin=SessionOrigin.RECORDED, + framework="mastra", + inputs={"mastra_memory_replay": {"version": 3, "complete": False}}, + metadata={"mastra_replay_state": "pending"}, + ), + actor=ACTOR, + ) + final_inputs = complete_mastra_memory_replay_inputs + await service.update_session( + created.id, + SessionUpdate( + status=SessionStatus.COMPLETED, + inputs=final_inputs, + metadata={"mastra_replay_state": "eligible"}, + ), + actor=ACTOR, + ) + stored = await service.get_session(created.id, actor=ACTOR) + assert stored.status == SessionStatus.COMPLETED + assert stored.inputs is not None and stored.inputs.value == final_inputs + assert stored.metadata["mastra_replay_state"] == "eligible" + + +@pytest.mark.parametrize("invalid_field", ["raw_input", "file_hash"]) +async def test_mastra_finalization_rejects_invalid_replay_prerequisite( + service: SessionService, + complete_mastra_memory_replay_inputs: dict[str, Any], + invalid_field: str, +) -> None: + """An eligible marker cannot accompany an incomplete v3 envelope.""" + created = await service.create_session( + SessionCreate( + agent_id=uuid.uuid4(), + origin=SessionOrigin.RECORDED, + framework="mastra", + inputs={"mastra_memory_replay": {"version": 3, "complete": False}}, + metadata={"mastra_replay_state": "pending"}, + ), + actor=ACTOR, + ) + invalid = deepcopy(complete_mastra_memory_replay_inputs) + if invalid_field == "raw_input": + del invalid["mastra_memory_replay"]["rawInput"] + else: + invalid["mastra_memory_replay"]["files"][0]["sha256"] = "0" * 64 + with pytest.raises(SessionReplayFinalizationInvalid): + await service.update_session( + created.id, + SessionUpdate( + status=SessionStatus.COMPLETED, + inputs=invalid, + metadata={"mastra_replay_state": "eligible"}, + ), + actor=ACTOR, + ) + stored = await service.get_session(created.id, actor=ACTOR) + assert stored.status == SessionStatus.IN_PROGRESS + assert stored.metadata["mastra_replay_state"] == "pending" + + +async def test_unrelated_session_cannot_replace_inputs( + service: SessionService, +) -> None: + """The replay finalization input patch is restricted to pending Mastra records.""" + created = await service.create_session( + SessionCreate( + agent_id=uuid.uuid4(), + origin=SessionOrigin.RECORDED, + framework="langgraph", + inputs={"original": True}, + ), + actor=ACTOR, + ) + with pytest.raises(SessionReplayFinalizationInvalid): + await service.update_session( + created.id, + SessionUpdate( + status=SessionStatus.COMPLETED, + inputs={"changed": True}, + metadata={"mastra_replay_state": "eligible"}, + ), + actor=ACTOR, + ) + stored = await service.get_session(created.id, actor=ACTOR) + assert stored.status == SessionStatus.IN_PROGRESS + assert stored.inputs is not None and stored.inputs.value == {"original": True} + + +async def test_pending_mastra_ineligibility_is_terminal( + service: SessionService, +) -> None: + """A failed recording cannot expose an eligible-looking pending input.""" + created = await service.create_session( + SessionCreate( + agent_id=uuid.uuid4(), + origin=SessionOrigin.RECORDED, + framework="mastra", + inputs={"mastra_memory_replay": {"version": 2, "complete": False}}, + metadata={"mastra_replay_state": "pending"}, + ), + actor=ACTOR, + ) + with pytest.raises(SessionReplayFinalizationInvalid): + await service.update_session( + created.id, + SessionUpdate(metadata={"mastra_replay_state": "ineligible"}), + actor=ACTOR, + ) + updated = await service.update_session( + created.id, + SessionUpdate( + status=SessionStatus.COMPLETED, + metadata={ + "mastra_replay_state": "ineligible", + "mastra_replay_reason": "capture_incomplete", + }, + ), + actor=ACTOR, + ) + assert updated.status == SessionStatus.COMPLETED + assert updated.metadata["mastra_replay_reason"] == "capture_incomplete" + + +async def test_mastra_om_cannot_be_finalized_without_result_tape( + service: SessionService, +) -> None: + """An OM session cannot claim eligibility before its tape is persisted.""" + created = await service.create_session( + SessionCreate( + agent_id=uuid.uuid4(), + origin=SessionOrigin.RECORDED, + framework="mastra", + inputs={"mastra_memory_replay": {"version": 3, "complete": False}}, + metadata={"mastra_replay_state": "pending"}, + ), + actor=ACTOR, + ) + with pytest.raises(SessionReplayFinalizationInvalid): + await service.update_session( + created.id, + SessionUpdate( + status=SessionStatus.COMPLETED, + inputs={ + "mastra_memory_replay": { + "version": 3, + "complete": True, + "configuration": { + "memoryConfig": {"observationalMemory": {"scope": "thread"}} + }, + } + }, + metadata={"mastra_replay_state": "eligible"}, + ), + actor=ACTOR, + ) + + async def test_update_session_omitted_fields_unchanged( service: SessionService, ) -> None: diff --git a/tests/server/test_sessions_api.py b/tests/server/test_sessions_api.py index d4f081772..7db14bc0a 100644 --- a/tests/server/test_sessions_api.py +++ b/tests/server/test_sessions_api.py @@ -17,6 +17,7 @@ import uuid from collections.abc import AsyncGenerator from datetime import UTC, datetime +from typing import Any import httpx import pytest @@ -643,6 +644,41 @@ async def test_update_session_clears_outputs_with_explicit_null( assert fetched["outputs"] is None +async def test_finalize_mastra_inputs_through_session_patch( + client: httpx.AsyncClient, + complete_mastra_memory_replay_inputs: dict[str, Any], +) -> None: + """Expose the final replay input and eligibility through one REST update.""" + created = ( + await client.post( + "/api/v1/sessions", + json=_session_body( + framework="mastra", + inputs={"mastra_memory_replay": {"version": 3, "complete": False}}, + metadata={"mastra_replay_state": "pending"}, + ), + ) + ).json() + final_inputs = complete_mastra_memory_replay_inputs + response = await client.patch( + f"/api/v1/sessions/{created['id']}", + json={ + "status": "completed", + "inputs": final_inputs, + "metadata": {"mastra_replay_state": "eligible"}, + }, + ) + assert response.status_code == 200 + fetched = (await client.get(f"/api/v1/sessions/{created['id']}")).json() + assert fetched["inputs"] == final_inputs + assert fetched["status"] == "completed" + assert fetched["metadata"]["mastra_replay_state"] == "eligible" + repeat = await client.patch( + f"/api/v1/sessions/{created['id']}", json={"inputs": final_inputs} + ) + assert repeat.status_code == 409 + + async def test_update_session_omitted_outputs_unchanged( client: httpx.AsyncClient, ) -> None: diff --git a/tests/server/test_sessions_api_pg.py b/tests/server/test_sessions_api_pg.py index 3c6f69e14..4a27f5d93 100644 --- a/tests/server/test_sessions_api_pg.py +++ b/tests/server/test_sessions_api_pg.py @@ -16,6 +16,7 @@ import json import uuid from collections.abc import AsyncGenerator +from typing import Any import httpx import pytest @@ -151,6 +152,87 @@ async def test_update_persists_across_requests( assert body["updated"] > created["updated"] +async def test_mastra_replay_input_finalization_persists_across_requests( + client: httpx.AsyncClient, + agent_id: str, + complete_mastra_memory_replay_inputs: dict[str, Any], +) -> None: + """Commit final inputs, eligibility, and status in the same database row.""" + created = ( + await client.post( + "/api/v1/sessions", + json=_session_body( + agent_id, + framework="mastra", + inputs={"mastra_memory_replay": {"version": 3, "complete": False}}, + metadata={"mastra_replay_state": "pending"}, + ), + ) + ).json() + final_inputs = complete_mastra_memory_replay_inputs + response = await client.patch( + f"/api/v1/sessions/{created['id']}", + json={ + "status": "completed", + "inputs": final_inputs, + "metadata": {"mastra_replay_state": "eligible"}, + }, + ) + assert response.status_code == 200 + fetched = (await client.get(f"/api/v1/sessions/{created['id']}")).json() + assert fetched["status"] == "completed" + assert fetched["metadata"]["mastra_replay_state"] == "eligible" + assert fetched["inputs"] == final_inputs + + +async def test_large_mastra_finalization_uses_payload_offload( + complete_mastra_memory_replay_inputs: dict[str, Any], +) -> None: + """Carry a final replay input above the old 1 MiB adapter gate as one update.""" + settings = db_settings(PAYLOAD_OFFLOAD_THRESHOLD_BYTES=64) + async with lifespan_client(settings) as client: + agent = (await client.post("/api/v1/agents", json={"name": "assistant"})).json() + created = ( + await client.post( + "/api/v1/sessions", + json=_session_body( + agent["id"], + framework="mastra", + inputs={"mastra_memory_replay": {"version": 3, "complete": False}}, + metadata={"mastra_replay_state": "pending"}, + ), + ) + ).json() + final_inputs = complete_mastra_memory_replay_inputs + final_inputs["mastra_memory_replay"]["rawInput"] = {"history": "x" * 1_200_000} + response = await client.patch( + f"/api/v1/sessions/{created['id']}", + json={ + "status": "completed", + "inputs": final_inputs, + "metadata": {"mastra_replay_state": "eligible"}, + }, + ) + assert response.status_code == 200 + fetched = (await client.get(f"/api/v1/sessions/{created['id']}")).json() + assert fetched["inputs"] == final_inputs + + engine = create_async_engine(DatabaseService.generate_database_uri(settings)) + try: + async with engine.connect() as connection: + inline, blob_id = ( + await connection.execute( + select(SessionORM.inputs, SessionORM.inputs_blob_id).where( + SessionORM.id == uuid.UUID(created["id"]) + ) + ) + ).one() + assert inline is None + assert blob_id is not None + finally: + await engine.dispose() + + async def test_delete_persists_across_requests( client: httpx.AsyncClient, agent_id: str ) -> None: From 5acefb3f3cf70ec0d7b24d00ce9b1ef763631e00 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 24 Sep 2026 07:53:38 +0200 Subject: [PATCH 019/115] Fix devtools import lint on Linux --- devtools/resilience.py | 2 +- devtools/seed.py | 2 +- pyproject.toml | 1 + 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/devtools/resilience.py b/devtools/resilience.py index f0bcf626b..bedeac88a 100644 --- a/devtools/resilience.py +++ b/devtools/resilience.py @@ -11,7 +11,6 @@ from dataclasses import dataclass, field from pathlib import Path -from fixtures import register_agent from seed import TERMINAL_JOB_STATUSES, poll_until from stack import ( RUN_DIR, @@ -23,6 +22,7 @@ wait_for_health, ) +from fixtures import register_agent from kitaru.api_models.v1.job import JobResponse, JobStatus from kitaru.api_models.v1.session_run import SessionRunCreateRequest from kitaru.api_models.v1.task import TaskResponse, TaskStatus diff --git a/devtools/seed.py b/devtools/seed.py index 94e539b38..a02d249b4 100644 --- a/devtools/seed.py +++ b/devtools/seed.py @@ -9,7 +9,6 @@ from collections.abc import Awaitable, Callable from typing import TypeVar -from fixtures import DEFAULT_IMPORTER_NAME, register_all from simulation import build_session_inputs from stack import ( RUN_DIR, @@ -28,6 +27,7 @@ malformed_line_count, ) +from fixtures import DEFAULT_IMPORTER_NAME, register_all from kitaru.api_models.v1.cohort import CohortCreateRequest from kitaru.api_models.v1.cohort_version import CohortVersionCreateRequest from kitaru.api_models.v1.experiment import ExperimentCreateRequest diff --git a/pyproject.toml b/pyproject.toml index a2f6721b6..2829063d6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -178,6 +178,7 @@ convention = "google" [tool.ruff.lint.isort] known-first-party = [ + "fixtures", "kitaru", "kitaru_openai_agents", "kitaru_pydantic_ai", From 81a614bc72b38d1a36618edcfb697edf78adb0da Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 24 Sep 2026 09:01:26 +0200 Subject: [PATCH 020/115] Guard shared Mastra resource replay --- docs/book/adapters/mastra.md | 4 +- packages/mastra/README.md | 4 +- packages/mastra/src/memory-binding.ts | 112 ++++++----- .../mastra/test/helpers/file-memory-access.ts | 188 +++++++++++------- packages/mastra/test/memory-binding.test.ts | 59 ++++++ .../mastra/test/memory-lease-process.test.ts | 58 +++++- 6 files changed, 297 insertions(+), 128 deletions(-) diff --git a/docs/book/adapters/mastra.md b/docs/book/adapters/mastra.md index 020283fdb..a243fbf71 100644 --- a/docs/book/adapters/mastra.md +++ b/docs/book/adapters/mastra.md @@ -256,9 +256,9 @@ Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KIT ### Source ownership and supported configuration -All writers to a source thread must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. Overlapping turns can continue natively, but their recordings become ineligible for replay. `settled()` joins pending work on the source `Memory` instance; it does not provide exclusive access. +All writers to a source thread or resource must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. Turns sharing either selector can continue natively, but their recordings become ineligible for replay. `settled()` joins pending work on the source `Memory` instance; it does not provide exclusive access. -A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must preserve ownership through the final eligible-session update, mark an overlapping or unowned write unsafe through `markUnsafeWrite()`, and retain unsafe state across process loss. Call `resetAfterQuiescence()` only after all possible stale writers have stopped. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. +A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must atomically reserve both thread and resource IDs, preserve ownership through the final eligible-session update, mark an overlapping or unowned write unsafe through `markUnsafeWrite()`, and retain unsafe state across process loss. Call `resetAfterQuiescence()` only after all possible stale writers for both selectors have stopped. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. Schema working memory requires explicit `scope: "thread"`. An observational-memory configuration object may omit `scope`, using Mastra's implicit thread scope, or set it to `"thread"`. Supply explicit observer/reflector model identities, either shared through `observationalMemory.model` or in the phase configuration. Resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` remain unsupported. diff --git a/packages/mastra/README.md b/packages/mastra/README.md index 7b3f691a9..b868addb6 100644 --- a/packages/mastra/README.md +++ b/packages/mastra/README.md @@ -220,9 +220,9 @@ Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KIT ### Source ownership and supported configuration -All writers to a source thread must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. Overlapping turns can continue natively, but their recordings become ineligible for replay. `settled()` joins pending work on the source `Memory` instance; it does not provide exclusive access. +All writers to a source thread or resource must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. Turns sharing either selector can continue natively, but their recordings become ineligible for replay. `settled()` joins pending work on the source `Memory` instance; it does not provide exclusive access. -A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must preserve ownership through the final eligible-session update, mark an overlapping or unowned write unsafe through `markUnsafeWrite()`, and retain unsafe state across process loss. Call `resetAfterQuiescence()` only after all possible stale writers have stopped. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. +A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must atomically reserve both thread and resource IDs, preserve ownership through the final eligible-session update, mark an overlapping or unowned write unsafe through `markUnsafeWrite()`, and retain unsafe state across process loss. Call `resetAfterQuiescence()` only after all possible stale writers for both selectors have stopped. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. Schema working memory requires explicit `scope: "thread"`. An observational-memory configuration object may omit `scope`, using Mastra's implicit thread scope, or set it to `"thread"`. Supply explicit observer/reflector model identities, either shared through `observationalMemory.model` or in the phase configuration. Resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` remain unsupported. diff --git a/packages/mastra/src/memory-binding.ts b/packages/mastra/src/memory-binding.ts index 998a5321e..50b0123db 100644 --- a/packages/mastra/src/memory-binding.ts +++ b/packages/mastra/src/memory-binding.ts @@ -30,10 +30,10 @@ export interface MastraMemoryLease { } /** - * Coordinate every writer of a source thread across all application processes. + * Coordinate every writer of a source thread or resource across all processes. * - * The implementation must atomically poison eligibility for the thread when a - * competing native turn proceeds without ownership or an owner loses its lease. + * The implementation must atomically poison eligibility for both selectors when + * a competing native turn proceeds without ownership or an owner loses its lease. * Poison must survive process loss and prevent a later acquisition from becoming * eligible until all possible stale writers have quiesced. A timeout or failed * coordination call must fail closed for replay eligibility, while native Mastra @@ -63,50 +63,75 @@ export interface MastraExclusiveMemoryAccess { */ export function createProcessLocalMemoryAccess(): MastraExclusiveMemoryAccess { type Turn = { onConflict?: () => void; invalidated: boolean }; - type ThreadState = { + type ScopeState = { turns: Set; poisoned: boolean; persistentLoss: boolean; }; - const threads = new Map(); + const scopes = new Map(); let unknownWriterPoisoned = false; - return { - async acquire({ threadId }, options = {}) { - if (options.signal?.aborted) - throw new Error("Exclusive source-thread ownership was cancelled."); - const state = threads.get(threadId) ?? { + + function keys({ threadId, resourceId }: MastraMemorySelector): string[] { + return [`thread:${threadId}`, `resource:${resourceId}`]; + } + + function getState(key: string): ScopeState { + let state = scopes.get(key); + if (!state) { + state = { turns: new Set(), poisoned: false, persistentLoss: false, }; - threads.set(threadId, state); - if (state.turns.size > 0) { - state.poisoned = true; - for (const turn of state.turns) { - turn.invalidated = true; - turn.onConflict?.(); - } - } + scopes.set(key, state); + } + return state; + } + + function poison(states: readonly ScopeState[], persistent: boolean): void { + const affected = new Set(); + for (const state of states) { + state.poisoned = true; + if (persistent) state.persistentLoss = true; + for (const turn of state.turns) affected.add(turn); + } + for (const turn of affected) { + turn.invalidated = true; + turn.onConflict?.(); + } + } + + return { + async acquire(selector, options = {}) { + if (options.signal?.aborted) + throw new Error("Exclusive source-thread ownership was cancelled."); + const scopeKeys = keys(selector); + const states = scopeKeys.map(getState); + const occupied = states.filter((state) => state.turns.size > 0); + if (occupied.length > 0) poison(occupied, false); const owner = { onConflict: options.onConflict, invalidated: - state.poisoned || state.persistentLoss || unknownWriterPoisoned, + states.some((state) => state.poisoned || state.persistentLoss) || + unknownWriterPoisoned, }; - state.turns.add(owner); + for (const state of states) state.turns.add(owner); let released = false; const release = async () => { - if (!released) state.turns.delete(owner); + if (!released) for (const state of states) state.turns.delete(owner); released = true; - if (state.turns.size === 0 && !state.persistentLoss) - threads.delete(threadId); + for (const key of scopeKeys) { + const state = scopes.get(key); + if (state && state.turns.size === 0 && !state.persistentLoss) + scopes.delete(key); + } }; return Object.assign(release, { async verifyEligibility() { return ( !released && !owner.invalidated && - !state.poisoned && - !state.persistentLoss && + states.every((state) => !state.poisoned && !state.persistentLoss) && !unknownWriterPoisoned ); }, @@ -115,46 +140,25 @@ export function createProcessLocalMemoryAccess(): MastraExclusiveMemoryAccess { async markUnsafeWrite(selector) { if (!selector) { unknownWriterPoisoned = true; - for (const state of threads.values()) { - state.poisoned = true; - state.persistentLoss = true; - for (const turn of state.turns) { - turn.invalidated = true; - turn.onConflict?.(); - } - } + poison([...scopes.values()], true); return; } - const { threadId } = selector; - const state = threads.get(threadId) ?? { - turns: new Set(), - poisoned: false, - persistentLoss: false, - }; - threads.set(threadId, state); - state.poisoned = true; // The unsafe writer may outlive every current lease. Do not let the // final release erase the conflict before that writer quiesces. - state.persistentLoss = true; - for (const turn of state.turns) { - turn.invalidated = true; - turn.onConflict?.(); - } + poison(keys(selector).map(getState), true); }, async resetAfterQuiescence(selector) { if (!selector) { - if ([...threads.values()].some((state) => state.turns.size > 0)) + if ([...scopes.values()].some((state) => state.turns.size > 0)) throw new Error("Source-thread writers are still active."); - threads.clear(); + scopes.clear(); unknownWriterPoisoned = false; return; } - const { threadId } = selector; - const state = threads.get(threadId); - if (!state) return; - if (state.turns.size > 0) - throw new Error("Source-thread writers are still active."); - threads.delete(threadId); + const scopeKeys = keys(selector); + if (scopeKeys.some((key) => (scopes.get(key)?.turns.size ?? 0) > 0)) + throw new Error("Source-thread or resource writers are still active."); + for (const key of scopeKeys) scopes.delete(key); }, }; } diff --git a/packages/mastra/test/helpers/file-memory-access.ts b/packages/mastra/test/helpers/file-memory-access.ts index f482ce3dd..a17216130 100644 --- a/packages/mastra/test/helpers/file-memory-access.ts +++ b/packages/mastra/test/helpers/file-memory-access.ts @@ -27,44 +27,70 @@ function isExists(error: unknown): boolean { ); } -/** Test-only shared lease: atomic directory creation and persistent poison. */ +/** Test-only shared lease: atomic thread/resource coordination and poison. */ export function createFileMemoryAccess( root: string, ): MastraExclusiveMemoryAccess & { simulateLeaseLoss(selector: MastraMemorySelector): Promise; } { - function directory(selector: MastraMemorySelector): string { + function directory(kind: "thread" | "resource", id: string): string { return join( root, - createHash("sha256").update(selector.threadId).digest("hex"), + `${kind}-${createHash("sha256").update(id).digest("hex")}`, ); } const globalPoison = join(root, "unknown-writer-poison"); - async function withMutex( - selector: MastraMemorySelector, - run: (dir: string) => Promise, + async function withMutexes( + mutexes: string[], + run: () => Promise, ): Promise { - const dir = directory(selector); - await mkdir(dir, { recursive: true }); - const mutex = join(dir, "mutex"); + const held: string[] = []; const deadline = Date.now() + 1000; - while (true) { - try { - await mkdir(mutex); - break; - } catch (error) { - if (!isExists(error) || Date.now() >= deadline) throw error; - await pause(2); - } - } try { - return await run(dir); + for (const mutex of mutexes) { + while (true) { + try { + await mkdir(mutex); + held.push(mutex); + break; + } catch (error) { + if (!isExists(error) || Date.now() >= deadline) throw error; + await pause(2); + } + } + } + return await run(); } finally { - await rm(mutex, { recursive: true, force: true }); + for (const mutex of held.reverse()) + await rm(mutex, { recursive: true, force: true }); } } + async function withGlobalMutex(run: () => Promise): Promise { + await mkdir(root, { recursive: true }); + return withMutexes([join(root, "coordination-mutex")], run); + } + + async function withScopes( + selector: MastraMemorySelector, + run: (dirs: string[]) => Promise, + ): Promise { + return withGlobalMutex(async () => { + const dirs = [ + directory("thread", selector.threadId), + directory("resource", selector.resourceId), + ].sort(); + for (const dir of dirs) await mkdir(dir, { recursive: true }); + // All processes acquire both scope locks in the same order. The global + // mutex also makes unknown-selector poison and reset atomic with them. + return withMutexes( + dirs.map((dir) => join(dir, "mutex")), + () => run(dirs), + ); + }); + } + async function exists(path: string): Promise { try { await readFile(path); @@ -105,34 +131,38 @@ export function createFileMemoryAccess( let released = false; const release = async () => { if (released) return; - await withMutex(selector, async (dir) => { - if (owns) { - const current = await readFile(join(dir, "owner"), "utf8").catch( - () => undefined, - ); - if (current === token) await rm(join(dir, "owner"), { force: true }); - } else { - await rm(join(dir, `turn-${token}`), { force: true }); + await withScopes(selector, async (dirs) => { + for (const dir of dirs) { + if (owns) { + const current = await readFile(join(dir, "owner"), "utf8").catch( + () => undefined, + ); + if (current === token) + await rm(join(dir, "owner"), { force: true }); + } else { + await rm(join(dir, `turn-${token}`), { force: true }); + } + if ( + !(await exists(join(dir, "owner"))) && + (await activeTurns(dir)).length === 0 && + !(await exists(join(dir, "persistent-loss"))) + ) + await rm(join(dir, "poison"), { force: true }); } - if ( - !(await exists(join(dir, "owner"))) && - (await activeTurns(dir)).length === 0 && - !(await exists(join(dir, "persistent-loss"))) - ) - await rm(join(dir, "poison"), { force: true }); }); released = true; }; return Object.assign(release, { async verifyEligibility() { if (released || !owns) return false; - return withMutex(selector, async (dir) => { - if ( - (await exists(globalPoison)) || - (await exists(join(dir, "poison"))) - ) - return false; - return (await readFile(join(dir, "owner"), "utf8")) === token; + return withScopes(selector, async (dirs) => { + if (await exists(globalPoison)) return false; + for (const dir of dirs) { + if (await exists(join(dir, "poison"))) return false; + if ((await readFile(join(dir, "owner"), "utf8")) !== token) + return false; + } + return true; }); }, }); @@ -146,24 +176,35 @@ export function createFileMemoryAccess( while (true) { if (options.signal?.aborted) throw new Error("Source-thread lease wait cancelled."); - const status = await withMutex(selector, async (dir) => { + const status = await withScopes(selector, async (dirs) => { if ( (await exists(globalPoison)) || - (await exists(join(dir, "poison"))) + ( + await Promise.all(dirs.map((dir) => exists(join(dir, "poison")))) + ).some(Boolean) ) { - await writeFile(join(dir, `turn-${token}`), ""); + for (const dir of dirs) + await writeFile(join(dir, `turn-${token}`), ""); return "denied"; } - if (await exists(join(dir, "owner"))) return "busy"; - await writeFile(join(dir, "owner"), token, { flag: "wx" }); + if ( + ( + await Promise.all(dirs.map((dir) => exists(join(dir, "owner")))) + ).some(Boolean) + ) + return "busy"; + for (const dir of dirs) + await writeFile(join(dir, "owner"), token, { flag: "wx" }); return "owned"; }); if (status === "owned") return makeLease(selector, token, true); if (status === "denied") return makeLease(selector, token, false); if (Date.now() >= deadline) { - await withMutex(selector, async (dir) => { - await poison(dir, false); - await writeFile(join(dir, `turn-${token}`), ""); + await withScopes(selector, async (dirs) => { + for (const dir of dirs) { + await poison(dir, false); + await writeFile(join(dir, `turn-${token}`), ""); + } }); return makeLease(selector, token, false); } @@ -172,41 +213,50 @@ export function createFileMemoryAccess( }, async markUnsafeWrite(selector) { if (!selector) { - await writeFile(globalPoison, "unsafe write"); + await withGlobalMutex(() => writeFile(globalPoison, "unsafe write")); return; } - await withMutex(selector, async (dir) => { - await poison(dir, true); + await withScopes(selector, async (dirs) => { + for (const dir of dirs) await poison(dir, true); }); }, async resetAfterQuiescence(selector) { if (!selector) { - for (const name of await readdir(root)) { - if (name === "unknown-writer-poison") continue; - const dir = join(root, name); + await withGlobalMutex(async () => { + for (const name of await readdir(root)) { + if (!name.startsWith("thread-") && !name.startsWith("resource-")) + continue; + const dir = join(root, name); + if ( + (await exists(join(dir, "owner"))) || + (await activeTurns(dir)).length > 0 + ) + throw new Error("Source-thread writers are still active."); + } + await rm(globalPoison, { force: true }); + }); + return; + } + await withScopes(selector, async (dirs) => { + for (const dir of dirs) { if ( (await exists(join(dir, "owner"))) || (await activeTurns(dir)).length > 0 ) throw new Error("Source-thread writers are still active."); } - await rm(globalPoison, { force: true }); - return; - } - await withMutex(selector, async (dir) => { - if ( - (await exists(join(dir, "owner"))) || - (await activeTurns(dir)).length > 0 - ) - throw new Error("Source-thread writers are still active."); - await rm(join(dir, "persistent-loss"), { force: true }); - await rm(join(dir, "poison"), { force: true }); + for (const dir of dirs) { + await rm(join(dir, "persistent-loss"), { force: true }); + await rm(join(dir, "poison"), { force: true }); + } }); }, async simulateLeaseLoss(selector) { - await withMutex(selector, async (dir) => { - await poison(dir, true); - await rm(join(dir, "owner"), { force: true }); + await withScopes(selector, async (dirs) => { + for (const dir of dirs) { + await poison(dir, true); + await rm(join(dir, "owner"), { force: true }); + } }); }, }; diff --git a/packages/mastra/test/memory-binding.test.ts b/packages/mastra/test/memory-binding.test.ts index be0db6987..fbdc36ae1 100644 --- a/packages/mastra/test/memory-binding.test.ts +++ b/packages/mastra/test/memory-binding.test.ts @@ -73,6 +73,65 @@ it("rejects shared-thread overlap while allowing independent source threads", as await two.binding.release(); }); +it("invalidates recordings on different threads sharing a resource", async () => { + const access = createProcessLocalMemoryAccess(); + const first = await access.acquire({ + threadId: "first-thread", + resourceId: "shared-resource", + }); + expect(await first.verifyEligibility()).toBe(true); + const second = await access.acquire({ + threadId: "second-thread", + resourceId: "shared-resource", + }); + expect(await first.verifyEligibility()).toBe(false); + expect(await second.verifyEligibility()).toBe(false); + const independent = await access.acquire({ + threadId: "third-thread", + resourceId: "other-resource", + }); + expect(await independent.verifyEligibility()).toBe(true); + await first(); + await second(); + await independent(); + const recovered = await access.acquire({ + threadId: "third-thread", + resourceId: "shared-resource", + }); + expect(await recovered.verifyEligibility()).toBe(true); + await recovered(); +}); + +it("keeps a shared resource unsafe after an unowned native write", async () => { + const access = createProcessLocalMemoryAccess(); + const first = await access.acquire({ + threadId: "first-thread", + resourceId: "shared-resource", + }); + await access.markUnsafeWrite({ + threadId: "second-thread", + resourceId: "shared-resource", + }); + expect(await first.verifyEligibility()).toBe(false); + await first(); + const later = await access.acquire({ + threadId: "third-thread", + resourceId: "shared-resource", + }); + expect(await later.verifyEligibility()).toBe(false); + await later(); + await access.resetAfterQuiescence({ + threadId: "second-thread", + resourceId: "shared-resource", + }); + const recovered = await access.acquire({ + threadId: "third-thread", + resourceId: "shared-resource", + }); + expect(await recovered.verifyEligibility()).toBe(true); + await recovered(); +}); + it("shows why separate process-local helpers cannot qualify for multi-server replay", async () => { const first = createProcessLocalMemoryAccess(); const second = createProcessLocalMemoryAccess(); diff --git a/packages/mastra/test/memory-lease-process.test.ts b/packages/mastra/test/memory-lease-process.test.ts index edc404564..8e6663a15 100644 --- a/packages/mastra/test/memory-lease-process.test.ts +++ b/packages/mastra/test/memory-lease-process.test.ts @@ -9,7 +9,13 @@ const workerPath = fileURLToPath( new URL("./helpers/file-memory-worker.mjs", import.meta.url), ); const thread = { threadId: "shared-thread", resourceId: "resource" }; -const other = { threadId: "other-thread", resourceId: "resource" }; +const sameResource = { threadId: "other-thread", resourceId: "resource" }; +const sameThread = { threadId: "shared-thread", resourceId: "other-resource" }; +const other = { threadId: "other-thread", resourceId: "other-resource" }; +const unrelated = { + threadId: "unrelated-thread", + resourceId: "unrelated-resource", +}; let root: string | undefined; const children: ChildProcess[] = []; let nextId = 0; @@ -91,6 +97,56 @@ it("invalidates both processes on bounded same-thread contention, then recovers await second.request({ action: "release", name: "next", selector: thread }); }); +it.each([ + ["resource", sameResource], + ["thread", sameThread], +] as const)( + "invalidates both processes when selectors share a %s", + async (_scope, competing) => { + const first = await worker(); + const second = await worker(); + expect( + await first.request({ + action: "acquire", + name: "first", + selector: thread, + }), + ).toBe(true); + expect( + await second.request({ + action: "acquire", + name: "competing", + selector: competing, + waitMs: 40, + }), + ).toBe(false); + expect(await first.request({ action: "verify", name: "first" })).toBe( + false, + ); + // Poison is scoped to the conflicting thread and resource, not all turns. + expect( + await second.request({ + action: "acquire", + name: "other", + selector: unrelated, + }), + ).toBe(true); + await Promise.all([ + first.request({ action: "release", name: "first" }), + second.request({ action: "release", name: "competing" }), + second.request({ action: "release", name: "other" }), + ]); + expect( + await second.request({ + action: "acquire", + name: "next", + selector: thread, + }), + ).toBe(true); + await second.request({ action: "release", name: "next" }); + }, +); + it("poisons a lost holder before a stale native write and blocks a successor", async () => { const first = await worker(); const second = await worker(); From 0f9527e4bd6c8b4dec6f4a15d3eaa948d0920159 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 24 Sep 2026 12:55:16 +0200 Subject: [PATCH 021/115] Read ISO-string memory dates at capture PostgreSQL and LibSQL stores keep buffered observation chunks in a JSON column, so reading a record back returns each chunk's `createdAt` and `lastObservedAt` as ISO strings. The snapshot validator requires Dates, so every turn that started after a buffered observation was marked ineligible, and older generations keep their chunks, so the thread never recovered. `normalizeStoredMemoryDates` converts known date fields back to Dates before validation. It accepts only an exact ISO rendering, so any other string still fails validation. --- packages/mastra/src/memory-binding.ts | 5 +- packages/mastra/src/memory-snapshot.ts | 57 +++++++++++++++++++++ packages/mastra/test/memory-binding.test.ts | 55 ++++++++++++++++++++ 3 files changed, 116 insertions(+), 1 deletion(-) diff --git a/packages/mastra/src/memory-binding.ts b/packages/mastra/src/memory-binding.ts index 50b0123db..d930cfb34 100644 --- a/packages/mastra/src/memory-binding.ts +++ b/packages/mastra/src/memory-binding.ts @@ -5,6 +5,7 @@ import { decodeMemoryValue, encodeMemoryValue, type MastraMemorySnapshot, + normalizeStoredMemoryDates, validateMemorySnapshot, } from "./memory-snapshot.js"; @@ -598,7 +599,9 @@ export function createMemoryCaptureBinding( records, }; // No storage-owned objects or Dates escape the explicit codec. - const copy = decodeMemoryValue(encodeMemoryValue(snapshot)); + const copy = normalizeStoredMemoryDates( + decodeMemoryValue(encodeMemoryValue(snapshot)), + ); validateMemorySnapshot(copy); return copy; } finally { diff --git a/packages/mastra/src/memory-snapshot.ts b/packages/mastra/src/memory-snapshot.ts index a250c7efb..0625c2b5c 100644 --- a/packages/mastra/src/memory-snapshot.ts +++ b/packages/mastra/src/memory-snapshot.ts @@ -257,6 +257,63 @@ export function decodeMemoryValue(value: JsonValue): unknown { return visit(converted); } +function readStoredDates(value: unknown, keys: readonly string[]): unknown { + if (!isRecord(value)) return value; + const copy = { ...value }; + for (const key of keys) { + const stored = copy[key]; + if (typeof stored !== "string") continue; + const date = new Date(stored); + // Only an exact ISO rendering is a serialized Date; other strings stay + // unchanged so validation still rejects them. + if (Number.isFinite(date.getTime()) && date.toISOString() === stored) + copy[key] = date; + } + return copy; +} + +/** + * Convert ISO-string dates at known snapshot fields back to Dates. + * + * SQL stores keep buffered observation chunks in a JSON column, so reading a + * record back returns each chunk's `createdAt` and `lastObservedAt` as strings. + */ +export function normalizeStoredMemoryDates(snapshot: unknown): unknown { + if (!isRecord(snapshot)) return snapshot; + const timestamps = ["createdAt", "updatedAt"]; + const records = Array.isArray(snapshot.records) + ? snapshot.records.map((value) => { + const record = readStoredDates(value, [ + ...timestamps, + "lastObservedAt", + "lastBufferedAtTime", + ]); + if ( + !isRecord(record) || + !Array.isArray(record.bufferedObservationChunks) + ) + return record; + return { + ...record, + bufferedObservationChunks: record.bufferedObservationChunks.map( + (chunk) => readStoredDates(chunk, ["createdAt", "lastObservedAt"]), + ), + }; + }) + : snapshot.records; + return { + ...snapshot, + thread: readStoredDates(snapshot.thread, timestamps), + resource: readStoredDates(snapshot.resource, timestamps), + messages: Array.isArray(snapshot.messages) + ? snapshot.messages.map((message) => + readStoredDates(message, ["createdAt"]), + ) + : snapshot.messages, + records, + }; +} + /** Validate the complete native state before an isolated store receives any writes. */ export function validateMemorySnapshot( value: unknown, diff --git a/packages/mastra/test/memory-binding.test.ts b/packages/mastra/test/memory-binding.test.ts index fbdc36ae1..73e396f9e 100644 --- a/packages/mastra/test/memory-binding.test.ts +++ b/packages/mastra/test/memory-binding.test.ts @@ -381,6 +381,61 @@ it("serializes overlapping native mutations and preserves original storage error await binding.release(); }); +async function readChunksLikeDatabase( + runtime: Awaited>["runtime"], + createdAt: string, +) { + const record = required( + await runtime.domain.getObservationalMemory(THREAD, RESOURCE), + ); + await runtime.domain.updateBufferedObservations({ + id: record.id, + chunk: { + cycleId: "cycle", + observations: "buffer", + tokenCount: 3, + messageIds: ["historical-message"], + messageTokens: 20, + lastObservedAt: new Date("2026-02-01T10:00:00.000Z"), + }, + }); + const history = runtime.domain.getObservationalMemoryHistory.bind( + runtime.domain, + ); + // @mastra/pg and LibSQL JSON.parse this column on every read. + vi.spyOn(runtime.domain, "getObservationalMemoryHistory").mockImplementation( + async (...args) => + (await history(...args)).map((value) => ({ + ...value, + bufferedObservationChunks: value.bufferedObservationChunks?.map( + (chunk) => ({ + ...JSON.parse(JSON.stringify(chunk)), + createdAt, + }), + ), + })), + ); +} + +it("captures buffered chunk dates that a database returns as ISO strings", async () => { + const { runtime, binding } = await fixture(); + await readChunksLikeDatabase(runtime, "2026-02-01T10:00:01.000Z"); + const initial = await binding.captureInitial(runtime.memory); + expect(binding.incompleteReasons).toEqual([]); + const chunk = required(initial?.records[0]?.bufferedObservationChunks?.[0]); + expect(chunk.createdAt).toEqual(new Date("2026-02-01T10:00:01.000Z")); + expect(chunk.lastObservedAt).toEqual(new Date("2026-02-01T10:00:00.000Z")); + await binding.release(); +}); + +it("still refuses buffered chunk dates that are not ISO timestamps", async () => { + const { runtime, binding } = await fixture(); + await readChunksLikeDatabase(runtime, "Sun Feb 01 2026"); + expect(await binding.captureInitial(runtime.memory)).toBeUndefined(); + expect(binding.incompleteReasons.join()).toMatch(/Initial memory capture/); + await binding.release(); +}); + it("records OM flags, buffers, config, activation and working memory with stable request identity", async () => { const { runtime, binding, recordMutation } = await fixture(); await binding.captureInitial(runtime.memory); From f512fdf1c902dc00018e50fab5e0afe7daeb8611 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 24 Sep 2026 12:55:16 +0200 Subject: [PATCH 022/115] Match source store semantics in memory replay Replay rebuilds memory in Mastra's InMemoryStore, which behaves unlike database stores in ways observational memory depends on. Its `createReflectionGeneration` fills a missing `lastObservedAt` with the current time where PostgreSQL and LibSQL copy it through, so after a replayed reflection Mastra treated the turn's messages as observed, skipped a recorded observer call and the replay diverged. It also hands out live record objects instead of copies, drops reflection metadata, keeps omitted observed-message ids, and activates caller-supplied buffered chunks. The baseline now records whether the source domain is an InMemoryStore. Unless it was, `createIsolatedMemoryReplay` gives the isolated store database record semantics for those methods. --- packages/mastra/src/memory-replay.ts | 97 ++++++++++++++++++++++ packages/mastra/src/stateful-agent.ts | 10 +++ packages/mastra/test/implicit-om.test.ts | 3 + packages/mastra/test/memory-replay.test.ts | 96 +++++++++++++++++++++ 4 files changed, 206 insertions(+) diff --git a/packages/mastra/src/memory-replay.ts b/packages/mastra/src/memory-replay.ts index efa77aebf..df6972247 100644 --- a/packages/mastra/src/memory-replay.ts +++ b/packages/mastra/src/memory-replay.ts @@ -1,6 +1,7 @@ import { createRequire } from "node:module"; import type { MastraModelConfig } from "@mastra/core/llm"; import type { MemoryConfigInternal } from "@mastra/core/memory"; +import type { MemoryStorage } from "@mastra/core/storage"; import { createMemoryCaptureBinding, createProcessLocalMemoryAccess, @@ -174,9 +175,103 @@ export async function bindOMResultModels( return { ...config, observationalMemory: bound } as MemoryConfigInternal; } +/** + * How a source memory store hands out observational-memory records. + * + * `in-memory` is Mastra's InMemoryStore, which returns its stored objects. + * `persistent` covers database stores such as PostgreSQL and LibSQL, which + * return a fresh copy on every read. + */ +export type MastraMemoryStoreSemantics = "in-memory" | "persistent"; + +/** Classify a source memory domain by the record semantics it implements. */ +export async function getMemoryStoreSemantics( + domain: MemoryStorage, +): Promise { + const { InMemoryMemory } = await import("@mastra/core/storage"); + return domain instanceof InMemoryMemory ? "in-memory" : "persistent"; +} + +/** Copy plain objects and arrays; Dates and class instances stay shared. */ +function copyStoredValue(value: T): T { + if (Array.isArray(value)) return value.map(copyStoredValue) as T; + if ( + !record(value) || + ![Object.prototype, null].includes(Object.getPrototypeOf(value)) + ) + return value; + return Object.fromEntries( + Object.entries(value).map(([key, item]) => [key, copyStoredValue(item)]), + ) as T; +} + +/** + * Give an InMemoryStore memory domain the record semantics of a database store. + * + * Observational memory decides what to observe from the record it read at the + * start of a step. InMemoryStore returns live objects, so later writes change + * that record under Mastra; a database returns a copy that stays stale. Its + * `createReflectionGeneration` also substitutes the current time for a missing + * `lastObservedAt`, which makes Mastra treat unobserved messages as observed. + */ +function applyPersistentStoreSemantics( + domain: MemoryStorage, + selector: Pick, +): void { + const read = domain.getObservationalMemory.bind(domain); + const history = domain.getObservationalMemoryHistory.bind(domain); + const initialize = domain.initializeObservationalMemory.bind(domain); + const insert = domain.insertObservationalMemoryRecord.bind(domain); + const observe = domain.updateActiveObservations.bind(domain); + const activate = domain.swapBufferedToActive.bind(domain); + const reflect = domain.createReflectionGeneration.bind(domain); + // The isolated store holds one thread, and InMemoryStore's history returns + // the stored objects themselves, so assignments below change the store. + const getStoredRecord = async (id: string) => + (await history(selector.threadId, selector.resourceId)).find( + (value) => value.id === id, + ); + // Instance properties shadow the prototype, so InMemoryStore's own internal + // calls, such as swapBufferedReflectionToActive, also use these versions. + domain.getObservationalMemory = async (...args) => + copyStoredValue(await read(...args)); + domain.getObservationalMemoryHistory = async (...args) => + copyStoredValue(await history(...args)); + domain.initializeObservationalMemory = async (input) => { + const created = await initialize(input); + const stored = await getStoredRecord(created.id); + if (stored) stored.metadata = undefined; + return copyStoredValue({ ...created, metadata: undefined }); + }; + domain.insertObservationalMemoryRecord = (value) => + insert(copyStoredValue(value)); + domain.updateActiveObservations = async (input) => { + await observe(input); + // A database overwrites the column, so omitted identities become unset. + if (input.observedMessageIds) return; + const stored = await getStoredRecord(input.id); + if (stored) stored.observedMessageIds = undefined; + }; + domain.swapBufferedToActive = (input) => + // A database activates its persisted chunks, never the caller's copy. + activate({ ...input, bufferedChunks: undefined }); + domain.createReflectionGeneration = async (input) => { + const created = await reflect(input); + const stored = await getStoredRecord(created.id); + const copied = { + lastObservedAt: input.currentRecord.lastObservedAt, + metadata: input.currentRecord.metadata, + }; + if (stored) Object.assign(stored, copied); + return copyStoredValue({ ...created, ...copied }); + }; +} + export interface IsolatedMemoryReplayOptions { invocationId: string; initialSnapshot: MastraMemorySnapshot; + /** Defaults to `persistent`, the semantics of database-backed stores. */ + storeSemantics?: MastraMemoryStoreSemantics; configuration: Record; resolveModel: (id: string) => Promise | MastraModelConfig; recordMutation: MastraMemoryCaptureOptions["recordMutation"]; @@ -211,6 +306,8 @@ export async function createIsolatedMemoryReplay( const store = new InMemoryStore(); const domain = store.stores.memory; if (!domain) return unsupported("Native in-memory storage is unavailable."); + if (options.storeSemantics !== "in-memory") + applyPersistentStoreSemantics(domain, snapshot); try { if (snapshot.thread) await domain.saveThread({ thread: structuredClone(snapshot.thread) }); diff --git a/packages/mastra/src/stateful-agent.ts b/packages/mastra/src/stateful-agent.ts index a62758884..b799ad3ed 100644 --- a/packages/mastra/src/stateful-agent.ts +++ b/packages/mastra/src/stateful-agent.ts @@ -27,6 +27,8 @@ import { bindOMResultModels, createIsolatedMemoryReplay, getMemoryModelId, + getMemoryStoreSemantics, + type MastraMemoryStoreSemantics, serializeMemoryConfiguration, } from "./memory-replay.js"; import { @@ -492,10 +494,16 @@ export function createMemoryReplayAgent( : baselineFiles?.evidenceSanitizer(markUnknownCredentialUrl); if (!sanitizer) throw new Error("Controlled evidence sanitizer was not initialized."); + let memoryStore: MastraMemoryStoreSemantics; if (historical) { + memoryStore = + historical.configuration.memoryStore === "in-memory" + ? "in-memory" + : "persistent"; runtime = await createIsolatedMemoryReplay({ invocationId, initialSnapshot: historical.initialSnapshot, + storeSemantics: memoryStore, configuration: requireRecord( historical.configuration.memoryConfig, "memory configuration", @@ -508,6 +516,7 @@ export function createMemoryReplayAgent( }); } else { const source = await options.sourceMemory(); + memoryStore = await getMemoryStoreSemantics(source.domain); const binding = createMemoryCaptureBinding({ invocationId, ...selector, @@ -693,6 +702,7 @@ export function createMemoryReplayAgent( memoryConfig: historical?.configuration.memoryConfig ?? serializeMemoryConfiguration(runtime.memory.getMergedThreadConfig()), + memoryStore, ...(workspace ? { workspaceManifest: workspace.manifest } : {}), }; let recordedRawInput = invocationInput; diff --git a/packages/mastra/test/implicit-om.test.ts b/packages/mastra/test/implicit-om.test.ts index 4a80d2fcc..f5747aa79 100644 --- a/packages/mastra/test/implicit-om.test.ts +++ b/packages/mastra/test/implicit-om.test.ts @@ -105,6 +105,9 @@ it("records a first turn with implicit thread-scoped OM and string models", asyn )?.body?.inputs as Record; expect(input.mastra_memory_replay?.version).toBe(3); expect(input.mastra_memory_replay?.omTape.length).toBeGreaterThan(0); + expect(JSON.stringify(input.mastra_memory_replay)).toContain( + '"memoryStore":"in-memory"', + ); const observerCalls = observe.mock.calls.length; const reflectorCalls = reflect.mock.calls.length; vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); diff --git a/packages/mastra/test/memory-replay.test.ts b/packages/mastra/test/memory-replay.test.ts index 34cf2102e..377985c11 100644 --- a/packages/mastra/test/memory-replay.test.ts +++ b/packages/mastra/test/memory-replay.test.ts @@ -259,3 +259,99 @@ it("waits for native background observation and evidence before releasing replay await runtime.finish(); } }, 20000); + +it("replays on the record semantics of a database store unless the source was in memory", async () => { + const { source, initialSnapshot, configuration } = await baseline(); + const historical = initialSnapshot.records[0]; + if (!historical) throw new Error("Missing native memory fixture"); + // A database copies an unobserved generation's cursor into its reflection. + historical.lastObservedAt = undefined; + historical.metadata = { origin: "database" }; + const options = { + initialSnapshot, + configuration, + resolveModel: async (id: string) => + id.endsWith("observer") ? source.observer.model : source.reflector.model, + recordMutation: async () => {}, + }; + const persistent = await createIsolatedMemoryReplay({ + ...options, + invocationId: "persistent", + }); + const domain = persistent.binding.domain; + const current = await domain.getObservationalMemory(THREAD, RESOURCE); + if (!current) throw new Error("Missing restored observational memory"); + current.activeObservations = "CHANGED BY THE CALLER"; + expect( + (await domain.getObservationalMemory(THREAD, RESOURCE))?.activeObservations, + ).toBe(historical.activeObservations); + expect( + await domain.swapBufferedToActive({ + id: current.id, + activationRatio: 1, + messageTokensThreshold: 100, + currentPendingTokens: 50, + bufferedChunks: [ + { + id: "caller-chunk", + cycleId: "caller-cycle", + observations: "NOT PERSISTED", + tokenCount: 3, + messageIds: ["historical-message"], + messageTokens: 50, + lastObservedAt: new Date(), + createdAt: new Date(), + }, + ], + }), + ).toMatchObject({ chunksActivated: 0 }); + await domain.createReflectionGeneration({ + currentRecord: current, + reflection: "REFLECTED", + tokenCount: 3, + }); + const reflected = await domain.getObservationalMemory(THREAD, RESOURCE); + if (!reflected) throw new Error("Missing reflection generation"); + expect(reflected).toMatchObject({ + originType: "reflection", + metadata: { origin: "database" }, + }); + expect(reflected.lastObservedAt).toBeUndefined(); + const observation = { + id: reflected.id, + observations: "OBSERVED", + tokenCount: 3, + lastObservedAt: new Date("2026-01-02T00:00:00Z"), + }; + await domain.updateActiveObservations({ + ...observation, + observedMessageIds: ["historical-message"], + }); + await domain.updateActiveObservations(observation); + expect( + (await domain.getObservationalMemory(THREAD, RESOURCE))?.observedMessageIds, + ).toBeUndefined(); + await persistent.finish(); + + const inMemory = await createIsolatedMemoryReplay({ + ...options, + invocationId: "in-memory", + storeSemantics: "in-memory", + }); + const nativeDomain = inMemory.binding.domain; + const nativeCurrent = await nativeDomain.getObservationalMemory( + THREAD, + RESOURCE, + ); + if (!nativeCurrent) throw new Error("Missing restored observational memory"); + await nativeDomain.createReflectionGeneration({ + currentRecord: nativeCurrent, + reflection: "REFLECTED", + tokenCount: 3, + }); + expect( + (await nativeDomain.getObservationalMemory(THREAD, RESOURCE)) + ?.lastObservedAt, + ).toBeInstanceOf(Date); + await inMemory.finish(); +}); From 59f6039216006a4b6ec9d3e4a46e3c62b6f3d5c9 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 24 Sep 2026 12:55:16 +0200 Subject: [PATCH 023/115] Test Mastra memory replay on `@mastra/pg` Add `@mastra/pg` 1.25.0, the newest release compatible with `@mastra/core` 1.67.0, as a dev dependency and a test that records six turns on a real PostgreSQL store, including a buffered observation and a reflection. It requires every turn to be eligible and replays the reflection turn from the recorded tape. The test runs when `KITARU_TEST_MASTRA_POSTGRES_URL` is set and uses its own schema. The TypeScript CI job points it at its PostgreSQL service. --- .github/workflows/ci.yml | 7 +- packages/mastra/package.json | 1 + .../test/postgres-memory-replay.test.ts | 202 ++++++++++++++++++ pnpm-lock.yaml | 126 +++++++++++ 4 files changed, 335 insertions(+), 1 deletion(-) create mode 100644 packages/mastra/test/postgres-memory-replay.test.ts diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 352094aa3..65fafcc30 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -72,7 +72,12 @@ jobs: - run: pnpm run generate:check - run: pnpm run lint - run: pnpm run typecheck - - run: pnpm run test:built + - name: Test TypeScript packages, including Mastra memory on PostgreSQL + env: + POSTGRES_PASSWORD: password + run: >- + KITARU_TEST_MASTRA_POSTGRES_URL="postgres://postgres:${POSTGRES_PASSWORD}@localhost:5433/postgres" + pnpm run test:built - run: pnpm run pack:check:built - run: uv sync --frozen --extra server --extra cli --extra worker - name: Test compiled TypeScript adapters through Worker diff --git a/packages/mastra/package.json b/packages/mastra/package.json index cc79da4c7..1d380a3bf 100644 --- a/packages/mastra/package.json +++ b/packages/mastra/package.json @@ -56,6 +56,7 @@ "devDependencies": { "@mastra/core": "1.67.0", "@mastra/memory": "1.30.0", + "@mastra/pg": "1.25.0", "zod": "3.25.76" }, "scripts": { diff --git a/packages/mastra/test/postgres-memory-replay.test.ts b/packages/mastra/test/postgres-memory-replay.test.ts new file mode 100644 index 000000000..5eb923061 --- /dev/null +++ b/packages/mastra/test/postgres-memory-replay.test.ts @@ -0,0 +1,202 @@ +import { randomUUID } from "node:crypto"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { Memory } from "@mastra/memory"; +import { PostgresStore } from "@mastra/pg"; +import { afterEach, expect, it, vi } from "vitest"; +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, +} from "../src/memory.js"; +import { decodeMemoryReplayEnvelope } from "../src/memory-snapshot.js"; +import { textStream } from "./helpers/memory-agent.js"; +import { + AGENT_ID, + type ApiCall, + installTestApi, + REPLAY_ID, +} from "./helpers.js"; + +// A PostgreSQL connection URL, such as the repository's local test database +// on localhost:5433. The test creates and drops its own schema. +const POSTGRES_URL = process.env.KITARU_TEST_MASTRA_POSTGRES_URL; +const THREAD = "postgres-thread"; +const RESOURCE = "postgres-resource"; + +afterEach(() => { + vi.unstubAllEnvs(); + vi.unstubAllGlobals(); +}); + +function omModel(kind: "observer" | "reflector") { + const doStream = vi.fn(async () => + textStream( + kind === "observer" + ? `\n${"The user keeps sharing travel plans for the spring trip. ".repeat(12)}\n\nContinue.` + : "\nREFLECTED: the user is planning a spring trip.\n", + ), + ); + return { + doStream, + model: new MastraLanguageModelV2Mock({ + modelId: kind, + provider: "fixture", + doStream, + }), + }; +} + +function lastUpdate(calls: ApiCall[], sessionId: string) { + return calls.findLast( + (call) => call.method === "PATCH" && call.path.endsWith(sessionId), + )?.body; +} + +it.skipIf(!POSTGRES_URL)( + "records every turn on @mastra/pg and replays a reflection from the tape", + async () => { + const schemaName = `kitaru_mastra_${randomUUID().replaceAll("-", "")}`; + const store = new PostgresStore({ + id: "kitaru-postgres-replay", + connectionString: POSTGRES_URL as string, + schemaName, + }); + try { + await store.init(); + const domain = await store.getStore("memory"); + if (!domain) throw new Error("Missing PostgreSQL memory domain"); + const source = new Memory({ + storage: store, + options: { + lastMessages: 20, + observationalMemory: { + observation: { + model: "fixture/observer", + messageTokens: 300, + bufferTokens: 0.2, + bufferActivation: 1, + blockAfter: 1.1, + }, + reflection: { + model: "fixture/reflector", + observationTokens: 200, + bufferActivation: 1, + }, + }, + }, + }); + const observer = omModel("observer"); + const reflector = omModel("reflector"); + const actor = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("noted"), + }); + const api = installTestApi(); + const agent = createMemoryReplayAgent( + ({ memory }) => ({ + id: "postgres-memory", + name: "Postgres memory", + instructions: "Answer briefly.", + model: actor, + memory, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => source.settled(), + domain, + configuration: source.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: (id) => + id === "fixture/observer" + ? observer.model + : id === "fixture/reflector" + ? reflector.model + : actor, + }, + ); + + const turns: { sessionId: string; inputs: Record }[] = + []; + for (let turn = 1; turn <= 6; turn += 1) { + const result = await agent.stream( + `Turn ${turn}: ${"Here are more details about the spring trip. ".repeat(20)}`, + { memory: { thread: THREAD, resource: RESOURCE } }, + ); + await result.consumeStream(); + expect(await result.text).toBe("noted"); + const sessionId = api.sessionIds[turn - 1] as string; + await vi.waitFor( + () => + expect(lastUpdate(api.calls, sessionId)?.metadata).toHaveProperty( + "mastra_replay_state", + ), + { timeout: 10_000 }, + ); + const body = lastUpdate(api.calls, sessionId); + expect(body?.metadata, `turn ${turn}`).toMatchObject({ + mastra_replay_state: "eligible", + }); + turns.push({ + sessionId, + inputs: body?.inputs as Record, + }); + } + + const envelopes = turns.map(({ inputs }) => + decodeMemoryReplayEnvelope(inputs.mastra_memory_replay), + ); + expect(envelopes[0]?.configuration.memoryStore).toBe("persistent"); + // A later turn must start from a buffered chunk that PostgreSQL returned + // from its JSON column. + expect( + envelopes.some(({ initialSnapshot }) => + initialSnapshot.records.some( + (record) => (record.bufferedObservationChunks?.length ?? 0) > 0, + ), + ), + ).toBe(true); + const reflectionTurn = envelopes.findIndex(({ omTape }) => + omTape?.some( + (entry) => + (entry as { phase?: unknown } | null)?.phase === "reflector", + ), + ); + expect(reflectionTurn).toBeGreaterThan(0); + expect( + envelopes[reflectionTurn]?.omTape?.some( + (entry) => + (entry as { phase?: unknown } | null)?.phase === "observer", + ), + ).toBe(true); + + const observerCalls = observer.doStream.mock.calls.length; + const reflectorCalls = reflector.doStream.mock.calls.length; + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv( + "KITARU_TASK_INPUTS", + JSON.stringify(turns[reflectionTurn]?.inputs), + ); + const replay = await agent.stream("ignored"); + await replay.consumeStream(); + expect(await replay.text).toBe("noted"); + const replaySession = api.sessionIds.at(-1) as string; + await vi.waitFor( + () => + expect(lastUpdate(api.calls, replaySession)).toMatchObject({ + status: "completed", + }), + { timeout: 10_000 }, + ); + expect(observer.doStream).toHaveBeenCalledTimes(observerCalls); + expect(reflector.doStream).toHaveBeenCalledTimes(reflectorCalls); + } finally { + await store.pool.query(`DROP SCHEMA IF EXISTS "${schemaName}" CASCADE`); + await store.close(); + } + }, + 60_000, +); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 3037fe930..0986d28de 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -104,6 +104,9 @@ importers: '@mastra/memory': specifier: 1.30.0 version: 1.30.0(@mastra/core@1.67.0(ai@7.0.107(zod@3.25.76))(express@5.2.1)(zod@3.25.76)) + '@mastra/pg': + specifier: 1.25.0 + version: 1.25.0(@mastra/core@1.67.0(ai@7.0.107(zod@3.25.76))(express@5.2.1)(zod@3.25.76)) zod: specifier: 3.25.76 version: 3.25.76 @@ -303,6 +306,12 @@ packages: peerDependencies: '@mastra/core': '>=1.4.1-0 <2.0.0-0' + '@mastra/pg@1.25.0': + resolution: {integrity: sha512-RRgIRMBIKeCT/i5Mr8zFB1Pu6h9nLGui5Sd4C0wY454rOflexuHp1pUwlGsLRc3sB/79PhpAzPzyPkOoOV9ufQ==} + engines: {node: '>=22.13.0'} + peerDependencies: + '@mastra/core': '>=1.63.1-0 <2.0.0-0' + '@mastra/schema-compat@1.3.10': resolution: {integrity: sha512-ZrjvutWy3QJoJROA5PeSFdeFwDgyfGN1sOOObFirex9h9c/R4ReqRllCojNHSN0cKsKxkRlH6zBbWVTiybamUw==} engines: {node: '>=22.13.0'} @@ -1244,6 +1253,40 @@ packages: pathe@2.0.3: resolution: {integrity: sha512-WUjGcAqP1gQacoQe+OBJsFA7Ld4DyXuUIjZ5cc75cLHvJ7dtNsTugphxIADwspS+AraAUePCKrSVtPLFj/F88w==} + pg-cloudflare@1.4.0: + resolution: {integrity: sha512-Vo7z/6rrQYxpNRylp4Tlob2elzbh+N/MOQbxFVWCxS7oEx6jF53GTJFxK2WWpKuBRkmiin4Mt+xofFDjx09R0A==} + + pg-connection-string@2.14.0: + resolution: {integrity: sha512-XwWDGcLRGCXAR8F/AM5bG7Q+A3Wm2s6QeEjlOKZLlH3UYcguiqCWKyWXVag5TLTIjR7oOJUY8kcADaZgWPyLeg==} + + pg-int8@1.0.1: + resolution: {integrity: sha512-WCtabS6t3c8SkpDBUlb1kjOs7l66xsGdKpIPZsg4wR+B3+u9UAum2odSsF9tnvxg80h4ZxLWMy4pRjOsFIqQpw==} + engines: {node: '>=4.0.0'} + + pg-pool@3.14.0: + resolution: {integrity: sha512-gKtPkFdQPU3DksooVLi9LsjZxrsBUZIpa+7aVx+LV5pNh0KzP4Zleud2po+ConrxbuXGBJ6Hfer6hdgpIBpBaw==} + peerDependencies: + pg: '>=8.0' + + pg-protocol@1.16.0: + resolution: {integrity: sha512-sILXutLVjCLjcDuOmvhX5e2Z4cS5qG/6Bu3VkpFwdf/633ElGLpEh9bgmuI5I4sqKqkifQiGyiCcx1HdtrK7tg==} + + pg-types@2.2.0: + resolution: {integrity: sha512-qTAAlrEsl8s4OiEQY69wDvcMIdQN6wdz5ojQiOy6YRMuynxenON0O5oCpJI6lshc6scgAY8qvJ2On/p+CXY0GA==} + engines: {node: '>=4'} + + pg@8.23.0: + resolution: {integrity: sha512-Ip2EQCngowJLGOfCwkFhPXU7/ljlhn6Rxlmy4XYfL2Y+vyRM59+8uR2xqRWKdYmbXmxCFOAmKxBuSUCdF34qLg==} + engines: {node: '>= 16.0.0'} + peerDependencies: + pg-native: '>=3.0.1' + peerDependenciesMeta: + pg-native: + optional: true + + pgpass@1.0.5: + resolution: {integrity: sha512-FdW9r/jQZhSeohs1Z3sI1yxFQNFvMcnmfuj4WBMUTxOrAyLMaTcE1aAMBiTlbMNaXvBCQuVi0R7hd8udDSP7ug==} + picocolors@1.1.1: resolution: {integrity: sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA==} @@ -1259,6 +1302,22 @@ packages: resolution: {integrity: sha512-RRuzqDtt5Y9h3quz5hWhK+TPnsmVs6WwSU6LkJMeY4HstUEDuYTG8UJSdawMRzmzAtV+KEoG8N3Qg2qLy5vM/A==} engines: {node: ^10 || ^12 || >=14} + postgres-array@2.0.0: + resolution: {integrity: sha512-VpZrUqU5A69eQyW2c5CA1jtLecCsN2U/bD6VilrFDWq5+5UIEVO7nazS3TEcHf1zuPYO/sqGvUvW62g86RXZuA==} + engines: {node: '>=4'} + + postgres-bytea@1.0.1: + resolution: {integrity: sha512-5+5HqXnsZPE65IJZSMkZtURARZelel2oXUEO8rH83VS/hxH5vv1uHquPg5wZs8yMAfdv971IU+kcPUczi7NVBQ==} + engines: {node: '>=0.10.0'} + + postgres-date@1.0.7: + resolution: {integrity: sha512-suDmjLVQg78nMK2UZ454hAG+OAW+HQPZ6n++TNDUX+L0+uUlLywnoxJKDou51Zm+zTCjrCl0Nq6J9C5hP9vK/Q==} + engines: {node: '>=0.10.0'} + + postgres-interval@1.2.0: + resolution: {integrity: sha512-9ZhXKM/rw350N1ovuWHbGxnGh/SNJ4cnxHiM0rxE4VN41wsg8P8zWn9hv/buK00RP4WvlOyr/RBDiptyxVbkZQ==} + engines: {node: '>=0.10.0'} + posthog-node@5.46.1: resolution: {integrity: sha512-WjCqExq44pBdyg9MSsH6UAE0tNZ88p4aIuVFicgqhjf2Fbws6IhS4ioYUa4aBrbUPS9EDRXtBTtF5DpP1ml8Pw==} engines: {node: ^20.20.0 || >=22.22.0} @@ -1377,6 +1436,10 @@ packages: resolution: {integrity: sha512-UXWMKhLOwVKb728IUtQPXxfYU+usdybtUrK/8uGE8CQMvrhOpwvzDBwj0QhSL7MQc7vIsISBG8VQ8+IDQxpfQA==} engines: {node: '>=0.10.0'} + split2@4.2.0: + resolution: {integrity: sha512-UcjcJOWknrNkF6PLX83qcHM6KHgVKNkV62Y8a5uYDVv9ydGQVwAHMKqHdJje1VTWpljG0WYpCDhrCdAOYH4TWg==} + engines: {node: '>= 10.x'} + sprintf-js@1.0.3: resolution: {integrity: sha512-D9cPgkvLlV3t3IzL0D0YLvGA9Ahk4PcvVwUbN0dSGr1aP0Nrt4AEnTUbuGvquEC0mA64Gqt1fzirlRs5ibXx8g==} @@ -1602,6 +1665,10 @@ packages: utf-8-validate: optional: true + xtend@4.0.2: + resolution: {integrity: sha512-LKYU1iAXJXUgAXn9URjiu+MWhyUXHsvfp7mcuYm9dSUKK0/CjtrUwFAxD82/mCWbtLsGjFIad0wIsod4zrTAEQ==} + engines: {node: '>=0.4'} + xxhash-wasm@1.1.0: resolution: {integrity: sha512-147y/6YNh+tlp6nd/2pWq38i9h6mz/EuQ6njIrmW8D1BS5nCqs0P6DG+m6zTGnNz5I+uhZ0SHxBs9BsPrwcKDA==} @@ -1906,6 +1973,16 @@ snapshots: transitivePeerDependencies: - supports-color + '@mastra/pg@1.25.0(@mastra/core@1.67.0(ai@7.0.107(zod@3.25.76))(express@5.2.1)(zod@3.25.76))': + dependencies: + '@mastra/core': 1.67.0(ai@7.0.107(zod@3.25.76))(express@5.2.1)(zod@3.25.76) + async-mutex: 0.5.0 + pg: 8.23.0 + pg-connection-string: 2.14.0 + xxhash-wasm: 1.1.0 + transitivePeerDependencies: + - pg-native + '@mastra/schema-compat@1.3.10(zod@3.25.76)': dependencies: json-schema-to-zod: 2.8.1 @@ -3003,6 +3080,41 @@ snapshots: pathe@2.0.3: {} + pg-cloudflare@1.4.0: + optional: true + + pg-connection-string@2.14.0: {} + + pg-int8@1.0.1: {} + + pg-pool@3.14.0(pg@8.23.0): + dependencies: + pg: 8.23.0 + + pg-protocol@1.16.0: {} + + pg-types@2.2.0: + dependencies: + pg-int8: 1.0.1 + postgres-array: 2.0.0 + postgres-bytea: 1.0.1 + postgres-date: 1.0.7 + postgres-interval: 1.2.0 + + pg@8.23.0: + dependencies: + pg-connection-string: 2.14.0 + pg-pool: 3.14.0(pg@8.23.0) + pg-protocol: 1.16.0 + pg-types: 2.2.0 + pgpass: 1.0.5 + optionalDependencies: + pg-cloudflare: 1.4.0 + + pgpass@1.0.5: + dependencies: + split2: 4.2.0 + picocolors@1.1.1: {} picomatch@4.0.7: {} @@ -3015,6 +3127,16 @@ snapshots: picocolors: 1.1.1 source-map-js: 1.2.1 + postgres-array@2.0.0: {} + + postgres-bytea@1.0.1: {} + + postgres-date@1.0.7: {} + + postgres-interval@1.2.0: + dependencies: + xtend: 4.0.2 + posthog-node@5.46.1: dependencies: '@posthog/core': 1.45.1 @@ -3201,6 +3323,8 @@ snapshots: source-map-js@1.2.1: {} + split2@4.2.0: {} + sprintf-js@1.0.3: {} stackback@0.0.2: {} @@ -3360,6 +3484,8 @@ snapshots: ws@8.21.1: {} + xtend@4.0.2: {} + xxhash-wasm@1.1.0: {} yaml-ast-parser@0.0.43: {} From b6051acc61f7f3e67e69d94e76470e6305090a19 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 24 Sep 2026 13:24:27 +0200 Subject: [PATCH 024/115] Keep tape models out of production OM records The baseline gives Mastra tape-instrumented observer and reflector models, and Mastra persists its resolved OM configuration into the record it initializes. The source store therefore received Kitaru's wrapper instead of the configured model id. A JSON store serialized the wrapped provider client, so a `ModelRouterLanguageModel` grew each record's config from about 4.7 KB to about 464 KB and pushed later reflection evidence over the replay item limit. InMemoryStore kept the wrapper object itself, which the snapshot codec rejects, so no turn after the first was replayable. The capture binding now hands `initializeObservationalMemory` the model values from the source configuration, and the wrapper serializes as that value if it reaches any other JSON path. Initial snapshots and every mutation that carries OM records now apply the identity projection that initialization evidence already used: models become identities and Mastra's built-in extractors become slugs. Stores that return live Extractor objects can now record later turns. --- packages/mastra/src/memory-binding.ts | 120 ++++++++-- packages/mastra/src/memory-replay.ts | 2 +- packages/mastra/src/om-result-tape.ts | 30 ++- packages/mastra/test/helpers/memory-agent.ts | 35 ++- .../mastra/test/om-production-row.test.ts | 222 ++++++++++++++++++ packages/mastra/test/om-result-tape.test.ts | 17 ++ .../test/postgres-memory-replay.test.ts | 91 ++++--- 7 files changed, 455 insertions(+), 62 deletions(-) create mode 100644 packages/mastra/test/om-production-row.test.ts diff --git a/packages/mastra/src/memory-binding.ts b/packages/mastra/src/memory-binding.ts index d930cfb34..e1729eb5c 100644 --- a/packages/mastra/src/memory-binding.ts +++ b/packages/mastra/src/memory-binding.ts @@ -8,6 +8,7 @@ import { normalizeStoredMemoryDates, validateMemorySnapshot, } from "./memory-snapshot.js"; +import { getNativeOMModel } from "./om-result-tape.js"; export interface MastraMemorySelector { threadId: string; @@ -273,37 +274,105 @@ function modelIdentity(model: unknown): string { throw new Error("Observational-memory model has no stable identity."); } +const BUILTIN_EXTRACTORS = new Set([ + "current-task", + "suggested-response", + "thread-title", +]); + +function extractorIdentity(extractor: unknown): { + mastraBuiltinExtractor: string; +} { + if ( + record(extractor) && + typeof extractor.mastraBuiltinExtractor === "string" && + BUILTIN_EXTRACTORS.has(extractor.mastraBuiltinExtractor) + ) + return { mastraBuiltinExtractor: extractor.mastraBuiltinExtractor }; + if ( + record(extractor) && + extractor.internal === true && + typeof extractor.slug === "string" && + BUILTIN_EXTRACTORS.has(extractor.slug) + ) + return { mastraBuiltinExtractor: extractor.slug }; + throw new Error("Unsupported observational-memory extractor."); +} + +/** Replace OM models with identities and built-in extractors with their slugs. */ +function projectOMConfig( + config: Record, +): Record { + const projected = { ...config }; + if (projected.model !== undefined) + projected.model = modelIdentity(projected.model); + for (const name of ["observation", "reflection"]) { + const phase = projected[name]; + if (!record(phase)) continue; + const copy = { ...phase }; + if (copy.model !== undefined) copy.model = modelIdentity(copy.model); + if (Array.isArray(copy.extractors)) + copy.extractors = copy.extractors.map(extractorIdentity); + projected[name] = copy; + } + return projected; +} + +/** + * Project the OM configuration held by stored records. + * + * Mastra keeps its resolved OM configuration, including model and Extractor + * objects, in every record. Stores that keep objects in memory return them + * as they are, and those objects have no stable encoding. + */ +function projectOMRecords(value: unknown): unknown { + if (Array.isArray(value)) return value.map(projectOMRecords); + if (!record(value)) return value; + const projected = { ...value }; + if (record(value.config)) projected.config = projectOMConfig(value.config); + if (record(value.currentRecord)) + projected.currentRecord = projectOMRecords(value.currentRecord); + return projected; +} + +const OM_RECORD_METHODS = new Set([ + "initializeObservationalMemory", + "insertObservationalMemoryRecord", + "createReflectionGeneration", + "swapBufferedReflectionToActive", +]); + /** Keep provider clients out of evidence without changing native storage calls. */ function mutationEvidenceValue(method: PropertyKey, value: unknown): unknown { - if (method !== "initializeObservationalMemory") return value; - function project(input: unknown): unknown { + return OM_RECORD_METHODS.has(method) ? projectOMRecords(value) : value; +} + +/** + * Give the source store the OM models from the source configuration. + * + * Mastra persists its OM configuration into the record it initializes. Tape + * instrumented models must not reach production rows, where they would add + * Kitaru objects or a serialized provider client. + */ +function nativeStorageArguments( + method: PropertyKey, + args: unknown[], +): unknown[] { + if (method !== "initializeObservationalMemory") return args; + return args.map((input) => { if (!record(input) || !record(input.config)) return input; + let changed = false; const config = { ...input.config }; - if (config.model !== undefined) config.model = modelIdentity(config.model); for (const name of ["observation", "reflection"]) { const phase = config[name]; if (!record(phase)) continue; - const projected = { ...phase }; - if (projected.model !== undefined) - projected.model = modelIdentity(projected.model); - if (Array.isArray(projected.extractors)) - projected.extractors = projected.extractors.map((extractor) => { - if ( - record(extractor) && - extractor.internal === true && - typeof extractor.slug === "string" && - ["current-task", "suggested-response", "thread-title"].includes( - extractor.slug, - ) - ) - return { mastraBuiltinExtractor: extractor.slug }; - throw new Error("Unsupported observational-memory extractor."); - }); - config[name] = projected; + const native = getNativeOMModel(phase.model); + if (native === phase.model) continue; + config[name] = { ...phase, model: native }; + changed = true; } - return { ...input, config }; - } - return Array.isArray(value) ? value.map(project) : project(value); + return changed ? { ...input, config } : input; + }); } /** Capture one native invocation without changing the shared source domain or Agent. */ @@ -342,7 +411,8 @@ export function createMemoryCaptureBinding( methods.set(property, bound); return bound; } - const bound = (...args: unknown[]): Promise => { + const bound = (...callerArgs: unknown[]): Promise => { + const args = nativeStorageArguments(property, callerArgs); let encodedArguments: JsonValue = null; let complete = true; let requestId: string | undefined; @@ -596,7 +666,7 @@ export function createMemoryCaptureBinding( thread, resource, messages, - records, + records: projectOMRecords(records), }; // No storage-owned objects or Dates escape the explicit codec. const copy = normalizeStoredMemoryDates( diff --git a/packages/mastra/src/memory-replay.ts b/packages/mastra/src/memory-replay.ts index df6972247..e192c1518 100644 --- a/packages/mastra/src/memory-replay.ts +++ b/packages/mastra/src/memory-replay.ts @@ -169,7 +169,7 @@ export async function bindOMResultModels( : modelIdentity; if (!record(model) || typeof model.doStream !== "function") unsupported("OM model resolution did not return a stream-capable model."); - settings.model = tape.instrument(model, phase); + settings.model = tape.instrument(model, phase, modelIdentity); bound[name] = settings; } return { ...config, observationalMemory: bound } as MemoryConfigInternal; diff --git a/packages/mastra/src/om-result-tape.ts b/packages/mastra/src/om-result-tape.ts index d311f1da2..306a27fda 100644 --- a/packages/mastra/src/om-result-tape.ts +++ b/packages/mastra/src/om-result-tape.ts @@ -52,6 +52,21 @@ interface ModelLike { doStream?: (input: unknown) => Promise; } +const nativeModels = new WeakMap(); + +/** Return the configured model value that an instrumented OM model replaced. */ +export function getNativeOMModel(model: unknown): unknown { + return typeof model === "object" && model !== null && nativeModels.has(model) + ? nativeModels.get(model) + : model; +} + +function serializeNativeModel(native: unknown): unknown { + if (typeof native !== "object" || native === null) return native; + const toJSON: unknown = Reflect.get(native, "toJSON"); + return typeof toJSON === "function" ? toJSON.call(native) : native; +} + /** Intercept only OM model calls; the actor model remains untouched. */ export function createOMResultTape( recorded: readonly OMResultEntry[] | undefined, @@ -93,9 +108,21 @@ export function createOMResultTape( return entry; } - function instrument(model: T, phase: OMPhase): T { + /** + * Wrap an OM model so its calls go through the tape. + * + * `native` is the value from the source memory configuration, such as a + * model id string. Mastra persists the OM configuration into the record, so + * a store that serializes the wrapper writes this value instead. + */ + function instrument( + model: T, + phase: OMPhase, + native: unknown = model, + ): T { const proxy = new Proxy(model, { get(target, key) { + if (key === "toJSON") return () => serializeNativeModel(native); const value = Reflect.get(target, key, target); if (key !== "doGenerate" && key !== "doStream") return typeof value === "function" ? value.bind(target) : value; @@ -188,6 +215,7 @@ export function createOMResultTape( }; }, }); + nativeModels.set(proxy, native); return proxy; } diff --git a/packages/mastra/test/helpers/memory-agent.ts b/packages/mastra/test/helpers/memory-agent.ts index e97e3ab1a..780f1252e 100644 --- a/packages/mastra/test/helpers/memory-agent.ts +++ b/packages/mastra/test/helpers/memory-agent.ts @@ -1,7 +1,7 @@ import { Agent } from "@mastra/core/agent"; import type { InputProcessor } from "@mastra/core/processors"; import { RequestContext } from "@mastra/core/request-context"; -import { InMemoryStore } from "@mastra/core/storage"; +import { InMemoryStore, type MemoryStorage } from "@mastra/core/storage"; import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; import { createTool } from "@mastra/core/tools"; import { LocalSkillSource, Workspace } from "@mastra/core/workspace"; @@ -75,6 +75,39 @@ export function memoryModel( }; } +/** Run one short native Mastra turn and return the OM configuration it stores. */ +export async function getNativeOMRecordConfig( + options: NonNullable[0]>["options"], + prepare?: (domain: MemoryStorage) => void, +): Promise { + const store = new InMemoryStore(); + const domain = store.stores.memory; + if (!domain) throw new Error("Missing native memory domain"); + prepare?.(domain); + const agent = new Agent({ + id: "native-om-record", + name: "Native OM record", + instructions: "Answer briefly.", + model: new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("noted"), + }), + memory: new Memory({ storage: store, options }), + }); + try { + const result = await agent.stream("Hello.", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + await result.consumeStream(); + const [row] = await domain.getObservationalMemoryHistory(THREAD, RESOURCE); + if (!row) throw new Error("Native Mastra did not create an OM record"); + return row.config; + } finally { + await store.close(); + } +} + export function createMemoryRuntime( options: { observerWait?: () => Promise; messageTokens?: number } = {}, ) { diff --git a/packages/mastra/test/om-production-row.test.ts b/packages/mastra/test/om-production-row.test.ts new file mode 100644 index 000000000..2e802ea40 --- /dev/null +++ b/packages/mastra/test/om-production-row.test.ts @@ -0,0 +1,222 @@ +import { ModelRouterLanguageModel } from "@mastra/core/llm"; +import type { MemoryStorage } from "@mastra/core/storage"; +import { InMemoryStore } from "@mastra/core/storage"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { Memory } from "@mastra/memory"; +import { afterEach, expect, it, vi } from "vitest"; +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, +} from "../src/memory.js"; +import { decodeMemoryReplayEnvelope } from "../src/memory-snapshot.js"; +import { getNativeOMRecordConfig, textStream } from "./helpers/memory-agent.js"; +import { + AGENT_ID, + type ApiCall, + installTestApi, + REPLAY_ID, +} from "./helpers.js"; + +const THREAD = "row-thread"; +const RESOURCE = "row-resource"; +const TURNS = 5; + +// Default continuation hints put Mastra's built-in extractors into the record. +const OPTIONS = { + lastMessages: 20, + observationalMemory: { + observation: { + model: "fixture/observer", + messageTokens: 300, + bufferTokens: 0.2, + bufferActivation: 1, + blockAfter: 1.1, + }, + reflection: { + model: "fixture/reflector", + observationTokens: 200, + bufferActivation: 1, + }, + }, +}; + +afterEach(() => { + vi.unstubAllEnvs(); + vi.unstubAllGlobals(); +}); + +function omOutput(kind: "observer" | "reflector") { + return textStream( + kind === "observer" + ? `\n${"The user keeps sharing travel plans for the spring trip. ".repeat(12)}\n\nContinue.` + : "\nREFLECTED: the user is planning a spring trip.\n", + ); +} + +function omModel(kind: "observer" | "reflector", router: boolean) { + const doStream = vi.fn(async () => omOutput(kind)); + if (!router) + return { + doStream, + model: new MastraLanguageModelV2Mock({ + modelId: kind, + provider: "fixture", + doStream, + }), + }; + // The router model carries the gateway catalog as enumerable state. + const model = new ModelRouterLanguageModel("openai/gpt-5-nano"); + Object.assign(model, { doStream }); + return { doStream, model }; +} + +// Store `config` the way a JSON column does, as PostgreSQL and LibSQL do. +function storeConfigAsJson(domain: MemoryStorage): void { + const initialize = domain.initializeObservationalMemory.bind(domain); + const reflect = domain.createReflectionGeneration.bind(domain); + domain.initializeObservationalMemory = (input) => + initialize({ ...input, config: JSON.parse(JSON.stringify(input.config)) }); + domain.createReflectionGeneration = (input) => + reflect({ + ...input, + currentRecord: { + ...input.currentRecord, + config: JSON.parse(JSON.stringify(input.currentRecord.config)), + }, + }); +} + +function lastUpdate(calls: ApiCall[], sessionId: string) { + return calls.findLast( + (call) => call.method === "PATCH" && call.path.endsWith(sessionId), + )?.body; +} + +it.each([ + { store: "in-memory", model: "mock" }, + { store: "in-memory", model: "router" }, + { store: "json", model: "mock" }, + { store: "json", model: "router" }, +] as const)( + "keeps Kitaru objects out of production OM records ($store store, $model model)", + async ({ store: storeKind, model: modelKind }) => { + vi.stubEnv("OPENAI_API_KEY", "sk-test-placeholder"); + const nativeConfig = JSON.stringify( + await getNativeOMRecordConfig( + OPTIONS, + storeKind === "json" ? storeConfigAsJson : undefined, + ), + ); + const store = new InMemoryStore(); + const domain = store.stores.memory as MemoryStorage; + if (storeKind === "json") storeConfigAsJson(domain); + const source = new Memory({ storage: store, options: OPTIONS }); + const observer = omModel("observer", modelKind === "router"); + const reflector = omModel("reflector", modelKind === "router"); + const actor = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => textStream("noted"), + }); + const api = installTestApi(); + const agent = createMemoryReplayAgent( + ({ memory }) => ({ + id: "om-row", + name: "OM row", + instructions: "Answer briefly.", + model: actor, + memory, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + requestedModelId: "fixture/actor", + sourceMemory: () => ({ + settled: () => source.settled(), + domain, + configuration: source.getMergedThreadConfig(), + exclusiveAccess: createProcessLocalMemoryAccess(), + }), + resolveModel: (id) => + id === "fixture/observer" + ? observer.model + : id === "fixture/reflector" + ? reflector.model + : actor, + }, + ); + try { + const turns: Record[] = []; + for (let turn = 1; turn <= TURNS; turn += 1) { + const result = await agent.stream( + `Turn ${turn}: ${"Here are more details about the spring trip. ".repeat(20)}`, + { memory: { thread: THREAD, resource: RESOURCE } }, + ); + await result.consumeStream(); + expect(await result.text).toBe("noted"); + const sessionId = api.sessionIds[turn - 1] as string; + await vi.waitFor( + () => + expect(lastUpdate(api.calls, sessionId)?.metadata).toHaveProperty( + "mastra_replay_state", + ), + { timeout: 10_000 }, + ); + const body = lastUpdate(api.calls, sessionId); + expect(body?.metadata, `turn ${turn}`).toMatchObject({ + mastra_replay_state: "eligible", + }); + turns.push(body?.inputs as Record); + + const rows = await domain.getObservationalMemoryHistory( + THREAD, + RESOURCE, + ); + expect(rows.length).toBeGreaterThan(0); + for (const row of rows) { + expect(row.config.observation).toMatchObject({ + model: "fixture/observer", + }); + expect(row.config.reflection).toMatchObject({ + model: "fixture/reflector", + }); + expect(JSON.stringify(row.config)).toBe(nativeConfig); + } + } + + const envelopes = turns.map((inputs) => + decodeMemoryReplayEnvelope(inputs.mastra_memory_replay), + ); + const reflectionTurn = envelopes.findIndex(({ omTape }) => + omTape?.some( + (entry) => + (entry as { phase?: unknown } | null)?.phase === "reflector", + ), + ); + expect(reflectionTurn).toBeGreaterThan(0); + const observerCalls = observer.doStream.mock.calls.length; + const reflectorCalls = reflector.doStream.mock.calls.length; + for (const index of [reflectionTurn, TURNS - 1]) { + vi.stubEnv("KITARU_REPLAY_ID", REPLAY_ID); + vi.stubEnv("KITARU_TASK_INPUTS", JSON.stringify(turns[index])); + const replay = await agent.stream("ignored"); + await replay.consumeStream(); + expect(await replay.text).toBe("noted"); + const replaySession = api.sessionIds.at(-1) as string; + await vi.waitFor( + () => + expect(lastUpdate(api.calls, replaySession)).toMatchObject({ + status: "completed", + }), + { timeout: 10_000 }, + ); + } + expect(observer.doStream).toHaveBeenCalledTimes(observerCalls); + expect(reflector.doStream).toHaveBeenCalledTimes(reflectorCalls); + } finally { + await source.settled(); + await store.close(); + } + }, + 60_000, +); diff --git a/packages/mastra/test/om-result-tape.test.ts b/packages/mastra/test/om-result-tape.test.ts index 01a695db4..5143e780c 100644 --- a/packages/mastra/test/om-result-tape.test.ts +++ b/packages/mastra/test/om-result-tape.test.ts @@ -2,6 +2,7 @@ import { expect, it, vi } from "vitest"; import { decodeMemoryValue } from "../src/memory-snapshot.js"; import { createOMResultTape, + getNativeOMModel, MastraOMDivergenceError, } from "../src/om-result-tape.js"; @@ -57,6 +58,22 @@ it("reuses ordered recorded OM output without calling the live model", async () await replay.finish(); }); +it("serializes an instrumented OM model as its configured value", () => { + const tape = createOMResultTape(undefined, () => {}); + const fromId = tape.instrument(model(), "observer", "fixture/observer"); + expect(JSON.stringify({ model: fromId })).toBe( + '{"model":"fixture/observer"}', + ); + expect(getNativeOMModel(fromId)).toBe("fixture/observer"); + const configured = model(); + const fromObject = tape.instrument(configured, "reflector"); + expect(JSON.parse(JSON.stringify(fromObject))).toEqual( + JSON.parse(JSON.stringify(configured)), + ); + expect(getNativeOMModel(fromObject)).toBe(configured); + expect(getNativeOMModel(configured)).toBe(configured); +}); + it("rejects an extra or missing OM call", async () => { const capture = createOMResultTape(undefined, () => {}); const output = await capture.instrument(model(), "reflector").doStream({}); diff --git a/packages/mastra/test/postgres-memory-replay.test.ts b/packages/mastra/test/postgres-memory-replay.test.ts index 5eb923061..5a05340a0 100644 --- a/packages/mastra/test/postgres-memory-replay.test.ts +++ b/packages/mastra/test/postgres-memory-replay.test.ts @@ -1,4 +1,5 @@ import { randomUUID } from "node:crypto"; +import { ModelRouterLanguageModel } from "@mastra/core/llm"; import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; import { Memory } from "@mastra/memory"; import { PostgresStore } from "@mastra/pg"; @@ -8,7 +9,7 @@ import { createProcessLocalMemoryAccess, } from "../src/memory.js"; import { decodeMemoryReplayEnvelope } from "../src/memory-snapshot.js"; -import { textStream } from "./helpers/memory-agent.js"; +import { getNativeOMRecordConfig, textStream } from "./helpers/memory-agent.js"; import { AGENT_ID, type ApiCall, @@ -27,7 +28,25 @@ afterEach(() => { vi.unstubAllGlobals(); }); -function omModel(kind: "observer" | "reflector") { +const OPTIONS = { + lastMessages: 20, + observationalMemory: { + observation: { + model: "fixture/observer", + messageTokens: 300, + bufferTokens: 0.2, + bufferActivation: 1, + blockAfter: 1.1, + }, + reflection: { + model: "fixture/reflector", + observationTokens: 200, + bufferActivation: 1, + }, + }, +}; + +function omModel(kind: "observer" | "reflector", router: boolean) { const doStream = vi.fn(async () => textStream( kind === "observer" @@ -35,14 +54,19 @@ function omModel(kind: "observer" | "reflector") { : "\nREFLECTED: the user is planning a spring trip.\n", ), ); - return { - doStream, - model: new MastraLanguageModelV2Mock({ - modelId: kind, - provider: "fixture", + if (!router) + return { doStream, - }), - }; + model: new MastraLanguageModelV2Mock({ + modelId: kind, + provider: "fixture", + doStream, + }), + }; + // The router model carries the gateway catalog as enumerable state. + const model = new ModelRouterLanguageModel("openai/gpt-5-nano"); + Object.assign(model, { doStream }); + return { doStream, model }; } function lastUpdate(calls: ApiCall[], sessionId: string) { @@ -51,9 +75,13 @@ function lastUpdate(calls: ApiCall[], sessionId: string) { )?.body; } -it.skipIf(!POSTGRES_URL)( - "records every turn on @mastra/pg and replays a reflection from the tape", - async () => { +it.skipIf(!POSTGRES_URL).each(["mock", "router"] as const)( + "records every turn on @mastra/pg and replays a reflection from the tape (%s OM models)", + async (modelKind) => { + vi.stubEnv("OPENAI_API_KEY", "sk-test-placeholder"); + const nativeConfigBytes = JSON.stringify( + await getNativeOMRecordConfig(OPTIONS), + ).length; const schemaName = `kitaru_mastra_${randomUUID().replaceAll("-", "")}`; const store = new PostgresStore({ id: "kitaru-postgres-replay", @@ -64,28 +92,9 @@ it.skipIf(!POSTGRES_URL)( await store.init(); const domain = await store.getStore("memory"); if (!domain) throw new Error("Missing PostgreSQL memory domain"); - const source = new Memory({ - storage: store, - options: { - lastMessages: 20, - observationalMemory: { - observation: { - model: "fixture/observer", - messageTokens: 300, - bufferTokens: 0.2, - bufferActivation: 1, - blockAfter: 1.1, - }, - reflection: { - model: "fixture/reflector", - observationTokens: 200, - bufferActivation: 1, - }, - }, - }, - }); - const observer = omModel("observer"); - const reflector = omModel("reflector"); + const source = new Memory({ storage: store, options: OPTIONS }); + const observer = omModel("observer", modelKind === "router"); + const reflector = omModel("reflector", modelKind === "router"); const actor = new MastraLanguageModelV2Mock({ modelId: "actor", provider: "fixture", @@ -146,6 +155,20 @@ it.skipIf(!POSTGRES_URL)( }); } + // PostgreSQL reorders JSONB keys, so compare sizes rather than text. + for (const row of await domain.getObservationalMemoryHistory( + THREAD, + RESOURCE, + )) { + expect(row.config.observation).toMatchObject({ + model: "fixture/observer", + }); + expect(row.config.reflection).toMatchObject({ + model: "fixture/reflector", + }); + expect(JSON.stringify(row.config).length).toBe(nativeConfigBytes); + } + const envelopes = turns.map(({ inputs }) => decodeMemoryReplayEnvelope(inputs.mastra_memory_replay), ); From 1be648f67d729eb61caf69874b3b3e3135139d50 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 24 Sep 2026 14:21:18 +0200 Subject: [PATCH 025/115] Scope Mastra lease poison to the turns it affects The source-thread lease marked a thread and its resource as permanently unsafe whenever the adapter itself wrote outside an eligible lease: a buffered reflection finishing after release, a reply sent while the previous turn was still finalizing, a native fallback after a Kitaru error, or one call without memory selectors. Every later turn for that user then stayed unreplayable until restart, and forever under a shared lease that keeps markers across process loss. - A write outside an eligible lease now registers through `acquire(selector, { waitMs: 0 })` for the length of the write. An overlapping holder loses eligibility, and nothing outlives the write. `markUnsafeWrite` is left for writes that cannot register. - Native fallbacks register each write under the caller's or the factory's default selectors. A call with no selectors marks nothing unless it actually writes. - Finalization joins Mastra's buffered observation and reflection through `waitForBuffering` before the drain, bounded by the new `finalizationWaitMs` (60 s). A turn that misses it is ineligible, and the lease is released before the final session update. - The process-local lease waits up to `waitMs` for a releasing holder. - The OM tape tracks a call from its start, so a slow provider no longer leaves an empty slot. - A source `memory` option lets the application's `settled()` join each recorded turn's memory work, and capture then waits only for the recorded thread's buffering. --- docs/book/adapters/mastra.md | 5 +- packages/mastra/README.md | 5 +- packages/mastra/src/memory-binding.ts | 389 +++++++++++---- packages/mastra/src/memory-replay.ts | 11 +- packages/mastra/src/om-result-tape.ts | 29 +- packages/mastra/src/stateful-agent.ts | 139 ++++-- packages/mastra/src/stream-recording.ts | 5 + packages/mastra/test/memory-binding.test.ts | 49 +- .../test/memory-lease-lifecycle.test.ts | 448 ++++++++++++++++++ .../mastra/test/memory-replay-safety.test.ts | 13 +- packages/mastra/test/processor-replay.test.ts | 8 +- .../mastra/test/stateful-overrides.test.ts | 34 +- 12 files changed, 971 insertions(+), 164 deletions(-) create mode 100644 packages/mastra/test/memory-lease-lifecycle.test.ts diff --git a/docs/book/adapters/mastra.md b/docs/book/adapters/mastra.md index a243fbf71..e1fb6ae93 100644 --- a/docs/book/adapters/mastra.md +++ b/docs/book/adapters/mastra.md @@ -235,6 +235,7 @@ const recorded = createMemoryReplayAgent( domain: store.stores.memory!, configuration: sourceMemory.getMergedThreadConfig(), settled: () => sourceMemory.settled(), + memory: sourceMemory, exclusiveAccess, }), resolveModel: (id) => { @@ -256,9 +257,9 @@ Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KIT ### Source ownership and supported configuration -All writers to a source thread or resource must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. Turns sharing either selector can continue natively, but their recordings become ineligible for replay. `settled()` joins pending work on the source `Memory` instance; it does not provide exclusive access. +All writers to a source thread or resource must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. A new turn waits up to 100 ms for an earlier turn on the same thread or resource to release it. A recorded turn holds both selectors until its memory writes, including delayed observational-memory work, have settled, or until `finalizationWaitMs` passes (60 seconds by default, after which the turn is ineligible). Kitaru then releases the selectors and sends the final session update. Turns that overlap on either selector still answer natively; only the overlapping turns become ineligible for replay, and later turns are unaffected. Pass the source `Memory` as `memory` so that its `settled()` also waits for the observational-memory work of recorded turns before you close storage. `settled()` does not provide exclusive access. -A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must atomically reserve both thread and resource IDs, preserve ownership through the final eligible-session update, mark an overlapping or unowned write unsafe through `markUnsafeWrite()`, and retain unsafe state across process loss. Call `resetAfterQuiescence()` only after all possible stale writers for both selectors have stopped. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. +A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must atomically reserve both thread and resource IDs. When either ID is still held after `waitMs`, it must invalidate the current holders and return a lease that is not eligible but holds both IDs until it is released; that invalidation ends once every overlapping lease has been released. `waitMs: 0` must not wait: Kitaru registers each write it makes outside its own lease this way and releases the registration when the write finishes. `markUnsafeWrite()` is only for a write that could not register, because coordination failed or its selector is unknown, and that marker must survive process loss. Kitaru never calls `resetAfterQuiescence()`. Your application calls it for the marked selectors, or with no selector after an unknown-selector marker, once every process that might have written without registering has stopped or restarted. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. Schema working memory requires explicit `scope: "thread"`. An observational-memory configuration object may omit `scope`, using Mastra's implicit thread scope, or set it to `"thread"`. Supply explicit observer/reflector model identities, either shared through `observationalMemory.model` or in the phase configuration. Resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` remain unsupported. diff --git a/packages/mastra/README.md b/packages/mastra/README.md index b868addb6..76b50e7e1 100644 --- a/packages/mastra/README.md +++ b/packages/mastra/README.md @@ -199,6 +199,7 @@ const recorded = createMemoryReplayAgent( domain: store.stores.memory!, configuration: sourceMemory.getMergedThreadConfig(), settled: () => sourceMemory.settled(), + memory: sourceMemory, exclusiveAccess, }), resolveModel: (id) => { @@ -220,9 +221,9 @@ Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KIT ### Source ownership and supported configuration -All writers to a source thread or resource must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. Turns sharing either selector can continue natively, but their recordings become ineligible for replay. `settled()` joins pending work on the source `Memory` instance; it does not provide exclusive access. +All writers to a source thread or resource must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. A new turn waits up to 100 ms for an earlier turn on the same thread or resource to release it. A recorded turn holds both selectors until its memory writes, including delayed observational-memory work, have settled, or until `finalizationWaitMs` passes (60 seconds by default, after which the turn is ineligible). Kitaru then releases the selectors and sends the final session update. Turns that overlap on either selector still answer natively; only the overlapping turns become ineligible for replay, and later turns are unaffected. Pass the source `Memory` as `memory` so that its `settled()` also waits for the observational-memory work of recorded turns before you close storage. `settled()` does not provide exclusive access. -A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must atomically reserve both thread and resource IDs, preserve ownership through the final eligible-session update, mark an overlapping or unowned write unsafe through `markUnsafeWrite()`, and retain unsafe state across process loss. Call `resetAfterQuiescence()` only after all possible stale writers for both selectors have stopped. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. +A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must atomically reserve both thread and resource IDs. When either ID is still held after `waitMs`, it must invalidate the current holders and return a lease that is not eligible but holds both IDs until it is released; that invalidation ends once every overlapping lease has been released. `waitMs: 0` must not wait: Kitaru registers each write it makes outside its own lease this way and releases the registration when the write finishes. `markUnsafeWrite()` is only for a write that could not register, because coordination failed or its selector is unknown, and that marker must survive process loss. Kitaru never calls `resetAfterQuiescence()`. Your application calls it for the marked selectors, or with no selector after an unknown-selector marker, once every process that might have written without registering has stopped or restarted. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. Schema working memory requires explicit `scope: "thread"`. An observational-memory configuration object may omit `scope`, using Mastra's implicit thread scope, or set it to `"thread"`. Supply explicit observer/reflector model identities, either shared through `observationalMemory.model` or in the phase configuration. Resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` remain unsupported. diff --git a/packages/mastra/src/memory-binding.ts b/packages/mastra/src/memory-binding.ts index e1729eb5c..36b1ba2a5 100644 --- a/packages/mastra/src/memory-binding.ts +++ b/packages/mastra/src/memory-binding.ts @@ -16,7 +16,10 @@ export interface MastraMemorySelector { } export interface MastraMemoryLeaseOptions { - /** Bound acquisition so a live answer does not wait for a delayed reflection. */ + /** + * Bound acquisition so a live answer does not wait for a delayed reflection. + * Zero means do not wait: invalidate any current holder at once. + */ waitMs?: number; /** Cancel a waiting acquisition without releasing another writer's lease. */ signal?: AbortSignal; @@ -34,13 +37,23 @@ export interface MastraMemoryLease { /** * Coordinate every writer of a source thread or resource across all processes. * - * The implementation must atomically poison eligibility for both selectors when - * a competing native turn proceeds without ownership or an owner loses its lease. - * Poison must survive process loss and prevent a later acquisition from becoming - * eligible until all possible stale writers have quiesced. A timeout or failed - * coordination call must fail closed for replay eligibility, while native Mastra - * storage writes still run. Keep ownership through the final eligible-session - * update, then release only after no delayed source write remains possible. + * `acquire` atomically reserves both selectors. When either selector is still + * held after `waitMs`, the implementation invalidates every current holder and + * returns a lease that is not eligible but still occupies both selectors until + * it is released. That invalidation ends once every overlapping lease has been + * released. Kitaru holds a turn's lease until the turn's memory writes, + * including delayed observational-memory work, have settled or reached their + * finalization deadline, and until the final eligibility check has passed. + * A write Kitaru makes outside that lease registers through + * `acquire(selector, { waitMs: 0 })` and releases immediately afterwards. + * + * `markUnsafeWrite` is only for a write that could not register: its selector + * is unknown or coordination failed. That marker must survive process loss and + * keep both selectors ineligible until `resetAfterQuiescence`. Kitaru never + * calls `resetAfterQuiescence`; the application calls it after every process + * that could have written without registering has stopped or restarted. A + * timeout or failed coordination call must fail closed for replay eligibility, + * while native Mastra storage writes still run. * * The process-local helper below is valid only when every writer shares one * instance in one process. A production multi-server application must provide @@ -51,7 +64,7 @@ export interface MastraExclusiveMemoryAccess { selector: MastraMemorySelector, options?: MastraMemoryLeaseOptions, ): Promise; - /** Persist an unsafe-write marker before an unowned native write proceeds. + /** Persist an unsafe-write marker before an unregistered native write proceeds. * An unknown selector poisons every thread until global quiescence is proven. */ markUnsafeWrite(selector?: MastraMemorySelector): Promise; @@ -71,6 +84,7 @@ export function createProcessLocalMemoryAccess(): MastraExclusiveMemoryAccess { persistentLoss: boolean; }; const scopes = new Map(); + const releaseWaiters = new Set<() => void>(); let unknownWriterPoisoned = false; function keys({ threadId, resourceId }: MastraMemorySelector): string[] { @@ -103,11 +117,39 @@ export function createProcessLocalMemoryAccess(): MastraExclusiveMemoryAccess { } } + function isOccupied(scopeKeys: readonly string[]): boolean { + return scopeKeys.some((key) => (scopes.get(key)?.turns.size ?? 0) > 0); + } + + async function waitForRelease( + scopeKeys: readonly string[], + waitMs: number, + signal: AbortSignal | undefined, + ): Promise { + const deadline = Date.now() + waitMs; + while (isOccupied(scopeKeys) && !signal?.aborted) { + const remaining = deadline - Date.now(); + if (remaining <= 0) return; + await new Promise((resolve) => { + const done = () => { + clearTimeout(timer); + releaseWaiters.delete(done); + signal?.removeEventListener("abort", done); + resolve(); + }; + const timer = setTimeout(done, remaining); + releaseWaiters.add(done); + signal?.addEventListener("abort", done, { once: true }); + }); + } + } + return { async acquire(selector, options = {}) { + const scopeKeys = keys(selector); + await waitForRelease(scopeKeys, options.waitMs ?? 0, options.signal); if (options.signal?.aborted) throw new Error("Exclusive source-thread ownership was cancelled."); - const scopeKeys = keys(selector); const states = scopeKeys.map(getState); const occupied = states.filter((state) => state.turns.size > 0); if (occupied.length > 0) poison(occupied, false); @@ -127,6 +169,7 @@ export function createProcessLocalMemoryAccess(): MastraExclusiveMemoryAccess { if (state && state.turns.size === 0 && !state.persistentLoss) scopes.delete(key); } + for (const notify of [...releaseWaiters]) notify(); }; return Object.assign(release, { async verifyEligibility() { @@ -201,11 +244,33 @@ export interface MastraMemoryCaptureBinding { }): Promise; markIncomplete(reason: string): void; drain(): Promise; + /** + * Join the invocation's memory work, including buffered observation and + * reflection, and record its evidence. Without `waitMs` this waits until the + * work settles; with it, resolve false once `waitMs` passes first. A failed + * join rejects either way. + */ + settle(memory: MastraSettlingMemory, waitMs?: number): Promise; /** Check shared ownership immediately before persisting eligible inputs. */ verifyEligibility(): Promise; release(): Promise; } +/** A Memory whose background observational-memory work can be joined. */ +export interface MastraSettlingMemory { + settled(): Promise; + readonly omEngine: Promise<{ + waitForBuffering( + threadId: string, + resourceId: string, + timeoutMs?: number, + ): Promise; + } | null>; +} + +const BUFFERING_WAIT_MS = 30_000; +const ACQUIRE_RESPONSE_MARGIN_MS = 25; + async function boundedCoordination( operation: Promise, waitMs: number, @@ -224,6 +289,147 @@ async function boundedCoordination( } } +/** + * Acquire within `boundWaitMs`, even from a backend that ignores cancellation. + * + * A lease that arrives after the caller has given up is released at once so it + * cannot keep the selectors occupied. + */ +async function acquireWithin( + access: MastraExclusiveMemoryAccess, + selector: MastraMemorySelector, + options: MastraMemoryLeaseOptions, + boundWaitMs: number, + onLateReleaseFailure: () => void, +): Promise { + const timeout = AbortSignal.timeout(boundWaitMs); + const signal = options.signal + ? AbortSignal.any([timeout, options.signal]) + : timeout; + const attempted = access.acquire(selector, { ...options, signal }); + let accepted = false; + void attempted.then( + async (lateLease) => { + if (accepted || !signal.aborted) return; + try { + await lateLease(); + } catch { + onLateReleaseFailure(); + } + }, + () => undefined, + ); + const lease = await Promise.race([ + attempted, + new Promise((_resolve, reject) => { + if (signal.aborted) reject(signal.reason); + else + signal.addEventListener("abort", () => reject(signal.reason), { + once: true, + }); + }), + ]); + accepted = true; + return lease; +} + +/** + * Register storage writes that have no eligible lease behind them. + * + * Each registration overlaps any current holder, so that turn becomes + * ineligible, and it ends with its write. Only a write that cannot register + * leaves a persistent marker; an unknown selector cannot register at all. That + * marker already covers every later write, so later writes skip coordination + * instead of delaying the native call again. + */ +function createWriteRegistrar( + access: MastraExclusiveMemoryAccess, + getSelector: () => MastraMemorySelector | undefined, + waitMs: number, + onFailure: (reason: string) => void, +): () => Promise { + let markedUnsafe = false; + return async () => { + if (markedUnsafe) return undefined; + const selector = getSelector(); + if (selector) { + try { + return await acquireWithin( + access, + selector, + { waitMs: 0 }, + waitMs, + () => onFailure("Late source-thread lease release failed."), + ); + } catch { + // Fall through to the persistent marker below. + } + } + markedUnsafe = true; + try { + await boundedCoordination(access.markUnsafeWrite(selector), waitMs); + } catch { + onFailure("Unsafe memory write could not be fenced."); + } + return undefined; + }; +} + +async function releaseRegistration( + registration: MastraMemoryLease | undefined, + waitMs: number, + onFailure: (reason: string) => void, +): Promise { + if (!registration) return; + try { + await boundedCoordination(registration(), waitMs); + } catch { + onFailure("Source-thread write registration release failed."); + } +} + +/** + * Register each native write to a source domain for the duration of the write. + * + * `getSelector` is read at write time, so a caller can resolve its selector + * after constructing the Memory that owns this domain. Coordination failures + * are reported and never fail the native write. + */ +export function createRegisteredWriteDomain( + domain: MemoryStorage, + access: MastraExclusiveMemoryAccess, + getSelector: () => MastraMemorySelector | undefined, + onFailure: (reason: string) => void, + waitMs = 100, +): MemoryStorage { + const methods = new Map(); + const registerWrite = createWriteRegistrar( + access, + getSelector, + waitMs, + onFailure, + ); + return new Proxy(domain, { + get(target, property) { + const value: unknown = Reflect.get(target, property, target); + if (typeof value !== "function") return value; + if (methods.has(property)) return methods.get(property); + const bound = MUTATIONS.has(property as keyof MemoryStorage) + ? async (...args: unknown[]): Promise => { + const registration = await registerWrite(); + try { + return await Reflect.apply(value, target, args); + } finally { + await releaseRegistration(registration, waitMs, onFailure); + } + } + : value.bind(target); + methods.set(property, bound); + return bound; + }, + }); +} + // The pinned public MemoryStorage mutation inventory. Delegation binds `this` to // the original domain, so a native method's own helper calls record only once. const MUTATIONS = new Set([ @@ -390,6 +596,17 @@ export function createMemoryCaptureBinding( const reasons: string[] = []; const methods = new Map(); const waitMs = options.leaseWaitMs ?? 100; + const selector = { + threadId: options.threadId, + resourceId: options.resourceId, + }; + let settling: Promise | undefined; + const registerWrite = createWriteRegistrar( + options.exclusiveAccess, + () => selector, + waitMs, + markIncomplete, + ); function markIncomplete(reason: string): void { if (reasons.includes(reason)) return; @@ -401,6 +618,25 @@ export function createMemoryCaptureBinding( } } + async function holdsEligibleLease(): Promise { + if (released || !lease) { + markIncomplete( + "Memory mutation occurred without source-thread ownership.", + ); + return false; + } + try { + if (await boundedCoordination(lease.verifyEligibility(), waitMs)) + return true; + markIncomplete("Exclusive source-thread ownership was lost."); + } catch { + markIncomplete( + "Exclusive source-thread ownership could not be verified.", + ); + } + return false; + } + const domain = new Proxy(options.domain, { get(target, property) { const value: unknown = Reflect.get(target, property, target); @@ -441,51 +677,16 @@ export function createMemoryCaptureBinding( const duringCapture = capturing; const result = mutations.then(async () => { let output: unknown; + let registration: MastraMemoryLease | undefined; try { - if (released || !lease) { - markIncomplete( - "Memory mutation occurred without source-thread ownership.", - ); - try { - await boundedCoordination( - options.exclusiveAccess.markUnsafeWrite(options), - waitMs, - ); - } catch { - markIncomplete("Unsafe memory write could not be fenced."); - } - } else { - try { - if ( - !(await boundedCoordination( - lease.verifyEligibility(), - waitMs, - )) - ) { - markIncomplete("Exclusive source-thread ownership was lost."); - await boundedCoordination( - options.exclusiveAccess.markUnsafeWrite(options), - waitMs, - ); - } - } catch { - markIncomplete( - "Exclusive source-thread ownership could not be verified.", - ); - try { - await boundedCoordination( - options.exclusiveAccess.markUnsafeWrite(options), - waitMs, - ); - } catch { - markIncomplete("Unsafe memory write could not be fenced."); - } - } - } + if (!(await holdsEligibleLease())) + registration = await registerWrite(); output = await Reflect.apply(value, target, args); } catch (error) { markIncomplete("Native memory storage mutation failed."); throw error; + } finally { + await releaseRegistration(registration, waitMs, markIncomplete); } // Joined work from a previous turn belongs to the initial snapshot. if (duringCapture) return output; @@ -562,6 +763,26 @@ export function createMemoryCaptureBinding( } } + async function joinMemoryWork(memory: MastraSettlingMemory): Promise { + const engine = await memory.omEngine; + // Mastra's settled() does not join a buffered reflection, and settled + // work can start more buffering. Repeat until a round records no write. + while (true) { + const before = revision; + await memory.settled(); + const waitStarted = Date.now(); + await engine?.waitForBuffering( + options.threadId, + options.resourceId, + BUFFERING_WAIT_MS, + ); + await drain(); + // waitForBuffering resolves, rather than rejects, when it times out. + if (Date.now() - waitStarted >= BUFFERING_WAIT_MS) continue; + if (revision === before) return; + } + } + async function verifyEligibility(): Promise { if (released || !lease) { markIncomplete("Exclusive source-thread ownership is unavailable."); @@ -597,44 +818,22 @@ export function createMemoryCaptureBinding( capturing = true; try { try { - const timeout = AbortSignal.timeout(waitMs); - const signal = options.leaseSignal - ? AbortSignal.any([timeout, options.leaseSignal]) - : timeout; - const attempted = options.exclusiveAccess.acquire(options, { - waitMs, - signal, - onConflict: () => - markIncomplete( - "Exclusive source-thread ownership was invalidated by an overlapping invocation.", - ), - }); - let accepted = false; - // A backend that ignores cancellation must not retain ownership if - // its acquire resolves after the caller has resumed natively. - void attempted.then( - async (lateLease) => { - if (!accepted && signal.aborted) { - try { - await lateLease(); - } catch { - markIncomplete("Late source-thread lease release failed."); - } - } + lease = await acquireWithin( + options.exclusiveAccess, + selector, + { + // Leave the backend time to answer with a lease that registers + // the overlap before this invocation stops waiting for it. + waitMs: Math.max(0, waitMs - ACQUIRE_RESPONSE_MARGIN_MS), + signal: options.leaseSignal, + onConflict: () => + markIncomplete( + "Exclusive source-thread ownership was invalidated by an overlapping invocation.", + ), }, - () => undefined, + waitMs, + () => markIncomplete("Late source-thread lease release failed."), ); - lease = await Promise.race([ - attempted, - new Promise((_resolve, reject) => { - if (signal.aborted) reject(signal.reason); - else - signal.addEventListener("abort", () => reject(signal.reason), { - once: true, - }); - }), - ]); - accepted = true; await verifyEligibility(); if (reasons.length) return undefined; } catch { @@ -698,6 +897,24 @@ export function createMemoryCaptureBinding( } }, drain, + async settle(memory, settleWaitMs) { + settling ??= joinMemoryWork(memory); + if (settleWaitMs === undefined) { + await settling; + return true; + } + let timer: ReturnType | undefined; + try { + return await Promise.race([ + settling.then(() => true), + new Promise((resolve) => { + timer = setTimeout(() => resolve(false), settleWaitMs); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } + }, verifyEligibility, async release() { if (released) return; diff --git a/packages/mastra/src/memory-replay.ts b/packages/mastra/src/memory-replay.ts index e192c1518..f9bcd6a29 100644 --- a/packages/mastra/src/memory-replay.ts +++ b/packages/mastra/src/memory-replay.ts @@ -342,16 +342,15 @@ export async function createIsolatedMemoryReplay( "Restored memory did not produce a coherent initial snapshot.", ); } - let finished: Promise | undefined; + let finished: Promise | undefined; return { memory, binding, initialSnapshot, - finish(): Promise { + finish(): Promise { finished ??= (async () => { try { - await memory.settled(); - await binding.drain(); + return await binding.settle(memory); } finally { await binding.release(); await store.close(); @@ -359,8 +358,8 @@ export async function createIsolatedMemoryReplay( })(); return finished; }, - release(): Promise { - return finished ?? Promise.resolve(); + async release(): Promise { + await finished; }, }; } catch (error) { diff --git a/packages/mastra/src/om-result-tape.ts b/packages/mastra/src/om-result-tape.ts index 306a27fda..8e23025bf 100644 --- a/packages/mastra/src/om-result-tape.ts +++ b/packages/mastra/src/om-result-tape.ts @@ -147,11 +147,25 @@ export function createOMResultTape( return async (input: unknown) => { const ordinal = next++; const inputFingerprint = fingerprint(input); - const result = await Reflect.apply( - value as (input: unknown) => Promise, - target, - [input], - ); + // Track the call from its start: finish() must wait for a call that + // is still waiting on the provider, not report its slot as missing. + let settleCall!: () => void; + const call = new Promise((resolve) => { + settleCall = resolve; + }); + pending.add(call); + void call.finally(() => pending.delete(call)); + let result: unknown; + try { + result = await Reflect.apply( + value as (input: unknown) => Promise, + target, + [input], + ); + } catch (error) { + settleCall(); + throw error; + } if (method === "doGenerate") { try { entries[ordinal] = { @@ -164,12 +178,14 @@ export function createOMResultTape( } catch { failCapture(); } + settleCall(); return result; } const stream = (result as { stream?: ReadableStream }) ?.stream; if (!(stream instanceof ReadableStream)) { failCapture(); + settleCall(); return result; } const [native, capture] = stream.tee(); @@ -209,8 +225,7 @@ export function createOMResultTape( reader.releaseLock(); } })(); - pending.add(work); - void work.finally(() => pending.delete(work)); + void work.finally(settleCall); return { ...(result as object), stream: native }; }; }, diff --git a/packages/mastra/src/stateful-agent.ts b/packages/mastra/src/stateful-agent.ts index b799ad3ed..a5603d6e2 100644 --- a/packages/mastra/src/stateful-agent.ts +++ b/packages/mastra/src/stateful-agent.ts @@ -18,9 +18,11 @@ import { } from "@zenml-io/kitaru/adapter"; import { createMemoryCaptureBinding, + createRegisteredWriteDomain, type MastraExclusiveMemoryAccess, type MastraMemoryCaptureBinding, type MastraMemoryMutation, + type MastraMemorySelector, } from "./memory-binding.js"; import { assertMemoryReplayVersions, @@ -68,6 +70,12 @@ import type { KitaruAgentOptions, RuntimeStreamOptions } from "./types.js"; interface MastraMemorySource { settled(): Promise; + /** + * The source Memory instance. When supplied, its `settled()` also joins the + * memory work each recorded turn runs on its own Memory instance, and Kitaru + * joins only the recorded thread's buffered work before capturing a turn. + */ + memory?: Memory; domain: MemoryStorage; configuration: MemoryConfigInternal; exclusiveAccess: MastraExclusiveMemoryAccess; @@ -113,6 +121,12 @@ export interface MemoryReplayAgentOptions extends KitaruAgentOptions { ) => Promise<{ bytes: Uint8Array; mediaType: string }>; skillsDirectory?: string; resolveModel: (id: string) => MastraModelConfig | Promise; + /** + * How long a baseline waits after its stream closes for observational-memory + * work before it releases the source lease. A turn whose work does not settle + * in time is recorded as ineligible. Defaults to 60 seconds. + */ + finalizationWaitMs?: number; } export interface MemoryReplayAgentBindings { @@ -125,6 +139,9 @@ export type MemoryReplayAgentFactory = ( bindings: MemoryReplayAgentBindings, ) => AgentConfig | Promise; +const DEFAULT_FINALIZATION_WAIT_MS = 60_000; +const CAPTURE_BUFFERING_WAIT_MS = 5_000; + function record(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); } @@ -259,34 +276,19 @@ export function createMemoryReplayAgent( const { MastraCompositeStore } = await import("@mastra/core/storage"); const { Memory } = await import("@mastra/memory"); const source = await options.sourceMemory(); - let unsafeSelector: { threadId: string; resourceId: string } | undefined; - try { - unsafeSelector = getSelector(callerOptions); - } catch { - // An implicit native selector can come from factory defaults. - } - try { - let timeout: ReturnType | undefined; - try { - await Promise.race([ - source.exclusiveAccess.markUnsafeWrite(unsafeSelector), - new Promise((_resolve, reject) => { - timeout = setTimeout( - () => reject(new Error("Source-thread unsafe marker timed out.")), - 100, - ); - }), - ]); - } finally { - if (timeout) clearTimeout(timeout); - } - } catch (error) { - reportLocalRecordingError(error); - } + // Mastra writes only once the stream runs, after the factory has resolved + // the default memory selectors below. + let selector: MastraMemorySelector | undefined; + const domain = createRegisteredWriteDomain( + source.domain, + source.exclusiveAccess, + () => selector, + (reason) => reportLocalRecordingError(new Error(reason)), + ); const memory = new Memory({ storage: new MastraCompositeStore({ id: `kitaru-native-${globalThis.crypto.randomUUID()}`, - domains: { memory: source.domain }, + domains: { memory: domain }, }), options: source.configuration, }); @@ -302,12 +304,41 @@ export function createMemoryReplayAgent( }), workspace: workspace?.workspace, }); + selector = await getNativeSelector(config, callerOptions); const native = new Agent({ ...config, memory }) as unknown as { stream(input: unknown, options: RuntimeStreamOptions): Promise; }; return native.stream(rawInput, callerOptions); } + /** Resolve the selectors a native call writes, or undefined when unknown. */ + async function getNativeSelector( + config: AgentConfig, + callerOptions: RuntimeStreamOptions, + ): Promise { + try { + const defaults = requireRecord( + typeof config.defaultOptions === "function" + ? await config.defaultOptions({ + requestContext: + callerOptions.requestContext ?? new RequestContext(), + mastra: options.mastra, + }) + : (config.defaultOptions ?? {}), + "default options", + ); + const { deepMerge } = await import("@mastra/core/utils"); + return getSelector( + deepMerge( + record(defaults.memory) ? { memory: defaults.memory } : {}, + record(callerOptions.memory) ? { memory: callerOptions.memory } : {}, + ), + ); + } catch { + return undefined; + } + } + function reportLocalRecordingError(error: unknown): void { if (options.onRecordingError) { void Promise.resolve() @@ -479,7 +510,9 @@ export function createMemoryReplayAgent( memory: Memory; binding: MastraMemoryCaptureBinding; initialSnapshot: MastraMemorySnapshot | undefined; - finish(): Promise; + beginFinalization?(): void; + /** Resolve false when memory work missed the finalization deadline. */ + finish(): Promise; release(): Promise; }; let unsafeEvidenceReason: UnsafeEvidenceReason | undefined; @@ -539,16 +572,50 @@ export function createMemoryReplayAgent( omTape, ), }); - const initialSnapshot = await binding.captureInitial(source); - let finished: Promise | undefined; + const sourceEngine = source.memory + ? await source.memory.omEngine.catch(() => null) + : null; + // Once the source Memory joins every turn's work, its settled() also + // waits for turns on other threads. Join only this thread's buffering. + const initialSnapshot = await binding.captureInitial( + source.memory + ? { + settled: async () => { + await (await memory.omEngine)?.waitForBuffering( + selector.threadId, + selector.resourceId, + CAPTURE_BUFFERING_WAIT_MS, + ); + }, + } + : source, + ); + let tracked = false; + const trackSourceWork = () => { + if (tracked || !sourceEngine) return; + tracked = true; + void sourceEngine + .trackBackgroundWork(binding.settle(memory)) + .catch(() => undefined); + }; + let finished: Promise | undefined; runtime = { memory, binding, initialSnapshot, + beginFinalization: trackSourceWork, finish() { finished ??= (async () => { - await memory.settled(); - await binding.drain(); + trackSourceWork(); + const settled = await binding.settle( + memory, + options.finalizationWaitMs ?? DEFAULT_FINALIZATION_WAIT_MS, + ); + if (!settled) + binding.markIncomplete( + "Observational-memory work did not settle before the finalization deadline.", + ); + return settled; })(); return finished; }, @@ -874,11 +941,19 @@ export function createMemoryReplayAgent( ); return evidence; }, + beginFinalization() { + runtime.beginFinalization?.(); + }, async finish() { - await runtime.finish(); + const settled = await runtime.finish(); await capture.drain(); if (!historical) await runtime.binding.verifyEligibility(); - const omResults = await omTape.finish(); + // The turn's writes are settled and checked, so the next turn on + // this thread can acquire while the session update is sent. + await runtime.release(); + // An OM call past the deadline may never return; the turn is + // already ineligible, so do not wait for its tape entry. + const omResults = settled ? await omTape.finish() : []; for (const reason of omCaptureErrors) runtime.binding.markIncomplete(reason); if (omMismatches.length) diff --git a/packages/mastra/src/stream-recording.ts b/packages/mastra/src/stream-recording.ts index c6e829f43..a4d82e482 100644 --- a/packages/mastra/src/stream-recording.ts +++ b/packages/mastra/src/stream-recording.ts @@ -55,6 +55,8 @@ export interface StatefulStreamRecording { sanitizeEvidence?: (value: T) => T; initialize(state: AdapterRunState): void; takeRequest(): RequestEvidence | undefined; + /** Start joining the invocation's background memory work without waiting. */ + beginFinalization?(): void; finish(): Promise; release(): Promise; } @@ -328,6 +330,9 @@ class StreamLifecycle { private async finalize(complete: boolean, result?: unknown): Promise { this.#finalizerPromise ??= (async () => { try { + // Start before the first await so a caller that joins background work + // right after the stream closes already sees this invocation's work. + this.stateful?.beginFinalization?.(); await this.#stepTail; let finalInput: JsonValue | undefined; try { diff --git a/packages/mastra/test/memory-binding.test.ts b/packages/mastra/test/memory-binding.test.ts index 73e396f9e..226abd099 100644 --- a/packages/mastra/test/memory-binding.test.ts +++ b/packages/mastra/test/memory-binding.test.ts @@ -212,7 +212,7 @@ it("invalidates the first recording when a conflicting invocation cannot get its await next.binding.release(); }); -it("preserves a native write after lost ownership and blocks eligibility until quiescence reset", async () => { +it("keeps a native write after lost ownership and recovers once overlapping turns release", async () => { const access = createProcessLocalMemoryAccess(); const first = await fixture("first", access); const second = await fixture("second", access); @@ -230,28 +230,45 @@ it("preserves a native write after lost ownership and blocks eligibility until q await first.binding.release(); await second.binding.release(); const next = await fixture("after", access); - expect( - await next.binding.captureInitial(next.runtime.memory), - ).toBeUndefined(); + expect(await next.binding.captureInitial(next.runtime.memory)).toBeDefined(); await next.binding.release(); - await access.resetAfterQuiescence({ threadId: THREAD, resourceId: RESOURCE }); - const recovered = await fixture("recovered", access); - expect( - await recovered.binding.captureInitial(recovered.runtime.memory), - ).toBeDefined(); - await recovered.binding.release(); }); -it("keeps an untracked late write unsafe until an explicit quiescence reset", async () => { +it("registers a late write for its duration instead of poisoning later turns", async () => { const access = createProcessLocalMemoryAccess(); + const unsafeWrite = vi.spyOn(access, "markUnsafeWrite"); const first = await fixture("first", access); await first.binding.captureInitial(first.runtime.memory); await first.binding.release(); + const holder = await fixture("holder", access); + expect( + await holder.binding.captureInitial(holder.runtime.memory), + ).toBeDefined(); const native = await first.binding.domain.updateThread({ id: THREAD, title: "late native write", }); expect(native.title).toBe("late native write"); + // The late write overlapped the holder, so only the holder's turn is lost. + await holder.binding.verifyEligibility(); + expect(holder.binding.incompleteReasons.join()).toMatch(/overlapping/); + await holder.binding.release(); + expect(unsafeWrite).not.toHaveBeenCalled(); + const next = await fixture("next", access); + expect(await next.binding.captureInitial(next.runtime.memory)).toBeDefined(); + await next.binding.release(); +}); + +it("keeps a thread unsafe when a late write cannot register", async () => { + const access = createProcessLocalMemoryAccess(); + const first = await fixture("first", access); + await first.binding.captureInitial(first.runtime.memory); + await first.binding.release(); + const acquire = vi + .spyOn(access, "acquire") + .mockRejectedValueOnce(new Error("coordination unavailable")); + await first.binding.domain.updateThread({ id: THREAD, title: "late" }); + acquire.mockRestore(); const next = await fixture("next", access); expect( await next.binding.captureInitial(next.runtime.memory), @@ -265,6 +282,16 @@ it("keeps an untracked late write unsafe until an explicit quiescence reset", as await recovered.binding.release(); }); +it("waits within waitMs for a releasing holder instead of invalidating it", async () => { + const access = createProcessLocalMemoryAccess(); + const selector = { threadId: THREAD, resourceId: RESOURCE }; + const first = await access.acquire(selector); + setTimeout(() => void first(), 20); + const second = await access.acquire(selector, { waitMs: 200 }); + expect(await second.verifyEligibility()).toBe(true); + await second(); +}); + it("bounds a hung acquisition and releases a lease returned after cancellation", async () => { const release = Object.assign( vi.fn(async () => {}), diff --git a/packages/mastra/test/memory-lease-lifecycle.test.ts b/packages/mastra/test/memory-lease-lifecycle.test.ts new file mode 100644 index 000000000..0e4b45285 --- /dev/null +++ b/packages/mastra/test/memory-lease-lifecycle.test.ts @@ -0,0 +1,448 @@ +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { InMemoryStore } from "@mastra/core/storage"; +import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; +import { createTool } from "@mastra/core/tools"; +import { Memory } from "@mastra/memory"; +import { afterEach, expect, it, vi } from "vitest"; +import { z } from "zod/v4"; +import { + createMemoryReplayAgent, + createProcessLocalMemoryAccess, + type MastraExclusiveMemoryAccess, +} from "../src/memory.js"; +import { createFileMemoryAccess } from "./helpers/file-memory-access.js"; +import { + RESOURCE, + seedMemory, + streamParts, + THREAD, + textStream, +} from "./helpers/memory-agent.js"; +import { AGENT_ID, installTestApi, type TestApi } from "./helpers.js"; + +const roots: string[] = []; +const stores: InMemoryStore[] = []; + +afterEach(async () => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); + for (const store of stores.splice(0)) await store.close(); + for (const root of roots.splice(0)) + await rm(root, { recursive: true, force: true, maxRetries: 5 }); +}); + +const LONG_MESSAGE = `Observe this: ${"The user prefers green for every report. ".repeat(300)}`; + +const pause = (ms: number) => + new Promise((resolve) => setTimeout(resolve, ms)); + +const LEASES: Array<[string, () => Promise]> = [ + ["process-local", async () => createProcessLocalMemoryAccess()], + [ + "shared file", + async () => { + const root = await mkdtemp(join(tmpdir(), "kitaru-lease-lifecycle-")); + roots.push(root); + return createFileMemoryAccess(root); + }, + ], +]; + +function gate() { + let open!: () => void; + const opened = new Promise((resolve) => { + open = resolve; + }); + return { open, opened }; +} + +function omModel(kind: "observer" | "reflector", wait?: () => Promise) { + const text = + kind === "observer" + ? `OBSERVED: ${"The user changed the preference to green. ".repeat(15)}` + : "REFLECTED: the user prefers green."; + return new MastraLanguageModelV2Mock({ + modelId: kind, + provider: "fixture", + doStream: async () => { + await wait?.(); + return textStream( + `\n${text}\n\nContinue.`, + ); + }, + }); +} + +/** A source Memory whose first evidence-heavy turn starts a buffered reflection. */ +async function setup( + access: MastraExclusiveMemoryAccess, + options: { + observerWait?: () => Promise; + reflectorWait?: () => Promise; + observe?: boolean; + /** Observe in the background: buffer without reaching the threshold. */ + backgroundObservation?: boolean; + finalizationWaitMs?: number; + joinSourceMemory?: boolean; + /** + * Use a fresh thread. Mastra keeps buffering state per thread for the whole + * process, so background observation needs a thread no other test used. + */ + thread?: string; + /** Wrap the test API before the adapter's client captures fetch. */ + wrapFetch?: (fetch: typeof globalThis.fetch) => typeof globalThis.fetch; + } = {}, +) { + const api = installTestApi(); + if (options.wrapFetch) + vi.stubGlobal("fetch", options.wrapFetch(globalThis.fetch)); + const observer = omModel("observer", options.observerWait); + const reflector = omModel("reflector", options.reflectorWait); + const store = new InMemoryStore(); + stores.push(store); + const domain = store.stores.memory; + if (!domain) throw new Error("Missing native memory domain"); + const memory = new Memory({ + storage: store, + options: { + lastMessages: 20, + semanticRecall: false, + workingMemory: { + enabled: true, + scope: "thread", + schema: z.object({ preference: z.string() }), + }, + observationalMemory: { + scope: "thread", + observation: { + model: observer, + messageTokens: options.backgroundObservation + ? 10_000 + : options.observe + ? 300 + : 1_000_000, + bufferTokens: 0.2, + bufferActivation: 1, + blockAfter: 1.1, + }, + reflection: { + model: reflector, + observationTokens: 200, + bufferActivation: 0.5, + blockAfter: 1.1, + }, + }, + }, + }); + const thread = options.thread ?? THREAD; + if (options.thread) + await memory.createThread({ threadId: thread, resourceId: RESOURCE }); + else + await seedMemory({ + store, + domain, + memory, + observer: { calls: [], model: observer }, + reflector: { calls: [], model: reflector }, + }); + let actorCalls = 0; + const actor = new MastraLanguageModelV2Mock({ + modelId: "actor", + provider: "fixture", + doStream: async () => { + actorCalls += 1; + if (!options.observe || actorCalls % 2 === 0) return textStream("done"); + return streamParts( + [ + { + type: "tool-call", + toolCallId: `call-${actorCalls}`, + toolName: "evidence", + input: "{}", + }, + ], + "tool-calls", + ); + }, + }); + const adapter = createMemoryReplayAgent( + ({ memory }) => ({ + id: "lease-lifecycle", + name: "Lease lifecycle", + instructions: "Answer", + memory, + model: actor, + tools: { + evidence: createTool({ + id: "evidence", + description: "Read evidence", + inputSchema: z.object({}), + execute: async () => "The user now prefers green. ".repeat(60), + }), + }, + }), + { + agentId: AGENT_ID, + apiUrl: "https://kitaru.invalid", + apiKey: "fixture", + requestedModelId: "fixture/actor", + finalizationWaitMs: options.finalizationWaitMs, + onRecordingError: () => undefined, + sourceMemory: () => ({ + settled: () => memory.settled(), + ...(options.joinSourceMemory ? { memory } : {}), + domain, + configuration: memory.getMergedThreadConfig(), + exclusiveAccess: access, + }), + resolveModel: (id) => + id.endsWith("observer") + ? observer + : id.endsWith("reflector") + ? reflector + : actor, + }, + ); + async function turn( + message: string, + selector: { thread: string; resource: string } | null = { + thread, + resource: RESOURCE, + }, + ) { + const output = await adapter.stream(message, { + maxSteps: 5, + ...(selector ? { memory: selector } : {}), + }); + await output.consumeStream(); + } + return { api, memory, domain, turn }; +} + +/** The final replay state of every session, in creation order. */ +function outcomes(api: TestApi): string[] { + const created = api.calls.filter( + (call) => call.method === "POST" && call.path === "/api/v1/sessions", + ); + return api.sessionIds.map((id, index) => { + const initial = created[index]?.body?.metadata as + | Record + | undefined; + // A setup failure reports its own session after the native answer. + if (initial?.mastra_replay_state === "ineligible") return "setup-failure"; + const update = api.calls + .filter( + (call) => + call.method === "PATCH" && + call.path.endsWith(id) && + call.body?.status !== "in_progress", + ) + .at(-1); + if (!update) return "pending"; + const metadata = update.body?.metadata as + | Record + | undefined; + return `${String(update.body?.status)}/${String(metadata?.mastra_replay_state)}`; + }); +} + +function memoryMethods(api: TestApi, sessionId: string): unknown[] { + return api + .nodeBatches(sessionId) + .flat() + .filter((node) => node.name === "memory_mutation") + .map((node) => (node.attributes as Record).memory_method); +} + +it.each(LEASES)( + "records two turns sent 100 ms apart on one thread as eligible (%s lease)", + async (_kind, createAccess) => { + const { api, turn } = await setup(await createAccess(), { + // A slow session update must not keep the thread leased. + wrapFetch: (recorded) => async (input, init) => { + if (init?.method === "PATCH") await pause(300); + return recorded(input, init); + }, + }); + await turn("Turn one."); + await pause(100); + await turn("Turn two."); + await vi.waitFor( + () => + expect(outcomes(api)).toEqual([ + "completed/eligible", + "completed/eligible", + ]), + { timeout: 3000 }, + ); + }, +); + +it.each(LEASES)( + "records a late buffered reflection in its own turn and keeps the next turn eligible (%s lease)", + async (_kind, createAccess) => { + const reflection = gate(); + let reflectorStarted = false; + const { api, turn } = await setup(await createAccess(), { + observe: true, + reflectorWait: async () => { + reflectorStarted = true; + await reflection.opened; + }, + }); + await turn("Remember the evidence."); + await vi.waitFor(() => expect(reflectorStarted).toBe(true)); + // The stream has closed; the reflection finishes well after it. + await pause(300); + expect(outcomes(api)).toEqual(["pending"]); + reflection.open(); + await vi.waitFor( + () => expect(outcomes(api)).toEqual(["completed/eligible"]), + { + timeout: 3000, + }, + ); + const [first] = api.sessionIds; + expect(memoryMethods(api, String(first))).toContain( + "updateBufferedReflection", + ); + const finalInputs = JSON.stringify( + api.calls.find( + (call) => call.method === "PATCH" && call.body?.status === "completed", + )?.body?.inputs, + ); + expect(finalInputs).toContain('"phase":"reflector"'); + await turn("Next turn."); + await vi.waitFor( + () => + expect(outcomes(api)).toEqual([ + "completed/eligible", + "completed/eligible", + ]), + { timeout: 3000 }, + ); + }, +); + +it.each(LEASES)( + "keeps turns after a native fallback turn eligible (%s lease)", + async (_kind, createAccess) => { + let failed = false; + const { api, turn } = await setup(await createAccess(), { + wrapFetch: (recorded) => async (input, init) => { + if ( + !failed && + init?.method === "POST" && + new URL(String(input)).pathname === "/api/v1/sessions" + ) { + failed = true; + throw new TypeError("Kitaru is unavailable"); + } + return recorded(input, init); + }, + }); + await turn("Answered natively."); + await turn("Same thread."); + await turn("Other thread.", { thread: "other-thread", resource: RESOURCE }); + await vi.waitFor( + () => + expect(outcomes(api)).toEqual([ + "completed/eligible", + "completed/eligible", + ]), + { timeout: 3000 }, + ); + }, +); + +it.each(LEASES)( + "does not let a call without memory selectors affect other threads (%s lease)", + async (_kind, createAccess) => { + const { api, turn } = await setup(await createAccess()); + await turn("Summarize this without memory.", null); + await turn("A normal turn."); + await vi.waitFor( + () => + expect(outcomes(api).sort()).toEqual([ + "completed/eligible", + "setup-failure", + ]), + { timeout: 3000 }, + ); + }, +); + +it.each(LEASES)( + "keeps a thread ineligible after an unregistered writer until quiescence reset (%s lease)", + async (_kind, createAccess) => { + const access = await createAccess(); + const { api, turn } = await setup(access); + await access.markUnsafeWrite({ threadId: THREAD, resourceId: RESOURCE }); + await turn("After a foreign write."); + await turn("Other thread, same resource.", { + thread: "other-thread", + resource: RESOURCE, + }); + await vi.waitFor( + () => + expect(outcomes(api)).toEqual([ + "failed/ineligible", + "failed/ineligible", + ]), + { timeout: 3000 }, + ); + await access.resetAfterQuiescence({ + threadId: THREAD, + resourceId: RESOURCE, + }); + await turn("After the application reset."); + await vi.waitFor(() => + expect(outcomes(api).at(-1)).toBe("completed/eligible"), + ); + }, +); + +it("releases the lease and closes the session when OM work misses the finalization deadline", async () => { + const access = createProcessLocalMemoryAccess(); + const hung = gate(); + const { api, turn } = await setup(access, { + thread: "hung-observer-thread", + backgroundObservation: true, + finalizationWaitMs: 200, + observerWait: () => hung.opened, + }); + try { + await turn(LONG_MESSAGE); + await vi.waitFor( + () => expect(outcomes(api)).toEqual(["failed/ineligible"]), + { timeout: 3000 }, + ); + const probe = await access.acquire({ + threadId: "other-thread", + resourceId: RESOURCE, + }); + expect(await probe.verifyEligibility()).toBe(true); + await probe(); + } finally { + hung.open(); + } +}); + +it("lets the source Memory's settled() join a turn's observational-memory work", async () => { + let observed = false; + const { memory, turn } = await setup(createProcessLocalMemoryAccess(), { + thread: "source-settled-thread", + backgroundObservation: true, + joinSourceMemory: true, + observerWait: async () => { + await pause(300); + observed = true; + }, + }); + await turn(LONG_MESSAGE); + expect(observed).toBe(false); + await memory.settled(); + expect(observed).toBe(true); +}); diff --git a/packages/mastra/test/memory-replay-safety.test.ts b/packages/mastra/test/memory-replay-safety.test.ts index 44a76f08f..df47ea827 100644 --- a/packages/mastra/test/memory-replay-safety.test.ts +++ b/packages/mastra/test/memory-replay-safety.test.ts @@ -94,6 +94,15 @@ it("stores captured file references and refuses signed source URLs", () => { ); }); +/** The native fallback registers each write; no recording leased the thread. */ +function expectOnlyWriteRegistrations(acquire: { + mock: { calls: unknown[][] }; +}): void { + expect(acquire.mock.calls.length).toBeGreaterThan(0); + for (const [, options] of acquire.mock.calls) + expect(options).toMatchObject({ waitMs: 0 }); +} + function fixture( factory?: MemoryReplayAgentFactory, overrides: Partial = {}, @@ -209,7 +218,7 @@ it.each([MASTRA_THREAD_ID_KEY, MASTRA_RESOURCE_ID_KEY])( }); await result.consumeStream(); expect(await result.text).toBe("done"); - expect(acquire).not.toHaveBeenCalled(); + expectOnlyWriteRegistrations(acquire); expect(modelCall).toHaveBeenCalledTimes(1); await vi.waitFor(() => expect( @@ -237,7 +246,7 @@ it("does not make mismatched middleware selectors replayable through selective c }); await result.consumeStream(); expect(await result.text).toBe("done"); - expect(acquire).not.toHaveBeenCalled(); + expectOnlyWriteRegistrations(acquire); await vi.waitFor(() => expect( api.calls.find( diff --git a/packages/mastra/test/processor-replay.test.ts b/packages/mastra/test/processor-replay.test.ts index 7c7c17253..d42913940 100644 --- a/packages/mastra/test/processor-replay.test.ts +++ b/packages/mastra/test/processor-replay.test.ts @@ -472,11 +472,13 @@ it.each(["late", "memory-name-spoof", "copied-memory-id"])( replaying = true; const result = await adapter.stream("ignored"); await result.consumeStream(); + // The baseline finalizes in the background, so its update can land last. await vi.waitFor(() => expect( - api.calls.filter((call) => call.method === "PATCH").at(-1)?.body - ?.status, - ).toBe("failed"), + api.calls + .filter((call) => call.method === "PATCH") + .map((call) => call.body?.status), + ).toContain("failed"), ); expect(execute).not.toHaveBeenCalled(); expect( diff --git a/packages/mastra/test/stateful-overrides.test.ts b/packages/mastra/test/stateful-overrides.test.ts index c83a21ee2..3e20d0ef2 100644 --- a/packages/mastra/test/stateful-overrides.test.ts +++ b/packages/mastra/test/stateful-overrides.test.ts @@ -228,6 +228,7 @@ it("executes a second-step memory tool during replay", async () => { import { createMemoryReplayAgent, createProcessLocalMemoryAccess, + type MastraMemoryLease, MEMORY_REPLAY_KEY, } from "../src/memory.js"; import { @@ -654,23 +655,22 @@ it("runs natively when request context cannot be captured safely", async () => { }), ); expect(JSON.stringify(api.calls)).not.toContain("secret"); - expect(unsafeWrite).toHaveBeenCalledWith({ - threadId: THREAD, - resourceId: RESOURCE, - }); + // The fallback registered each write, so nothing outlives its native turn. + expect(unsafeWrite).not.toHaveBeenCalled(); const next = await access.acquire({ threadId: THREAD, resourceId: RESOURCE }); - expect(await next.verifyEligibility()).toBe(false); + expect(await next.verifyEligibility()).toBe(true); await next(); await runtime.memory.settled(); await runtime.store.close(); }); -it("poisons all source threads when a native fallback has only an implicit selector", async () => { +it("registers native fallback writes under an implicit default selector", async () => { const runtime = createMemoryRuntime({ messageTokens: 10000 }); await seedMemory(runtime); installTestApi(); const access = createProcessLocalMemoryAccess(); const unsafeWrite = vi.spyOn(access, "markUnsafeWrite"); + const acquire = vi.spyOn(access, "acquire"); const model = new MastraLanguageModelV2Mock({ modelId: "actor", provider: "fixture", @@ -704,12 +704,16 @@ it("poisons all source threads when a native fallback has only an implicit selec const output = await adapter.stream("Hello"); await output.consumeStream(); expect(await output.text).toBe("native answer"); - expect(unsafeWrite).toHaveBeenCalledWith(undefined); + expect(unsafeWrite).not.toHaveBeenCalled(); + expect(acquire).toHaveBeenCalledWith( + { threadId: THREAD, resourceId: RESOURCE }, + expect.objectContaining({ waitMs: 0 }), + ); const next = await access.acquire({ threadId: THREAD, resourceId: RESOURCE, }); - expect(await next.verifyEligibility()).toBe(false); + expect(await next.verifyEligibility()).toBe(true); await next(); } finally { await runtime.store.close(); @@ -723,6 +727,7 @@ it("keeps the native answer moving when unsafe-write coordination hangs", async const reported = vi.fn(); const access = { ...createProcessLocalMemoryAccess(), + acquire: vi.fn(() => new Promise(() => undefined)), markUnsafeWrite: vi.fn(() => new Promise(() => undefined)), }; const model = new MastraLanguageModelV2Mock({ @@ -759,13 +764,16 @@ it("keeps the native answer moving when unsafe-write coordination hangs", async }); await output.consumeStream(); expect(await output.text).toBe("native answer"); + // One persistent marker covers the fallback, so later writes skip waiting. expect(access.markUnsafeWrite).toHaveBeenCalledOnce(); - expect(reported).toHaveBeenCalledWith( - expect.objectContaining({ - error: expect.objectContaining({ - message: "Source-thread unsafe marker timed out.", + await vi.waitFor(() => + expect(reported).toHaveBeenCalledWith( + expect.objectContaining({ + error: expect.objectContaining({ + message: "Unsafe memory write could not be fenced.", + }), }), - }), + ), ); } finally { await runtime.store.close(); From 4a2503148501b5f2e67ac567ab0f202789193837 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 24 Sep 2026 14:53:10 +0200 Subject: [PATCH 026/115] Honor context selectors and release before uploads A native fallback whose thread or resource came from Mastra's reserved `mastra__threadId`/`mastra__resourceId` request-context keys was treated as selector-unknown, so its writes set the global unknown-writer marker and every thread of every user stopped being replayable. `getNativeSelector` now follows Mastra's precedence: the context keys override the merged default and caller memory options. A turn also kept its source lease while its memory evidence uploaded to Kitaru after the stream closed. With a slow Kitaru API, a reply sent shortly afterwards overlapped the lease and both turns became ineligible. The binding now releases once its storage writes settle and the eligibility check passes, and `drain()` waits for evidence after. The finalization deadline now bounds every wait on OM work: the tracked work that the source Memory's `settled()` joins, the join loop that kept polling a hung buffering operation, and replay, which could wait indefinitely on a buffering operation left hung on the same thread and now fails instead. The lease contract and docs now require a shared backend to end a lease whose holder process died, since Kitaru never renews one. --- docs/book/adapters/mastra.md | 4 +- packages/mastra/README.md | 4 +- packages/mastra/src/memory-binding.ts | 39 ++++++-- packages/mastra/src/memory-replay.ts | 14 ++- packages/mastra/src/stateful-agent.ts | 56 +++++++---- .../test/memory-lease-lifecycle.test.ts | 96 ++++++++++++++++++- packages/mastra/test/memory-replay.test.ts | 63 ++++++++++++ 7 files changed, 245 insertions(+), 31 deletions(-) diff --git a/docs/book/adapters/mastra.md b/docs/book/adapters/mastra.md index e1fb6ae93..c1730db51 100644 --- a/docs/book/adapters/mastra.md +++ b/docs/book/adapters/mastra.md @@ -257,9 +257,9 @@ Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KIT ### Source ownership and supported configuration -All writers to a source thread or resource must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. A new turn waits up to 100 ms for an earlier turn on the same thread or resource to release it. A recorded turn holds both selectors until its memory writes, including delayed observational-memory work, have settled, or until `finalizationWaitMs` passes (60 seconds by default, after which the turn is ineligible). Kitaru then releases the selectors and sends the final session update. Turns that overlap on either selector still answer natively; only the overlapping turns become ineligible for replay, and later turns are unaffected. Pass the source `Memory` as `memory` so that its `settled()` also waits for the observational-memory work of recorded turns before you close storage. `settled()` does not provide exclusive access. +All writers to a source thread or resource must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. A new turn waits up to 100 ms for an earlier turn on the same thread or resource to release it. A recorded turn holds both selectors until its memory writes, including delayed observational-memory work, have settled, or until `finalizationWaitMs` passes (60 seconds by default, after which the turn is ineligible). Kitaru then releases the selectors before it uploads the remaining evidence and the final session update. Turns that overlap on either selector still answer natively; only the overlapping turns become ineligible for replay, and later turns are unaffected. For example, a reply sent while the previous turn's buffered reflection is still running makes both of those turns ineligible. The selectors Kitaru coordinates on are the ones Mastra uses, including the reserved `mastra__threadId` and `mastra__resourceId` request-context keys. Pass the source `Memory` as `memory` so that its `settled()` also waits, for up to `finalizationWaitMs`, for the observational-memory work of recorded turns before you close storage. `settled()` does not provide exclusive access. -A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must atomically reserve both thread and resource IDs. When either ID is still held after `waitMs`, it must invalidate the current holders and return a lease that is not eligible but holds both IDs until it is released; that invalidation ends once every overlapping lease has been released. `waitMs: 0` must not wait: Kitaru registers each write it makes outside its own lease this way and releases the registration when the write finishes. `markUnsafeWrite()` is only for a write that could not register, because coordination failed or its selector is unknown, and that marker must survive process loss. Kitaru never calls `resetAfterQuiescence()`. Your application calls it for the marked selectors, or with no selector after an unknown-selector marker, once every process that might have written without registering has stopped or restarted. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. +A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must atomically reserve both thread and resource IDs. When either ID is still held after `waitMs`, it must invalidate the current holders and return a lease that is not eligible but holds both IDs until it is released; that invalidation ends once every overlapping lease has been released. `waitMs: 0` must not wait: Kitaru registers each write it makes outside its own lease this way and releases the registration when the write finishes. Kitaru never renews a lease, so the backend must also end a lease whose holder process died without releasing it: give each lease a time-to-live longer than your longest turn plus `finalizationWaitMs`, or tie it to a liveness check of the holder. `verifyEligibility()` must return false once a lease has expired. Without this, a server that stops mid-turn, for example during a rolling deploy, leaves every later turn on that thread and resource ineligible. `markUnsafeWrite()` is only for a write that could not register, because coordination failed or its selector is unknown, and that marker must survive process loss. Kitaru never calls `resetAfterQuiescence()`. Your application calls it for the marked selectors, or with no selector after an unknown-selector marker, once every process that might have written without registering has stopped or restarted. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. Schema working memory requires explicit `scope: "thread"`. An observational-memory configuration object may omit `scope`, using Mastra's implicit thread scope, or set it to `"thread"`. Supply explicit observer/reflector model identities, either shared through `observationalMemory.model` or in the phase configuration. Resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` remain unsupported. diff --git a/packages/mastra/README.md b/packages/mastra/README.md index 76b50e7e1..8fc93c880 100644 --- a/packages/mastra/README.md +++ b/packages/mastra/README.md @@ -221,9 +221,9 @@ Run this entrypoint with `KITARU_API_URL`, a Kitaru credential, an existing `KIT ### Source ownership and supported configuration -All writers to a source thread or resource must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. A new turn waits up to 100 ms for an earlier turn on the same thread or resource to release it. A recorded turn holds both selectors until its memory writes, including delayed observational-memory work, have settled, or until `finalizationWaitMs` passes (60 seconds by default, after which the turn is ineligible). Kitaru then releases the selectors and sends the final session update. Turns that overlap on either selector still answer natively; only the overlapping turns become ineligible for replay, and later turns are unaffected. Pass the source `Memory` as `memory` so that its `settled()` also waits for the observational-memory work of recorded turns before you close storage. `settled()` does not provide exclusive access. +All writers to a source thread or resource must participate in the same `MastraExclusiveMemoryAccess` implementation. The process-local helper works only when every writer shares that instance in one process. A new turn waits up to 100 ms for an earlier turn on the same thread or resource to release it. A recorded turn holds both selectors until its memory writes, including delayed observational-memory work, have settled, or until `finalizationWaitMs` passes (60 seconds by default, after which the turn is ineligible). Kitaru then releases the selectors before it uploads the remaining evidence and the final session update. Turns that overlap on either selector still answer natively; only the overlapping turns become ineligible for replay, and later turns are unaffected. For example, a reply sent while the previous turn's buffered reflection is still running makes both of those turns ineligible. The selectors Kitaru coordinates on are the ones Mastra uses, including the reserved `mastra__threadId` and `mastra__resourceId` request-context keys. Pass the source `Memory` as `memory` so that its `settled()` also waits, for up to `finalizationWaitMs`, for the observational-memory work of recorded turns before you close storage. `settled()` does not provide exclusive access. -A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must atomically reserve both thread and resource IDs. When either ID is still held after `waitMs`, it must invalidate the current holders and return a lease that is not eligible but holds both IDs until it is released; that invalidation ends once every overlapping lease has been released. `waitMs: 0` must not wait: Kitaru registers each write it makes outside its own lease this way and releases the registration when the write finishes. `markUnsafeWrite()` is only for a write that could not register, because coordination failed or its selector is unknown, and that marker must survive process loss. Kitaru never calls `resetAfterQuiescence()`. Your application calls it for the marked selectors, or with no selector after an unknown-selector marker, once every process that might have written without registering has stopped or restarted. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. +A multi-process or multi-server deployment must supply a backend using shared atomic storage; Kitaru does not include a production distributed lease backend. `acquire()` returns a callable release function with `verifyEligibility()`. The backend must atomically reserve both thread and resource IDs. When either ID is still held after `waitMs`, it must invalidate the current holders and return a lease that is not eligible but holds both IDs until it is released; that invalidation ends once every overlapping lease has been released. `waitMs: 0` must not wait: Kitaru registers each write it makes outside its own lease this way and releases the registration when the write finishes. Kitaru never renews a lease, so the backend must also end a lease whose holder process died without releasing it: give each lease a time-to-live longer than your longest turn plus `finalizationWaitMs`, or tie it to a liveness check of the holder. `verifyEligibility()` must return false once a lease has expired. Without this, a server that stops mid-turn, for example during a rolling deploy, leaves every later turn on that thread and resource ineligible. `markUnsafeWrite()` is only for a write that could not register, because coordination failed or its selector is unknown, and that marker must survive process loss. Kitaru never calls `resetAfterQuiescence()`. Your application calls it for the marked selectors, or with no selector after an unknown-selector marker, once every process that might have written without registering has stopped or restarted. Coordination failure must prevent replay eligibility even when native writes continue. Validate these guarantees against your actual storage, deployment topology, and failure recovery before enabling production replay; the process-local example does not establish customer deployment readiness. Schema working memory requires explicit `scope: "thread"`. An observational-memory configuration object may omit `scope`, using Mastra's implicit thread scope, or set it to `"thread"`. Supply explicit observer/reflector model identities, either shared through `observationalMemory.model` or in the phase configuration. Resource-scoped state, semantic recall, automatic title generation, and per-call `memory.options` remain unsupported. diff --git a/packages/mastra/src/memory-binding.ts b/packages/mastra/src/memory-binding.ts index 36b1ba2a5..101e504c7 100644 --- a/packages/mastra/src/memory-binding.ts +++ b/packages/mastra/src/memory-binding.ts @@ -47,6 +47,12 @@ export interface MastraMemoryLease { * A write Kitaru makes outside that lease registers through * `acquire(selector, { waitMs: 0 })` and releases immediately afterwards. * + * Kitaru never renews a lease. A shared implementation must still end a lease + * whose holder process died without releasing it, through a time-to-live + * longer than the longest turn plus its finalization wait, or a liveness check + * of the holder. `verifyEligibility` must return false once the lease has + * expired, so a turn that outlives its lease is ineligible rather than unsafe. + * * `markUnsafeWrite` is only for a write that could not register: its selector * is unknown or coordination failed. That marker must survive process loss and * keep both selectors ineligible until `resetAfterQuiescence`. Kitaru never @@ -243,16 +249,21 @@ export interface MastraMemoryCaptureBinding { settled(): Promise; }): Promise; markIncomplete(reason: string): void; + /** Wait for every storage write and its evidence upload to finish. */ drain(): Promise; /** * Join the invocation's memory work, including buffered observation and - * reflection, and record its evidence. Without `waitMs` this waits until the - * work settles; with it, resolve false once `waitMs` passes first. A failed - * join rejects either way. + * reflection, until its storage writes settle. Without `waitMs` this waits + * until the work settles; with it, resolve false once `waitMs` passes first + * and stop joining further rounds. A failed join rejects either way. */ settle(memory: MastraSettlingMemory, waitMs?: number): Promise; /** Check shared ownership immediately before persisting eligible inputs. */ verifyEligibility(): Promise; + /** + * Release ownership once storage writes have settled. Evidence uploads can + * still be running; `drain()` before reading `incompleteReasons`. + */ release(): Promise; } @@ -601,6 +612,9 @@ export function createMemoryCaptureBinding( resourceId: options.resourceId, }; let settling: Promise | undefined; + // A hung buffering operation stays in Mastra's process-wide map; stop + // polling it once a bounded caller has given up on this invocation. + let joinAbandoned = false; const registerWrite = createWriteRegistrar( options.exclusiveAccess, () => selector, @@ -751,6 +765,14 @@ export function createMemoryCaptureBinding( }, }); + async function settleMutations(): Promise { + while (true) { + const current = mutations; + await current; + if (current === mutations) return; + } + } + async function drain(): Promise { // Evidence can grow while a storage operation settles; follow both tails. while (true) { @@ -767,7 +789,8 @@ export function createMemoryCaptureBinding( const engine = await memory.omEngine; // Mastra's settled() does not join a buffered reflection, and settled // work can start more buffering. Repeat until a round records no write. - while (true) { + // Evidence uploads do not change storage, so they are not joined here. + while (!joinAbandoned) { const before = revision; await memory.settled(); const waitStarted = Date.now(); @@ -776,7 +799,7 @@ export function createMemoryCaptureBinding( options.resourceId, BUFFERING_WAIT_MS, ); - await drain(); + await settleMutations(); // waitForBuffering resolves, rather than rejects, when it times out. if (Date.now() - waitStarted >= BUFFERING_WAIT_MS) continue; if (revision === before) return; @@ -905,12 +928,14 @@ export function createMemoryCaptureBinding( } let timer: ReturnType | undefined; try { - return await Promise.race([ + const settled = await Promise.race([ settling.then(() => true), new Promise((resolve) => { timer = setTimeout(() => resolve(false), settleWaitMs); }), ]); + if (!settled) joinAbandoned = true; + return settled; } finally { if (timer) clearTimeout(timer); } @@ -918,7 +943,7 @@ export function createMemoryCaptureBinding( verifyEligibility, async release() { if (released) return; - await drain(); + await settleMutations(); released = true; try { await lease?.(); diff --git a/packages/mastra/src/memory-replay.ts b/packages/mastra/src/memory-replay.ts index f9bcd6a29..09ebae84d 100644 --- a/packages/mastra/src/memory-replay.ts +++ b/packages/mastra/src/memory-replay.ts @@ -278,6 +278,8 @@ export interface IsolatedMemoryReplayOptions { getRequestId?: MastraMemoryCaptureOptions["getRequestId"]; onIncomplete?: MastraMemoryCaptureOptions["onIncomplete"]; omTape?: ReturnType; + /** Fail the replay when its memory work has not settled after this long. */ + finalizationWaitMs: number; } /** Restore historical state into a fresh store; no production store is accepted. */ @@ -350,7 +352,17 @@ export async function createIsolatedMemoryReplay( finish(): Promise { finished ??= (async () => { try { - return await binding.settle(memory); + // A buffering operation that another turn left hung on the same + // thread would otherwise stall this replay indefinitely. + const settled = await binding.settle( + memory, + options.finalizationWaitMs, + ); + if (!settled) + binding.markIncomplete( + "Observational-memory work did not settle before the finalization deadline.", + ); + return settled; } finally { await binding.release(); await store.close(); diff --git a/packages/mastra/src/stateful-agent.ts b/packages/mastra/src/stateful-agent.ts index a5603d6e2..485067877 100644 --- a/packages/mastra/src/stateful-agent.ts +++ b/packages/mastra/src/stateful-agent.ts @@ -4,7 +4,11 @@ import type { MastraModelConfig } from "@mastra/core/llm"; import type { Mastra } from "@mastra/core/mastra"; import type { MemoryConfigInternal } from "@mastra/core/memory"; import type { InputProcessor } from "@mastra/core/processors"; -import { RequestContext } from "@mastra/core/request-context"; +import { + MASTRA_RESOURCE_ID_KEY, + MASTRA_THREAD_ID_KEY, + RequestContext, +} from "@mastra/core/request-context"; import type { MemoryStorage } from "@mastra/core/storage"; import type { Memory } from "@mastra/memory"; import { KitaruClient, type SessionNodeCreateRequest } from "@zenml-io/kitaru"; @@ -122,9 +126,10 @@ export interface MemoryReplayAgentOptions extends KitaruAgentOptions { skillsDirectory?: string; resolveModel: (id: string) => MastraModelConfig | Promise; /** - * How long a baseline waits after its stream closes for observational-memory - * work before it releases the source lease. A turn whose work does not settle - * in time is recorded as ineligible. Defaults to 60 seconds. + * How long a turn waits after its stream closes for observational-memory + * work. A baseline then releases the source lease, and a turn whose work did + * not settle in time is recorded as ineligible; a replay whose work did not + * settle fails. Defaults to 60 seconds. */ finalizationWaitMs?: number; } @@ -264,6 +269,8 @@ export function createMemoryReplayAgent( ...supplied, recordingLimits: normalizeRecordingLimits(supplied.recordingLimits), }; + const finalizationWaitMs = + supplied.finalizationWaitMs ?? DEFAULT_FINALIZATION_WAIT_MS; const client = new KitaruClient({ apiKey: options.apiKey, apiUrl: options.apiUrl, @@ -311,29 +318,42 @@ export function createMemoryReplayAgent( return native.stream(rawInput, callerOptions); } - /** Resolve the selectors a native call writes, or undefined when unknown. */ + /** + * Resolve the selectors a native call writes, or undefined when unknown. + * + * Follows Mastra's precedence: the reserved request-context keys override + * the merged default and caller memory options. + */ async function getNativeSelector( config: AgentConfig, callerOptions: RuntimeStreamOptions, ): Promise { try { + const requestContext = + callerOptions.requestContext ?? new RequestContext(); const defaults = requireRecord( typeof config.defaultOptions === "function" ? await config.defaultOptions({ - requestContext: - callerOptions.requestContext ?? new RequestContext(), + requestContext, mastra: options.mastra, }) : (config.defaultOptions ?? {}), "default options", ); const { deepMerge } = await import("@mastra/core/utils"); - return getSelector( - deepMerge( - record(defaults.memory) ? { memory: defaults.memory } : {}, - record(callerOptions.memory) ? { memory: callerOptions.memory } : {}, - ), + const merged = deepMerge( + record(defaults.memory) ? defaults.memory : {}, + record(callerOptions.memory) ? callerOptions.memory : {}, ); + const contextThread = requestContext.get(MASTRA_THREAD_ID_KEY); + const contextResource = requestContext.get(MASTRA_RESOURCE_ID_KEY); + return getSelector({ + memory: { + ...merged, + ...(contextThread ? { thread: contextThread } : {}), + ...(contextResource ? { resource: contextResource } : {}), + }, + }); } catch { return undefined; } @@ -546,6 +566,7 @@ export function createMemoryReplayAgent( onIncomplete, getRequestId: () => requestCapture?.currentRequestId, omTape, + finalizationWaitMs, }); } else { const source = await options.sourceMemory(); @@ -595,7 +616,7 @@ export function createMemoryReplayAgent( if (tracked || !sourceEngine) return; tracked = true; void sourceEngine - .trackBackgroundWork(binding.settle(memory)) + .trackBackgroundWork(binding.settle(memory, finalizationWaitMs)) .catch(() => undefined); }; let finished: Promise | undefined; @@ -607,10 +628,7 @@ export function createMemoryReplayAgent( finish() { finished ??= (async () => { trackSourceWork(); - const settled = await binding.settle( - memory, - options.finalizationWaitMs ?? DEFAULT_FINALIZATION_WAIT_MS, - ); + const settled = await binding.settle(memory, finalizationWaitMs); if (!settled) binding.markIncomplete( "Observational-memory work did not settle before the finalization deadline.", @@ -949,8 +967,10 @@ export function createMemoryReplayAgent( await capture.drain(); if (!historical) await runtime.binding.verifyEligibility(); // The turn's writes are settled and checked, so the next turn on - // this thread can acquire while the session update is sent. + // this thread can acquire while evidence uploads and the session + // update are sent. await runtime.release(); + await runtime.binding.drain(); // An OM call past the deadline may never return; the turn is // already ineligible, so do not wait for its tape entry. const omResults = settled ? await omTape.finish() : []; diff --git a/packages/mastra/test/memory-lease-lifecycle.test.ts b/packages/mastra/test/memory-lease-lifecycle.test.ts index 0e4b45285..9ed664733 100644 --- a/packages/mastra/test/memory-lease-lifecycle.test.ts +++ b/packages/mastra/test/memory-lease-lifecycle.test.ts @@ -1,6 +1,11 @@ import { mkdtemp, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { + MASTRA_RESOURCE_ID_KEY, + MASTRA_THREAD_ID_KEY, + RequestContext, +} from "@mastra/core/request-context"; import { InMemoryStore } from "@mastra/core/storage"; import { MastraLanguageModelV2Mock } from "@mastra/core/test-utils/llm-mock"; import { createTool } from "@mastra/core/tools"; @@ -207,14 +212,16 @@ async function setup( ); async function turn( message: string, - selector: { thread: string; resource: string } | null = { + selector: { thread: string; resource?: string } | null = { thread, resource: RESOURCE, }, + requestContext?: RequestContext, ) { const output = await adapter.stream(message, { maxSteps: 5, ...(selector ? { memory: selector } : {}), + ...(requestContext ? { requestContext } : {}), }); await output.consumeStream(); } @@ -374,6 +381,72 @@ it.each(LEASES)( }, ); +it.each(LEASES)( + "records two turns 100 ms apart as eligible while every Kitaru request is slow (%s lease)", + async (_kind, createAccess) => { + const { api, turn } = await setup(await createAccess(), { + // Evidence uploads after the stream closes must not keep the thread leased. + wrapFetch: (recorded) => async (input, init) => { + await pause(250); + return recorded(input, init); + }, + }); + await turn("Turn one."); + await pause(100); + await turn("Turn two."); + await vi.waitFor( + () => + expect(outcomes(api)).toEqual([ + "completed/eligible", + "completed/eligible", + ]), + { timeout: 5000 }, + ); + }, + 15_000, +); + +it.each(LEASES)( + "does not let a call whose selectors come from the request context affect other threads (%s lease)", + async (_kind, createAccess) => { + const access = await createAccess(); + const markUnsafeWrite = vi.spyOn(access, "markUnsafeWrite"); + const { api, domain, turn } = await setup(access); + const resourceOnly = new RequestContext(); + resourceOnly.set(MASTRA_RESOURCE_ID_KEY, "context-resource"); + await turn( + "Thread in options, resource in context.", + { thread: "context-thread-a" }, + resourceOnly, + ); + const both = new RequestContext(); + both.set(MASTRA_THREAD_ID_KEY, "context-thread-b"); + both.set(MASTRA_RESOURCE_ID_KEY, "context-resource"); + await turn("Both selectors in context.", null, both); + await turn("An unrelated user.", { + thread: "unrelated-thread", + resource: "unrelated-resource", + }); + await vi.waitFor( + () => + expect(outcomes(api).sort()).toEqual([ + "completed/eligible", + "setup-failure", + "setup-failure", + ]), + { timeout: 3000 }, + ); + expect(markUnsafeWrite).not.toHaveBeenCalled(); + for (const threadId of ["context-thread-a", "context-thread-b"]) { + const { messages } = await domain.listMessages({ threadId }); + expect(messages.length).toBeGreaterThan(0); + expect( + messages.every((message) => message.resourceId === "context-resource"), + ).toBe(true); + } + }, +); + it.each(LEASES)( "keeps a thread ineligible after an unregistered writer until quiescence reset (%s lease)", async (_kind, createAccess) => { @@ -446,3 +519,24 @@ it("lets the source Memory's settled() join a turn's observational-memory work", await memory.settled(); expect(observed).toBe(true); }); + +it("bounds the source Memory's settled() by the finalization deadline when OM work hangs", async () => { + const hung = gate(); + const { memory, turn } = await setup(createProcessLocalMemoryAccess(), { + thread: "hung-settled-thread", + backgroundObservation: true, + joinSourceMemory: true, + finalizationWaitMs: 200, + observerWait: () => hung.opened, + }); + try { + await turn(LONG_MESSAGE); + const outcome = await Promise.race([ + memory.settled().then(() => "settled"), + pause(3000).then(() => "still waiting"), + ]); + expect(outcome).toBe("settled"); + } finally { + hung.open(); + } +}); diff --git a/packages/mastra/test/memory-replay.test.ts b/packages/mastra/test/memory-replay.test.ts index 377985c11..2b755ed8a 100644 --- a/packages/mastra/test/memory-replay.test.ts +++ b/packages/mastra/test/memory-replay.test.ts @@ -45,6 +45,7 @@ it("restores historical memory with original timestamps into independent native resolveModel: async (id: string) => id.endsWith("observer") ? source.observer.model : source.reflector.model, recordMutation: async () => {}, + finalizationWaitMs: 60_000, }; const one = await createIsolatedMemoryReplay({ ...options, @@ -105,6 +106,7 @@ it("rejects unsupported memory dependencies and unjoined work before model resol throw new Error("must not resolve"); }, recordMutation: async () => {}, + finalizationWaitMs: 60_000, }), ).rejects.toThrow(/Unjoined/); expect(resolved).toBe(false); @@ -129,6 +131,7 @@ it("runs native observation, reflection and working-memory changes with ordered recordMutation: async (event) => { changes.push(event); }, + finalizationWaitMs: 60_000, }); const calls: ModelCall[] = []; const model = new MastraLanguageModelV2Mock({ @@ -227,6 +230,7 @@ it("waits for native background observation and evidence before releasing replay resolveModel: (id) => id.endsWith("observer") ? source.observer.model : source.reflector.model, recordMutation: async () => {}, + finalizationWaitMs: 60_000, }); try { const agent = new Agent({ @@ -260,6 +264,64 @@ it("waits for native background observation and evidence before releasing replay } }, 20000); +it("fails a replay whose background observation misses the finalization deadline", async () => { + let release!: () => void; + let signalStarted!: () => void; + const blocked = new Promise((resolve) => { + release = resolve; + }); + const started = new Promise((resolve) => { + signalStarted = resolve; + }); + const source = createMemoryRuntime({ + messageTokens: 10000, + observerWait: async () => { + signalStarted(); + await blocked; + }, + }); + await seedMemory(source); + const runtime = await createIsolatedMemoryReplay({ + invocationId: "hung-background", + initialSnapshot: { + ...(await snapshotMemory(source, true)), + threadId: THREAD, + resourceId: RESOURCE, + }, + configuration: serializeMemoryConfiguration( + source.memory.getMergedThreadConfig(), + ), + resolveModel: (id) => + id.endsWith("observer") ? source.observer.model : source.reflector.model, + recordMutation: async () => {}, + finalizationWaitMs: 200, + }); + try { + const agent = new Agent({ + id: "hung-background", + name: "Hung background", + instructions: "Answer", + model: new MastraLanguageModelV2Mock({ + doStream: async () => textStream("done"), + }), + memory: runtime.memory, + }); + const result = await agent.stream("Remember this preference.", { + memory: { thread: THREAD, resource: RESOURCE }, + }); + for await (const _ of result.textStream) { + /* Consume native stream. */ + } + await started; + expect(await runtime.finish()).toBe(false); + expect(runtime.binding.incompleteReasons).toContain( + "Observational-memory work did not settle before the finalization deadline.", + ); + } finally { + release(); + } +}, 20000); + it("replays on the record semantics of a database store unless the source was in memory", async () => { const { source, initialSnapshot, configuration } = await baseline(); const historical = initialSnapshot.records[0]; @@ -273,6 +335,7 @@ it("replays on the record semantics of a database store unless the source was in resolveModel: async (id: string) => id.endsWith("observer") ? source.observer.model : source.reflector.model, recordMutation: async () => {}, + finalizationWaitMs: 60_000, }; const persistent = await createIsolatedMemoryReplay({ ...options, From 978d26df67de84ed7c7d5e3167a0bba2b139749f Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 24 Sep 2026 15:31:31 +0200 Subject: [PATCH 027/115] Tolerate OM call drift in Mastra memory replay Replay matched recorded observer and reflector results by call order and count only. Mastra's number of OM calls depends on timing: a slow production observer merges buffering rounds that an instant replay makes separately, so an unchanged replay failed as `mastra_om_call_order`. - Match each replay call to an unused recorded call with the same phase, method and input fingerprint, else the next unused one of its phase. - Serve surplus calls without a provider: an empty observation for the observer, the last reflection for the reflector. Report input mismatches, surplus calls and unused results in an `om_call_divergence` span and `mastra_om_divergence` session metadata. - Fail closed only when a phase has no recorded result at all. - Record failed OM attempts, so a baseline whose observer succeeded after a retry stays eligible and replays the successful result. - Fingerprint OM inputs without message times, dates, ids and per-part creation timestamps, so `om_input_mismatch` reports only real input drift. - Close the session when a memory processor tripwire ends the run, since Mastra then calls neither `onFinish` nor `onError`. --- changelog.d/mastra-memory-replay.added.md | 2 +- docs/book/adapters/mastra.md | 2 +- packages/mastra/README.md | 2 +- packages/mastra/src/om-result-tape.ts | 345 ++++++++++++++---- packages/mastra/src/stateful-agent.ts | 52 ++- packages/mastra/src/stateful-tools.ts | 55 +++ packages/mastra/src/stream-recording.ts | 36 +- .../mastra/test/om-replay-tolerance.test.ts | 321 ++++++++++++++++ packages/mastra/test/om-result-tape.test.ts | 303 +++++++++++++-- 9 files changed, 1014 insertions(+), 104 deletions(-) create mode 100644 packages/mastra/test/om-replay-tolerance.test.ts diff --git a/changelog.d/mastra-memory-replay.added.md b/changelog.d/mastra-memory-replay.added.md index 2b6b2a1cc..cfbc52b94 100644 --- a/changelog.d/mastra-memory-replay.added.md +++ b/changelog.d/mastra-memory-replay.added.md @@ -1 +1 @@ -- Add opt-in Mastra streamed replay with isolated native thread memory. `createMemoryReplayAgent()` from `@zenml-io/kitaru-mastra/memory` records starting working and observational memory, ordered observation and reflection outputs, effective actor requests, controlled file bytes, and memory changes. Replay reuses those recorded outputs without reading or writing production memory. Requires Mastra core 1.67.0 and memory 1.30.0; existing history-only recordings must be recorded again to use this path. +- Add opt-in Mastra streamed replay with isolated native thread memory. `createMemoryReplayAgent()` from `@zenml-io/kitaru-mastra/memory` records starting working and observational memory, observation and reflection outputs including failed attempts, effective actor requests, controlled file bytes, and memory changes. Replay reuses those recorded outputs, matched to calls by phase and input, without reading or writing production memory or calling the observer or reflector model. Requires Mastra core 1.67.0 and memory 1.30.0; existing history-only recordings must be recorded again to use this path. diff --git a/docs/book/adapters/mastra.md b/docs/book/adapters/mastra.md index c1730db51..b6b1480c0 100644 --- a/docs/book/adapters/mastra.md +++ b/docs/book/adapters/mastra.md @@ -191,7 +191,7 @@ pnpm add @zenml-io/kitaru-mastra @mastra/core@1.67.0 @mastra/memory@1.30.0 zod The factory creates a fresh native agent for each invocation. A baseline uses your source storage and records the starting state before recall, then records the observer and reflector model outputs produced during that invocation. Replay restores the starting state into a separate in-memory store, runs the actor again, and supplies those recorded outputs to native observational memory (OM). Working-memory tools and native OM storage updates therefore run against isolated state. Replay never calls `sourceMemory()` or makes fresh observer/reflector provider calls. -Recorded OM outputs match by call order, phase (observer or reflector), and model method. A missing, extra, reordered, or unused call fails replay with `KITARU_REPLAY_DIVERGED:mastra_om_call_order`. Changed OM input alone does not fail replay: Kitaru reuses the recorded output and records an `om_input_mismatch` span. This lets you compare actor instruction/model changes with recorded OM results, but does not measure how a fresh observer or reflector would respond to the changed conversation. +Each replay OM call takes the unused recorded output with the same phase (observer or reflector), model method, and input. The input comparison ignores message times, dates, and generated ids. A call whose input matches no unused output takes the next unused output of its phase, and replay records an `om_input_mismatch` span. Mastra's number of OM calls depends on timing: a slow production observer merges buffering rounds that an instant replay makes separately. After a phase's recorded outputs are used up, further observer calls receive an empty observation and further reflector calls repeat the last recorded reflection. Neither reaches a provider. Replay reports these departures in an `om_call_divergence` span and in the session's `metadata.mastra_om_divergence` counts (`input_mismatches`, `surplus_calls`, `unused_results`). Replay fails with `KITARU_REPLAY_DIVERGED:mastra_om_call_order` only when a call's phase has no recorded output at all. A baseline also records failed OM attempts, so a turn whose observer succeeded after Mastra retried it stays eligible, and its replay serves the successful output directly. When a failed blocking observation ends the Mastra stream with a tripwire, the session still closes: a baseline becomes `ineligible` and a replay fails. Reusing recorded outputs lets you compare actor instruction/model changes, but does not measure how a fresh observer or reflector would respond to the changed conversation. The following binding uses a process-local store. Supply your existing public memory storage domain and its complete configuration for a persistent application: diff --git a/packages/mastra/README.md b/packages/mastra/README.md index 8fc93c880..0b9ee4971 100644 --- a/packages/mastra/README.md +++ b/packages/mastra/README.md @@ -155,7 +155,7 @@ pnpm add @zenml-io/kitaru-mastra @mastra/core@1.67.0 @mastra/memory@1.30.0 zod The factory creates a fresh native agent for each invocation. A baseline uses your source storage and records the starting state before recall, then records the observer and reflector model outputs produced during that invocation. Replay restores the starting state into a separate in-memory store, runs the actor again, and supplies those recorded outputs to native observational memory (OM). Working-memory tools and native OM storage updates therefore run against isolated state. Replay never calls `sourceMemory()` or makes fresh observer/reflector provider calls. -Recorded OM outputs match by call order, phase (observer or reflector), and model method. A missing, extra, reordered, or unused call fails replay with `KITARU_REPLAY_DIVERGED:mastra_om_call_order`. Changed OM input alone does not fail replay: Kitaru reuses the recorded output and records an `om_input_mismatch` span. This lets you compare actor instruction/model changes with recorded OM results, but does not measure how a fresh observer or reflector would respond to the changed conversation. +Each replay OM call takes the unused recorded output with the same phase (observer or reflector), model method, and input. The input comparison ignores message times, dates, and generated ids. A call whose input matches no unused output takes the next unused output of its phase, and replay records an `om_input_mismatch` span. Mastra's number of OM calls depends on timing: a slow production observer merges buffering rounds that an instant replay makes separately. After a phase's recorded outputs are used up, further observer calls receive an empty observation and further reflector calls repeat the last recorded reflection. Neither reaches a provider. Replay reports these departures in an `om_call_divergence` span and in the session's `metadata.mastra_om_divergence` counts (`input_mismatches`, `surplus_calls`, `unused_results`). Replay fails with `KITARU_REPLAY_DIVERGED:mastra_om_call_order` only when a call's phase has no recorded output at all. A baseline also records failed OM attempts, so a turn whose observer succeeded after Mastra retried it stays eligible, and its replay serves the successful output directly. When a failed blocking observation ends the Mastra stream with a tripwire, the session still closes: a baseline becomes `ineligible` and a replay fails. Reusing recorded outputs lets you compare actor instruction/model changes, but does not measure how a fresh observer or reflector would respond to the changed conversation. The following binding uses a process-local store. Supply your existing public memory storage domain and its complete configuration for a persistent application: diff --git a/packages/mastra/src/om-result-tape.ts b/packages/mastra/src/om-result-tape.ts index 8e23025bf..152921ad2 100644 --- a/packages/mastra/src/om-result-tape.ts +++ b/packages/mastra/src/om-result-tape.ts @@ -15,6 +15,23 @@ export interface OMResultEntry { method: OMMethod; inputFingerprint: string; output: JsonValue; + /** The provider call failed, so `output` is null. Mastra may have retried it. */ + failed?: true; +} + +/** How a replay's OM calls departed from the recorded calls without a live call. */ +export interface OMReplayDivergence { + /** Calls whose input matched no unused recorded call and took the next one. */ + inputMismatches: number; + /** Calls after every recorded result of their phase was used. */ + surplusCalls: number; + /** Recorded results that the replay never requested. */ + unusedResults: number; +} + +export interface OMTapeResult { + entries: OMResultEntry[]; + divergence: OMReplayDivergence; } export class MastraOMDivergenceError extends Error { @@ -25,9 +42,62 @@ export class MastraOMDivergenceError extends Error { } } -function fingerprint(input: unknown): string { +/** Replays a provider failure that every recorded attempt of a call ended in. */ +export class MastraOMRecordedFailureError extends Error { + readonly code = "mastra_om_recorded_failure"; + + constructor() { + super("The recorded observational-memory model call failed."); + } +} + +const VOLATILE_KEYS = new Set(["createdAt", "updatedAt", "abortSignal"]); +const VOLATILE_TEXT: ReadonlyArray = [ + [ + /\b\d{4}-\d{2}-\d{2}T\d{2}:\d{2}(?::\d{2}(?:\.\d+)?)?(?:Z|[+-]\d{2}:?\d{2})?/g, + "", + ], + [ + /\b(?:(?:Mon|Tue|Wed|Thu|Fri|Sat|Sun)[a-z]*,? )?(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\.? \d{1,2},? \d{4}\b/g, + "", + ], + [/\b\d{1,2}:\d{2}(?::\d{2})?\s?[AP]M\b/gi, "