@mastra/pg 0.10.2-alpha.0 → 0.10.2-alpha.2
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/.turbo/turbo-build.log +7 -7
- package/CHANGELOG.md +20 -0
- package/dist/_tsup-dts-rollup.d.cts +33 -40
- package/dist/_tsup-dts-rollup.d.ts +33 -40
- package/dist/index.cjs +286 -225
- package/dist/index.js +286 -225
- package/package.json +4 -4
- package/src/storage/index.test.ts +308 -191
- package/src/storage/index.ts +348 -345
|
@@ -1,6 +1,16 @@
|
|
|
1
1
|
import { randomUUID } from 'crypto';
|
|
2
|
-
import {
|
|
3
|
-
|
|
2
|
+
import {
|
|
3
|
+
createSampleEval,
|
|
4
|
+
createSampleTraceForDB,
|
|
5
|
+
createSampleThread,
|
|
6
|
+
createSampleMessageV1,
|
|
7
|
+
createSampleMessageV2,
|
|
8
|
+
createSampleWorkflowSnapshot,
|
|
9
|
+
resetRole,
|
|
10
|
+
checkWorkflowSnapshot,
|
|
11
|
+
} from '@internal/storage-test-utils';
|
|
12
|
+
import type { MastraMessageV1, MastraMessageV2, StorageThreadType } from '@mastra/core/memory';
|
|
13
|
+
import type { StorageColumn, TABLE_NAMES } from '@mastra/core/storage';
|
|
4
14
|
import {
|
|
5
15
|
TABLE_WORKFLOW_SNAPSHOT,
|
|
6
16
|
TABLE_MESSAGES,
|
|
@@ -27,65 +37,6 @@ const connectionString = `postgresql://${TEST_CONFIG.user}:${TEST_CONFIG.passwor
|
|
|
27
37
|
|
|
28
38
|
vi.setConfig({ testTimeout: 60_000, hookTimeout: 60_000 });
|
|
29
39
|
|
|
30
|
-
// Sample test data factory functions
|
|
31
|
-
const createSampleThread = () => ({
|
|
32
|
-
id: `thread-${randomUUID()}`,
|
|
33
|
-
resourceId: `resource-${randomUUID()}`,
|
|
34
|
-
title: 'Test Thread',
|
|
35
|
-
createdAt: new Date(),
|
|
36
|
-
updatedAt: new Date(),
|
|
37
|
-
metadata: { key: 'value' },
|
|
38
|
-
});
|
|
39
|
-
|
|
40
|
-
let role: 'user' | 'assistant' = 'assistant';
|
|
41
|
-
const getRole = () => {
|
|
42
|
-
if (role === `user`) role = `assistant`;
|
|
43
|
-
else role = `user`;
|
|
44
|
-
return role;
|
|
45
|
-
};
|
|
46
|
-
const createSampleMessage = (threadId: string): MastraMessageV1 => ({
|
|
47
|
-
id: `msg-${randomUUID()}`,
|
|
48
|
-
resourceId: `resource-${randomUUID()}`,
|
|
49
|
-
role: getRole(),
|
|
50
|
-
type: 'text',
|
|
51
|
-
threadId,
|
|
52
|
-
content: [{ type: 'text', text: 'Hello' }],
|
|
53
|
-
createdAt: new Date(),
|
|
54
|
-
});
|
|
55
|
-
|
|
56
|
-
const createSampleWorkflowSnapshot = (status: WorkflowRunState['context'][string]['status'], createdAt?: Date) => {
|
|
57
|
-
const runId = `run-${randomUUID()}`;
|
|
58
|
-
const stepId = `step-${randomUUID()}`;
|
|
59
|
-
const timestamp = createdAt || new Date();
|
|
60
|
-
const snapshot = {
|
|
61
|
-
result: { success: true },
|
|
62
|
-
value: {},
|
|
63
|
-
context: {
|
|
64
|
-
[stepId]: {
|
|
65
|
-
status,
|
|
66
|
-
payload: {},
|
|
67
|
-
error: undefined,
|
|
68
|
-
startedAt: timestamp.getTime(),
|
|
69
|
-
endedAt: new Date(timestamp.getTime() + 15000).getTime(),
|
|
70
|
-
},
|
|
71
|
-
input: {},
|
|
72
|
-
},
|
|
73
|
-
serializedStepGraph: [],
|
|
74
|
-
activePaths: [],
|
|
75
|
-
suspendedPaths: {},
|
|
76
|
-
runId,
|
|
77
|
-
timestamp: timestamp.getTime(),
|
|
78
|
-
} as unknown as WorkflowRunState;
|
|
79
|
-
return { snapshot, runId, stepId };
|
|
80
|
-
};
|
|
81
|
-
|
|
82
|
-
const checkWorkflowSnapshot = (snapshot: WorkflowRunState | string, stepId: string, status: string) => {
|
|
83
|
-
if (typeof snapshot === 'string') {
|
|
84
|
-
throw new Error('Expected WorkflowRunState, got string');
|
|
85
|
-
}
|
|
86
|
-
expect(snapshot.context?.[stepId]?.status).toBe(status);
|
|
87
|
-
};
|
|
88
|
-
|
|
89
40
|
describe('PostgresStore', () => {
|
|
90
41
|
let store: PostgresStore;
|
|
91
42
|
|
|
@@ -210,7 +161,7 @@ describe('PostgresStore', () => {
|
|
|
210
161
|
await store.saveThread({ thread });
|
|
211
162
|
|
|
212
163
|
// Add some messages
|
|
213
|
-
const messages = [
|
|
164
|
+
const messages = [createSampleMessageV1({ threadId: thread.id }), createSampleMessageV1({ threadId: thread.id })];
|
|
214
165
|
await store.saveMessages({ messages });
|
|
215
166
|
|
|
216
167
|
await store.deleteThread({ threadId: thread.id });
|
|
@@ -229,7 +180,7 @@ describe('PostgresStore', () => {
|
|
|
229
180
|
const thread = createSampleThread();
|
|
230
181
|
await store.saveThread({ thread });
|
|
231
182
|
|
|
232
|
-
const messages = [
|
|
183
|
+
const messages = [createSampleMessageV1({ threadId: thread.id }), createSampleMessageV1({ threadId: thread.id })];
|
|
233
184
|
|
|
234
185
|
// Save messages
|
|
235
186
|
const savedMessages = await store.saveMessages({ messages });
|
|
@@ -254,24 +205,18 @@ describe('PostgresStore', () => {
|
|
|
254
205
|
const thread = createSampleThread();
|
|
255
206
|
await store.saveThread({ thread });
|
|
256
207
|
|
|
257
|
-
const
|
|
258
|
-
{ ...createSampleMessage(thread.id), content: [{ type: 'text', text: 'First' }] },
|
|
259
|
-
{
|
|
260
|
-
...createSampleMessage(thread.id),
|
|
261
|
-
content: [{ type: 'text', text: 'Second' }],
|
|
262
|
-
},
|
|
263
|
-
{ ...createSampleMessage(thread.id), content: [{ type: 'text', text: 'Third' }] },
|
|
264
|
-
] satisfies MastraMessageV1[];
|
|
208
|
+
const messageContent = ['First', 'Second', 'Third'];
|
|
265
209
|
|
|
266
|
-
|
|
210
|
+
const messages = messageContent.map(content => createSampleMessageV2({ threadId: thread.id, content }));
|
|
267
211
|
|
|
268
|
-
|
|
212
|
+
await store.saveMessages({ messages, format: 'v2' });
|
|
213
|
+
|
|
214
|
+
const retrievedMessages = await store.getMessages({ threadId: thread.id, format: 'v2' });
|
|
269
215
|
expect(retrievedMessages).toHaveLength(3);
|
|
270
216
|
|
|
271
217
|
// Verify order is maintained
|
|
272
218
|
retrievedMessages.forEach((msg, idx) => {
|
|
273
|
-
|
|
274
|
-
expect(msg.content[0].text).toBe(messages[idx].content[0].text);
|
|
219
|
+
expect((msg.content.parts[0] as any).text).toEqual(messageContent[idx]);
|
|
275
220
|
});
|
|
276
221
|
});
|
|
277
222
|
|
|
@@ -280,8 +225,8 @@ describe('PostgresStore', () => {
|
|
|
280
225
|
await store.saveThread({ thread });
|
|
281
226
|
|
|
282
227
|
const messages = [
|
|
283
|
-
|
|
284
|
-
{ ...
|
|
228
|
+
createSampleMessageV1({ threadId: thread.id }),
|
|
229
|
+
{ ...createSampleMessageV1({ threadId: thread.id }), id: null } as any, // This will cause an error
|
|
285
230
|
];
|
|
286
231
|
|
|
287
232
|
await expect(store.saveMessages({ messages })).rejects.toThrow();
|
|
@@ -291,41 +236,108 @@ describe('PostgresStore', () => {
|
|
|
291
236
|
expect(savedMessages).toHaveLength(0);
|
|
292
237
|
});
|
|
293
238
|
|
|
294
|
-
it('should
|
|
295
|
-
const thread = createSampleThread();
|
|
239
|
+
it('should retrieve messages w/ next/prev messages by message id + resource id', async () => {
|
|
240
|
+
const thread = createSampleThread({ id: 'thread-one' });
|
|
296
241
|
await store.saveThread({ thread });
|
|
297
|
-
const now = new Date();
|
|
298
|
-
const yesterday = new Date(now.getTime() - 24 * 60 * 60 * 1000);
|
|
299
|
-
const dayBeforeYesterday = new Date(now.getTime() - 48 * 60 * 60 * 1000);
|
|
300
242
|
|
|
301
|
-
const
|
|
302
|
-
|
|
303
|
-
|
|
304
|
-
|
|
305
|
-
|
|
306
|
-
|
|
307
|
-
|
|
308
|
-
|
|
309
|
-
|
|
310
|
-
|
|
311
|
-
|
|
312
|
-
|
|
313
|
-
|
|
314
|
-
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
319
|
-
|
|
243
|
+
const thread2 = createSampleThread({ id: 'thread-two' });
|
|
244
|
+
await store.saveThread({ thread: thread2 });
|
|
245
|
+
|
|
246
|
+
const thread3 = createSampleThread({ id: 'thread-three' });
|
|
247
|
+
await store.saveThread({ thread: thread3 });
|
|
248
|
+
|
|
249
|
+
const messages: MastraMessageV2[] = [
|
|
250
|
+
createSampleMessageV2({ threadId: 'thread-one', content: 'First', resourceId: 'cross-thread-resource' }),
|
|
251
|
+
createSampleMessageV2({ threadId: 'thread-one', content: 'Second', resourceId: 'cross-thread-resource' }),
|
|
252
|
+
createSampleMessageV2({ threadId: 'thread-one', content: 'Third', resourceId: 'cross-thread-resource' }),
|
|
253
|
+
|
|
254
|
+
createSampleMessageV2({ threadId: 'thread-two', content: 'Fourth', resourceId: 'cross-thread-resource' }),
|
|
255
|
+
createSampleMessageV2({ threadId: 'thread-two', content: 'Fifth', resourceId: 'cross-thread-resource' }),
|
|
256
|
+
createSampleMessageV2({ threadId: 'thread-two', content: 'Sixth', resourceId: 'cross-thread-resource' }),
|
|
257
|
+
|
|
258
|
+
createSampleMessageV2({ threadId: 'thread-three', content: 'Seventh', resourceId: 'other-resource' }),
|
|
259
|
+
createSampleMessageV2({ threadId: 'thread-three', content: 'Eighth', resourceId: 'other-resource' }),
|
|
260
|
+
];
|
|
261
|
+
|
|
262
|
+
await store.saveMessages({ messages: messages, format: 'v2' });
|
|
263
|
+
|
|
264
|
+
const retrievedMessages = await store.getMessages({ threadId: 'thread-one', format: 'v2' });
|
|
265
|
+
expect(retrievedMessages).toHaveLength(3);
|
|
266
|
+
expect(retrievedMessages.map((m: any) => m.content.parts[0].text)).toEqual(['First', 'Second', 'Third']);
|
|
267
|
+
|
|
268
|
+
const retrievedMessages2 = await store.getMessages({ threadId: 'thread-two', format: 'v2' });
|
|
269
|
+
expect(retrievedMessages2).toHaveLength(3);
|
|
270
|
+
expect(retrievedMessages2.map((m: any) => m.content.parts[0].text)).toEqual(['Fourth', 'Fifth', 'Sixth']);
|
|
271
|
+
|
|
272
|
+
const retrievedMessages3 = await store.getMessages({ threadId: 'thread-three', format: 'v2' });
|
|
273
|
+
expect(retrievedMessages3).toHaveLength(2);
|
|
274
|
+
expect(retrievedMessages3.map((m: any) => m.content.parts[0].text)).toEqual(['Seventh', 'Eighth']);
|
|
275
|
+
|
|
276
|
+
const crossThreadMessages: MastraMessageV2[] = await store.getMessages({
|
|
277
|
+
threadId: 'thread-doesnt-exist',
|
|
278
|
+
format: 'v2',
|
|
279
|
+
selectBy: {
|
|
280
|
+
last: 0,
|
|
281
|
+
include: [
|
|
282
|
+
{
|
|
283
|
+
id: messages[1].id,
|
|
284
|
+
threadId: 'thread-one',
|
|
285
|
+
withNextMessages: 2,
|
|
286
|
+
withPreviousMessages: 2,
|
|
287
|
+
},
|
|
288
|
+
{
|
|
289
|
+
id: messages[4].id,
|
|
290
|
+
threadId: 'thread-two',
|
|
291
|
+
withPreviousMessages: 2,
|
|
292
|
+
withNextMessages: 2,
|
|
293
|
+
},
|
|
294
|
+
],
|
|
295
|
+
},
|
|
320
296
|
});
|
|
321
|
-
expect(resultPage.total).toBe(5);
|
|
322
|
-
expect(resultPage.messages).toHaveLength(3);
|
|
323
297
|
|
|
324
|
-
expect(
|
|
325
|
-
expect(
|
|
326
|
-
expect(
|
|
327
|
-
|
|
328
|
-
|
|
298
|
+
expect(crossThreadMessages).toHaveLength(6);
|
|
299
|
+
expect(crossThreadMessages.filter(m => m.threadId === `thread-one`)).toHaveLength(3);
|
|
300
|
+
expect(crossThreadMessages.filter(m => m.threadId === `thread-two`)).toHaveLength(3);
|
|
301
|
+
|
|
302
|
+
const crossThreadMessages2: MastraMessageV2[] = await store.getMessages({
|
|
303
|
+
threadId: 'thread-one',
|
|
304
|
+
format: 'v2',
|
|
305
|
+
selectBy: {
|
|
306
|
+
last: 0,
|
|
307
|
+
include: [
|
|
308
|
+
{
|
|
309
|
+
id: messages[4].id,
|
|
310
|
+
threadId: 'thread-two',
|
|
311
|
+
withPreviousMessages: 1,
|
|
312
|
+
withNextMessages: 1,
|
|
313
|
+
},
|
|
314
|
+
],
|
|
315
|
+
},
|
|
316
|
+
});
|
|
317
|
+
|
|
318
|
+
expect(crossThreadMessages2).toHaveLength(3);
|
|
319
|
+
expect(crossThreadMessages2.filter(m => m.threadId === `thread-one`)).toHaveLength(0);
|
|
320
|
+
expect(crossThreadMessages2.filter(m => m.threadId === `thread-two`)).toHaveLength(3);
|
|
321
|
+
|
|
322
|
+
const crossThreadMessages3: MastraMessageV2[] = await store.getMessages({
|
|
323
|
+
threadId: 'thread-two',
|
|
324
|
+
format: 'v2',
|
|
325
|
+
selectBy: {
|
|
326
|
+
last: 0,
|
|
327
|
+
include: [
|
|
328
|
+
{
|
|
329
|
+
id: messages[1].id,
|
|
330
|
+
threadId: 'thread-one',
|
|
331
|
+
withNextMessages: 1,
|
|
332
|
+
withPreviousMessages: 1,
|
|
333
|
+
},
|
|
334
|
+
],
|
|
335
|
+
},
|
|
336
|
+
});
|
|
337
|
+
|
|
338
|
+
expect(crossThreadMessages3).toHaveLength(3);
|
|
339
|
+
expect(crossThreadMessages3.filter(m => m.threadId === `thread-one`)).toHaveLength(3);
|
|
340
|
+
expect(crossThreadMessages3.filter(m => m.threadId === `thread-two`)).toHaveLength(0);
|
|
329
341
|
});
|
|
330
342
|
});
|
|
331
343
|
|
|
@@ -887,6 +899,96 @@ describe('PostgresStore', () => {
|
|
|
887
899
|
});
|
|
888
900
|
});
|
|
889
901
|
|
|
902
|
+
describe('alterTable', () => {
|
|
903
|
+
const TEST_TABLE = 'test_alter_table';
|
|
904
|
+
const BASE_SCHEMA = {
|
|
905
|
+
id: { type: 'integer', primaryKey: true, nullable: false },
|
|
906
|
+
name: { type: 'text', nullable: true },
|
|
907
|
+
} as Record<string, StorageColumn>;
|
|
908
|
+
|
|
909
|
+
beforeEach(async () => {
|
|
910
|
+
await store.createTable({ tableName: TEST_TABLE as TABLE_NAMES, schema: BASE_SCHEMA });
|
|
911
|
+
});
|
|
912
|
+
|
|
913
|
+
afterEach(async () => {
|
|
914
|
+
await store.clearTable({ tableName: TEST_TABLE as TABLE_NAMES });
|
|
915
|
+
});
|
|
916
|
+
|
|
917
|
+
it('adds a new column to an existing table', async () => {
|
|
918
|
+
await store.alterTable({
|
|
919
|
+
tableName: TEST_TABLE as TABLE_NAMES,
|
|
920
|
+
schema: { ...BASE_SCHEMA, age: { type: 'integer', nullable: true } },
|
|
921
|
+
ifNotExists: ['age'],
|
|
922
|
+
});
|
|
923
|
+
|
|
924
|
+
await store.insert({
|
|
925
|
+
tableName: TEST_TABLE as TABLE_NAMES,
|
|
926
|
+
record: { id: 1, name: 'Alice', age: 42 },
|
|
927
|
+
});
|
|
928
|
+
|
|
929
|
+
const row = await store.load<{ id: string; name: string; age?: number }>({
|
|
930
|
+
tableName: TEST_TABLE as TABLE_NAMES,
|
|
931
|
+
keys: { id: '1' },
|
|
932
|
+
});
|
|
933
|
+
expect(row?.age).toBe(42);
|
|
934
|
+
});
|
|
935
|
+
|
|
936
|
+
it('is idempotent when adding an existing column', async () => {
|
|
937
|
+
await store.alterTable({
|
|
938
|
+
tableName: TEST_TABLE as TABLE_NAMES,
|
|
939
|
+
schema: { ...BASE_SCHEMA, foo: { type: 'text', nullable: true } },
|
|
940
|
+
ifNotExists: ['foo'],
|
|
941
|
+
});
|
|
942
|
+
// Add the column again (should not throw)
|
|
943
|
+
await expect(
|
|
944
|
+
store.alterTable({
|
|
945
|
+
tableName: TEST_TABLE as TABLE_NAMES,
|
|
946
|
+
schema: { ...BASE_SCHEMA, foo: { type: 'text', nullable: true } },
|
|
947
|
+
ifNotExists: ['foo'],
|
|
948
|
+
}),
|
|
949
|
+
).resolves.not.toThrow();
|
|
950
|
+
});
|
|
951
|
+
|
|
952
|
+
it('should add a default value to a column when using not null', async () => {
|
|
953
|
+
await store.insert({
|
|
954
|
+
tableName: TEST_TABLE as TABLE_NAMES,
|
|
955
|
+
record: { id: 1, name: 'Bob' },
|
|
956
|
+
});
|
|
957
|
+
|
|
958
|
+
await expect(
|
|
959
|
+
store.alterTable({
|
|
960
|
+
tableName: TEST_TABLE as TABLE_NAMES,
|
|
961
|
+
schema: { ...BASE_SCHEMA, text_column: { type: 'text', nullable: false } },
|
|
962
|
+
ifNotExists: ['text_column'],
|
|
963
|
+
}),
|
|
964
|
+
).resolves.not.toThrow();
|
|
965
|
+
|
|
966
|
+
await expect(
|
|
967
|
+
store.alterTable({
|
|
968
|
+
tableName: TEST_TABLE as TABLE_NAMES,
|
|
969
|
+
schema: { ...BASE_SCHEMA, timestamp_column: { type: 'timestamp', nullable: false } },
|
|
970
|
+
ifNotExists: ['timestamp_column'],
|
|
971
|
+
}),
|
|
972
|
+
).resolves.not.toThrow();
|
|
973
|
+
|
|
974
|
+
await expect(
|
|
975
|
+
store.alterTable({
|
|
976
|
+
tableName: TEST_TABLE as TABLE_NAMES,
|
|
977
|
+
schema: { ...BASE_SCHEMA, bigint_column: { type: 'bigint', nullable: false } },
|
|
978
|
+
ifNotExists: ['bigint_column'],
|
|
979
|
+
}),
|
|
980
|
+
).resolves.not.toThrow();
|
|
981
|
+
|
|
982
|
+
await expect(
|
|
983
|
+
store.alterTable({
|
|
984
|
+
tableName: TEST_TABLE as TABLE_NAMES,
|
|
985
|
+
schema: { ...BASE_SCHEMA, jsonb_column: { type: 'jsonb', nullable: false } },
|
|
986
|
+
ifNotExists: ['jsonb_column'],
|
|
987
|
+
}),
|
|
988
|
+
).resolves.not.toThrow();
|
|
989
|
+
});
|
|
990
|
+
});
|
|
991
|
+
|
|
890
992
|
describe('Schema Support', () => {
|
|
891
993
|
const customSchema = 'mastra_test';
|
|
892
994
|
let customSchemaStore: PostgresStore;
|
|
@@ -1022,11 +1124,11 @@ describe('PostgresStore', () => {
|
|
|
1022
1124
|
});
|
|
1023
1125
|
await Promise.all(evalPromises);
|
|
1024
1126
|
|
|
1025
|
-
const result = await store.getEvals({ agentName,
|
|
1127
|
+
const result = await store.getEvals({ agentName, perPage: 5, page: 2 });
|
|
1026
1128
|
expect(result.evals).toHaveLength(5);
|
|
1027
1129
|
expect(result.total).toBe(15);
|
|
1028
|
-
expect(result.page).
|
|
1029
|
-
expect(result.perPage).
|
|
1130
|
+
expect(result.page).toBe(2);
|
|
1131
|
+
expect(result.perPage).toBe(5);
|
|
1030
1132
|
expect(result.hasMore).toBe(false);
|
|
1031
1133
|
});
|
|
1032
1134
|
|
|
@@ -1117,7 +1219,7 @@ describe('PostgresStore', () => {
|
|
|
1117
1219
|
createEvalAtDate(now),
|
|
1118
1220
|
]);
|
|
1119
1221
|
|
|
1120
|
-
const fromYesterday = await store.getEvals({ agentName,
|
|
1222
|
+
const fromYesterday = await store.getEvals({ agentName, dateRange: { start: yesterday }, page: 0, perPage: 3 });
|
|
1121
1223
|
expect(fromYesterday.total).toBe(7); // 3 yesterday + 4 now
|
|
1122
1224
|
expect(fromYesterday.evals).toHaveLength(3);
|
|
1123
1225
|
// Evals are sorted DESC, so first 3 are from 'now'
|
|
@@ -1127,7 +1229,9 @@ describe('PostgresStore', () => {
|
|
|
1127
1229
|
|
|
1128
1230
|
const onlyDayBefore = await store.getEvals({
|
|
1129
1231
|
agentName,
|
|
1130
|
-
|
|
1232
|
+
dateRange: {
|
|
1233
|
+
end: new Date(yesterday.getTime() - 1),
|
|
1234
|
+
},
|
|
1131
1235
|
page: 0,
|
|
1132
1236
|
perPage: 5,
|
|
1133
1237
|
});
|
|
@@ -1143,11 +1247,10 @@ describe('PostgresStore', () => {
|
|
|
1143
1247
|
);
|
|
1144
1248
|
await Promise.all(tracePromises);
|
|
1145
1249
|
|
|
1146
|
-
const page1 = await store.
|
|
1250
|
+
const page1 = await store.getTracesPaginated({
|
|
1147
1251
|
scope: 'pg-test-scope',
|
|
1148
1252
|
page: 0,
|
|
1149
1253
|
perPage: 8,
|
|
1150
|
-
returnPaginationResults: true,
|
|
1151
1254
|
});
|
|
1152
1255
|
expect(page1.traces).toHaveLength(8);
|
|
1153
1256
|
expect(page1.total).toBe(18);
|
|
@@ -1155,11 +1258,10 @@ describe('PostgresStore', () => {
|
|
|
1155
1258
|
expect(page1.perPage).toBe(8);
|
|
1156
1259
|
expect(page1.hasMore).toBe(true);
|
|
1157
1260
|
|
|
1158
|
-
const page3 = await store.
|
|
1261
|
+
const page3 = await store.getTracesPaginated({
|
|
1159
1262
|
scope: 'pg-test-scope',
|
|
1160
1263
|
page: 2,
|
|
1161
1264
|
perPage: 8,
|
|
1162
|
-
returnPaginationResults: true,
|
|
1163
1265
|
});
|
|
1164
1266
|
expect(page3.traces).toHaveLength(2);
|
|
1165
1267
|
expect(page3.total).toBe(18);
|
|
@@ -1181,12 +1283,11 @@ describe('PostgresStore', () => {
|
|
|
1181
1283
|
);
|
|
1182
1284
|
await Promise.all([...tracesWithAttr, ...tracesWithoutAttr]);
|
|
1183
1285
|
|
|
1184
|
-
const prodTraces = await store.
|
|
1286
|
+
const prodTraces = await store.getTracesPaginated({
|
|
1185
1287
|
scope: 'pg-attr-scope',
|
|
1186
1288
|
attributes: { environment: 'prod' },
|
|
1187
1289
|
page: 0,
|
|
1188
1290
|
perPage: 5,
|
|
1189
|
-
returnPaginationResults: true,
|
|
1190
1291
|
});
|
|
1191
1292
|
expect(prodTraces.traces).toHaveLength(5);
|
|
1192
1293
|
expect(prodTraces.total).toBe(8);
|
|
@@ -1210,12 +1311,13 @@ describe('PostgresStore', () => {
|
|
|
1210
1311
|
store.insert({ tableName: TABLE_TRACES, record: createSampleTraceForDB('t5', scope, undefined, now) }),
|
|
1211
1312
|
]);
|
|
1212
1313
|
|
|
1213
|
-
const fromYesterday = await store.
|
|
1314
|
+
const fromYesterday = await store.getTracesPaginated({
|
|
1214
1315
|
scope,
|
|
1215
|
-
|
|
1316
|
+
dateRange: {
|
|
1317
|
+
start: yesterday,
|
|
1318
|
+
},
|
|
1216
1319
|
page: 0,
|
|
1217
1320
|
perPage: 2,
|
|
1218
|
-
returnPaginationResults: true,
|
|
1219
1321
|
});
|
|
1220
1322
|
expect(fromYesterday.total).toBe(4); // 2 yesterday + 2 now
|
|
1221
1323
|
expect(fromYesterday.traces).toHaveLength(2);
|
|
@@ -1223,97 +1325,110 @@ describe('PostgresStore', () => {
|
|
|
1223
1325
|
expect(new Date(t.createdAt).getTime()).toBeGreaterThanOrEqual(yesterday.getTime()),
|
|
1224
1326
|
);
|
|
1225
1327
|
|
|
1226
|
-
const onlyNow = await store.
|
|
1328
|
+
const onlyNow = await store.getTracesPaginated({
|
|
1227
1329
|
scope,
|
|
1228
|
-
|
|
1229
|
-
|
|
1330
|
+
dateRange: {
|
|
1331
|
+
start: now,
|
|
1332
|
+
end: now,
|
|
1333
|
+
},
|
|
1230
1334
|
page: 0,
|
|
1231
1335
|
perPage: 5,
|
|
1232
|
-
returnPaginationResults: true,
|
|
1233
1336
|
});
|
|
1234
1337
|
expect(onlyNow.total).toBe(2);
|
|
1235
1338
|
expect(onlyNow.traces).toHaveLength(2);
|
|
1236
1339
|
});
|
|
1237
|
-
|
|
1238
|
-
it('should return array when returnPaginationResults is false or undefined', async () => {
|
|
1239
|
-
await store.insert({ tableName: TABLE_TRACES, record: createSampleTraceForDB('trace-arr', 'pg-array-scope') });
|
|
1240
|
-
const tracesArray = await store.getTraces({ scope: 'pg-array-scope', page: 0, perPage: 5 }); // returnPaginationResults is undefined
|
|
1241
|
-
expect(Array.isArray(tracesArray)).toBe(true);
|
|
1242
|
-
expect(tracesArray.length).toBe(1);
|
|
1243
|
-
});
|
|
1244
1340
|
});
|
|
1245
1341
|
|
|
1246
1342
|
describe('getMessages with pagination', () => {
|
|
1247
1343
|
it('should return paginated messages with total count', async () => {
|
|
1248
1344
|
const thread = createSampleThread();
|
|
1249
1345
|
await store.saveThread({ thread });
|
|
1250
|
-
|
|
1251
|
-
|
|
1252
|
-
|
|
1253
|
-
|
|
1254
|
-
|
|
1255
|
-
|
|
1346
|
+
// Reset role to 'assistant' before creating messages
|
|
1347
|
+
resetRole();
|
|
1348
|
+
// Create messages sequentially to ensure unique timestamps
|
|
1349
|
+
for (let i = 0; i < 15; i++) {
|
|
1350
|
+
const message = createSampleMessageV1({ threadId: thread.id, content: `Message ${i + 1}` });
|
|
1351
|
+
await store.saveMessages({
|
|
1352
|
+
messages: [message],
|
|
1353
|
+
});
|
|
1354
|
+
await new Promise(r => setTimeout(r, 5));
|
|
1355
|
+
}
|
|
1256
1356
|
|
|
1257
|
-
const page1 = await store.
|
|
1357
|
+
const page1 = await store.getMessagesPaginated({
|
|
1358
|
+
threadId: thread.id,
|
|
1359
|
+
selectBy: { pagination: { page: 0, perPage: 5 } },
|
|
1360
|
+
format: 'v2',
|
|
1361
|
+
});
|
|
1362
|
+
console.log(page1);
|
|
1258
1363
|
expect(page1.messages).toHaveLength(5);
|
|
1259
1364
|
expect(page1.total).toBe(15);
|
|
1260
1365
|
expect(page1.page).toBe(0);
|
|
1261
1366
|
expect(page1.perPage).toBe(5);
|
|
1262
1367
|
expect(page1.hasMore).toBe(true);
|
|
1263
1368
|
|
|
1264
|
-
const page3 = await store.
|
|
1369
|
+
const page3 = await store.getMessagesPaginated({
|
|
1370
|
+
threadId: thread.id,
|
|
1371
|
+
selectBy: { pagination: { page: 2, perPage: 5 } },
|
|
1372
|
+
format: 'v2',
|
|
1373
|
+
});
|
|
1265
1374
|
expect(page3.messages).toHaveLength(5);
|
|
1266
1375
|
expect(page3.total).toBe(15);
|
|
1267
1376
|
expect(page3.hasMore).toBe(false);
|
|
1268
1377
|
});
|
|
1269
1378
|
|
|
1270
1379
|
it('should filter by date with pagination for getMessages', async () => {
|
|
1271
|
-
|
|
1272
|
-
|
|
1380
|
+
resetRole();
|
|
1381
|
+
const threadData = createSampleThread();
|
|
1382
|
+
const thread = await store.saveThread({ thread: threadData as StorageThreadType });
|
|
1273
1383
|
const now = new Date();
|
|
1274
|
-
const yesterday = new Date(
|
|
1275
|
-
|
|
1276
|
-
|
|
1277
|
-
|
|
1278
|
-
|
|
1279
|
-
|
|
1280
|
-
|
|
1281
|
-
|
|
1282
|
-
|
|
1283
|
-
|
|
1284
|
-
|
|
1285
|
-
|
|
1286
|
-
|
|
1287
|
-
|
|
1288
|
-
|
|
1384
|
+
const yesterday = new Date(
|
|
1385
|
+
now.getFullYear(),
|
|
1386
|
+
now.getMonth(),
|
|
1387
|
+
now.getDate() - 1,
|
|
1388
|
+
now.getHours(),
|
|
1389
|
+
now.getMinutes(),
|
|
1390
|
+
now.getSeconds(),
|
|
1391
|
+
);
|
|
1392
|
+
const dayBeforeYesterday = new Date(
|
|
1393
|
+
now.getFullYear(),
|
|
1394
|
+
now.getMonth(),
|
|
1395
|
+
now.getDate() - 2,
|
|
1396
|
+
now.getHours(),
|
|
1397
|
+
now.getMinutes(),
|
|
1398
|
+
now.getSeconds(),
|
|
1399
|
+
);
|
|
1289
1400
|
|
|
1290
|
-
|
|
1401
|
+
// Ensure timestamps are distinct for reliable sorting by creating them with a slight delay for testing clarity
|
|
1402
|
+
const messagesToSave: MastraMessageV1[] = [];
|
|
1403
|
+
messagesToSave.push(createSampleMessageV1({ threadId: thread.id, createdAt: dayBeforeYesterday }));
|
|
1404
|
+
await new Promise(r => setTimeout(r, 5));
|
|
1405
|
+
messagesToSave.push(createSampleMessageV1({ threadId: thread.id, createdAt: dayBeforeYesterday }));
|
|
1406
|
+
await new Promise(r => setTimeout(r, 5));
|
|
1407
|
+
messagesToSave.push(createSampleMessageV1({ threadId: thread.id, createdAt: yesterday }));
|
|
1408
|
+
await new Promise(r => setTimeout(r, 5));
|
|
1409
|
+
messagesToSave.push(createSampleMessageV1({ threadId: thread.id, createdAt: yesterday }));
|
|
1410
|
+
await new Promise(r => setTimeout(r, 5));
|
|
1411
|
+
messagesToSave.push(createSampleMessageV1({ threadId: thread.id, createdAt: now }));
|
|
1412
|
+
await new Promise(r => setTimeout(r, 5));
|
|
1413
|
+
messagesToSave.push(createSampleMessageV1({ threadId: thread.id, createdAt: now }));
|
|
1414
|
+
|
|
1415
|
+
await store.saveMessages({ messages: messagesToSave, format: 'v1' });
|
|
1416
|
+
// Total 6 messages: 2 now, 2 yesterday, 2 dayBeforeYesterday (oldest to newest)
|
|
1417
|
+
|
|
1418
|
+
const fromYesterday = await store.getMessagesPaginated({
|
|
1291
1419
|
threadId: thread.id,
|
|
1292
|
-
|
|
1293
|
-
|
|
1294
|
-
perPage: 3,
|
|
1295
|
-
format: 'v1',
|
|
1420
|
+
selectBy: { pagination: { page: 0, perPage: 3, dateRange: { start: yesterday } } },
|
|
1421
|
+
format: 'v2',
|
|
1296
1422
|
});
|
|
1297
|
-
expect(
|
|
1298
|
-
expect(
|
|
1299
|
-
|
|
1300
|
-
expect(new Date(
|
|
1301
|
-
|
|
1302
|
-
|
|
1303
|
-
|
|
1304
|
-
|
|
1305
|
-
|
|
1306
|
-
|
|
1307
|
-
it('should maintain backward compatibility for getMessages (no pagination params)', async () => {
|
|
1308
|
-
const thread = createSampleThread();
|
|
1309
|
-
await store.saveThread({ thread });
|
|
1310
|
-
await store.saveMessages({ messages: [createSampleMessage(thread.id)] });
|
|
1311
|
-
|
|
1312
|
-
const messages = await store.getMessages({ threadId: thread.id, format: 'v1' });
|
|
1313
|
-
expect(Array.isArray(messages)).toBe(true);
|
|
1314
|
-
expect(messages.length).toBe(1);
|
|
1315
|
-
// @ts-expect-error - messages should not have pagination properties
|
|
1316
|
-
expect(messages.total).toBeUndefined();
|
|
1423
|
+
expect(fromYesterday.total).toBe(4);
|
|
1424
|
+
expect(fromYesterday.messages).toHaveLength(3);
|
|
1425
|
+
const firstMessageTime = new Date((fromYesterday.messages[0] as MastraMessageV1).createdAt).getTime();
|
|
1426
|
+
expect(firstMessageTime).toBeGreaterThanOrEqual(new Date(yesterday.toISOString()).getTime());
|
|
1427
|
+
if (fromYesterday.messages.length > 0) {
|
|
1428
|
+
expect(new Date((fromYesterday.messages[0] as MastraMessageV1).createdAt).toISOString().slice(0, 10)).toEqual(
|
|
1429
|
+
yesterday.toISOString().slice(0, 10),
|
|
1430
|
+
);
|
|
1431
|
+
}
|
|
1317
1432
|
});
|
|
1318
1433
|
});
|
|
1319
1434
|
|
|
@@ -1325,28 +1440,30 @@ describe('PostgresStore', () => {
|
|
|
1325
1440
|
);
|
|
1326
1441
|
await Promise.all(threadPromises);
|
|
1327
1442
|
|
|
1328
|
-
const page1 = await store.
|
|
1443
|
+
const page1 = await store.getThreadsByResourceIdPaginated({ resourceId, page: 0, perPage: 7 });
|
|
1329
1444
|
expect(page1.threads).toHaveLength(7);
|
|
1330
1445
|
expect(page1.total).toBe(17);
|
|
1331
1446
|
expect(page1.page).toBe(0);
|
|
1332
1447
|
expect(page1.perPage).toBe(7);
|
|
1333
1448
|
expect(page1.hasMore).toBe(true);
|
|
1334
1449
|
|
|
1335
|
-
const page3 = await store.
|
|
1450
|
+
const page3 = await store.getThreadsByResourceIdPaginated({ resourceId, page: 2, perPage: 7 });
|
|
1336
1451
|
expect(page3.threads).toHaveLength(3); // 17 total, 7 per page, 3rd page has 17 - 2*7 = 3
|
|
1337
1452
|
expect(page3.total).toBe(17);
|
|
1338
1453
|
expect(page3.hasMore).toBe(false);
|
|
1339
1454
|
});
|
|
1340
1455
|
|
|
1341
|
-
it('should return
|
|
1456
|
+
it('should return paginated results when no pagination params for getThreadsByResourceId', async () => {
|
|
1342
1457
|
const resourceId = `pg-non-paginated-resource-${randomUUID()}`;
|
|
1343
1458
|
await store.saveThread({ thread: { ...createSampleThread(), resourceId } });
|
|
1344
1459
|
|
|
1345
|
-
const
|
|
1346
|
-
expect(Array.isArray(threads)).toBe(true);
|
|
1347
|
-
expect(threads.length).toBe(1);
|
|
1348
|
-
|
|
1349
|
-
expect(
|
|
1460
|
+
const results = await store.getThreadsByResourceIdPaginated({ resourceId });
|
|
1461
|
+
expect(Array.isArray(results.threads)).toBe(true);
|
|
1462
|
+
expect(results.threads.length).toBe(1);
|
|
1463
|
+
expect(results.total).toBe(1);
|
|
1464
|
+
expect(results.page).toBe(0);
|
|
1465
|
+
expect(results.perPage).toBe(100);
|
|
1466
|
+
expect(results.hasMore).toBe(false);
|
|
1350
1467
|
});
|
|
1351
1468
|
});
|
|
1352
1469
|
});
|