refactor(instance-ai): decouple builder memory compaction

This commit is contained in:
Oleg Ivaniv 2026-05-05 10:58:47 +02:00
parent 3323e173f0
commit 8d4a1fce6a
No known key found for this signature in database
2 changed files with 70 additions and 26 deletions

View File

@ -1,11 +1,24 @@
import type { MastraDBMessage } from '@mastra/core/agent';
import type { MastraCompositeStore } from '@mastra/core/storage';
import type { Workspace } from '@mastra/core/workspace';
import { BuilderSandboxSessionRegistry } from '../../../runtime/builder-sandbox-session-registry';
import { compactBuilderMemoryThread } from '../builder-memory-compaction';
function makeMessage(id: string, text: string): MastraDBMessage {
type CompactionInput = Parameters<typeof compactBuilderMemoryThread>[0];
interface TestBuilderMemoryMessage {
id: string;
role: string;
threadId: string;
resourceId: string;
createdAt: Date;
type?: string;
content: {
format: number;
parts: Array<{ type: string; text: string }>;
content: string;
metadata?: Record<string, unknown>;
};
}
function makeMessage(id: string, text: string): TestBuilderMemoryMessage {
return {
id,
role: 'assistant',
@ -20,19 +33,17 @@ function makeMessage(id: string, text: string): MastraDBMessage {
};
}
function makeStorage(memoryStore: unknown): MastraCompositeStore {
function makeStorage(memoryStore: unknown): CompactionInput['context']['storage'] {
return {
getStore: jest.fn(async (storeName: string) => {
await Promise.resolve();
return storeName === 'memory' ? memoryStore : undefined;
}),
} as unknown as MastraCompositeStore;
};
}
type CompactionInput = Parameters<typeof compactBuilderMemoryThread>[0];
function makeCompactionInput(
storage: MastraCompositeStore,
storage: CompactionInput['context']['storage'],
overrides: Partial<CompactionInput> = {},
): CompactionInput {
return {
@ -85,10 +96,12 @@ describe('compactBuilderMemoryThread', () => {
deleteMessages: jest.fn(async () => {
await Promise.resolve();
}),
saveMessages: jest.fn(async ({ messages: saved }: { messages: MastraDBMessage[] }) => {
await Promise.resolve();
return { messages: saved };
}),
saveMessages: jest.fn(
async ({ messages: saved }: { messages: TestBuilderMemoryMessage[] }) => {
await Promise.resolve();
return { messages: saved };
},
),
};
const result = await compactBuilderMemoryThread(makeCompactionInput(makeStorage(memoryStore)));
@ -123,7 +136,7 @@ describe('compactBuilderMemoryThread', () => {
builderThreadId: 'builder-thread-1',
builderResourceId: 'user-1:workflow-builder',
builderWorkspace: {
workspace: {} as Workspace,
workspace: {} as never,
cleanup,
},
root: '/home/daytona/workspace',
@ -143,7 +156,7 @@ describe('compactBuilderMemoryThread', () => {
deleteMessages: jest.fn(async () => {
await Promise.resolve();
}),
saveMessages: jest.fn(async ({ messages }: { messages: MastraDBMessage[] }) => {
saveMessages: jest.fn(async ({ messages }: { messages: TestBuilderMemoryMessage[] }) => {
await Promise.resolve();
return { messages };
}),
@ -178,7 +191,7 @@ describe('compactBuilderMemoryThread', () => {
await Promise.resolve();
storedMessages = storedMessages.filter((message) => !messageIds.includes(message.id));
}),
saveMessages: jest.fn(async ({ messages }: { messages: MastraDBMessage[] }) => {
saveMessages: jest.fn(async ({ messages }: { messages: TestBuilderMemoryMessage[] }) => {
await Promise.resolve();
storedMessages.push(...messages);
return { messages };

View File

@ -1,8 +1,5 @@
import type { MastraDBMessage } from '@mastra/core/agent';
import type { MastraCompositeStore, MemoryStorage } from '@mastra/core/storage';
import { randomUUID } from 'node:crypto';
import type { OrchestrationContext } from '../../types';
import type { WorkflowBuildOutcome } from '../../workflow-loop';
const BUILDER_MEMORY_SUMMARY_TYPE = 'builder-memory-summary';
@ -12,8 +9,19 @@ interface BuilderMemoryBinding {
thread: string;
}
interface BuilderMemoryStorageProvider {
getStore(storeName: string): Promise<unknown> | unknown;
}
interface BuilderMemoryCompactionContext {
storage: BuilderMemoryStorageProvider;
threadId: string;
runId: string;
messageGroupId?: string;
}
interface BuilderMemoryCompactionInput {
context: Pick<OrchestrationContext, 'storage' | 'threadId' | 'runId' | 'messageGroupId'>;
context: BuilderMemoryCompactionContext;
binding: BuilderMemoryBinding;
sessionId?: string;
workflowId?: string;
@ -38,10 +46,33 @@ export interface BuilderMemoryCompactionResult {
compactedTokenEstimate: number;
}
interface BuilderMemoryMessage {
id: string;
role: string;
threadId: string;
resourceId: string;
createdAt: Date;
type?: string;
content: unknown;
}
interface BuilderMemoryListResult {
messages: BuilderMemoryMessage[];
total?: number;
page?: number;
perPage?: number | false;
hasMore?: boolean;
}
interface BuilderMemoryStore {
listMessages: MemoryStorage['listMessages'];
saveMessages: MemoryStorage['saveMessages'];
deleteMessages: MemoryStorage['deleteMessages'];
listMessages: (args: {
threadId: string;
resourceId: string;
perPage: false;
orderBy: { field: 'createdAt'; direction: 'ASC' };
}) => Promise<BuilderMemoryListResult>;
saveMessages: (args: { messages: BuilderMemoryMessage[] }) => Promise<unknown>;
deleteMessages: (messageIds: string[]) => Promise<unknown>;
}
function estimateTokens(value: string): number {
@ -68,7 +99,7 @@ function hasBuilderMemoryStore(value: unknown): value is BuilderMemoryStore {
}
async function getBuilderMemoryStore(
storage: MastraCompositeStore,
storage: BuilderMemoryStorageProvider,
): Promise<BuilderMemoryStore | undefined> {
const store = await storage.getStore('memory');
return hasBuilderMemoryStore(store) ? store : undefined;
@ -154,7 +185,7 @@ function buildSummaryContent(input: BuilderMemoryCompactionInput): string {
function buildSummaryMessage(
input: BuilderMemoryCompactionInput,
content: string,
): MastraDBMessage {
): BuilderMemoryMessage {
return {
id: `${BUILDER_MEMORY_SUMMARY_TYPE}-${randomUUID()}`,
role: 'assistant',