Skip to content

Commit e146aad

Browse files
abhiaiyer91Mastra Code (anthropic/claude-opus-4-6)
andauthored
fix(core): merge thread metadata and working memory on partial updates (#16846)
Co-authored-by: Mastra Code (anthropic/claude-opus-4-6) <noreply@mastra.ai>
1 parent 865a28e commit e146aad

6 files changed

Lines changed: 313 additions & 6 deletions

File tree

.changeset/perky-sloths-shake.md

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
---
2+
'@mastra/core': patch
3+
---
4+
5+
Fixed thread metadata updates to merge with existing fields instead of replacing them. Previously, updating a thread's metadata would silently drop any fields not included in the update. Now existing metadata fields are preserved when updating.
6+
7+
Fixed MockMemory working memory tool to support partial JSON updates when using schema-based working memory. Previously, sending a partial update would overwrite all existing data. Now for schema-based configs, unchanged fields are preserved automatically (matching @mastra/memory behavior).
8+
9+
Fixed MockMemory constructor to preserve workingMemory config options (like schema) when enableWorkingMemory is true.

packages/core/src/agent/__tests__/memory-metadata.test.ts

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,51 @@ function memoryMetadataTests(version: 'v1' | 'v2') {
186186
expect(thread?.metadata).toEqual({ client: 'updated' });
187187
});
188188

189+
it('should merge metadata with existing fields instead of replacing', async () => {
190+
const mockMemory = new MockMemory();
191+
const initialThread: StorageThreadType = {
192+
id: 'thread-1',
193+
resourceId: 'user-1',
194+
metadata: { existingField: 'should-persist', client: 'initial' },
195+
createdAt: new Date(),
196+
updatedAt: new Date(),
197+
};
198+
await mockMemory.saveThread({ thread: initialThread });
199+
200+
const agent = new Agent({
201+
id: 'test-agent',
202+
name: 'Test Agent',
203+
instructions: 'test',
204+
model: dummyModel,
205+
memory: mockMemory,
206+
});
207+
208+
if (version === 'v1') {
209+
await agent.generateLegacy('hello', {
210+
memory: {
211+
resource: 'user-1',
212+
thread: {
213+
id: 'thread-1',
214+
metadata: { client: 'updated' },
215+
},
216+
},
217+
});
218+
} else {
219+
await agent.generate('hello', {
220+
memory: {
221+
resource: 'user-1',
222+
thread: {
223+
id: 'thread-1',
224+
metadata: { client: 'updated' },
225+
},
226+
},
227+
});
228+
}
229+
230+
const thread = await mockMemory.getThreadById({ threadId: 'thread-1' });
231+
expect(thread?.metadata).toEqual({ existingField: 'should-persist', client: 'updated' });
232+
});
233+
189234
it('should not update metadata if it is the same using generate', async () => {
190235
const mockMemory = new MockMemory();
191236
const initialThread: StorageThreadType = {

packages/core/src/agent/agent-legacy.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -398,7 +398,7 @@ export class AgentLegacyHandler {
398398
(thread.metadata && !deepEqual(existingThread.metadata, thread.metadata))
399399
) {
400400
threadObject = await memory.saveThread({
401-
thread: { ...existingThread, metadata: thread.metadata },
401+
thread: { ...existingThread, metadata: { ...(existingThread.metadata ?? {}), ...thread.metadata } },
402402
memoryConfig,
403403
});
404404
} else {

packages/core/src/agent/workflows/prepare-stream/prepare-memory-step.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ export function createPrepareMemoryStep<OUTPUT = undefined>({
151151
(thread.metadata && !deepEqual(existingThread.metadata, thread.metadata))
152152
) {
153153
threadObject = await memory.saveThread({
154-
thread: { ...existingThread, metadata: thread.metadata },
154+
thread: { ...existingThread, metadata: { ...(existingThread.metadata ?? {}), ...thread.metadata } },
155155
memoryConfig,
156156
});
157157
} else {
Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
import { describe, it, expect } from 'vitest';
2+
import { z } from 'zod/v4';
3+
import { MockMemory } from './mock';
4+
5+
describe('MockMemory working memory merge semantics', () => {
6+
const threadId = 'thread-1';
7+
const resourceId = 'resource-1';
8+
9+
async function setupMemory(useSchema: boolean) {
10+
const options: ConstructorParameters<typeof MockMemory>[0] = {
11+
enableWorkingMemory: true,
12+
};
13+
14+
if (useSchema) {
15+
options.options = {
16+
workingMemory: {
17+
enabled: true,
18+
schema: z.object({
19+
name: z.string().optional(),
20+
age: z.number().optional(),
21+
location: z.string().optional(),
22+
}),
23+
},
24+
};
25+
}
26+
27+
const memory = new MockMemory(options);
28+
29+
// Create a thread so the tool doesn't error
30+
await memory.createThread({ threadId, resourceId });
31+
32+
return memory;
33+
}
34+
35+
async function callUpdateTool(memory: MockMemory, input: string) {
36+
const config = (memory as any).getMergedThreadConfig();
37+
const tools = memory.listTools(config);
38+
const tool = tools.updateWorkingMemory;
39+
if (!tool) throw new Error('updateWorkingMemory tool not found');
40+
41+
await (tool as any).execute({ memory: input }, { agent: { threadId, resourceId }, memory });
42+
}
43+
44+
it('replaces working memory entirely for template-based (no schema)', async () => {
45+
const memory = await setupMemory(false);
46+
47+
await callUpdateTool(memory, JSON.stringify({ name: 'Alice', age: 30, location: 'NYC' }));
48+
await callUpdateTool(memory, JSON.stringify({ location: 'LA' }));
49+
50+
const wm = await memory.getWorkingMemory({ threadId, resourceId });
51+
const parsed = JSON.parse(wm!);
52+
expect(parsed).toEqual({ location: 'LA' });
53+
expect(parsed.name).toBeUndefined();
54+
});
55+
56+
it('merges working memory for schema-based configs', async () => {
57+
const memory = await setupMemory(true);
58+
59+
await callUpdateTool(memory, JSON.stringify({ name: 'Alice', age: 30, location: 'NYC' }));
60+
await callUpdateTool(memory, JSON.stringify({ location: 'LA' }));
61+
62+
const wm = await memory.getWorkingMemory({ threadId, resourceId });
63+
const parsed = JSON.parse(wm!);
64+
expect(parsed).toEqual({ name: 'Alice', age: 30, location: 'LA' });
65+
});
66+
67+
it('overwrites fields in schema-based merge when explicitly provided', async () => {
68+
const memory = await setupMemory(true);
69+
70+
await callUpdateTool(memory, JSON.stringify({ name: 'Alice', age: 30 }));
71+
await callUpdateTool(memory, JSON.stringify({ name: 'Bob', age: 25 }));
72+
73+
const wm = await memory.getWorkingMemory({ threadId, resourceId });
74+
const parsed = JSON.parse(wm!);
75+
expect(parsed).toEqual({ name: 'Bob', age: 25 });
76+
});
77+
78+
it('handles first write with no existing data in schema mode', async () => {
79+
const memory = await setupMemory(true);
80+
81+
await callUpdateTool(memory, JSON.stringify({ name: 'Alice' }));
82+
83+
const wm = await memory.getWorkingMemory({ threadId, resourceId });
84+
const parsed = JSON.parse(wm!);
85+
expect(parsed).toEqual({ name: 'Alice' });
86+
});
87+
88+
it('deep-merges nested objects in schema mode', async () => {
89+
const memory = new MockMemory({
90+
enableWorkingMemory: true,
91+
options: {
92+
workingMemory: {
93+
enabled: true,
94+
schema: z.object({
95+
user: z.object({
96+
name: z.string().optional(),
97+
address: z
98+
.object({
99+
city: z.string().optional(),
100+
state: z.string().optional(),
101+
})
102+
.optional(),
103+
}),
104+
}),
105+
},
106+
},
107+
});
108+
await memory.createThread({ threadId, resourceId });
109+
110+
await callUpdateTool(memory, JSON.stringify({ user: { name: 'Alice', address: { city: 'NYC', state: 'NY' } } }));
111+
await callUpdateTool(memory, JSON.stringify({ user: { address: { city: 'LA' } } }));
112+
113+
const wm = await memory.getWorkingMemory({ threadId, resourceId });
114+
const parsed = JSON.parse(wm!);
115+
// Deep merge: name preserved, state preserved, only city changed
116+
expect(parsed).toEqual({ user: { name: 'Alice', address: { city: 'LA', state: 'NY' } } });
117+
});
118+
119+
it('deletes keys set to null in schema mode', async () => {
120+
const memory = await setupMemory(true);
121+
122+
await callUpdateTool(memory, JSON.stringify({ name: 'Alice', age: 30, location: 'NYC' }));
123+
await callUpdateTool(memory, JSON.stringify({ age: null }));
124+
125+
const wm = await memory.getWorkingMemory({ threadId, resourceId });
126+
const parsed = JSON.parse(wm!);
127+
// null deletes the key
128+
expect(parsed).toEqual({ name: 'Alice', location: 'NYC' });
129+
expect(parsed.age).toBeUndefined();
130+
});
131+
132+
it('replaces arrays entirely in schema mode', async () => {
133+
const memory = new MockMemory({
134+
enableWorkingMemory: true,
135+
options: {
136+
workingMemory: {
137+
enabled: true,
138+
schema: z.object({
139+
tags: z.array(z.string()).optional(),
140+
count: z.number().optional(),
141+
}),
142+
},
143+
},
144+
});
145+
await memory.createThread({ threadId, resourceId });
146+
147+
await callUpdateTool(memory, JSON.stringify({ tags: ['a', 'b', 'c'], count: 3 }));
148+
await callUpdateTool(memory, JSON.stringify({ tags: ['x'] }));
149+
150+
const wm = await memory.getWorkingMemory({ threadId, resourceId });
151+
const parsed = JSON.parse(wm!);
152+
// Arrays replace, count preserved
153+
expect(parsed).toEqual({ tags: ['x'], count: 3 });
154+
});
155+
});

packages/core/src/memory/mock.ts

Lines changed: 102 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,54 @@ import type {
2424
SharedMemoryConfig,
2525
} from './types';
2626

27+
/**
28+
* Deep-merge working memory objects.
29+
* Matches the semantics of `deepMergeWorkingMemory` in `@mastra/memory`:
30+
* - `null` values delete the corresponding key
31+
* - Arrays are replaced entirely (not merged element-by-element)
32+
* - Nested plain objects are merged recursively
33+
* - Primitives and new keys are set directly
34+
*/
35+
function deepMergeWorkingMemory(
36+
existing: Record<string, unknown> | null | undefined,
37+
update: Record<string, unknown> | null | undefined,
38+
): Record<string, unknown> {
39+
if (!update || typeof update !== 'object' || Object.keys(update).length === 0) {
40+
return existing && typeof existing === 'object' ? { ...existing } : {};
41+
}
42+
if (!existing || typeof existing !== 'object') {
43+
return update;
44+
}
45+
46+
const result: Record<string, unknown> = { ...existing };
47+
48+
for (const key of Object.keys(update)) {
49+
const updateValue = update[key];
50+
const existingValue = result[key];
51+
52+
if (updateValue === null) {
53+
delete result[key];
54+
} else if (Array.isArray(updateValue)) {
55+
result[key] = updateValue;
56+
} else if (
57+
typeof updateValue === 'object' &&
58+
updateValue !== null &&
59+
typeof existingValue === 'object' &&
60+
existingValue !== null &&
61+
!Array.isArray(existingValue)
62+
) {
63+
result[key] = deepMergeWorkingMemory(
64+
existingValue as Record<string, unknown>,
65+
updateValue as Record<string, unknown>,
66+
);
67+
} else {
68+
result[key] = updateValue;
69+
}
70+
}
71+
72+
return result;
73+
}
74+
2775
export class MockMemory extends MastraMemory {
2876
constructor({
2977
storage,
@@ -44,7 +92,11 @@ export class MockMemory extends MastraMemory {
4492
options: {
4593
...options,
4694
workingMemory: enableWorkingMemory
47-
? ({ enabled: true, template: workingMemoryTemplate } as WorkingMemory)
95+
? ({
96+
...options?.workingMemory,
97+
enabled: true,
98+
...(workingMemoryTemplate !== undefined ? { template: workingMemoryTemplate } : {}),
99+
} as WorkingMemory)
48100
: options?.workingMemory,
49101
lastMessages: enableMessageHistory ? (options?.lastMessages ?? 10) : options?.lastMessages,
50102
},
@@ -175,10 +227,15 @@ export class MockMemory extends MastraMemory {
175227
return {};
176228
}
177229

230+
const usesMergeSemantics = Boolean(mergedConfig.workingMemory?.schema);
231+
const description = usesMergeSemantics
232+
? `Update the working memory with new information. Data is merged with existing memory - only include fields you want to add or update.`
233+
: `Update the working memory with new information. Any data not included will be overwritten.`;
234+
178235
return {
179236
updateWorkingMemory: createTool({
180237
id: 'update-working-memory',
181-
description: `Update the working memory with new information. Any data not included will be overwritten.`,
238+
description,
182239
inputSchema: z.object({ memory: z.string() }),
183240
execute: async (inputData, context) => {
184241
const threadId = context?.agent?.threadId;
@@ -218,8 +275,49 @@ export class MockMemory extends MastraMemory {
218275
}
219276
}
220277

221-
const workingMemory =
222-
typeof inputData.memory === 'string' ? inputData.memory : JSON.stringify(inputData.memory);
278+
let workingMemory: string;
279+
280+
if (usesMergeSemantics) {
281+
const existingRaw = await memory.getWorkingMemory({
282+
threadId,
283+
resourceId,
284+
memoryConfig: _config,
285+
});
286+
287+
let existingData: Record<string, unknown> | null = null;
288+
if (existingRaw) {
289+
try {
290+
existingData = typeof existingRaw === 'string' ? JSON.parse(existingRaw) : existingRaw;
291+
} catch {
292+
existingData = null;
293+
}
294+
}
295+
296+
const memoryInput = inputData.memory;
297+
let newData: unknown;
298+
if (typeof memoryInput === 'string') {
299+
try {
300+
newData = JSON.parse(memoryInput);
301+
} catch {
302+
newData = memoryInput;
303+
}
304+
} else {
305+
newData = memoryInput;
306+
}
307+
308+
if (newData && typeof newData === 'object' && !Array.isArray(newData)) {
309+
workingMemory = JSON.stringify(
310+
deepMergeWorkingMemory(
311+
existingData as Record<string, unknown> | null,
312+
newData as Record<string, unknown>,
313+
),
314+
);
315+
} else {
316+
workingMemory = typeof newData === 'string' ? newData : JSON.stringify(newData);
317+
}
318+
} else {
319+
workingMemory = typeof inputData.memory === 'string' ? inputData.memory : JSON.stringify(inputData.memory);
320+
}
223321

224322
await memory.updateWorkingMemory({
225323
threadId,

0 commit comments

Comments
 (0)