diff --git a/packages/junior-evals/evals/memory/actors.eval.ts b/packages/junior-evals/evals/memory/actors.eval.ts index e4aa97178a..2a584350e7 100644 --- a/packages/junior-evals/evals/memory/actors.eval.ts +++ b/packages/junior-evals/evals/memory/actors.eval.ts @@ -1,12 +1,6 @@ import { expect } from "vitest"; import { describeEval } from "vitest-evals"; -import { getDb } from "@/chat/db"; import { readActorIdentity } from "@/chat/plugins/viewer"; -import type { MemoryDb } from "@sentry/junior-memory"; -import { - juniorMemoryEmbeddings, - juniorMemoryMemories, -} from "../../../junior-memory/src/db/schema"; import { mention, rubric, @@ -14,6 +8,7 @@ import { steer, threadMessage, } from "../../src/helpers"; +import { clearMemories, readMemories, type MemoryThread } from "./helpers"; /** * Passive memory learning when a run has more than one Actor. @@ -46,34 +41,6 @@ const CAROL = { full_name: "Carol Example", }; -interface MemoryThread { - channel_type?: "channel" | "group" | "im" | "mpim"; - channel_id: string; - id: string; - thread_ts: string; -} - -function memoryDb(): MemoryDb { - return getDb() as unknown as MemoryDb; -} - -function memorySourceKey(thread: MemoryThread): string { - return `slack:${memoryTeamId}:${thread.channel_id}:${thread.thread_ts}`; -} - -async function readMemories(thread: MemoryThread) { - const rows = await memoryDb() - .select() - .from(juniorMemoryMemories) - .orderBy(juniorMemoryMemories.createdAtMs, juniorMemoryMemories.id); - return rows.filter((memory) => memory.sourceKey === memorySourceKey(thread)); -} - -async function clearMemories() { - await memoryDb().delete(juniorMemoryEmbeddings); - await memoryDb().delete(juniorMemoryMemories); -} - async function memoriesForActor( rows: Awaited>, slackUserId: string, diff --git a/packages/junior-evals/evals/memory/helpers.ts b/packages/junior-evals/evals/memory/helpers.ts index a18747db98..cfc9988891 100644 --- a/packages/junior-evals/evals/memory/helpers.ts +++ b/packages/junior-evals/evals/memory/helpers.ts @@ -4,12 +4,14 @@ import { getDb, getSqlExecutor } from "@/chat/db"; import { upsertIdentity } from "@/chat/identities/sql"; import { completeText, resolveGatewayModel } from "@/chat/pi/client"; import { createPluginEmbedder } from "@/chat/plugins/model"; -import { createMemoryStore, type MemoryDb } from "@sentry/junior-memory"; -import { createSlackSource } from "@sentry/junior-plugin-api"; import { - juniorMemoryEmbeddings, - juniorMemoryMemories, -} from "../../../junior-memory/src/db/schema"; + clearAll, + countEmbeddings, + createMemory, + listBySource, + type MemoryDb, +} from "@sentry/junior-memory/testing"; +import { createSlackSource } from "@sentry/junior-plugin-api"; import { TEST_USER_ID } from "@junior-tests/fixtures/slack/factories/ids"; export const memoryPluginOverrides = { @@ -47,37 +49,34 @@ export async function seedMemory(args: { if (!identity.userId) { throw new Error("Eval memory Actor did not resolve to a User"); } - const store = createMemoryStore( - memoryDb(), - { - conversationId: `slack:${args.thread.channel_id}:${args.thread.thread_ts}`, - actor: { - platform: "slack", - teamId: memoryTeamId, - userId: actorUserId, - }, - source: createSlackSource({ - channelId: args.thread.channel_id, - messageTs: args.thread.thread_ts, - teamId: memoryTeamId, - threadTs: args.thread.thread_ts, - visibility: - args.thread.channel_type === "channel" ? "public" : "private", - }), - userId: identity.userId, + const context = { + conversationId: `slack:${args.thread.channel_id}:${args.thread.thread_ts}`, + actor: { + platform: "slack" as const, + teamId: memoryTeamId, + userId: actorUserId, }, - { embedder: evalMemoryEmbedder }, - ); + source: createSlackSource({ + channelId: args.thread.channel_id, + messageTs: args.thread.thread_ts, + teamId: memoryTeamId, + threadTs: args.thread.thread_ts, + visibility: args.thread.channel_type === "channel" ? "public" : "private", + }), + userId: identity.userId, + }; const input = { content: args.content, idempotencyKey: args.idempotencyKey, kind: args.kind ?? "preference", }; - if (args.subject === "conversation") { - await store.createConversationMemory(input); - return; - } - await store.createMemory(input); + await createMemory({ + context, + db: memoryDb(), + embedder: evalMemoryEmbedder, + input, + subjectType: args.subject === "conversation" ? "conversation" : "user", + }); } function memoryDb(): MemoryDb { @@ -89,24 +88,16 @@ function memorySourceKey(thread: MemoryThread): string { } export async function readMemories(thread: MemoryThread) { - const rows = await memoryDb() - .select() - .from(juniorMemoryMemories) - .orderBy(juniorMemoryMemories.createdAtMs, juniorMemoryMemories.id); - return rows.filter((memory) => memory.sourceKey === memorySourceKey(thread)); + return listBySource(memoryDb(), memorySourceKey(thread)); } -/** Count vector rows for memories seeded in one eval thread. */ +/** Count embeddings for memories from one eval thread. */ export async function countMemoryEmbeddings(thread: MemoryThread) { const memories = await readMemories(thread); - if (memories.length === 0) { - return 0; - } - const memoryIds = new Set(memories.map((memory) => memory.id)); - const rows = await memoryDb() - .select({ memoryId: juniorMemoryEmbeddings.memoryId }) - .from(juniorMemoryEmbeddings); - return rows.filter((row) => memoryIds.has(row.memoryId)).length; + return countEmbeddings( + memoryDb(), + memories.map((memory) => memory.id), + ); } /** Read the durable memories currently eligible for recall in one eval thread. */ @@ -124,8 +115,7 @@ export async function readActiveMemories( } export async function clearMemories() { - await memoryDb().delete(juniorMemoryEmbeddings); - await memoryDb().delete(juniorMemoryMemories); + await clearAll(memoryDb()); } export function visibleAssistantText(result: { diff --git a/packages/junior-evals/vitest.evals.behavioral.config.ts b/packages/junior-evals/vitest.evals.behavioral.config.ts index f53a802d15..e1f9fd0b25 100644 --- a/packages/junior-evals/vitest.evals.behavioral.config.ts +++ b/packages/junior-evals/vitest.evals.behavioral.config.ts @@ -47,6 +47,10 @@ export default defineConfig({ resolve: { alias: { "@": path.resolve(juniorPackageRoot, "src"), + "@sentry/junior-memory/testing": path.resolve( + memoryPackageRoot, + "src/testing.ts", + ), "@sentry/junior-memory": path.resolve(memoryPackageRoot, "src/index.ts"), "@sentry/junior-plugin-api": path.resolve( pluginApiPackageRoot, diff --git a/packages/junior-memory/package.json b/packages/junior-memory/package.json index e4633ee0fb..97409460d3 100644 --- a/packages/junior-memory/package.json +++ b/packages/junior-memory/package.json @@ -15,6 +15,10 @@ ".": { "types": "./src/index.ts", "default": "./dist/index.js" + }, + "./testing": { + "types": "./src/testing.ts", + "default": "./dist/testing.js" } }, "files": [ diff --git a/packages/junior-memory/src/agent.ts b/packages/junior-memory/src/agent.ts index 3421a64b72..a9ffc18b03 100644 --- a/packages/junior-memory/src/agent.ts +++ b/packages/junior-memory/src/agent.ts @@ -9,7 +9,7 @@ import { memorySupersessionInputSchema, type MemorySupersessionDecision, type MemorySupersessionInput, -} from "./store"; +} from "./create"; import { MEMORY_KINDS, memoryRuntimeContextSchema, diff --git a/packages/junior-memory/src/api.ts b/packages/junior-memory/src/api.ts index 5c12be69c0..f2841bc142 100644 --- a/packages/junior-memory/src/api.ts +++ b/packages/junior-memory/src/api.ts @@ -10,7 +10,7 @@ import { type PluginRouteApp, type User, } from "@sentry/junior-plugin-api"; -import type { MemoryDb } from "./store"; +import type { MemoryDb } from "./memories"; import { archiveMemory, getMemory, diff --git a/packages/junior-memory/src/cli/search.ts b/packages/junior-memory/src/cli/search.ts index 41bba8cfa7..15c551c6e1 100644 --- a/packages/junior-memory/src/cli/search.ts +++ b/packages/junior-memory/src/cli/search.ts @@ -5,7 +5,7 @@ import type { PluginCliHost, } from "@sentry/junior-plugin-api"; import { juniorMemoryMemories } from "../db/schema"; -import type { MemoryDb } from "../store"; +import type { MemoryDb } from "../memories"; import { MEMORY_SCOPES, type MemoryScope } from "../types"; import { formatMemory } from "./format"; diff --git a/packages/junior-memory/src/cli/show.ts b/packages/junior-memory/src/cli/show.ts index 7d2ec6a9d9..97b304f864 100644 --- a/packages/junior-memory/src/cli/show.ts +++ b/packages/junior-memory/src/cli/show.ts @@ -5,7 +5,7 @@ import type { } from "@sentry/junior-plugin-api"; import { eq } from "drizzle-orm"; import { juniorMemoryMemories } from "../db/schema"; -import type { MemoryDb } from "../store"; +import type { MemoryDb } from "../memories"; import { formatMemory } from "./format"; async function runShow( diff --git a/packages/junior-memory/src/create.ts b/packages/junior-memory/src/create.ts new file mode 100644 index 0000000000..c3475501fd --- /dev/null +++ b/packages/junior-memory/src/create.ts @@ -0,0 +1,702 @@ +/** + * Memory creation owns access, idempotency, exact duplicates, and preference + * replacement. It adds the embedding after commit. An embedding failure does + * not roll back the memory. + */ +import { createHash, randomUUID } from "node:crypto"; +import { + and, + asc, + desc, + eq, + gt, + inArray, + isNotNull, + isNull, + or, + sql, + type SQL, +} from "drizzle-orm"; +import { cosineDistance } from "drizzle-orm/sql/functions"; +import { z } from "zod"; +import { getSourceKey } from "@sentry/junior-plugin-api"; +import { juniorMemoryEmbeddings, juniorMemoryMemories } from "./db/schema"; +import { + EMBEDDING_METRIC, + embedMemoryText, + hashEmbeddedContent, + normalizeMemoryContent, + type MemoryEmbedding, + type MemoryEmbeddingProvider, +} from "./embeddings"; +import { + archiveExpiredMemoryBatch, + parseMemoryRow, + type MemoryDb, + type Memory, +} from "./memories"; +import { + deriveMemoryScope, + deriveMemorySubject, + type ResolvedMemoryScope, + type ResolvedMemorySubject, +} from "./scope"; +import { + MEMORY_EMBEDDING_DIMENSIONS, + MEMORY_KINDS, + memoryRuntimeContextSchema, + type MemoryRuntimeContext, +} from "./types"; + +const PREFERENCE_ADJUDICATION_CANDIDATE_LIMIT = 10; +const PREFERENCE_ADJUDICATION_VECTOR_LIMIT = 5; +const MAX_MEMORY_CONTENT_CHARS = 4_000; + +const nonEmptyStringSchema = z.string().min(1); +const numberSchema = z.number().finite(); +const memoryContentSchema = z + .string() + .refine((content) => content.trim().length > 0, { + message: "Memory content is required.", + }); +const createMemoryInputSchema = z + .object({ + content: memoryContentSchema, + expiresAtMs: z.number().finite().optional(), + idempotencyKey: nonEmptyStringSchema, + kind: z.enum(MEMORY_KINDS), + }) + .strict(); +const memorySupersessionCandidateSchema = z + .object({ content: z.string().min(1), id: z.string().min(1) }) + .strict(); +const memorySupersessionCandidatesSchema = z + .array(memorySupersessionCandidateSchema) + .min(1) + .max(PREFERENCE_ADJUDICATION_CANDIDATE_LIMIT); +const supersededIdsSchema = z + .array(z.string().min(1)) + .min(1) + .max(PREFERENCE_ADJUDICATION_CANDIDATE_LIMIT); + +/** Validated preference comparison input supplied to a supersession decider. */ +export const memorySupersessionInputSchema = z + .object({ + candidate: z + .object({ content: z.string().min(1), kind: z.literal("preference") }) + .strict(), + existingMemories: memorySupersessionCandidatesSchema, + runtimeContext: memoryRuntimeContextSchema, + }) + .strict(); + +/** Validated preference decision limited to supplied existing memories. */ +export const memorySupersessionDecisionSchema = z.discriminatedUnion( + "decision", + [ + z + .object({ + decision: z.literal("duplicate"), + duplicateId: z.string().min(1), + }) + .strict(), + z + .object({ + decision: z.literal("supersedes_old"), + supersededIds: supersededIdsSchema, + }) + .strict(), + z.object({ decision: z.enum(["distinct", "uncertain"]) }).strict(), + ], +); + +export type CreateMemoryInput = z.output; +export interface CreateMemoryResult { + created: boolean; + /** True when this call found the memory written for the same input identity. */ + idempotent?: true; + memory: Memory; + /** Memory ids made inactive by this write. */ + supersededIds?: string[]; +} +export type MemorySupersessionInput = z.output< + typeof memorySupersessionInputSchema +>; +export type MemorySupersessionDecision = z.output< + typeof memorySupersessionDecisionSchema +>; +export interface MemorySupersessionDecider { + /** Classify a new preference against related active preferences. */ + adjudicateSupersession( + input: MemorySupersessionInput, + ): Promise | MemorySupersessionDecision; +} + +function idempotencyAliasId(args: { + idempotencyKey: string; + scope: ResolvedMemoryScope; + targetId: string; +}): string { + return `alias:${createHash("sha256") + .update(args.scope.scope) + .update("\0") + .update(args.scope.scopeKey) + .update("\0") + .update(args.idempotencyKey) + .update("\0") + .update(args.targetId) + .digest("hex")}`; +} + +function sourceKey(ctx: MemoryRuntimeContext): string { + const key = getSourceKey(ctx.source); + if (!key) { + throw new Error("Memory Source has no stable key."); + } + return key; +} + +/** Add the embedding without failing memory creation. */ +async function storeMemoryEmbedding(args: { + content: string; + db: MemoryDb; + embedder?: MemoryEmbeddingProvider; + embedding?: MemoryEmbedding; + memoryId: string; + nowMs: number; +}): Promise { + if (!args.embedder && !args.embedding) return; + try { + const existing = await args.db + .select({ memoryId: juniorMemoryEmbeddings.memoryId }) + .from(juniorMemoryEmbeddings) + .where(eq(juniorMemoryEmbeddings.memoryId, args.memoryId)) + .limit(1); + if (existing[0]) return; + } catch { + return; + } + let embedding: MemoryEmbedding; + if (args.embedding) { + embedding = args.embedding; + } else { + if (!args.embedder) return; + try { + embedding = await embedMemoryText(args.embedder, args.content); + } catch { + return; + } + } + try { + await args.db + .insert(juniorMemoryEmbeddings) + .values({ + contentHash: hashEmbeddedContent(args.content), + createdAtMs: args.nowMs, + dimensions: MEMORY_EMBEDDING_DIMENSIONS, + embedding: embedding.vector, + memoryId: args.memoryId, + metric: EMBEDDING_METRIC, + model: embedding.model, + provider: embedding.provider, + }) + .onConflictDoNothing(); + } catch { + return; + } +} + +function activeScopedSubjectPredicate(args: { + kind: Memory["kind"]; + nowMs: number; + scope: ResolvedMemoryScope; + subject: ResolvedMemorySubject; +}): SQL { + const predicate = and( + eq(juniorMemoryMemories.scope, args.scope.scope), + eq(juniorMemoryMemories.scopeKey, args.scope.scopeKey), + eq(juniorMemoryMemories.kind, args.kind), + eq(juniorMemoryMemories.subjectType, args.subject.subjectType), + eq(juniorMemoryMemories.subjectKey, args.subject.subjectKey), + isNull(juniorMemoryMemories.archivedAtMs), + isNull(juniorMemoryMemories.supersededAtMs), + isNull(juniorMemoryMemories.supersededById), + or( + isNull(juniorMemoryMemories.expiresAtMs), + gt(juniorMemoryMemories.expiresAtMs, args.nowMs), + ), + ); + if (!predicate) { + throw new Error("Memory duplicate predicate is empty."); + } + return predicate; +} + +type IdempotencyMatch = { + memory: Memory; + outcome: "created" | "duplicate"; +}; + +async function findByIdempotencyKey(args: { + db: MemoryDb; + idempotencyKey: string; + nowMs: number; + scope: ResolvedMemoryScope; +}): Promise { + const activeRows = await args.db + .select() + .from(juniorMemoryMemories) + .where( + and( + eq(juniorMemoryMemories.scope, args.scope.scope), + eq(juniorMemoryMemories.scopeKey, args.scope.scopeKey), + eq(juniorMemoryMemories.idempotencyKey, args.idempotencyKey), + isNull(juniorMemoryMemories.archivedAtMs), + isNull(juniorMemoryMemories.supersededAtMs), + isNull(juniorMemoryMemories.supersededById), + or( + isNull(juniorMemoryMemories.expiresAtMs), + gt(juniorMemoryMemories.expiresAtMs, args.nowMs), + ), + ), + ) + .limit(1); + if (activeRows[0]) { + return { memory: parseMemoryRow(activeRows[0]), outcome: "created" }; + } + + const aliases = await args.db + .select({ supersededById: juniorMemoryMemories.supersededById }) + .from(juniorMemoryMemories) + .where( + and( + eq(juniorMemoryMemories.scope, args.scope.scope), + eq(juniorMemoryMemories.scopeKey, args.scope.scopeKey), + eq(juniorMemoryMemories.idempotencyKey, args.idempotencyKey), + isNull(juniorMemoryMemories.archivedAtMs), + isNotNull(juniorMemoryMemories.supersededAtMs), + isNotNull(juniorMemoryMemories.supersededById), + or( + isNull(juniorMemoryMemories.expiresAtMs), + gt(juniorMemoryMemories.expiresAtMs, args.nowMs), + ), + ), + ) + .orderBy( + desc(juniorMemoryMemories.createdAtMs), + asc(juniorMemoryMemories.id), + ); + for (const alias of aliases) { + if (!alias.supersededById) continue; + const rows = await args.db + .select() + .from(juniorMemoryMemories) + .where( + and( + eq(juniorMemoryMemories.id, alias.supersededById), + eq(juniorMemoryMemories.scope, args.scope.scope), + eq(juniorMemoryMemories.scopeKey, args.scope.scopeKey), + isNull(juniorMemoryMemories.archivedAtMs), + isNull(juniorMemoryMemories.supersededAtMs), + isNull(juniorMemoryMemories.supersededById), + or( + isNull(juniorMemoryMemories.expiresAtMs), + gt(juniorMemoryMemories.expiresAtMs, args.nowMs), + ), + ), + ) + .limit(1); + if (rows[0]) { + return { memory: parseMemoryRow(rows[0]), outcome: "duplicate" }; + } + } + return undefined; +} + +async function findExactDuplicate(args: { + content: string; + db: MemoryDb; + kind: Memory["kind"]; + nowMs: number; + scope: ResolvedMemoryScope; + subject: ResolvedMemorySubject; +}): Promise { + const rows = await args.db + .select() + .from(juniorMemoryMemories) + .where( + and( + activeScopedSubjectPredicate(args), + eq(juniorMemoryMemories.content, args.content), + ), + ) + .orderBy( + desc(juniorMemoryMemories.createdAtMs), + asc(juniorMemoryMemories.id), + ) + .limit(1); + return rows[0] ? parseMemoryRow(rows[0]) : undefined; +} + +async function rememberDuplicateIdempotency(args: { + content: string; + db: MemoryDb; + duplicate: Memory; + idempotencyKey: string; + nowMs: number; + runtimeContext: MemoryRuntimeContext; + scope: ResolvedMemoryScope; + subject: ResolvedMemorySubject; +}): Promise { + await args.db + .insert(juniorMemoryMemories) + .values({ + content: args.content, + createdAtMs: args.nowMs, + expiresAtMs: args.duplicate.expiresAtMs, + id: idempotencyAliasId({ + idempotencyKey: args.idempotencyKey, + scope: args.scope, + targetId: args.duplicate.id, + }), + idempotencyKey: args.idempotencyKey, + locationId: args.runtimeContext.locationId, + observedAtMs: args.nowMs, + scope: args.scope.scope, + scopeKey: args.scope.scopeKey, + sourceKey: sourceKey(args.runtimeContext), + sourcePlatform: args.runtimeContext.source.platform, + subjectKey: args.subject.subjectKey, + subjectType: args.subject.subjectType, + supersededAtMs: args.nowMs, + supersededById: args.duplicate.id, + kind: args.duplicate.kind, + }) + .onConflictDoNothing(); +} + +async function listVectorPreferenceCandidates(args: { + db: MemoryDb; + embedding: MemoryEmbedding; + nowMs: number; + scope: ResolvedMemoryScope; + subject: ResolvedMemorySubject; +}): Promise { + const distance = cosineDistance( + juniorMemoryEmbeddings.embedding, + args.embedding.vector, + ); + const rows = await args.db + .select({ + contentHash: juniorMemoryEmbeddings.contentHash, + distance, + memory: juniorMemoryMemories, + }) + .from(juniorMemoryMemories) + .innerJoin( + juniorMemoryEmbeddings, + eq(juniorMemoryEmbeddings.memoryId, juniorMemoryMemories.id), + ) + .where( + and( + activeScopedSubjectPredicate({ ...args, kind: "preference" }), + eq(juniorMemoryEmbeddings.provider, args.embedding.provider), + eq(juniorMemoryEmbeddings.model, args.embedding.model), + eq(juniorMemoryEmbeddings.dimensions, MEMORY_EMBEDDING_DIMENSIONS), + eq(juniorMemoryEmbeddings.metric, EMBEDDING_METRIC), + ), + ) + .orderBy( + distance, + desc(juniorMemoryMemories.createdAtMs), + asc(juniorMemoryMemories.id), + ) + .limit(PREFERENCE_ADJUDICATION_VECTOR_LIMIT); + return rows.flatMap((row) => + hashEmbeddedContent(row.memory.content) === row.contentHash + ? [parseMemoryRow(row.memory)] + : [], + ); +} + +/** Find vector matches, then fill the candidate limit with recent memories. */ +async function listPreferenceCandidates(args: { + db: MemoryDb; + embedding: MemoryEmbedding | undefined; + nowMs: number; + scope: ResolvedMemoryScope; + subject: ResolvedMemorySubject; +}): Promise { + const vector = args.embedding + ? await listVectorPreferenceCandidates({ + ...args, + embedding: args.embedding, + }) + : []; + const recent = ( + await args.db + .select() + .from(juniorMemoryMemories) + .where(activeScopedSubjectPredicate({ ...args, kind: "preference" })) + .orderBy( + desc(juniorMemoryMemories.createdAtMs), + asc(juniorMemoryMemories.id), + ) + .limit(PREFERENCE_ADJUDICATION_CANDIDATE_LIMIT) + ).map(parseMemoryRow); + return [ + ...new Map( + [...vector, ...recent].map((memory) => [memory.id, memory]), + ).values(), + ].slice(0, PREFERENCE_ADJUDICATION_CANDIDATE_LIMIT); +} + +type PreferenceDecision = + | { decision: "create" } + | { decision: "duplicate"; memory: Memory } + | { decision: "supersede"; ids: [string, ...string[]] }; + +/** Keep old memories active when the model fails or returns unknown ids. */ +async function adjudicatePreference(args: { + candidates: Memory[]; + content: string; + decider: MemorySupersessionDecider; + runtimeContext: MemoryRuntimeContext; +}): Promise { + if (!args.candidates[0]) return { decision: "create" }; + const candidateIds = new Set(args.candidates.map(({ id }) => id)); + try { + const decision = await args.decider.adjudicateSupersession({ + candidate: { content: args.content, kind: "preference" }, + existingMemories: args.candidates.map(({ content, id }) => ({ + content, + id, + })), + runtimeContext: args.runtimeContext, + }); + if (decision.decision === "duplicate") { + const memory = args.candidates.find( + ({ id }) => id === decision.duplicateId, + ); + return memory + ? { decision: "duplicate", memory } + : { decision: "create" }; + } + if (decision.decision === "supersedes_old") { + const ids = decision.supersededIds.filter((id) => candidateIds.has(id)); + const [first, ...rest] = ids; + return first + ? { decision: "supersede", ids: [first, ...rest] } + : { decision: "create" }; + } + } catch { + return { decision: "create" }; + } + return { decision: "create" }; +} + +/** Create one memory with runtime-owned scope and subject authority. */ +export async function createMemory(args: { + context: MemoryRuntimeContext; + db: MemoryDb; + embedder?: MemoryEmbeddingProvider; + input: CreateMemoryInput; + now?: () => number; + subjectType: ResolvedMemorySubject["subjectType"]; + supersessionDecider?: MemorySupersessionDecider; +}): Promise { + const context = memoryRuntimeContextSchema.parse(args.context); + const input = createMemoryInputSchema.parse(args.input); + const nowMs = numberSchema.parse(args.now?.() ?? Date.now()); + const content = normalizeMemoryContent(input.content); + const scope = deriveMemoryScope(context); + const subject = deriveMemorySubject(context, args.subjectType); + if (content.length > MAX_MEMORY_CONTENT_CHARS) { + throw new Error("Memory content exceeds the maximum length."); + } + await archiveExpiredMemoryBatch({ db: args.db, nowMs, scopes: [scope] }); + await archiveExpiredMemoryBatch({ + db: args.db, + idempotencyKey: input.idempotencyKey, + limit: 1, + nowMs, + scopes: [scope], + }); + + const reuse = async (duplicate: Memory): Promise => { + await rememberDuplicateIdempotency({ + content, + db: args.db, + duplicate, + idempotencyKey: input.idempotencyKey, + nowMs, + runtimeContext: context, + scope, + subject, + }); + await storeMemoryEmbedding({ + content: duplicate.content, + db: args.db, + embedder: args.embedder, + memoryId: duplicate.id, + nowMs, + }); + return { created: false, memory: duplicate }; + }; + + const idempotent = await findByIdempotencyKey({ + db: args.db, + idempotencyKey: input.idempotencyKey, + nowMs, + scope, + }); + if (idempotent) { + await storeMemoryEmbedding({ + content: idempotent.memory.content, + db: args.db, + embedder: args.embedder, + memoryId: idempotent.memory.id, + nowMs, + }); + return idempotent.outcome === "created" + ? { created: false, idempotent: true, memory: idempotent.memory } + : { created: false, memory: idempotent.memory }; + } + + const exactDuplicate = await findExactDuplicate({ + content, + db: args.db, + kind: input.kind, + nowMs, + scope, + subject, + }); + if (exactDuplicate) return await reuse(exactDuplicate); + + let candidateEmbedding: MemoryEmbedding | undefined; + if (args.embedder) { + try { + candidateEmbedding = await embedMemoryText(args.embedder, content); + } catch { + candidateEmbedding = undefined; + } + } + let supersededIds: string[] = []; + if ( + args.subjectType === "user" && + input.kind === "preference" && + args.supersessionDecider && + (input.expiresAtMs === undefined || input.expiresAtMs > nowMs) + ) { + const candidates = await listPreferenceCandidates({ + db: args.db, + embedding: candidateEmbedding, + nowMs, + scope, + subject, + }); + const decision = await adjudicatePreference({ + candidates, + content, + decider: args.supersessionDecider, + runtimeContext: context, + }); + if (decision.decision === "duplicate") return await reuse(decision.memory); + if (decision.decision === "supersede") supersededIds = decision.ids; + } + + const id = randomUUID(); + const write = await args.db.transaction(async (tx) => { + const inserted = await tx + .insert(juniorMemoryMemories) + .values({ + content, + createdAtMs: nowMs, + expiresAtMs: input.expiresAtMs, + id, + idempotencyKey: input.idempotencyKey, + locationId: context.locationId, + observedAtMs: nowMs, + scope: scope.scope, + scopeKey: scope.scopeKey, + sourceKey: sourceKey(context), + sourcePlatform: context.source.platform, + subjectKey: subject.subjectKey, + subjectType: subject.subjectType, + kind: input.kind, + }) + .onConflictDoNothing({ + target: [ + juniorMemoryMemories.scope, + juniorMemoryMemories.scopeKey, + juniorMemoryMemories.idempotencyKey, + ], + where: sql`${juniorMemoryMemories.idempotencyKey} IS NOT NULL AND ${juniorMemoryMemories.archivedAtMs} IS NULL AND ${juniorMemoryMemories.supersededAtMs} IS NULL AND ${juniorMemoryMemories.supersededById} IS NULL`, + }) + .returning(); + const insertedMemory = inserted[0]; + if (!insertedMemory || supersededIds.length === 0) { + return { inserted, supersededIds: [] }; + } + const superseded = await tx + .update(juniorMemoryMemories) + .set({ supersededAtMs: nowMs, supersededById: insertedMemory.id }) + .where( + and( + inArray(juniorMemoryMemories.id, supersededIds), + activeScopedSubjectPredicate({ + kind: input.kind, + nowMs, + scope, + subject, + }), + ), + ) + .returning({ id: juniorMemoryMemories.id }); + const idsToClean = superseded.map(({ id }) => id); + if (idsToClean.length > 0) { + await tx + .delete(juniorMemoryEmbeddings) + .where(inArray(juniorMemoryEmbeddings.memoryId, idsToClean)); + } + return { inserted, supersededIds: idsToClean }; + }); + if (write.inserted[0]) { + const memory = parseMemoryRow(write.inserted[0]); + await storeMemoryEmbedding({ + content: memory.content, + db: args.db, + embedder: args.embedder, + embedding: candidateEmbedding, + memoryId: memory.id, + nowMs, + }); + const result: CreateMemoryResult = { + created: true, + memory, + }; + if (write.supersededIds.length > 0) { + result.supersededIds = write.supersededIds; + } + return result; + } + const conflict = await findByIdempotencyKey({ + db: args.db, + idempotencyKey: input.idempotencyKey, + nowMs, + scope, + }); + if (!conflict) { + throw new Error("Memory idempotency conflict did not resolve."); + } + await storeMemoryEmbedding({ + content: conflict.memory.content, + db: args.db, + embedder: args.embedder, + memoryId: conflict.memory.id, + nowMs, + }); + return conflict.outcome === "created" + ? { created: false, idempotent: true, memory: conflict.memory } + : { created: false, memory: conflict.memory }; +} diff --git a/packages/junior-memory/src/embeddings.ts b/packages/junior-memory/src/embeddings.ts new file mode 100644 index 0000000000..582ce6fc90 --- /dev/null +++ b/packages/junior-memory/src/embeddings.ts @@ -0,0 +1,69 @@ +/** Shared memory embedding normalization and provider validation. */ +import { createHash } from "node:crypto"; +import { z } from "zod"; +import { MEMORY_EMBEDDING_DIMENSIONS } from "./types"; + +const numberSchema = z.number().finite(); +const nonEmptyStringSchema = z.string().min(1); +export const EMBEDDING_METRIC = "cosine"; +const embeddingVectorSchema = z + .array(numberSchema) + .length(MEMORY_EMBEDDING_DIMENSIONS); +const embeddingResultSchema = z + .object({ + costUsd: z.number().finite().nonnegative().optional(), + dimensions: z.literal(MEMORY_EMBEDDING_DIMENSIONS), + model: nonEmptyStringSchema, + provider: nonEmptyStringSchema, + vectors: z.array(embeddingVectorSchema), + }) + .strict(); + +export interface MemoryEmbedding { + model: string; + provider: string; + vector: number[]; +} + +export interface MemoryEmbeddingProvider { + /** Embed normalized memory text for derived vector retrieval. */ + embedTexts(input: { texts: string[] }): Promise<{ + costUsd?: number; + dimensions: number; + model: string; + provider: string; + vectors: number[][]; + }>; +} + +/** Normalize memory content before comparison, storage, or embedding. */ +export function normalizeMemoryContent(content: string): string { + return content.replace(/\s+/g, " ").trim(); +} + +/** Hash the exact normalized content represented by a derived embedding. */ +export function hashEmbeddedContent(content: string): string { + return createHash("sha256").update(content, "utf8").digest("hex"); +} + +/** Embed one non-empty memory or query with a validated provider response. */ +export async function embedMemoryText( + embedder: MemoryEmbeddingProvider, + text: string, +): Promise { + const normalized = normalizeMemoryContent(text); + if (!normalized) { + throw new Error("Embedding text is required."); + } + const result = embeddingResultSchema.parse( + await embedder.embedTexts({ texts: [normalized] }), + ); + if (result.vectors.length !== 1) { + throw new Error("Embedding provider returned an unexpected vector count."); + } + return { + model: result.model, + provider: result.provider, + vector: result.vectors[0], + }; +} diff --git a/packages/junior-memory/src/events.ts b/packages/junior-memory/src/events.ts index 2fc437d874..b168578af4 100644 --- a/packages/junior-memory/src/events.ts +++ b/packages/junior-memory/src/events.ts @@ -1,7 +1,7 @@ import { defineConversationEvent } from "@sentry/junior-plugin-api"; import { z } from "zod"; import { MEMORY_KINDS, MEMORY_SCOPES } from "./types"; -import type { MemoryRecord } from "./store"; +import type { Memory } from "./memories"; const capturedMemoryFields = { content: z.string().min(1), @@ -96,7 +96,7 @@ export const memoriesRecalledEvent = defineConversationEvent({ }); /** Select the stable, safe memory fields retained in conversation history. */ -export function capturedMemory(memory: MemoryRecord) { +export function capturedMemory(memory: Memory) { return { content: memory.content, id: memory.id, diff --git a/packages/junior-memory/src/index.ts b/packages/junior-memory/src/index.ts index 96ef0fca86..1b0221dfa4 100644 --- a/packages/junior-memory/src/index.ts +++ b/packages/junior-memory/src/index.ts @@ -8,18 +8,5 @@ export { type MemoryListResponse, } from "./api"; export type { MemoryPluginOptions } from "./plugin"; -export { createMemoryStore } from "./store"; -export type { - ArchiveMemoryInput, - CreateMemoryInput, - CreateMemoryResult, - ListMemoriesInput, - MemoryDb, - MemoryEmbeddingProvider, - MemoryRecord, - MemoryStore, - MemoryStoreOptions, - SearchMemoriesInput, -} from "./store"; export { MEMORY_KINDS } from "./types"; export type { MemoryKind, MemoryRuntimeContext } from "./types"; diff --git a/packages/junior-memory/src/memories.ts b/packages/junior-memory/src/memories.ts new file mode 100644 index 0000000000..70c458c53c --- /dev/null +++ b/packages/junior-memory/src/memories.ts @@ -0,0 +1,227 @@ +/** Memory row validation, access filters, and expired cleanup. */ +import { + and, + asc, + eq, + gt, + inArray, + isNull, + lte, + or, + type SQL, +} from "drizzle-orm"; +import type { PgDatabase } from "drizzle-orm/pg-core"; +import type { PgQueryResultHKT } from "drizzle-orm/pg-core/session"; +import { z } from "zod"; +import * as memorySqlSchema from "./db/schema"; +import { juniorMemoryEmbeddings, juniorMemoryMemories } from "./db/schema"; +import type { ResolvedMemoryScope } from "./scope"; +import { + MEMORY_KINDS, + MEMORY_SCOPES, + MEMORY_SOURCE_PLATFORMS, + MEMORY_SUBJECT_TYPES, +} from "./types"; + +export type MemoryDb = PgDatabase; + +const numberSchema = z.number().finite(); +const optionalNumberSchema = z.preprocess( + (value) => (value === null ? undefined : value), + z.coerce.number().optional(), +); +const optionalStringSchema = z.preprocess( + (value) => (value === null ? undefined : value), + z.string().optional(), +); +const optionalNonEmptyStringSchema = z.preprocess( + (value) => (value === null ? undefined : value), + z.string().min(1).optional(), +); +const memoryContentSchema = z + .string() + .refine((content) => content.trim().length > 0, { + message: "Memory content is required.", + }); +const memoryRowSchema = z + .object({ + archivedAtMs: optionalNumberSchema, + archiveReason: optionalStringSchema, + content: memoryContentSchema, + createdAtMs: z.coerce.number(), + expiresAtMs: optionalNumberSchema, + id: z.string().min(1), + idempotencyKey: optionalStringSchema, + locationId: optionalNonEmptyStringSchema, + observedAtMs: z.coerce.number(), + searchVector: z.string().optional(), + scope: z.enum(MEMORY_SCOPES), + scopeKey: z.string().min(1), + sourceKey: z.string().min(1), + sourcePlatform: z.enum(MEMORY_SOURCE_PLATFORMS), + subjectKey: optionalNonEmptyStringSchema, + subjectType: z.enum(MEMORY_SUBJECT_TYPES), + supersededAtMs: optionalNumberSchema, + supersededById: optionalStringSchema, + kind: z.enum(MEMORY_KINDS), + }) + .strict() + .superRefine((row, ctx) => { + if (row.subjectType === "general") { + if (row.subjectKey !== undefined) { + ctx.addIssue({ + code: "custom", + message: "General-subject memory rows must not have a subject key.", + path: ["subjectKey"], + }); + } + return; + } + if (row.subjectKey === undefined) { + ctx.addIssue({ + code: "custom", + message: "User and conversation memory rows require a subject key.", + path: ["subjectKey"], + }); + } + }); + +const memorySchema = z + .object({ + archivedAtMs: numberSchema.optional(), + archiveReason: z.string().min(1).optional(), + content: memoryContentSchema, + createdAtMs: numberSchema, + expiresAtMs: numberSchema.optional(), + id: z.string().min(1), + observedAtMs: numberSchema, + scope: z.enum(MEMORY_SCOPES), + subjectType: z.enum(MEMORY_SUBJECT_TYPES), + supersededAtMs: numberSchema.optional(), + supersededById: z.string().min(1).optional(), + kind: z.enum(MEMORY_KINDS), + }) + .strict(); + +export type Memory = z.output; + +/** Parse one SQL row into the public memory projection. */ +export function parseMemoryRow(row: unknown): Memory { + const parsed = memoryRowSchema.parse(row); + const memory: z.input = { + id: parsed.id, + scope: parsed.scope, + kind: parsed.kind, + subjectType: parsed.subjectType, + content: parsed.content, + observedAtMs: parsed.observedAtMs, + createdAtMs: parsed.createdAtMs, + }; + if (parsed.expiresAtMs !== undefined) memory.expiresAtMs = parsed.expiresAtMs; + if (parsed.supersededAtMs !== undefined) { + memory.supersededAtMs = parsed.supersededAtMs; + } + if (parsed.supersededById) memory.supersededById = parsed.supersededById; + if (parsed.archivedAtMs !== undefined) { + memory.archivedAtMs = parsed.archivedAtMs; + } + if (parsed.archiveReason) memory.archiveReason = parsed.archiveReason; + return memorySchema.parse(memory); +} + +function visibleScopePredicate(scopes: ResolvedMemoryScope[]): SQL | undefined { + if (scopes.length === 0) { + return undefined; + } + return or( + ...scopes.map((scope) => + and( + eq(juniorMemoryMemories.scope, scope.scope), + eq(juniorMemoryMemories.scopeKey, scope.scopeKey), + ), + ), + ); +} + +/** Build the active-row predicate for already-authorized memory scopes. */ +export function activeVisiblePredicate(args: { + nowMs: number; + scopes: ResolvedMemoryScope[]; +}): SQL | undefined { + const scopePredicate = visibleScopePredicate(args.scopes); + if (!scopePredicate) { + return undefined; + } + return and( + scopePredicate, + isNull(juniorMemoryMemories.archivedAtMs), + isNull(juniorMemoryMemories.supersededAtMs), + isNull(juniorMemoryMemories.supersededById), + or( + isNull(juniorMemoryMemories.expiresAtMs), + gt(juniorMemoryMemories.expiresAtMs, args.nowMs), + ), + ); +} + +function boundedLimit(value: number | undefined, fallback: number): number { + if (typeof value !== "number" || !Number.isFinite(value)) { + return fallback; + } + return Math.min(200, Math.max(1, Math.floor(value))); +} + +/** Archive a bounded batch of expired rows and remove their derived vectors. */ +export async function archiveExpiredMemoryBatch(args: { + db: MemoryDb; + idempotencyKey?: string; + limit?: number; + nowMs: number; + scopes: ResolvedMemoryScope[]; +}): Promise<{ archivedCount: number }> { + const scopePredicate = visibleScopePredicate(args.scopes); + if (!scopePredicate) { + return { archivedCount: 0 }; + } + const predicates: SQL[] = [ + scopePredicate, + isNull(juniorMemoryMemories.archivedAtMs), + isNull(juniorMemoryMemories.supersededAtMs), + isNull(juniorMemoryMemories.supersededById), + lte(juniorMemoryMemories.expiresAtMs, args.nowMs), + ]; + if (args.idempotencyKey !== undefined) { + predicates.push( + eq(juniorMemoryMemories.idempotencyKey, args.idempotencyKey), + ); + } + + const archivedIds = await args.db.transaction(async (tx) => { + const expired = await tx + .select({ id: juniorMemoryMemories.id }) + .from(juniorMemoryMemories) + .where(and(...predicates)) + .orderBy( + asc(juniorMemoryMemories.expiresAtMs), + asc(juniorMemoryMemories.id), + ) + .limit(boundedLimit(args.limit, 100)); + const ids = expired.map((row) => row.id); + if (ids.length === 0) { + return []; + } + const archived = await tx + .update(juniorMemoryMemories) + .set({ archivedAtMs: args.nowMs, archiveReason: "expired" }) + .where(and(inArray(juniorMemoryMemories.id, ids), ...predicates)) + .returning({ id: juniorMemoryMemories.id }); + const idsToClean = archived.map((row) => row.id); + if (idsToClean.length > 0) { + await tx + .delete(juniorMemoryEmbeddings) + .where(inArray(juniorMemoryEmbeddings.memoryId, idsToClean)); + } + return idsToClean; + }); + return { archivedCount: archivedIds.length }; +} diff --git a/packages/junior-memory/src/operational-report.ts b/packages/junior-memory/src/operational-report.ts index d021df2dd5..b14b3f4400 100644 --- a/packages/junior-memory/src/operational-report.ts +++ b/packages/junior-memory/src/operational-report.ts @@ -5,7 +5,7 @@ import type { import { and, eq, gt, isNull, or, sql } from "drizzle-orm"; import { z } from "zod"; import { juniorMemoryEmbeddings, juniorMemoryMemories } from "./db/schema"; -import type { MemoryDb } from "./store"; +import type { MemoryDb } from "./memories"; const DAY_MS = 24 * 60 * 60 * 1_000; const WINDOWS = [7, 30, 90] as const; diff --git a/packages/junior-memory/src/plugin.ts b/packages/junior-memory/src/plugin.ts index 7d8a8cf945..0da1bf69f8 100644 --- a/packages/junior-memory/src/plugin.ts +++ b/packages/junior-memory/src/plugin.ts @@ -19,7 +19,7 @@ import { memoriesCapturedEventV1, memoriesRecalledEvent, } from "./events"; -import type { MemoryDb } from "./store"; +import type { MemoryDb } from "./memories"; import { createMemoryUserPage } from "./user-pages"; const MEMORY_MODEL_ENV = "AI_MEMORY_MODEL"; diff --git a/packages/junior-memory/src/process-session.ts b/packages/junior-memory/src/process-session.ts index efc743f055..abf9fbbc85 100644 --- a/packages/junior-memory/src/process-session.ts +++ b/packages/junior-memory/src/process-session.ts @@ -7,11 +7,10 @@ import { } from "@sentry/junior-plugin-api"; import { z } from "zod"; import { - createMemoryStore, + createMemory, type CreateMemoryInput, type CreateMemoryResult, - type MemoryDb, -} from "./store"; +} from "./create"; import { createMemoryAgent, parseExtractedMemory, @@ -20,6 +19,9 @@ import { } from "./agent"; import { MEMORY_KINDS, memoryRuntimeContextSchema } from "./types"; import { capturedMemory, memoriesCapturedEvent } from "./events"; +import { archiveExpiredMemoryBatch, type MemoryDb } from "./memories"; +import { retrieveMemories } from "./retrieval"; +import { deriveVisibleMemoryScopes } from "./scope"; const MEMORY_TOOL_NAMES = new Set([ "createMemory", @@ -260,15 +262,19 @@ export async function processMemorySession( ...(run.actorUserId ? { userId: run.actorUserId } : undefined), }); const agent = createMemoryAgent(context.model); - const store = createMemoryStore(context.db as MemoryDb, runtimeContext, { - embedder: context.embedder, - supersessionDecider: agent, + const db = context.db as MemoryDb; + await archiveExpiredMemoryBatch({ + db, + nowMs: Date.now(), + scopes: deriveVisibleMemoryScopes(runtimeContext), }); - await store.archiveExpiredMemories(); const extraction = await getTaskExtraction(context, async () => { - const existingMemories = await store.searchMemories({ - limit: 10, - query: evidenceText, + const existingMemories = await retrieveMemories({ + context: runtimeContext, + db, + embedder: context.embedder, + input: { limit: 10, query: evidenceText }, + mode: "search", }); return await agent.extractSessionMemories({ existingMemories: existingMemories.map((memory) => ({ @@ -290,12 +296,14 @@ export async function processMemorySession( continue; } const input = passiveInput(run.runId, memory, sourceKey, target); - if (target === "conversation") { - const result = await store.createConversationMemory(input); - recordCapturedMemory(captured, result); - continue; - } - const result = await store.createMemory(input); + const result = await createMemory({ + context: runtimeContext, + db, + embedder: context.embedder, + input, + subjectType: target, + supersessionDecider: agent, + }); recordCapturedMemory(captured, result); } await context.events.emit( diff --git a/packages/junior-memory/src/ranking.ts b/packages/junior-memory/src/ranking.ts index 67649cd33a..02e3c028ad 100644 --- a/packages/junior-memory/src/ranking.ts +++ b/packages/junior-memory/src/ranking.ts @@ -1,4 +1,4 @@ -import type { MemoryRecord } from "./store"; +import type { Memory } from "./memories"; const RECIPROCAL_RANK_FUSION_K = 60; const ONE_DAY_MS = 24 * 60 * 60 * 1000; @@ -8,7 +8,7 @@ export interface MemoryMatch { lexical?: { rank: number; }; - memory: MemoryRecord; + memory: Memory; vector?: { rank: number; }; @@ -32,7 +32,7 @@ function matchScore( ); } -function observedAgeRank(memory: MemoryRecord, nowMs: number): number { +function observedAgeRank(memory: Memory, nowMs: number): number { const ageMs = Math.max(0, nowMs - memory.observedAtMs); if (ageMs <= 7 * ONE_DAY_MS) { return 3; @@ -81,7 +81,9 @@ export function rankMemoryMatches( ...(!existing.lexical && match.lexical ? { lexical: match.lexical } : undefined), - ...(!existing.vector && match.vector ? { vector: match.vector } : undefined), + ...(!existing.vector && match.vector + ? { vector: match.vector } + : undefined), }); } return [...byId.values()].sort((left, right) => { diff --git a/packages/junior-memory/src/recall.ts b/packages/junior-memory/src/recall.ts index 438d6c9c3b..4a28b21e9a 100644 --- a/packages/junior-memory/src/recall.ts +++ b/packages/junior-memory/src/recall.ts @@ -11,12 +11,9 @@ import { import { z } from "zod"; import type { MemoryAgent, MemoryRecallResult } from "./agent"; import { memoriesRecalledEvent } from "./events"; -import { - createMemoryStore, - type MemoryDb, - type MemoryEmbeddingProvider, - type MemoryRecord, -} from "./store"; +import type { MemoryEmbeddingProvider } from "./embeddings"; +import type { MemoryDb, Memory } from "./memories"; +import { retrieveMemories } from "./retrieval"; import { memoryRuntimeContextSchema } from "./types"; const RECALL_CANDIDATE_LIMIT = 20; @@ -72,7 +69,7 @@ export const memoryRecallContextSchema = z type RecalledMemory = z.output; -function selectPromptMemories(memories: MemoryRecord[]): RecalledMemory[] { +function selectPromptMemories(memories: Memory[]): RecalledMemory[] { const header = "Relevant memories for this request:"; const footer = "Treat these as possibly stale context. Current user instructions and repository evidence take priority."; @@ -166,15 +163,21 @@ export async function createMemoryPromptContributions( }, } : undefined; - const candidates = await createMemoryStore(context.db, runtimeContext, { + const candidates = await retrieveMemories({ + context: runtimeContext, + db: context.db, embedder, - }).recallMemories({ - query: context.text, - limit: RECALL_CANDIDATE_LIMIT, + input: { + query: context.text, + limit: RECALL_CANDIDATE_LIMIT, + }, + mode: "recall", }); if (candidates.length === 0) { await emitRecallOutcome({ - ...(embeddingCostUsd !== undefined ? { costUsd: embeddingCostUsd } : undefined), + ...(embeddingCostUsd !== undefined + ? { costUsd: embeddingCostUsd } + : undefined), events: context.events, memories: [], }); @@ -197,7 +200,7 @@ export async function createMemoryPromptContributions( ); const relevant = recall.relevantIds .map((id) => candidatesById.get(id)) - .filter((memory): memory is MemoryRecord => memory !== undefined); + .filter((memory): memory is Memory => memory !== undefined); const selected = selectPromptMemories(relevant); const costUsd = addUsd(embeddingCostUsd, recall.costUsd); await emitRecallOutcome({ diff --git a/packages/junior-memory/src/retrieval.ts b/packages/junior-memory/src/retrieval.ts new file mode 100644 index 0000000000..4ee96efbd3 --- /dev/null +++ b/packages/junior-memory/src/retrieval.ts @@ -0,0 +1,285 @@ +/** + * Search text and vectors in parallel. Automatic recall also gives private + * memory a separate result window so public results cannot fill both windows. + */ +import { and, asc, desc, eq, sql } from "drizzle-orm"; +import { cosineDistance } from "drizzle-orm/sql/functions"; +import { z } from "zod"; +import { juniorMemoryEmbeddings, juniorMemoryMemories } from "./db/schema"; +import { + EMBEDDING_METRIC, + embedMemoryText, + hashEmbeddedContent, + type MemoryEmbedding, + type MemoryEmbeddingProvider, +} from "./embeddings"; +import { + activeVisiblePredicate, + archiveExpiredMemoryBatch, + parseMemoryRow, + type MemoryDb, + type Memory, +} from "./memories"; +import { rankMemoryMatches, type MemoryMatch } from "./ranking"; +import { deriveVisibleMemoryScopes, type ResolvedMemoryScope } from "./scope"; +import { + MEMORY_EMBEDDING_DIMENSIONS, + memoryRuntimeContextSchema, + type MemoryRuntimeContext, +} from "./types"; + +const DEFAULT_SEARCH_LIMIT = 10; +const SEARCH_RETRIEVAL_OVERFETCH = 4; +const RECALL_RETRIEVAL_OVERFETCH = 2; +const MAX_RETRIEVAL_LEG_CANDIDATES = 200; +const MAX_LEXICAL_RANK_CANDIDATES = 200; +const LEXICAL_RANK_WINDOW_MULTIPLIER = 4; +const MAX_RETRIEVAL_QUERY_CHARS = 1_500; +const RECALL_MAX_VECTOR_DISTANCE = 0.45; +const numberSchema = z.number().finite(); + +const retrieveMemoriesInputSchema = z + .object({ + limit: z.number().finite().optional(), + query: z.string().min(1), + }) + .strict(); +type RetrieveMemoriesInput = z.output; + +function boundedLimit(value: number | undefined, fallback: number): number { + if (typeof value !== "number" || !Number.isFinite(value)) return fallback; + return Math.min(200, Math.max(1, Math.floor(value))); +} + +function normalizeQuery(query: string): string { + const normalized = query.replace(/\s+/g, " ").trim(); + return normalized.length <= MAX_RETRIEVAL_QUERY_CHARS + ? normalized + : normalized.slice(0, MAX_RETRIEVAL_QUERY_CHARS).trimEnd(); +} + +function retrievalLegLimit(limit: number, overfetch: number): number { + const requested = Math.max(1, limit); + return Math.min( + MAX_RETRIEVAL_LEG_CANDIDATES, + Math.max(requested, requested * Math.max(1, overfetch)), + ); +} + +function denseRanks(values: T[], key: (value: T) => string | number) { + let previous: string | number | undefined; + let rank = 0; + return values.map((value, index) => { + const current = key(value); + if (index === 0 || current !== previous) { + rank = index + 1; + previous = current; + } + return rank; + }); +} + +async function searchLexical(args: { + db: MemoryDb; + limit: number; + nowMs: number; + query: string; + scopes: ResolvedMemoryScope[]; +}): Promise { + const predicate = activeVisiblePredicate(args); + const query = normalizeQuery(args.query); + if (!predicate || !query) return []; + const queryVector = sql`to_tsvector('english', ${query})`; + const tsquery = sql`( + SELECT COALESCE( + string_agg(quote_literal(term), ' | ')::tsquery, + ''::tsquery + ) + FROM unnest(tsvector_to_array(${queryVector})) AS query_terms(term) + )`; + const candidateLimit = Math.min( + MAX_LEXICAL_RANK_CANDIDATES, + args.limit * LEXICAL_RANK_WINDOW_MULTIPLIER, + ); + const candidates = args.db + .select() + .from(juniorMemoryMemories) + .where( + and(predicate, sql`${juniorMemoryMemories.searchVector} @@ ${tsquery}`), + ) + .orderBy( + desc(juniorMemoryMemories.observedAtMs), + asc(juniorMemoryMemories.id), + ) + .limit(candidateLimit) + .as("lexical_candidates"); + const textRank = sql`ts_rank_cd(${candidates.searchVector}, ${tsquery})`; + const rows = await args.db + .select({ + memory: { + archiveReason: candidates.archiveReason, + archivedAtMs: candidates.archivedAtMs, + content: candidates.content, + createdAtMs: candidates.createdAtMs, + expiresAtMs: candidates.expiresAtMs, + id: candidates.id, + idempotencyKey: candidates.idempotencyKey, + kind: candidates.kind, + observedAtMs: candidates.observedAtMs, + scope: candidates.scope, + scopeKey: candidates.scopeKey, + searchVector: candidates.searchVector, + sourceKey: candidates.sourceKey, + sourcePlatform: candidates.sourcePlatform, + subjectKey: candidates.subjectKey, + subjectType: candidates.subjectType, + supersededAtMs: candidates.supersededAtMs, + supersededById: candidates.supersededById, + }, + textRank, + }) + .from(candidates) + .orderBy(desc(textRank), desc(candidates.observedAtMs), asc(candidates.id)) + .limit(args.limit); + const ranks = denseRanks(rows, (row) => Number(row.textRank)); + return rows.map((row, index) => ({ + lexical: { rank: ranks[index] }, + memory: parseMemoryRow(row.memory), + })); +} + +async function searchVector(args: { + db: MemoryDb; + embedding: MemoryEmbedding; + limit: number; + maxDistance: number | undefined; + nowMs: number; + scopes: ResolvedMemoryScope[]; +}): Promise { + const predicate = activeVisiblePredicate(args); + if (!predicate) return []; + const distance = cosineDistance( + juniorMemoryEmbeddings.embedding, + args.embedding.vector, + ); + const distancePredicate = + args.maxDistance === undefined + ? undefined + : sql`${distance} <= ${args.maxDistance}`; + const rows = await args.db + .select({ + contentHash: juniorMemoryEmbeddings.contentHash, + distance, + memory: juniorMemoryMemories, + }) + .from(juniorMemoryMemories) + .innerJoin( + juniorMemoryEmbeddings, + eq(juniorMemoryEmbeddings.memoryId, juniorMemoryMemories.id), + ) + .where( + and( + predicate, + eq(juniorMemoryEmbeddings.provider, args.embedding.provider), + eq(juniorMemoryEmbeddings.model, args.embedding.model), + eq(juniorMemoryEmbeddings.dimensions, MEMORY_EMBEDDING_DIMENSIONS), + eq(juniorMemoryEmbeddings.metric, EMBEDDING_METRIC), + ...(distancePredicate ? [distancePredicate] : []), + ), + ) + .orderBy( + distance, + desc(juniorMemoryMemories.createdAtMs), + asc(juniorMemoryMemories.id), + ) + .limit(args.limit); + const ranks = denseRanks(rows, (row) => Number(row.distance)); + return rows.flatMap((row, index) => { + const distanceValue = Number(row.distance); + if ( + row.distance === null || + !Number.isFinite(distanceValue) || + hashEmbeddedContent(row.memory.content) !== row.contentHash + ) { + return []; + } + return [ + { + memory: parseMemoryRow(row.memory), + vector: { rank: ranks[index] }, + }, + ]; + }); +} + +/** Retrieve active memories with search or recall ranking rules. */ +export async function retrieveMemories(args: { + context: MemoryRuntimeContext; + db: MemoryDb; + embedder?: MemoryEmbeddingProvider; + input: RetrieveMemoriesInput; + mode: "recall" | "search"; + now?: () => number; +}): Promise { + const context = memoryRuntimeContextSchema.parse(args.context); + const input = retrieveMemoriesInputSchema.parse(args.input); + const nowMs = numberSchema.parse(args.now?.() ?? Date.now()); + const scopes = deriveVisibleMemoryScopes(context); + await archiveExpiredMemoryBatch({ db: args.db, nowMs, scopes }); + const limit = boundedLimit(input.limit, DEFAULT_SEARCH_LIMIT); + const recall = args.mode === "recall"; + const candidateLimit = retrievalLegLimit( + limit, + recall ? RECALL_RETRIEVAL_OVERFETCH : SEARCH_RETRIEVAL_OVERFETCH, + ); + const privateScopes = scopes.filter((scope) => scope.scope === "private"); + const probePrivate = recall && privateScopes.length > 0; + const query = normalizeQuery(input.query); + let queryEmbedding: MemoryEmbedding | undefined; + if (args.embedder && query) { + try { + queryEmbedding = await embedMemoryText(args.embedder, query); + } catch { + queryEmbedding = undefined; + } + } + const empty = Promise.resolve([] as MemoryMatch[]); + const lexicalArgs = { + db: args.db, + limit: candidateLimit, + nowMs, + query: input.query, + }; + const matches = await Promise.all([ + queryEmbedding + ? searchVector({ + db: args.db, + embedding: queryEmbedding, + limit: candidateLimit, + maxDistance: recall ? RECALL_MAX_VECTOR_DISTANCE : undefined, + nowMs, + scopes, + }) + : empty, + searchLexical({ ...lexicalArgs, scopes }), + queryEmbedding && probePrivate + ? searchVector({ + db: args.db, + embedding: queryEmbedding, + limit: candidateLimit, + maxDistance: RECALL_MAX_VECTOR_DISTANCE, + nowMs, + scopes: privateScopes, + }) + : empty, + probePrivate + ? searchLexical({ ...lexicalArgs, scopes: privateScopes }) + : empty, + ]); + const weights = recall + ? { lexicalWeight: 1, nowMs, vectorWeight: 0.85 } + : { nowMs }; + return rankMemoryMatches(matches.flat(), weights) + .slice(0, limit) + .map(({ memory }) => memory); +} diff --git a/packages/junior-memory/src/store.ts b/packages/junior-memory/src/store.ts deleted file mode 100644 index e413d79883..0000000000 --- a/packages/junior-memory/src/store.ts +++ /dev/null @@ -1,1560 +0,0 @@ -/** - * SQL-backed memory store boundary. - * - * This module owns row parsing plus visible create/list/search/archive - * operations. Visibility, expiration, and supersession are enforced before - * records leave the store. - */ -import { createHash, randomUUID } from "node:crypto"; -import { - and, - asc, - desc, - eq, - gt, - inArray, - isNull, - isNotNull, - like, - lte, - or, - sql, - type SQL, -} from "drizzle-orm"; -import { cosineDistance } from "drizzle-orm/sql/functions"; -import type { PgDatabase } from "drizzle-orm/pg-core"; -import type { PgQueryResultHKT } from "drizzle-orm/pg-core/session"; -import { z } from "zod"; -import { getSourceKey } from "@sentry/junior-plugin-api"; -import * as memorySqlSchema from "./db/schema"; -import { juniorMemoryEmbeddings, juniorMemoryMemories } from "./db/schema"; -import { rankMemoryMatches, type MemoryMatch } from "./ranking"; -import { - MEMORY_EMBEDDING_DIMENSIONS, - MEMORY_SCOPES, - MEMORY_SOURCE_PLATFORMS, - MEMORY_SUBJECT_TYPES, - MEMORY_KINDS, - memoryRuntimeContextSchema, - type MemoryRuntimeContext, -} from "./types"; -import { - deriveMemoryScope, - deriveMemorySubject, - type ResolvedMemorySubject, - deriveVisibleMemoryScopes, - type ResolvedMemoryScope, -} from "./scope"; - -const DEFAULT_LIST_LIMIT = 50; -const DEFAULT_SEARCH_LIMIT = 10; -const DEFAULT_EXPIRED_ARCHIVE_LIMIT = 100; -const PREFERENCE_ADJUDICATION_CANDIDATE_LIMIT = 10; -const PREFERENCE_ADJUDICATION_VECTOR_LIMIT = 5; -/** Explicit search overfetch: keep a wider fusion window for tool/CLI search. */ -const SEARCH_RETRIEVAL_OVERFETCH = 4; -/** - * Automatic recall overfetch. Recall already asks for ~20 candidates before the - * relevance gate, so each hybrid leg only needs a small top-k probe. - */ -const RECALL_RETRIEVAL_OVERFETCH = 2; -/** - * Absolute ceiling per retrieval leg. Matches the store limit ceiling so a - * single healthy leg can still fill the caller's requested result window. - */ -const MAX_RETRIEVAL_LEG_CANDIDATES = 200; -/** Cap ts_rank_cd work after GIN filtering; ranking is not indexable. */ -const MAX_LEXICAL_RANK_CANDIDATES = 200; -/** Expand the GIN match window before ts_rank_cd, still under the hard cap. */ -const LEXICAL_RANK_WINDOW_MULTIPLIER = 4; -/** Bound query text before embedding / FTS construction. */ -const MAX_RETRIEVAL_QUERY_CHARS = 1_500; -const MAX_MEMORY_CONTENT_CHARS = 4_000; -const EMBEDDING_METRIC = "cosine"; -/** - * Cosine distance cutoff for automatic recall only (not explicit search). - * Tuned for text-embedding-3-small; retune if the embedding model changes. - */ -const RECALL_MAX_VECTOR_DISTANCE = 0.45; - -export type MemoryDb = PgDatabase; - -interface MemoryEmbedding { - model: string; - provider: string; - vector: number[]; -} - -const nonEmptyStringSchema = z.string().min(1); -const memoryContentSchema = z - .string() - .refine((content) => content.trim().length > 0, { - message: "Memory content is required.", - }); -const numberSchema = z.number().finite(); -const createMemoryInputSchema = z - .object({ - content: memoryContentSchema, - expiresAtMs: numberSchema.optional(), - idempotencyKey: nonEmptyStringSchema, - kind: z.enum(MEMORY_KINDS), - }) - .strict(); -const listMemoriesInputSchema = z - .object({ - limit: numberSchema.optional(), - }) - .strict(); -const searchMemoriesInputSchema = z - .object({ - limit: numberSchema.optional(), - query: nonEmptyStringSchema, - }) - .strict(); -const archiveMemoryInputSchema = z - .object({ - id: nonEmptyStringSchema, - reason: nonEmptyStringSchema.optional(), - }) - .strict(); -const archiveExpiredMemoriesInputSchema = z - .object({ - limit: numberSchema.optional(), - }) - .strict(); -const clockSchema = z.function({ input: [], output: numberSchema }).optional(); -const memoryStoreOptionsSchema = z - .object({ - now: clockSchema, - }) - .strict(); -const optionalNumberSchema = z.preprocess( - (value) => (value === null ? undefined : value), - z.coerce.number().optional(), -); -const optionalStringSchema = z.preprocess( - (value) => (value === null ? undefined : value), - z.string().optional(), -); -const optionalNonEmptyStringSchema = z.preprocess( - (value) => (value === null ? undefined : value), - z.string().min(1).optional(), -); -const memoryRowSchema = z - .object({ - archivedAtMs: optionalNumberSchema, - archiveReason: optionalStringSchema, - content: memoryContentSchema, - createdAtMs: z.coerce.number(), - expiresAtMs: optionalNumberSchema, - id: z.string().min(1), - idempotencyKey: optionalStringSchema, - locationId: optionalNonEmptyStringSchema, - observedAtMs: z.coerce.number(), - searchVector: z.string().optional(), - scope: z.enum(MEMORY_SCOPES), - scopeKey: z.string().min(1), - sourceKey: z.string().min(1), - sourcePlatform: z.enum(MEMORY_SOURCE_PLATFORMS), - subjectKey: optionalNonEmptyStringSchema, - subjectType: z.enum(MEMORY_SUBJECT_TYPES), - supersededAtMs: optionalNumberSchema, - supersededById: optionalStringSchema, - kind: z.enum(MEMORY_KINDS), - }) - .strict() - .superRefine((row, ctx) => { - if (row.subjectType === "general") { - if (row.subjectKey !== undefined) { - ctx.addIssue({ - code: "custom", - message: "General-subject memory rows must not have a subject key.", - path: ["subjectKey"], - }); - } - return; - } - if (row.subjectKey === undefined) { - ctx.addIssue({ - code: "custom", - message: "User and conversation memory rows require a subject key.", - path: ["subjectKey"], - }); - } - }); - -const memoryRecordSchema = z - .object({ - archivedAtMs: numberSchema.optional(), - archiveReason: nonEmptyStringSchema.optional(), - content: memoryContentSchema, - createdAtMs: numberSchema, - expiresAtMs: numberSchema.optional(), - id: nonEmptyStringSchema, - observedAtMs: numberSchema, - scope: z.enum(MEMORY_SCOPES), - subjectType: z.enum(MEMORY_SUBJECT_TYPES), - supersededAtMs: numberSchema.optional(), - supersededById: nonEmptyStringSchema.optional(), - kind: z.enum(MEMORY_KINDS), - }) - .strict(); -const embeddingVectorSchema = z - .array(numberSchema) - .length(MEMORY_EMBEDDING_DIMENSIONS); -const embeddingResultSchema = z - .object({ - costUsd: z.number().finite().nonnegative().optional(), - dimensions: z.literal(MEMORY_EMBEDDING_DIMENSIONS), - model: nonEmptyStringSchema, - provider: nonEmptyStringSchema, - vectors: z.array(embeddingVectorSchema), - }) - .strict(); -const memorySupersessionCandidateSchema = z - .object({ - content: z.string().min(1), - id: z.string().min(1), - }) - .strict(); -const memorySupersessionCandidatesSchema = z - .array(memorySupersessionCandidateSchema) - .min(1) - .max(PREFERENCE_ADJUDICATION_CANDIDATE_LIMIT); -const supersededIdsSchema = z - .array(z.string().min(1)) - .min(1) - .max(PREFERENCE_ADJUDICATION_CANDIDATE_LIMIT); - -/** Validated preference comparison input supplied to a supersession decider. */ -export const memorySupersessionInputSchema = z - .object({ - candidate: z - .object({ - content: z.string().min(1), - kind: z.literal("preference"), - }) - .strict(), - existingMemories: memorySupersessionCandidatesSchema, - runtimeContext: memoryRuntimeContextSchema, - }) - .strict(); - -/** - * Validated preference decision whose referenced ids must come from the - * supplied existing memories. - */ -export const memorySupersessionDecisionSchema = z.discriminatedUnion( - "decision", - [ - z - .object({ - decision: z.literal("duplicate"), - duplicateId: z.string().min(1), - }) - .strict(), - z - .object({ - decision: z.literal("supersedes_old"), - supersededIds: supersededIdsSchema, - }) - .strict(), - z - .object({ - decision: z.enum(["distinct", "uncertain"]), - }) - .strict(), - ], -); - -export type MemoryRecord = z.output; -export type CreateMemoryInput = z.output; - -/** Result of a memory write after idempotency checks. */ -export interface CreateMemoryResult { - created: boolean; - /** True when this call found the memory previously written for the same input identity. */ - idempotent?: true; - memory: MemoryRecord; - /** Memory ids made inactive by this write. */ - supersededIds?: string[]; -} - -export type ListMemoriesInput = z.output; - -export type SearchMemoriesInput = z.output; - -export type ArchiveMemoryInput = z.output; - -export type ArchiveExpiredMemoriesInput = z.output< - typeof archiveExpiredMemoriesInputSchema ->; - -export interface ArchiveExpiredMemoriesResult { - archivedCount: number; -} - -export interface MemoryEmbeddingProvider { - /** Embed normalized memory text for derived vector retrieval. */ - embedTexts(input: { texts: string[] }): Promise<{ - costUsd?: number; - dimensions: number; - model: string; - provider: string; - vectors: number[][]; - }>; -} - -export type MemorySupersessionInput = z.output< - typeof memorySupersessionInputSchema ->; - -export type MemorySupersessionDecision = z.output< - typeof memorySupersessionDecisionSchema ->; - -export interface MemorySupersessionDecider { - /** Classify a new preference against related active preferences. */ - adjudicateSupersession( - input: MemorySupersessionInput, - ): Promise | MemorySupersessionDecision; -} - -export interface MemoryStoreOptions { - embedder?: MemoryEmbeddingProvider; - now?: () => number; - supersessionDecider?: MemorySupersessionDecider; -} - -/** Context-bound storage operations for visible long-term memories. */ -export interface MemoryStore { - /** Archive expired memories visible in the current runtime context. */ - archiveExpiredMemories( - input?: ArchiveExpiredMemoriesInput, - ): Promise; - /** Archive a visible memory in the current runtime context. */ - archiveMemory(input: ArchiveMemoryInput): Promise; - /** Store a memory about the current User. The Source sets access. */ - createMemory(input: CreateMemoryInput): Promise; - /** Store a memory about the current Conversation. The Source sets access. */ - createConversationMemory( - input: CreateMemoryInput, - ): Promise; - /** List active memories visible in the current runtime context. */ - listMemories(input: ListMemoriesInput): Promise; - /** - * Retrieve a broad relevance-ranked candidate window for automatic recall. - * Prompt admission remains owned by the recall boundary. - */ - recallMemories(input: SearchMemoriesInput): Promise; - /** Search active memories visible in the current runtime context. */ - searchMemories(input: SearchMemoriesInput): Promise; -} - -function normalizeContent(content: string): string { - return content.replace(/\s+/g, " ").trim(); -} - -function hashEmbeddedContent(content: string): string { - return createHash("sha256").update(content, "utf8").digest("hex"); -} - -function idempotencyAliasId(args: { - idempotencyKey: string; - scope: ResolvedMemoryScope; - targetId: string; -}): string { - return `alias:${createHash("sha256") - .update(args.scope.scope) - .update("\0") - .update(args.scope.scopeKey) - .update("\0") - .update(args.idempotencyKey) - .update("\0") - .update(args.targetId) - .digest("hex")}`; -} - -function boundedLimit(value: number | undefined, fallback: number): number { - if (typeof value !== "number" || !Number.isFinite(value)) { - return fallback; - } - return Math.min(200, Math.max(1, Math.floor(value))); -} - -/** Build the stored key for the Source. */ -function sourceKey(ctx: MemoryRuntimeContext): string { - const key = getSourceKey(ctx.source); - if (!key) { - throw new Error("Memory Source has no stable key."); - } - return key; -} - -/** Parse one SQL row into the public memory projection. */ -export function parseMemoryRow(row: unknown): MemoryRecord { - const parsed = memoryRowSchema.parse(row); - return memoryRecordSchema.parse({ - id: parsed.id, - scope: parsed.scope, - kind: parsed.kind, - subjectType: parsed.subjectType, - content: parsed.content, - observedAtMs: parsed.observedAtMs, - createdAtMs: parsed.createdAtMs, - ...(parsed.expiresAtMs !== undefined - ? { expiresAtMs: parsed.expiresAtMs } - : undefined), - ...(parsed.supersededAtMs !== undefined - ? { supersededAtMs: parsed.supersededAtMs } - : undefined), - ...(parsed.supersededById ? { supersededById: parsed.supersededById } : undefined), - ...(parsed.archivedAtMs !== undefined - ? { archivedAtMs: parsed.archivedAtMs } - : undefined), - ...(parsed.archiveReason ? { archiveReason: parsed.archiveReason } : undefined), - }); -} - -/** Build the scoped SQL predicate and ordered params for visible memory reads. */ -function visibleScopePredicate(scopes: ResolvedMemoryScope[]): SQL | undefined { - if (scopes.length === 0) { - return undefined; - } - return or( - ...scopes.map((scope) => - and( - eq(juniorMemoryMemories.scope, scope.scope), - eq(juniorMemoryMemories.scopeKey, scope.scopeKey), - ), - ), - ); -} - -/** Build the active-row predicate for already-authorized memory scopes. */ -export function activeVisiblePredicate(args: { - nowMs: number; - scopes: ResolvedMemoryScope[]; -}): SQL | undefined { - const scopePredicate = visibleScopePredicate(args.scopes); - if (!scopePredicate) { - return undefined; - } - return and( - scopePredicate, - isNull(juniorMemoryMemories.archivedAtMs), - isNull(juniorMemoryMemories.supersededAtMs), - isNull(juniorMemoryMemories.supersededById), - or( - isNull(juniorMemoryMemories.expiresAtMs), - gt(juniorMemoryMemories.expiresAtMs, args.nowMs), - ), - ); -} - -/** Resolve retry attempts for the same scoped write idempotency key. */ -interface IdempotencyMatch { - memory: MemoryRecord; - outcome: "created" | "duplicate"; -} - -async function findByIdempotencyKey(args: { - db: MemoryDb; - idempotencyKey: string; - nowMs: number; - scope: ResolvedMemoryScope; -}): Promise { - const activeRows = await args.db - .select() - .from(juniorMemoryMemories) - .where( - and( - eq(juniorMemoryMemories.scope, args.scope.scope), - eq(juniorMemoryMemories.scopeKey, args.scope.scopeKey), - eq(juniorMemoryMemories.idempotencyKey, args.idempotencyKey), - isNull(juniorMemoryMemories.archivedAtMs), - isNull(juniorMemoryMemories.supersededAtMs), - isNull(juniorMemoryMemories.supersededById), - or( - isNull(juniorMemoryMemories.expiresAtMs), - gt(juniorMemoryMemories.expiresAtMs, args.nowMs), - ), - ), - ) - .limit(1); - if (activeRows[0]) { - return { memory: parseMemoryRow(activeRows[0]), outcome: "created" }; - } - - const aliasRows = await args.db - .select({ supersededById: juniorMemoryMemories.supersededById }) - .from(juniorMemoryMemories) - .where( - and( - eq(juniorMemoryMemories.scope, args.scope.scope), - eq(juniorMemoryMemories.scopeKey, args.scope.scopeKey), - eq(juniorMemoryMemories.idempotencyKey, args.idempotencyKey), - isNull(juniorMemoryMemories.archivedAtMs), - isNotNull(juniorMemoryMemories.supersededAtMs), - isNotNull(juniorMemoryMemories.supersededById), - or( - isNull(juniorMemoryMemories.expiresAtMs), - gt(juniorMemoryMemories.expiresAtMs, args.nowMs), - ), - ), - ) - .orderBy( - desc(juniorMemoryMemories.createdAtMs), - asc(juniorMemoryMemories.id), - ); - for (const alias of aliasRows) { - if (!alias.supersededById) { - continue; - } - const rows = await args.db - .select() - .from(juniorMemoryMemories) - .where( - and( - eq(juniorMemoryMemories.id, alias.supersededById), - eq(juniorMemoryMemories.scope, args.scope.scope), - eq(juniorMemoryMemories.scopeKey, args.scope.scopeKey), - isNull(juniorMemoryMemories.archivedAtMs), - isNull(juniorMemoryMemories.supersededAtMs), - isNull(juniorMemoryMemories.supersededById), - or( - isNull(juniorMemoryMemories.expiresAtMs), - gt(juniorMemoryMemories.expiresAtMs, args.nowMs), - ), - ), - ) - .limit(1); - if (rows[0]) { - return { memory: parseMemoryRow(rows[0]), outcome: "duplicate" }; - } - } - return undefined; -} - -/** - * Archive a bounded batch of expired active rows and remove their derived vectors. - */ -export async function archiveExpiredMemoryBatch(args: { - db: MemoryDb; - idempotencyKey?: string; - limit?: number; - nowMs: number; - scopes: ResolvedMemoryScope[]; -}): Promise { - const scopePredicate = visibleScopePredicate(args.scopes); - if (!scopePredicate) { - return { archivedCount: 0 }; - } - const predicates: SQL[] = [ - scopePredicate, - isNull(juniorMemoryMemories.archivedAtMs), - isNull(juniorMemoryMemories.supersededAtMs), - isNull(juniorMemoryMemories.supersededById), - lte(juniorMemoryMemories.expiresAtMs, args.nowMs), - ]; - if (args.idempotencyKey !== undefined) { - predicates.push( - eq(juniorMemoryMemories.idempotencyKey, args.idempotencyKey), - ); - } - - const archivedIds = await args.db.transaction(async (tx) => { - const expired = await tx - .select({ id: juniorMemoryMemories.id }) - .from(juniorMemoryMemories) - .where(and(...predicates)) - .orderBy( - asc(juniorMemoryMemories.expiresAtMs), - asc(juniorMemoryMemories.id), - ) - .limit(boundedLimit(args.limit, DEFAULT_EXPIRED_ARCHIVE_LIMIT)); - const ids = expired.map((row) => row.id); - if (ids.length === 0) { - return []; - } - - const archived = await tx - .update(juniorMemoryMemories) - .set({ - archivedAtMs: args.nowMs, - archiveReason: "expired", - }) - .where(and(inArray(juniorMemoryMemories.id, ids), ...predicates)) - .returning({ id: juniorMemoryMemories.id }); - const idsToClean = archived.map((row) => row.id); - if (idsToClean.length > 0) { - await tx - .delete(juniorMemoryEmbeddings) - .where(inArray(juniorMemoryEmbeddings.memoryId, idsToClean)); - } - return idsToClean; - }); - return { archivedCount: archivedIds.length }; -} - -function denseRanks( - values: T[], - key: (value: T) => string | number, -): number[] { - let previous: string | number | undefined; - let rank = 0; - return values.map((value, index) => { - const current = key(value); - if (index === 0 || current !== previous) { - rank = index + 1; - previous = current; - } - return rank; - }); -} - -async function embedOne( - embedder: MemoryEmbeddingProvider, - text: string, -): Promise { - const normalized = normalizeContent(text); - if (!normalized) { - throw new Error("Embedding text is required."); - } - const result = embeddingResultSchema.parse( - await embedder.embedTexts({ texts: [normalized] }), - ); - if (result.vectors.length !== 1) { - throw new Error("Embedding provider returned an unexpected vector count."); - } - return { - model: result.model, - provider: result.provider, - vector: result.vectors[0], - }; -} - -/** Store the derived vector index; failures must not block memory persistence. */ -async function storeEmbedding(args: { - content: string; - db: MemoryDb; - embedder: MemoryEmbeddingProvider | undefined; - embedding?: MemoryEmbedding; - memoryId: string; - nowMs: number; -}): Promise { - if (!args.embedder && !args.embedding) { - return; - } - try { - const existing = await args.db - .select({ memoryId: juniorMemoryEmbeddings.memoryId }) - .from(juniorMemoryEmbeddings) - .where(eq(juniorMemoryEmbeddings.memoryId, args.memoryId)) - .limit(1); - if (existing[0]) { - return; - } - } catch { - return; - } - let embedding: Awaited>; - if (args.embedding) { - embedding = args.embedding; - } else { - const embedder = args.embedder; - if (!embedder) { - return; - } - try { - embedding = await embedOne(embedder, args.content); - } catch { - return; - } - } - try { - await args.db - .insert(juniorMemoryEmbeddings) - .values({ - contentHash: hashEmbeddedContent(args.content), - createdAtMs: args.nowMs, - dimensions: MEMORY_EMBEDDING_DIMENSIONS, - embedding: embedding.vector, - memoryId: args.memoryId, - metric: EMBEDDING_METRIC, - model: embedding.model, - provider: embedding.provider, - }) - .onConflictDoNothing(); - } catch { - return; - } -} - -function activeScopedSubjectPredicate(args: { - kind: MemoryRecord["kind"]; - nowMs: number; - scope: ResolvedMemoryScope; - subject: ResolvedMemorySubject; -}): SQL { - const predicate = and( - eq(juniorMemoryMemories.scope, args.scope.scope), - eq(juniorMemoryMemories.scopeKey, args.scope.scopeKey), - eq(juniorMemoryMemories.kind, args.kind), - eq(juniorMemoryMemories.subjectType, args.subject.subjectType), - eq(juniorMemoryMemories.subjectKey, args.subject.subjectKey), - isNull(juniorMemoryMemories.archivedAtMs), - isNull(juniorMemoryMemories.supersededAtMs), - isNull(juniorMemoryMemories.supersededById), - or( - isNull(juniorMemoryMemories.expiresAtMs), - gt(juniorMemoryMemories.expiresAtMs, args.nowMs), - ), - ); - if (!predicate) { - throw new Error("Memory duplicate predicate is empty."); - } - return predicate; -} - -async function findExactDuplicateMemory(args: { - content: string; - db: MemoryDb; - kind: MemoryRecord["kind"]; - nowMs: number; - scope: ResolvedMemoryScope; - subject: ResolvedMemorySubject; -}): Promise { - const rows = await args.db - .select() - .from(juniorMemoryMemories) - .where( - and( - activeScopedSubjectPredicate(args), - eq(juniorMemoryMemories.content, args.content), - ), - ) - .orderBy( - desc(juniorMemoryMemories.createdAtMs), - asc(juniorMemoryMemories.id), - ) - .limit(1); - return rows[0] ? parseMemoryRow(rows[0]) : undefined; -} - -async function rememberDuplicateIdempotency(args: { - content: string; - db: MemoryDb; - duplicate: MemoryRecord; - idempotencyKey?: string; - nowMs: number; - runtimeContext: MemoryRuntimeContext; - scope: ResolvedMemoryScope; - subject: ResolvedMemorySubject; -}): Promise { - if (args.idempotencyKey === undefined) { - return; - } - await args.db - .insert(juniorMemoryMemories) - .values({ - content: args.content, - createdAtMs: args.nowMs, - expiresAtMs: args.duplicate.expiresAtMs, - id: idempotencyAliasId({ - idempotencyKey: args.idempotencyKey, - scope: args.scope, - targetId: args.duplicate.id, - }), - idempotencyKey: args.idempotencyKey, - locationId: args.runtimeContext.locationId, - observedAtMs: args.nowMs, - scope: args.scope.scope, - scopeKey: args.scope.scopeKey, - sourceKey: sourceKey(args.runtimeContext), - sourcePlatform: args.runtimeContext.source.platform, - subjectKey: args.subject.subjectKey, - subjectType: args.subject.subjectType, - supersededAtMs: args.nowMs, - supersededById: args.duplicate.id, - kind: args.duplicate.kind, - }) - .onConflictDoNothing(); -} - -/** Select semantic preferences, then fill the window by recency for unembedded records. */ -async function listPreferenceAdjudicationCandidates(args: { - db: MemoryDb; - embedding?: MemoryEmbedding; - nowMs: number; - scope: ResolvedMemoryScope; - subject: ResolvedMemorySubject; -}): Promise { - const vectorCandidates = args.embedding - ? await listVectorPreferenceAdjudicationCandidates({ - db: args.db, - embedding: args.embedding, - nowMs: args.nowMs, - scope: args.scope, - subject: args.subject, - }) - : []; - const recentCandidates = ( - await args.db - .select() - .from(juniorMemoryMemories) - .where( - activeScopedSubjectPredicate({ - ...args, - kind: "preference", - }), - ) - .orderBy( - desc(juniorMemoryMemories.createdAtMs), - asc(juniorMemoryMemories.id), - ) - .limit(PREFERENCE_ADJUDICATION_CANDIDATE_LIMIT) - ).map(parseMemoryRow); - return [ - ...new Map( - [...vectorCandidates, ...recentCandidates].map((memory) => [ - memory.id, - memory, - ]), - ).values(), - ].slice(0, PREFERENCE_ADJUDICATION_CANDIDATE_LIMIT); -} - -async function listVectorPreferenceAdjudicationCandidates(args: { - db: MemoryDb; - embedding: MemoryEmbedding; - nowMs: number; - scope: ResolvedMemoryScope; - subject: ResolvedMemorySubject; -}): Promise { - const distance = cosineDistance( - juniorMemoryEmbeddings.embedding, - args.embedding.vector, - ); - const rows = await args.db - .select({ - contentHash: juniorMemoryEmbeddings.contentHash, - distance, - memory: juniorMemoryMemories, - }) - .from(juniorMemoryMemories) - .innerJoin( - juniorMemoryEmbeddings, - eq(juniorMemoryEmbeddings.memoryId, juniorMemoryMemories.id), - ) - .where( - and( - activeScopedSubjectPredicate({ ...args, kind: "preference" }), - eq(juniorMemoryEmbeddings.provider, args.embedding.provider), - eq(juniorMemoryEmbeddings.model, args.embedding.model), - eq(juniorMemoryEmbeddings.dimensions, MEMORY_EMBEDDING_DIMENSIONS), - eq(juniorMemoryEmbeddings.metric, EMBEDDING_METRIC), - ), - ) - .orderBy( - distance, - desc(juniorMemoryMemories.createdAtMs), - asc(juniorMemoryMemories.id), - ) - .limit(PREFERENCE_ADJUDICATION_VECTOR_LIMIT); - return rows.flatMap((row) => { - if (hashEmbeddedContent(row.memory.content) !== row.contentHash) { - return []; - } - return [parseMemoryRow(row.memory)]; - }); -} - -type PreferenceAdjudicationResult = - | { decision: "create" } - | { decision: "duplicate"; memory: MemoryRecord } - | { decision: "supersede"; ids: [string, ...string[]] }; - -/** - * Normalize a preference decision to known duplicate or supersession targets. - * Uncertainty, invalid ids, and model failure leave existing memories active. - */ -async function adjudicatePreferenceCandidate(args: { - candidates: MemoryRecord[]; - content: string; - decider: MemorySupersessionDecider; - runtimeContext: MemoryRuntimeContext; -}): Promise { - const [firstCandidate, ...remainingCandidates] = args.candidates; - if (!firstCandidate) { - return { decision: "create" }; - } - const existingMemories = [ - { content: firstCandidate.content, id: firstCandidate.id }, - ...remainingCandidates.map((memory) => ({ - content: memory.content, - id: memory.id, - })), - ]; - const candidateIds = new Set(args.candidates.map((memory) => memory.id)); - try { - const decision = await args.decider.adjudicateSupersession({ - candidate: { - content: args.content, - kind: "preference", - }, - existingMemories, - runtimeContext: args.runtimeContext, - }); - if (decision.decision === "duplicate") { - const memory = args.candidates.find( - (candidate) => candidate.id === decision.duplicateId, - ); - return memory - ? { decision: "duplicate", memory } - : { decision: "create" }; - } - if (decision.decision === "supersedes_old") { - const ids = decision.supersededIds.filter((id) => candidateIds.has(id)); - const [firstId, ...remainingIds] = ids; - return firstId - ? { decision: "supersede", ids: [firstId, ...remainingIds] } - : { decision: "create" }; - } - return { decision: "create" }; - } catch { - return { decision: "create" }; - } -} - -/** List active records for the runtime-derived visible scopes. */ -async function listVisibleMemories(args: { - db: MemoryDb; - limit?: number; - nowMs: number; - scopes: ResolvedMemoryScope[]; -}): Promise { - const predicate = activeVisiblePredicate(args); - if (!predicate) { - return []; - } - const limit = boundedLimit(args.limit, DEFAULT_LIST_LIMIT); - const rows = await args.db - .select() - .from(juniorMemoryMemories) - .where(predicate) - .orderBy( - desc(juniorMemoryMemories.createdAtMs), - asc(juniorMemoryMemories.id), - ) - .limit(limit); - return rows.map(parseMemoryRow); -} - -function normalizeRetrievalQuery(query: string): string { - const normalized = query.replace(/\s+/g, " ").trim(); - if (normalized.length <= MAX_RETRIEVAL_QUERY_CHARS) { - return normalized; - } - return normalized.slice(0, MAX_RETRIEVAL_QUERY_CHARS).trimEnd(); -} - -function retrievalLegLimit(limit: number, overfetch: number): number { - const requested = Math.max(1, limit); - const withOverfetch = requested * Math.max(1, overfetch); - // Never return fewer candidates than the caller asked for. A hard overfetch - // cap below `limit` under-fills when one modality is empty or both overlap. - return Math.min( - MAX_RETRIEVAL_LEG_CANDIDATES, - Math.max(requested, withOverfetch), - ); -} - -/** Search a bounded active candidate set with PostgreSQL full-text ranking. */ -async function searchVisibleLexicalMemories(args: { - db: MemoryDb; - limit: number; - nowMs: number; - query: string; - scopes: ResolvedMemoryScope[]; -}): Promise { - const predicate = activeVisiblePredicate(args); - if (!predicate) { - return []; - } - const query = normalizeRetrievalQuery(args.query); - if (!query) { - return []; - } - const queryVector = sql`to_tsvector('english', ${query})`; - const tsquery = sql`( - SELECT COALESCE( - string_agg(quote_literal(term), ' | ')::tsquery, - ''::tsquery - ) - FROM unnest(tsvector_to_array(${queryVector})) AS query_terms(term) - )`; - // GIN filter first, then rank only a bounded recent match window. - const candidateLimit = Math.min( - MAX_LEXICAL_RANK_CANDIDATES, - args.limit * LEXICAL_RANK_WINDOW_MULTIPLIER, - ); - const candidates = args.db - .select() - .from(juniorMemoryMemories) - .where( - and(predicate, sql`${juniorMemoryMemories.searchVector} @@ ${tsquery}`), - ) - .orderBy( - desc(juniorMemoryMemories.observedAtMs), - asc(juniorMemoryMemories.id), - ) - .limit(candidateLimit) - .as("lexical_candidates"); - const textRank = sql`ts_rank_cd(${candidates.searchVector}, ${tsquery})`; - const rows = await args.db - .select({ - memory: { - archiveReason: candidates.archiveReason, - archivedAtMs: candidates.archivedAtMs, - content: candidates.content, - createdAtMs: candidates.createdAtMs, - expiresAtMs: candidates.expiresAtMs, - id: candidates.id, - idempotencyKey: candidates.idempotencyKey, - kind: candidates.kind, - observedAtMs: candidates.observedAtMs, - scope: candidates.scope, - scopeKey: candidates.scopeKey, - searchVector: candidates.searchVector, - sourceKey: candidates.sourceKey, - sourcePlatform: candidates.sourcePlatform, - subjectKey: candidates.subjectKey, - subjectType: candidates.subjectType, - supersededAtMs: candidates.supersededAtMs, - supersededById: candidates.supersededById, - }, - textRank, - }) - .from(candidates) - .orderBy(desc(textRank), desc(candidates.observedAtMs), asc(candidates.id)) - .limit(args.limit); - const ranks = denseRanks(rows, (row) => Number(row.textRank)); - return rows.map((row, index) => ({ - lexical: { rank: ranks[index] }, - memory: parseMemoryRow(row.memory), - })); -} - -/** Search active visible records with pgvector cosine distance. */ -async function searchVisibleVectorMemories(args: { - db: MemoryDb; - embedding: MemoryEmbedding; - limit: number; - maxDistance?: number; - nowMs: number; - scopes: ResolvedMemoryScope[]; -}): Promise { - const predicate = activeVisiblePredicate(args); - if (!predicate) { - return []; - } - const embedding = args.embedding; - const distance = cosineDistance( - juniorMemoryEmbeddings.embedding, - embedding.vector, - ); - // Push distance cutoff into SQL so recall does not overfetch weak neighbors. - const distancePredicate = - args.maxDistance === undefined - ? undefined - : sql`${distance} <= ${args.maxDistance}`; - const rows = await args.db - .select({ - contentHash: juniorMemoryEmbeddings.contentHash, - distance, - memory: juniorMemoryMemories, - }) - .from(juniorMemoryMemories) - .innerJoin( - juniorMemoryEmbeddings, - eq(juniorMemoryEmbeddings.memoryId, juniorMemoryMemories.id), - ) - .where( - and( - predicate, - eq(juniorMemoryEmbeddings.provider, embedding.provider), - eq(juniorMemoryEmbeddings.model, embedding.model), - eq(juniorMemoryEmbeddings.dimensions, MEMORY_EMBEDDING_DIMENSIONS), - eq(juniorMemoryEmbeddings.metric, EMBEDDING_METRIC), - ...(distancePredicate ? [distancePredicate] : []), - ), - ) - .orderBy( - distance, - desc(juniorMemoryMemories.createdAtMs), - asc(juniorMemoryMemories.id), - ) - .limit(args.limit); - const ranks = denseRanks(rows, (row) => Number(row.distance)); - return rows.flatMap((row, index) => { - const distanceValue = Number(row.distance); - if ( - row.distance === null || - !Number.isFinite(distanceValue) || - hashEmbeddedContent(row.memory.content) !== row.contentHash - ) { - return []; - } - return [ - { - memory: parseMemoryRow(row.memory), - vector: { - rank: ranks[index], - }, - }, - ]; - }); -} - -/** Create a context-bound SQL-backed store for explicit memory operations. */ -export function createMemoryStore( - db: MemoryDb, - context: MemoryRuntimeContext, - options: MemoryStoreOptions = {}, -): MemoryStore { - const runtimeContext = memoryRuntimeContextSchema.parse(context); - const parsedOptions = memoryStoreOptionsSchema.parse({ now: options.now }); - const embedder = options.embedder; - const supersessionDecider = options.supersessionDecider; - const getNowMs = parsedOptions.now ?? Date.now; - - async function archiveExpiredVisibleMemories( - input: ArchiveExpiredMemoriesInput | undefined, - nowMs: number, - ): Promise { - input = archiveExpiredMemoriesInputSchema.parse(input ?? {}); - return await archiveExpiredMemoryBatch({ - db, - limit: input.limit, - nowMs, - scopes: deriveVisibleMemoryScopes(runtimeContext), - }); - } - - async function reuseDuplicateMemory(args: { - content: string; - duplicate: MemoryRecord; - idempotencyKey?: string; - nowMs: number; - scope: ResolvedMemoryScope; - subject: ResolvedMemorySubject; - }): Promise { - await rememberDuplicateIdempotency({ - ...args, - db, - runtimeContext, - }); - await storeEmbedding({ - content: args.duplicate.content, - db, - embedder, - memoryId: args.duplicate.id, - nowMs: args.nowMs, - }); - return { created: false, memory: args.duplicate }; - } - - /** Persist a memory under the plugin-derived scope and subject. */ - async function createScopedMemory( - rawInput: CreateMemoryInput, - subjectType: ResolvedMemorySubject["subjectType"], - ): Promise { - const input = createMemoryInputSchema.parse(rawInput); - const nowMs = getNowMs(); - const content = normalizeContent(input.content); - const scope = deriveMemoryScope(runtimeContext); - const subject = deriveMemorySubject(runtimeContext, subjectType); - if (content.length > MAX_MEMORY_CONTENT_CHARS) { - throw new Error("Memory content exceeds the maximum length."); - } - await archiveExpiredMemoryBatch({ - db, - nowMs, - scopes: [scope], - }); - await archiveExpiredMemoryBatch({ - db, - idempotencyKey: input.idempotencyKey, - limit: 1, - nowMs, - scopes: [scope], - }); - if (input.idempotencyKey !== undefined) { - const idempotent = await findByIdempotencyKey({ - db, - idempotencyKey: input.idempotencyKey, - nowMs, - scope, - }); - if (idempotent) { - await storeEmbedding({ - content: idempotent.memory.content, - db, - embedder, - memoryId: idempotent.memory.id, - nowMs, - }); - return idempotent.outcome === "created" - ? { created: false, idempotent: true, memory: idempotent.memory } - : { created: false, memory: idempotent.memory }; - } - } - - const exactDuplicate = await findExactDuplicateMemory({ - content, - db, - kind: input.kind, - nowMs, - scope, - subject, - }); - if (exactDuplicate) { - return await reuseDuplicateMemory({ - content, - duplicate: exactDuplicate, - idempotencyKey: input.idempotencyKey, - nowMs, - scope, - subject, - }); - } - - let candidateEmbedding: MemoryEmbedding | undefined; - if (embedder) { - try { - candidateEmbedding = await embedOne(embedder, content); - } catch { - candidateEmbedding = undefined; - } - } - let supersededIds: string[] = []; - if ( - subjectType === "user" && - input.kind === "preference" && - supersessionDecider && - (input.expiresAtMs === undefined || input.expiresAtMs > nowMs) - ) { - const preferenceCandidates = await listPreferenceAdjudicationCandidates({ - db, - ...(candidateEmbedding ? { embedding: candidateEmbedding } : undefined), - nowMs, - scope, - subject, - }); - const adjudication = await adjudicatePreferenceCandidate({ - candidates: preferenceCandidates, - content, - decider: supersessionDecider, - runtimeContext, - }); - if (adjudication.decision === "duplicate") { - return await reuseDuplicateMemory({ - content, - duplicate: adjudication.memory, - idempotencyKey: input.idempotencyKey, - nowMs, - scope, - subject, - }); - } - if (adjudication.decision === "supersede") { - supersededIds = adjudication.ids; - } - } - - const id = randomUUID(); - const write = await db.transaction(async (tx) => { - const inserted = await tx - .insert(juniorMemoryMemories) - .values({ - content, - createdAtMs: nowMs, - expiresAtMs: input.expiresAtMs, - id, - idempotencyKey: input.idempotencyKey, - locationId: runtimeContext.locationId, - observedAtMs: nowMs, - scope: scope.scope, - scopeKey: scope.scopeKey, - sourceKey: sourceKey(runtimeContext), - sourcePlatform: runtimeContext.source.platform, - subjectKey: subject.subjectKey, - subjectType: subject.subjectType, - kind: input.kind, - }) - .onConflictDoNothing({ - target: [ - juniorMemoryMemories.scope, - juniorMemoryMemories.scopeKey, - juniorMemoryMemories.idempotencyKey, - ], - where: sql`${juniorMemoryMemories.idempotencyKey} IS NOT NULL AND ${juniorMemoryMemories.archivedAtMs} IS NULL AND ${juniorMemoryMemories.supersededAtMs} IS NULL AND ${juniorMemoryMemories.supersededById} IS NULL`, - }) - .returning(); - const insertedMemory = inserted[0]; - if (!insertedMemory || supersededIds.length === 0) { - return { inserted, supersededIds: [] }; - } - const superseded = await tx - .update(juniorMemoryMemories) - .set({ - supersededAtMs: nowMs, - supersededById: insertedMemory.id, - }) - .where( - and( - inArray(juniorMemoryMemories.id, supersededIds), - activeScopedSubjectPredicate({ - kind: input.kind, - nowMs, - scope, - subject, - }), - ), - ) - .returning({ id: juniorMemoryMemories.id }); - const idsToClean = superseded.map((row) => row.id); - if (idsToClean.length > 0) { - await tx - .delete(juniorMemoryEmbeddings) - .where(inArray(juniorMemoryEmbeddings.memoryId, idsToClean)); - } - return { inserted, supersededIds: idsToClean }; - }); - if (write.inserted[0]) { - const memory = parseMemoryRow(write.inserted[0]); - await storeEmbedding({ - content: memory.content, - db, - embedder, - embedding: candidateEmbedding, - memoryId: memory.id, - nowMs, - }); - return { - created: true, - memory, - ...(write.supersededIds.length > 0 - ? { supersededIds: write.supersededIds } - : undefined), - }; - } - - const idempotent = await findByIdempotencyKey({ - db, - idempotencyKey: input.idempotencyKey, - nowMs, - scope, - }); - if (!idempotent) { - throw new Error("Memory idempotency conflict did not resolve."); - } - await storeEmbedding({ - content: idempotent.memory.content, - db, - embedder, - memoryId: idempotent.memory.id, - nowMs, - }); - return idempotent.outcome === "created" - ? { created: false, idempotent: true, memory: idempotent.memory } - : { created: false, memory: idempotent.memory }; - } - - /** - * Hybrid retrieval for both automatic recall and explicit search. - * - * Keep both legs parallel and fuse ranks with RRF. Never skip lexical when - * vectors already hit: that drops exact/token memories and serializes the - * miss path. Each leg is a hard-capped top-k probe so Postgres work stays - * bounded even on broad queries. - * - * Automatic recall also searches private memory by itself. This keeps newer - * public memory with common words from hiding older private memory. - */ - async function retrieveVisibleMemories( - rawInput: SearchMemoriesInput, - vectorMaxDistance: number | undefined, - ): Promise { - const input = searchMemoriesInputSchema.parse(rawInput); - const nowMs = getNowMs(); - const scopes = deriveVisibleMemoryScopes(runtimeContext); - await archiveExpiredMemoryBatch({ - db, - nowMs, - scopes, - }); - const limit = boundedLimit(input.limit, DEFAULT_SEARCH_LIMIT); - const overfetch = - vectorMaxDistance === undefined - ? SEARCH_RETRIEVAL_OVERFETCH - : RECALL_RETRIEVAL_OVERFETCH; - const candidateLimit = retrievalLegLimit(limit, overfetch); - const privateScopes = scopes.filter((scope) => scope.scope === "private"); - // Search private memory by itself during recall so public results cannot - // fill both search windows. - const probePrivate = - vectorMaxDistance !== undefined && privateScopes.length > 0; - const query = normalizeRetrievalQuery(input.query); - let queryEmbedding: MemoryEmbedding | undefined; - if (embedder && query) { - try { - queryEmbedding = await embedOne(embedder, query); - } catch { - queryEmbedding = undefined; - } - } - const emptyMatches = Promise.resolve([] as MemoryMatch[]); - const lexicalArgs = { - db, - limit: candidateLimit, - nowMs, - query: input.query, - }; - // Always run both legs in parallel. Conditional lexical skip is unsafe: - // one in-threshold vector distractor can hide a stronger lexical hit. - // Embed once up front; vector probes only run when that embedding exists. - const matches = await Promise.all([ - queryEmbedding - ? searchVisibleVectorMemories({ - db, - embedding: queryEmbedding, - limit: candidateLimit, - ...(vectorMaxDistance !== undefined - ? { maxDistance: vectorMaxDistance } - : undefined), - nowMs, - scopes, - }) - : emptyMatches, - searchVisibleLexicalMemories({ - ...lexicalArgs, - scopes, - }), - queryEmbedding && probePrivate - ? searchVisibleVectorMemories({ - db, - embedding: queryEmbedding, - limit: candidateLimit, - maxDistance: vectorMaxDistance, - nowMs, - scopes: privateScopes, - }) - : emptyMatches, - probePrivate - ? searchVisibleLexicalMemories({ - ...lexicalArgs, - scopes: privateScopes, - }) - : emptyMatches, - ]); - return rankMemoryMatches(matches.flat(), { - nowMs, - // Slight lexical preference protects exact ids/names/timezones on ties. - ...(vectorMaxDistance === undefined - ? undefined - : { lexicalWeight: 1, vectorWeight: 0.85 }), - }) - .slice(0, limit) - .map(({ memory }) => memory); - } - - return { - async archiveExpiredMemories(input) { - return await archiveExpiredVisibleMemories(input, getNowMs()); - }, - - async createMemory(input) { - return await createScopedMemory(input, "user"); - }, - - async createConversationMemory(input) { - return await createScopedMemory(input, "conversation"); - }, - - async listMemories(input) { - input = listMemoriesInputSchema.parse(input); - const nowMs = getNowMs(); - const scopes = deriveVisibleMemoryScopes(runtimeContext); - await archiveExpiredMemoryBatch({ - db, - nowMs, - scopes, - }); - return await listVisibleMemories({ - db, - limit: input.limit, - nowMs, - scopes, - }); - }, - - async recallMemories(input) { - return await retrieveVisibleMemories(input, RECALL_MAX_VECTOR_DISTANCE); - }, - - async searchMemories(input) { - return await retrieveVisibleMemories(input, undefined); - }, - - async archiveMemory(input) { - input = archiveMemoryInputSchema.parse(input); - const nowMs = getNowMs(); - // Public memory is shared and has no single user owner. - const scopes = deriveVisibleMemoryScopes(runtimeContext).filter( - (scope) => scope.scope === "private", - ); - const predicate = activeVisiblePredicate({ nowMs, scopes }); - const idPrefix = input.id.trim(); - if (!idPrefix) { - throw new Error("Memory id is required."); - } - const rows = predicate - ? await db - .select() - .from(juniorMemoryMemories) - .where( - and( - predicate, - or( - eq(juniorMemoryMemories.id, idPrefix), - like(juniorMemoryMemories.id, `${idPrefix}%`), - ), - ), - ) - .orderBy(asc(juniorMemoryMemories.id)) - .limit(2) - : []; - if (rows.length === 0) { - throw new Error("Memory was not found in the current context."); - } - if (rows.length > 1) { - throw new Error("Memory id prefix is ambiguous."); - } - const memory = parseMemoryRow(rows[0]); - const updated = await db - .update(juniorMemoryMemories) - .set({ - archivedAtMs: nowMs, - archiveReason: input.reason ?? "user_removed", - }) - .where(eq(juniorMemoryMemories.id, memory.id)) - .returning(); - await db - .delete(juniorMemoryEmbeddings) - .where(eq(juniorMemoryEmbeddings.memoryId, memory.id)); - return parseMemoryRow(updated[0]); - }, - }; -} diff --git a/packages/junior-memory/src/testing.ts b/packages/junior-memory/src/testing.ts new file mode 100644 index 0000000000..0ed92068ca --- /dev/null +++ b/packages/junior-memory/src/testing.ts @@ -0,0 +1,49 @@ +import { asc, eq, inArray } from "drizzle-orm"; +import { createMemory } from "./create"; +import { juniorMemoryEmbeddings, juniorMemoryMemories } from "./db/schema"; +import type { MemoryDb } from "./memories"; + +/** Remove all memories and their embeddings from the test database. */ +export async function clearAll(db: MemoryDb): Promise { + await db.delete(juniorMemoryEmbeddings); + await db.delete(juniorMemoryMemories); +} + +/** List memories learned from one Source, oldest first. */ +export async function listBySource(db: MemoryDb, sourceKey: string) { + return db + .select({ + archivedAtMs: juniorMemoryMemories.archivedAtMs, + content: juniorMemoryMemories.content, + expiresAtMs: juniorMemoryMemories.expiresAtMs, + id: juniorMemoryMemories.id, + kind: juniorMemoryMemories.kind, + scope: juniorMemoryMemories.scope, + scopeKey: juniorMemoryMemories.scopeKey, + subjectKey: juniorMemoryMemories.subjectKey, + subjectType: juniorMemoryMemories.subjectType, + supersededAtMs: juniorMemoryMemories.supersededAtMs, + supersededById: juniorMemoryMemories.supersededById, + }) + .from(juniorMemoryMemories) + .where(eq(juniorMemoryMemories.sourceKey, sourceKey)) + .orderBy( + asc(juniorMemoryMemories.createdAtMs), + asc(juniorMemoryMemories.id), + ); +} + +/** Count embeddings for the given memory IDs. */ +export async function countEmbeddings( + db: MemoryDb, + memoryIds: string[], +): Promise { + const rows = await db + .select({ memoryId: juniorMemoryEmbeddings.memoryId }) + .from(juniorMemoryEmbeddings) + .where(inArray(juniorMemoryEmbeddings.memoryId, memoryIds)); + return rows.length; +} + +export { createMemory }; +export type { MemoryDb }; diff --git a/packages/junior-memory/src/tools.ts b/packages/junior-memory/src/tools.ts index 2c212aca1c..e6310d8bb0 100644 --- a/packages/junior-memory/src/tools.ts +++ b/packages/junior-memory/src/tools.ts @@ -1,5 +1,6 @@ import { Type, type Static } from "@sinclair/typebox"; import { Value } from "@sinclair/typebox/value"; +import { and, asc, desc, eq, like, or } from "drizzle-orm"; import { definePluginTool, getSourceKey, @@ -13,13 +14,21 @@ import { } from "@sentry/junior-plugin-api"; import { z } from "zod"; import { - createMemoryStore, + createMemory, type CreateMemoryInput, - type MemoryEmbeddingProvider, - type MemoryDb, - type MemoryRecord, type MemorySupersessionDecider, -} from "./store"; +} from "./create"; +import type { MemoryEmbeddingProvider } from "./embeddings"; +import { juniorMemoryEmbeddings, juniorMemoryMemories } from "./db/schema"; +import { + activeVisiblePredicate, + archiveExpiredMemoryBatch, + parseMemoryRow, + type MemoryDb, + type Memory, +} from "./memories"; +import { retrieveMemories } from "./retrieval"; +import { deriveVisibleMemoryScopes } from "./scope"; import { parseCreateMemoryRequest, parseMemoryReview, @@ -100,19 +109,6 @@ async function memoryRuntimeContext( }); } -function memoryStore( - context: MemoryToolContext, - runtimeContext: MemoryRuntimeContext, - options: { supersessionDecider?: MemorySupersessionDecider } = {}, -) { - return createMemoryStore(context.db, runtimeContext, { - embedder: context.embedder, - ...(options.supersessionDecider - ? { supersessionDecider: options.supersessionDecider } - : undefined), - }); -} - function boundedLimit(value: number | undefined, fallback: number): number { if (typeof value !== "number" || !Number.isFinite(value)) { return fallback; @@ -369,15 +365,8 @@ function createInput( } satisfies CreateMemoryInput; } -function targetForKind(kind: MemoryKind): "actor" | "conversation" { - if (kind === "preference") { - return "actor"; - } - return "conversation"; -} - /** Return the model-visible projection without hidden ownership/source fields. */ -function compactMemory(memory: MemoryRecord): MemoryToolProjection { +function compactMemory(memory: Memory): MemoryToolProjection { return Value.Parse(memoryToolProjectionSchema, { id: memory.id, content: memory.content, @@ -419,9 +408,6 @@ export function createMemoryCreateTool(context: MemoryCreateToolContext) { const toolCallId = requireToolCallId(options.toolCallId); const requestedExpiresAtMs = parseExpiresAt(parsedInput.expires_at); const runtimeContext = await memoryRuntimeContext(context); - const store = memoryStore(context, runtimeContext, { - supersessionDecider: context.supersessionDecider, - }); const review = await (async () => { try { return parseMemoryReview( @@ -476,10 +462,14 @@ export function createMemoryCreateTool(context: MemoryCreateToolContext) { ); const result = await (async () => { try { - if (targetForKind(review.kind) === "conversation") { - return await store.createConversationMemory(memoryInput); - } - return await store.createMemory(memoryInput); + return await createMemory({ + context: runtimeContext, + db: context.db, + embedder: context.embedder, + input: memoryInput, + subjectType: review.kind === "preference" ? "user" : "conversation", + supersessionDecider: context.supersessionDecider, + }); } catch (error) { asToolInputError(error); } @@ -511,10 +501,48 @@ export function createMemoryRemoveTool(context: MemoryToolContext) { const runtimeContext = await memoryRuntimeContext(context); const memory = await (async () => { try { - return await memoryStore(context, runtimeContext).archiveMemory({ - id: parsedInput.id, - reason: "tool_removed", - }); + const nowMs = Date.now(); + const scopes = deriveVisibleMemoryScopes(runtimeContext).filter( + (scope) => scope.scope === "private", + ); + const predicate = activeVisiblePredicate({ nowMs, scopes }); + const id = parsedInput.id.trim(); + if (!id) throw new Error("Memory id is required."); + const rows = predicate + ? await context.db + .select() + .from(juniorMemoryMemories) + .where( + and( + predicate, + or( + eq(juniorMemoryMemories.id, id), + like(juniorMemoryMemories.id, `${id}%`), + ), + ), + ) + .orderBy(asc(juniorMemoryMemories.id)) + .limit(2) + : []; + if (rows.length === 0) { + throw new Error("Memory was not found in the current context."); + } + if (rows.length > 1) { + throw new Error("Memory id prefix is ambiguous."); + } + const found = parseMemoryRow(rows[0]); + const updated = await context.db + .update(juniorMemoryMemories) + .set({ + archivedAtMs: nowMs, + archiveReason: "tool_removed", + }) + .where(eq(juniorMemoryMemories.id, found.id)) + .returning(); + await context.db + .delete(juniorMemoryEmbeddings) + .where(eq(juniorMemoryEmbeddings.memoryId, found.id)); + return parseMemoryRow(updated[0]); } catch (error) { asToolInputError(error); } @@ -542,11 +570,23 @@ export function createMemoryListTool(context: MemoryToolContext) { execute: async (input) => { const parsedInput = parseMemoryToolInput(listMemoriesInputSchema, input); const runtimeContext = await memoryRuntimeContext(context); - const memories = await memoryStore(context, runtimeContext).listMemories({ - limit: boundedLimit(parsedInput.limit, DEFAULT_RESULT_LIMIT), - }); + const nowMs = Date.now(); + const scopes = deriveVisibleMemoryScopes(runtimeContext); + await archiveExpiredMemoryBatch({ db: context.db, nowMs, scopes }); + const predicate = activeVisiblePredicate({ nowMs, scopes }); + const rows = predicate + ? await context.db + .select() + .from(juniorMemoryMemories) + .where(predicate) + .orderBy( + desc(juniorMemoryMemories.createdAtMs), + asc(juniorMemoryMemories.id), + ) + .limit(boundedLimit(parsedInput.limit, DEFAULT_RESULT_LIMIT)) + : []; return memoryToolResult("listMemories", { - memories: memories.map(compactMemory), + memories: rows.map(parseMemoryRow).map(compactMemory), }); }, }); @@ -571,12 +611,15 @@ export function createMemorySearchTool(context: MemoryToolContext) { input, ); const runtimeContext = await memoryRuntimeContext(context); - const memories = await memoryStore( - context, - runtimeContext, - ).searchMemories({ - query: parsedInput.query, - limit: boundedLimit(parsedInput.limit, DEFAULT_SEARCH_LIMIT), + const memories = await retrieveMemories({ + context: runtimeContext, + db: context.db, + embedder: context.embedder, + input: { + query: parsedInput.query, + limit: boundedLimit(parsedInput.limit, DEFAULT_SEARCH_LIMIT), + }, + mode: "search", }); return memoryToolResult("searchMemories", { memories: memories.map(compactMemory), diff --git a/packages/junior-memory/src/user-pages.ts b/packages/junior-memory/src/user-pages.ts index 7dd21b6ca1..2a78aef2c5 100644 --- a/packages/junior-memory/src/user-pages.ts +++ b/packages/junior-memory/src/user-pages.ts @@ -1,7 +1,7 @@ /** Render memory in Junior's User page format. */ import type { PluginUserPageDefinition } from "@sentry/junior-plugin-api"; import { listMemories, type MemoryVisibility, type MemoryView } from "./viewer"; -import type { MemoryDb } from "./store"; +import type { MemoryDb } from "./memories"; function titleCase(value: string): string { return value.charAt(0).toUpperCase() + value.slice(1); diff --git a/packages/junior-memory/src/viewer.ts b/packages/junior-memory/src/viewer.ts index 5f338c5fff..61cdfe85ae 100644 --- a/packages/junior-memory/src/viewer.ts +++ b/packages/junior-memory/src/viewer.ts @@ -19,7 +19,7 @@ import { import { z } from "zod"; import { juniorMemoryEmbeddings, juniorMemoryMemories } from "./db/schema"; import { publicMemoryScope } from "./scope"; -import { parseMemoryRow, type MemoryDb, type MemoryRecord } from "./store"; +import { parseMemoryRow, type MemoryDb, type Memory } from "./memories"; import { MEMORY_KINDS, type MemorySourcePlatform } from "./types"; const DAY_MS = 24 * 60 * 60 * 1_000; @@ -52,7 +52,7 @@ const timelineDaysSchema = z.number().int().min(1).max(365); export type MemoryVisibility = z.output; /** Memory fields returned to an authenticated User. */ -export type MemoryView = MemoryRecord & { +export type MemoryView = Memory & { origin: "automatic" | "explicit" | "other"; sourcePlatform: MemorySourcePlatform; visibility: MemoryVisibility; diff --git a/packages/junior-memory/tests/memory-operations.ts b/packages/junior-memory/tests/memory-operations.ts new file mode 100644 index 0000000000..79f13318b1 --- /dev/null +++ b/packages/junior-memory/tests/memory-operations.ts @@ -0,0 +1,82 @@ +import { createMemory, type CreateMemoryInput } from "../src/create"; +import type { MemoryDb } from "../src/memories"; +import { retrieveMemories } from "../src/retrieval"; +import type { MemoryRuntimeContext } from "../src/types"; + +type RetrieveMemoriesInput = Parameters[0]["input"]; + +interface MemoryFixture { + context: MemoryRuntimeContext; + db: MemoryDb; + options: Pick< + Parameters[0], + "embedder" | "now" | "supersessionDecider" + >; +} + +/** Bind the database, memory context, and optional test controls. */ +export function memoryFixture( + db: MemoryDb, + context: MemoryRuntimeContext, + options: MemoryFixture["options"] = {}, +): MemoryFixture { + return { context, db, options }; +} + +/** Create a memory about the Conversation. */ +export async function createConversationMemory( + test: MemoryFixture, + input: CreateMemoryInput, +) { + return await createMemory({ + context: test.context, + db: test.db, + ...test.options, + input, + subjectType: "conversation", + }); +} + +/** Create a memory about the User. */ +export async function createUserMemory( + test: MemoryFixture, + input: CreateMemoryInput, +) { + return await createMemory({ + context: test.context, + db: test.db, + ...test.options, + input, + subjectType: "user", + }); +} + +/** Retrieve candidates with automatic recall ranking rules. */ +export async function recallMemories( + test: MemoryFixture, + input: RetrieveMemoriesInput, +) { + return await retrieveMemories({ + context: test.context, + db: test.db, + embedder: test.options.embedder, + input, + mode: "recall", + now: test.options.now, + }); +} + +/** Retrieve candidates with explicit search ranking rules. */ +export async function searchMemories( + test: MemoryFixture, + input: RetrieveMemoriesInput, +) { + return await retrieveMemories({ + context: test.context, + db: test.db, + embedder: test.options.embedder, + input, + mode: "search", + now: test.options.now, + }); +} diff --git a/packages/junior-memory/tests/operational-report.test.ts b/packages/junior-memory/tests/operational-report.test.ts index 52ffac6e8b..c9b581f558 100644 --- a/packages/junior-memory/tests/operational-report.test.ts +++ b/packages/junior-memory/tests/operational-report.test.ts @@ -12,7 +12,12 @@ import { describe, expect, it } from "vitest"; import * as memorySqlSchema from "../src/db/schema"; import { juniorMemoryMemories } from "../src/db/schema"; import { buildMemoryOperationalReport } from "../src/operational-report"; -import { createMemoryStore, type MemoryDb } from "../src/store"; +import type { MemoryDb } from "../src/memories"; +import { + createConversationMemory, + createUserMemory, + memoryFixture, +} from "./memory-operations"; const TEST_NOW_MS = Date.parse("2026-07-28T12:00:00.000Z"); const __dirname = dirname(fileURLToPath(import.meta.url)); @@ -177,16 +182,16 @@ describe("memory operational report", () => { const fixture = await createMemoryFixture(); try { const db = fixture.db(); - const store = createMemoryStore(db, localContext(), { + const test = memoryFixture(db, localContext(), { embedder: testEmbedder(), now: () => TEST_NOW_MS, }); - await store.createMemory({ + await createUserMemory(test, { content: "Use compact pull request summaries.", idempotencyKey: "report-private", kind: "preference", }); - await store.createConversationMemory({ + await createConversationMemory(test, { content: "The checkout runbook lives in the service repository.", idempotencyKey: "report-conversation", kind: "procedure", diff --git a/packages/junior-memory/tests/ranking.test.ts b/packages/junior-memory/tests/ranking.test.ts index cfdbf8ec1b..920c459b86 100644 --- a/packages/junior-memory/tests/ranking.test.ts +++ b/packages/junior-memory/tests/ranking.test.ts @@ -1,10 +1,10 @@ import { describe, expect, it } from "vitest"; import { rankMemoryMatches, type MemoryMatch } from "../src/ranking"; -import type { MemoryRecord } from "../src/store"; +import type { Memory } from "../src/memories"; const NOW_MS = Date.parse("2026-07-28T12:00:00.000Z"); -function memory(id: string, observedAtMs = NOW_MS): MemoryRecord { +function memory(id: string, observedAtMs = NOW_MS): Memory { return { content: `Memory ${id}`, createdAtMs: observedAtMs, diff --git a/packages/junior-memory/tests/retrieval-quality.test.ts b/packages/junior-memory/tests/retrieval-quality.test.ts index d48db07675..13521b9a72 100644 --- a/packages/junior-memory/tests/retrieval-quality.test.ts +++ b/packages/junior-memory/tests/retrieval-quality.test.ts @@ -10,11 +10,13 @@ import { import { createSlackSource } from "@sentry/junior-plugin-api"; import { describe, expect, it } from "vitest"; import * as memorySqlSchema from "../src/db/schema"; +import type { MemoryEmbeddingProvider } from "../src/embeddings"; +import type { MemoryDb } from "../src/memories"; import { - createMemoryStore, - type MemoryDb, - type MemoryEmbeddingProvider, -} from "../src/store"; + createConversationMemory, + memoryFixture, + searchMemories, +} from "./memory-operations"; const EMBEDDING_DIMENSIONS = 1536; const NOW_MS = Date.parse("2026-07-29T12:00:00.000Z"); @@ -104,7 +106,7 @@ describe("memory retrieval quality", () => { [lexicalQuery]: unitEmbedding(1), [semanticQuery]: unitEmbedding(2), }); - const store = createMemoryStore(fixture.db(), runtimeContext(), { + const test = memoryFixture(fixture.db(), runtimeContext(), { embedder, now: () => NOW_MS, }); @@ -117,7 +119,7 @@ describe("memory retrieval quality", () => { ["lexical", runbookContent], ["semantic", semanticContent], ] as const) { - const result = await store.createConversationMemory({ + const result = await createConversationMemory(test, { content, idempotencyKey: `retrieval-quality:${key}`, kind: "knowledge", @@ -132,7 +134,7 @@ describe("memory retrieval quality", () => { ]; const outcomes = await Promise.all( cases.map(async ({ expected, query }) => { - const ranked = await store.searchMemories({ limit: 5, query }); + const ranked = await searchMemories(test, { limit: 5, query }); const rank = ranked.findIndex((memory) => memory.id === expected) + 1; return { hitAt1: rank === 1, diff --git a/packages/junior-memory/tests/storage.test.ts b/packages/junior-memory/tests/storage.test.ts index a1c4b630a0..59d4f5a9f5 100644 --- a/packages/junior-memory/tests/storage.test.ts +++ b/packages/junior-memory/tests/storage.test.ts @@ -22,7 +22,7 @@ import { type Actor, } from "@sentry/junior-plugin-api"; import { Command, CommanderError } from "commander"; -import { eq } from "drizzle-orm"; +import { and, asc, desc, eq, gt, isNull, or } from "drizzle-orm"; import { describe, expect, it, vi } from "vitest"; import * as memorySqlSchema from "../src/db/schema"; import { @@ -43,11 +43,18 @@ import { type MemoryReviewer, } from "../src/tools"; import { listMemories } from "../src/viewer"; -import { createMemoryStore, type MemoryDb } from "../src/store"; import type { MemorySupersessionDecider, MemorySupersessionInput, -} from "../src/store"; +} from "../src/create"; +import type { MemoryDb } from "../src/memories"; +import { + createConversationMemory, + createUserMemory, + memoryFixture, + recallMemories, + searchMemories, +} from "./memory-operations"; const TEST_NOW_MS = Date.parse("2026-06-19T12:00:00.000Z"); const TEST_EMBEDDING_DIMENSIONS = 1536; const __dirname = dirname(fileURLToPath(import.meta.url)); @@ -105,6 +112,31 @@ function memoryDb(fixture: MemoryFixture): MemoryDb { return fixture.db(); } +/** Read active rows for storage assertions without applying access rules. */ +async function readActiveMemoryRows( + fixture: MemoryFixture, + nowMs = Date.now(), +) { + return memoryDb(fixture) + .select() + .from(memorySqlSchema.juniorMemoryMemories) + .where( + and( + isNull(memorySqlSchema.juniorMemoryMemories.archivedAtMs), + isNull(memorySqlSchema.juniorMemoryMemories.supersededAtMs), + isNull(memorySqlSchema.juniorMemoryMemories.supersededById), + or( + isNull(memorySqlSchema.juniorMemoryMemories.expiresAtMs), + gt(memorySqlSchema.juniorMemoryMemories.expiresAtMs, nowMs), + ), + ), + ) + .orderBy( + desc(memorySqlSchema.juniorMemoryMemories.createdAtMs), + asc(memorySqlSchema.juniorMemoryMemories.id), + ); +} + async function runMemoryCli(fixture: MemoryFixture, argv: string[]) { const stdout: string[] = []; const stderr: string[] = []; @@ -1009,12 +1041,12 @@ describe("memory plugin storage", () => { const fixture = await createMemoryFixture(); try { - const store = createMemoryStore( + const store = memoryFixture( memoryDb(fixture), slackContext({ channelId: "D123" }), { now: () => TEST_NOW_MS }, ); - const oldMemory = await store.createMemory({ + const oldMemory = await createUserMemory(store, { content: "Prefers Python for automation scripts.", kind: "preference", idempotencyKey: "memory-test:passive-supersession-old", @@ -1438,12 +1470,12 @@ describe("memory plugin storage", () => { const fixture = await createMemoryFixture(); try { const duplicateContent = "Deployment runbooks live in Notion."; - const store = createMemoryStore( + const store = memoryFixture( memoryDb(fixture), slackContext({ channelId: "D123" }), { now: () => TEST_NOW_MS }, ); - await store.createConversationMemory({ + await createConversationMemory(store, { content: duplicateContent, idempotencyKey: "memory-test:existing-conversation-fact", kind: "knowledge", @@ -1757,10 +1789,10 @@ describe("memory plugin storage", () => { try { await installViewerCoreTables(fixture); - const store = createMemoryStore(memoryDb(fixture), runtime, { + const store = memoryFixture(memoryDb(fixture), runtime, { now: () => TEST_NOW_MS, }); - const created = await store.createMemory({ + const created = await createUserMemory(store, { content: "Prefers short dashboard answers.", idempotencyKey: "tool:api:personal-scope", kind: "preference", @@ -2320,29 +2352,27 @@ describe("memory plugin storage", () => { } }, 15_000); - it("persists, recalls, and archives visible memories", async () => { + it("lists and archives visible memories through tools", async () => { const fixture = await createMemoryFixture(); + let nowMs = TEST_NOW_MS; + const now = vi.spyOn(Date, "now").mockImplementation(() => nowMs); try { - let nowMs = TEST_NOW_MS; - const publicStore = createMemoryStore(memoryDb(fixture), slackContext(), { + const publicStore = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, }); - const publicMemory = await publicStore.createConversationMemory({ + const publicMemory = await createConversationMemory(publicStore, { content: "Deploy runbooks live in Notion.", kind: "knowledge", idempotencyKey: "memory-test:public", }); nowMs += 1; const privateContext = slackContext({ channelId: "D123" }); - const privateStore = createMemoryStore( - memoryDb(fixture), - privateContext, - { - now: () => nowMs, - }, - ); - const privateMemory = await privateStore.createMemory({ + const privateStore = memoryFixture(memoryDb(fixture), privateContext, { + embedder: createTestEmbedder(), + now: () => nowMs, + }); + const privateMemory = await createUserMemory(privateStore, { content: "Prefers short PR summaries.", kind: "preference", idempotencyKey: "memory-test:private", @@ -2356,65 +2386,66 @@ describe("memory plugin storage", () => { scope: "private", subjectType: "user", }); - await expect(privateStore.listMemories({})).resolves.toEqual([ - expect.objectContaining({ id: privateMemory.memory.id }), - expect.objectContaining({ id: publicMemory.memory.id }), - ]); - const sameUserStore = createMemoryStore( - memoryDb(fixture), - slackContext({ - ownerUserId: privateContext.userId, - channelId: "D999", - userId: "U456", - }), - { now: () => nowMs }, - ); - await expect(sameUserStore.listMemories({})).resolves.toEqual([ - expect.objectContaining({ id: privateMemory.memory.id }), - expect.objectContaining({ id: publicMemory.memory.id }), - ]); - - const otherUserStore = createMemoryStore( - memoryDb(fixture), - slackContext({ - channelId: "D123", - threadTs: "1718800001.000000", - userId: "U456", - }), - { now: () => nowMs }, - ); - await expect(otherUserStore.listMemories({})).resolves.toEqual([ - expect.objectContaining({ id: publicMemory.memory.id }), - ]); - const linkedIdentityStore = createMemoryStore( - memoryDb(fixture), - slackContext({ - ownerUserId: privateContext.userId, - teamId: "T999", - userId: "U456", - }), - { now: () => nowMs }, - ); - await expect(linkedIdentityStore.listMemories({})).resolves.toEqual([ - expect.objectContaining({ id: privateMemory.memory.id }), - expect.objectContaining({ id: publicMemory.memory.id }), - ]); + const listFor = (runtime: ReturnType) => + createMemoryListTool({ + agent: allowMemory("actor"), + db: memoryDb(fixture), + ...runtime, + users: memoryUsers(runtime.userId), + }); + const removeFor = (runtime: ReturnType) => + createMemoryRemoveTool({ + agent: allowMemory("actor"), + db: memoryDb(fixture), + ...runtime, + users: memoryUsers(runtime.userId), + }); + const otherUserContext = slackContext({ + channelId: "D123", + threadTs: "1718800001.000000", + userId: "U456", + }); + await expect( + listFor(otherUserContext).execute({}, {}), + ).resolves.toMatchObject({ + memories: [expect.objectContaining({ id: publicMemory.memory.id })], + }); + const linkedIdentityContext = slackContext({ + ownerUserId: privateContext.userId, + teamId: "T999", + userId: "U456", + }); await expect( - otherUserStore.archiveMemory({ id: privateMemory.memory.id }), + listFor(linkedIdentityContext).execute({}, {}), + ).resolves.toMatchObject({ + memories: [ + expect.objectContaining({ id: privateMemory.memory.id }), + expect.objectContaining({ id: publicMemory.memory.id }), + ], + }); + await expect( + removeFor(otherUserContext).execute( + { id: privateMemory.memory.id }, + {}, + ), ).rejects.toThrow("Memory was not found in the current context."); await expect( - publicStore.archiveMemory({ id: publicMemory.memory.id }), + removeFor(slackContext()).execute({ id: publicMemory.memory.id }, {}), ).rejects.toThrow("Memory was not found in the current context."); nowMs += 1; - const archived = await publicStore.archiveMemory({ - id: privateMemory.memory.id.slice(0, 12), - }); + const archived = await removeFor(linkedIdentityContext).execute( + { id: privateMemory.memory.id.slice(0, 12) }, + {}, + ); expect(archived).toMatchObject({ - id: privateMemory.memory.id, - archivedAtMs: nowMs, + memory: { id: privateMemory.memory.id }, }); + await expect( + memoryDb(fixture).select().from(memorySqlSchema.juniorMemoryEmbeddings), + ).resolves.toEqual([]); } finally { + now.mockRestore(); await fixture.close(); } }, 15_000); @@ -2438,41 +2469,39 @@ describe("memory plugin storage", () => { teamId: "T123", userId: "U123", }); - const firstStore = createMemoryStore(memoryDb(fixture), firstContext, { + const firstStore = memoryFixture(memoryDb(fixture), firstContext, { now: () => TEST_NOW_MS, }); - const secondStore = createMemoryStore(memoryDb(fixture), secondContext, { + const secondStore = memoryFixture(memoryDb(fixture), secondContext, { now: () => TEST_NOW_MS + 1, }); - const hiddenStore = createMemoryStore(memoryDb(fixture), hiddenContext, { + const hiddenStore = memoryFixture(memoryDb(fixture), hiddenContext, { now: () => TEST_NOW_MS + 2, }); - const privateStore = createMemoryStore( - memoryDb(fixture), - privateContext, - { now: () => TEST_NOW_MS + 3 }, - ); - const first = await firstStore.createMemory({ + const privateStore = memoryFixture(memoryDb(fixture), privateContext, { + now: () => TEST_NOW_MS + 3, + }); + const first = await createUserMemory(firstStore, { content: "Prefers concise release notes.", idempotencyKey: "tool:api:first", kind: "preference", }); - const second = await secondStore.createMemory({ + const second = await createUserMemory(secondStore, { content: "Deploy runbooks live in Notion.", idempotencyKey: "session:api:second", kind: "knowledge", }); - const hidden = await hiddenStore.createMemory({ + const hidden = await createUserMemory(hiddenStore, { content: "Hidden viewer memory.", idempotencyKey: "api:hidden", kind: "knowledge", }); - const publicMemory = await firstStore.createConversationMemory({ + const publicMemory = await createConversationMemory(firstStore, { content: "Public workspace memory.", idempotencyKey: "session:api:public", kind: "knowledge", }); - const privateMemory = await privateStore.createConversationMemory({ + const privateMemory = await createConversationMemory(privateStore, { content: "Private conversation memory.", idempotencyKey: "session:api:private", kind: "knowledge", @@ -2687,21 +2716,21 @@ describe("memory plugin storage", () => { try { const context = localContext({ userId: "cli-user" }); - const store = createMemoryStore(memoryDb(fixture), context, { + const store = memoryFixture(memoryDb(fixture), context, { now: () => TEST_NOW_MS, }); - const created = await store.createMemory({ + const created = await createUserMemory(store, { content: "Prefers CLI memory QA with scoped search.", kind: "preference", idempotencyKey: "memory-test:cli-search", }); - const expired = await store.createMemory({ + const expired = await createUserMemory(store, { content: "Prefers expired CLI memory rows to stay hidden.", kind: "preference", expiresAtMs: Date.now() - 1, idempotencyKey: "memory-test:cli-search-expired", }); - const superseded = await store.createMemory({ + const superseded = await createUserMemory(store, { content: "Prefers superseded CLI memory rows to stay hidden.", kind: "preference", idempotencyKey: "memory-test:cli-search-superseded", @@ -2819,18 +2848,18 @@ WHERE id = '${superseded.memory.id}' "client rendering library": unitEmbedding(1), }); let nowMs = TEST_NOW_MS; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => nowMs, }); - const react = await store.createMemory({ + const react = await createUserMemory(store, { content: reactMemory, kind: "preference", idempotencyKey: "memory-test:embedding-react", }); nowMs += 1; - await store.createMemory({ + await createUserMemory(store, { content: mangoMemory, kind: "preference", idempotencyKey: "memory-test:embedding-mango", @@ -2851,7 +2880,7 @@ WHERE id = '${superseded.memory.id}' }), ]), ); - const results = await store.searchMemories({ + const results = await searchMemories(store, { query: "client rendering library", }); expect(results[0]).toEqual( @@ -2874,28 +2903,30 @@ WHERE id = '${superseded.memory.id}' [closeContent]: cosineEmbedding(0.8), [weakContent]: cosineEmbedding(0.5), }); - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => TEST_NOW_MS, }); - const close = await store.createMemory({ + const close = await createUserMemory(store, { content: closeContent, kind: "knowledge", idempotencyKey: "memory-test:recall-distance-close", }); - const weak = await store.createMemory({ + const weak = await createUserMemory(store, { content: weakContent, kind: "knowledge", idempotencyKey: "memory-test:recall-distance-weak", }); - await expect(store.recallMemories({ limit: 2, query })).resolves.toEqual([ - expect.objectContaining({ id: close.memory.id }), - ]); - await expect(store.searchMemories({ limit: 2, query })).resolves.toEqual([ - expect.objectContaining({ id: close.memory.id }), - expect.objectContaining({ id: weak.memory.id }), - ]); + await expect(recallMemories(store, { limit: 2, query })).resolves.toEqual( + [expect.objectContaining({ id: close.memory.id })], + ); + await expect(searchMemories(store, { limit: 2, query })).resolves.toEqual( + [ + expect.objectContaining({ id: close.memory.id }), + expect.objectContaining({ id: weak.memory.id }), + ], + ); } finally { await fixture.close(); } @@ -2914,16 +2945,16 @@ WHERE id = '${superseded.memory.id}' [vectorDistractor]: cosineEmbedding(0.8), [lexicalAnswer]: unitEmbedding(1), }); - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => TEST_NOW_MS, }); - const distractor = await store.createMemory({ + const distractor = await createUserMemory(store, { content: vectorDistractor, kind: "preference", idempotencyKey: "memory-test:recall-vector-distractor", }); - const answer = await store.createMemory({ + const answer = await createUserMemory(store, { content: lexicalAnswer, kind: "knowledge", idempotencyKey: "memory-test:recall-lexical-answer", @@ -2931,7 +2962,7 @@ WHERE id = '${superseded.memory.id}' // Hybrid recall must keep both legs. A close vector hit must not hide // the exact/token memory that only lexical retrieval surfaces. - await expect(store.recallMemories({ limit: 2, query })).resolves.toEqual( + await expect(recallMemories(store, { limit: 2, query })).resolves.toEqual( expect.arrayContaining([ expect.objectContaining({ id: distractor.memory.id }), expect.objectContaining({ id: answer.memory.id }), @@ -2953,7 +2984,7 @@ WHERE id = '${superseded.memory.id}' const preferenceContent = "Located in San Francisco and uses Pacific Time (PT)."; let nowMs = TEST_NOW_MS; - const privateStore = createMemoryStore( + const privateStore = memoryFixture( memoryDb(fixture), slackContext({ channelId: "D123" }), { @@ -2961,11 +2992,11 @@ WHERE id = '${superseded.memory.id}' now: () => nowMs, }, ); - const publicStore = createMemoryStore(memoryDb(fixture), slackContext(), { + const publicStore = memoryFixture(memoryDb(fixture), slackContext(), { // No embedder: force the pure lexical path that production noise hits. now: () => nowMs, }); - const preference = await privateStore.createMemory({ + const preference = await createUserMemory(privateStore, { content: preferenceContent, kind: "preference", idempotencyKey: "memory-test:recall-personal-timezone", @@ -2973,7 +3004,7 @@ WHERE id = '${superseded.memory.id}' for (let index = 0; index < 80; index += 1) { nowMs = TEST_NOW_MS + index + 1; - await publicStore.createConversationMemory({ + await createConversationMemory(publicStore, { content: `Recent workspace time note ${index} about deploy time windows`, kind: "knowledge", idempotencyKey: `memory-test:recall-time-noise-${index}`, @@ -2984,7 +3015,7 @@ WHERE id = '${superseded.memory.id}' // Public lexical recall alone would keep only the newest noise. The // The separate private search must still find the older preference. await expect( - privateStore.recallMemories({ limit: 5, query }), + recallMemories(privateStore, { limit: 5, query }), ).resolves.toEqual( expect.arrayContaining([ expect.objectContaining({ @@ -3018,34 +3049,30 @@ WHERE id = '${superseded.memory.id}' } const embedder = createTestEmbedder(vectors); let nowMs = TEST_NOW_MS; - const vectorStore = createMemoryStore(memoryDb(fixture), slackContext(), { + const vectorStore = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => nowMs, }); for (const [index, memory] of vectorMemories.entries()) { nowMs += 1; - await vectorStore.createMemory({ + await createUserMemory(vectorStore, { content: memory, kind: "preference", idempotencyKey: `memory-test:fusion-vector-${index}`, }); } nowMs += 1; - const lexicalStore = createMemoryStore( - memoryDb(fixture), - slackContext(), - { - now: () => nowMs, - }, - ); - const lexical = await lexicalStore.createMemory({ + const lexicalStore = memoryFixture(memoryDb(fixture), slackContext(), { + now: () => nowMs, + }); + const lexical = await createUserMemory(lexicalStore, { content: lexicalMemory, kind: "preference", idempotencyKey: "memory-test:fusion-lexical", }); await expect( - vectorStore.searchMemories({ limit: 1, query }), + searchMemories(vectorStore, { limit: 1, query }), ).resolves.toEqual([expect.objectContaining({ id: lexical.memory.id })]); } finally { await fixture.close(); @@ -3057,24 +3084,24 @@ WHERE id = '${superseded.memory.id}' try { let nowMs = TEST_NOW_MS - 120 * 24 * 60 * 60 * 1000; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, }); - const oldRelevant = await store.createMemory({ + const oldRelevant = await createUserMemory(store, { content: "Deploy checklist ownership escalation requires release approval.", kind: "knowledge", idempotencyKey: "memory-test:recency-old-relevant", }); nowMs = TEST_NOW_MS; - const newLessRelevant = await store.createMemory({ + const newLessRelevant = await createUserMemory(store, { content: "Deploy snacks live near the office checklist.", kind: "knowledge", idempotencyKey: "memory-test:recency-new-less-relevant", }); await expect( - store.searchMemories({ + searchMemories(store, { limit: 2, query: "deploy checklist ownership escalation approval", }), @@ -3092,23 +3119,23 @@ WHERE id = '${superseded.memory.id}' try { let nowMs = TEST_NOW_MS - 120 * 24 * 60 * 60 * 1000; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, }); - const oldMemory = await store.createMemory({ + const oldMemory = await createUserMemory(store, { content: "Deploy checklist lives in the legacy wiki.", kind: "knowledge", idempotencyKey: "memory-test:recency-old-equal", }); nowMs = TEST_NOW_MS; - const newMemory = await store.createMemory({ + const newMemory = await createUserMemory(store, { content: "Deploy checklist lives in Notion.", kind: "knowledge", idempotencyKey: "memory-test:recency-new-equal", }); await expect( - store.searchMemories({ limit: 2, query: "deploy checklist" }), + searchMemories(store, { limit: 2, query: "deploy checklist" }), ).resolves.toEqual([ expect.objectContaining({ id: newMemory.memory.id }), expect.objectContaining({ id: oldMemory.memory.id }), @@ -3131,26 +3158,28 @@ WHERE id = '${superseded.memory.id}' [newContent]: cosineEmbedding(0.9), }); let nowMs = TEST_NOW_MS - 120 * 24 * 60 * 60 * 1000; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => nowMs, }); - const oldMemory = await store.createMemory({ + const oldMemory = await createUserMemory(store, { content: oldContent, kind: "knowledge", idempotencyKey: "memory-test:recency-vector-old", }); nowMs = TEST_NOW_MS; - const newMemory = await store.createMemory({ + const newMemory = await createUserMemory(store, { content: newContent, kind: "knowledge", idempotencyKey: "memory-test:recency-vector-new", }); - await expect(store.searchMemories({ limit: 2, query })).resolves.toEqual([ - expect.objectContaining({ id: oldMemory.memory.id }), - expect.objectContaining({ id: newMemory.memory.id }), - ]); + await expect(searchMemories(store, { limit: 2, query })).resolves.toEqual( + [ + expect.objectContaining({ id: oldMemory.memory.id }), + expect.objectContaining({ id: newMemory.memory.id }), + ], + ); } finally { await fixture.close(); } @@ -3161,18 +3190,18 @@ WHERE id = '${superseded.memory.id}' try { const embedder = createTestEmbedder(); - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => TEST_NOW_MS, }); - const created = await store.createMemory({ + const created = await createUserMemory(store, { content: "Prefers duplicate-safe vector writes.", kind: "preference", idempotencyKey: "memory-test:embedding-idempotent", }); await expect( - store.createMemory({ + createUserMemory(store, { content: "Changed retry content should not be re-embedded.", kind: "preference", idempotencyKey: "memory-test:embedding-idempotent", @@ -3196,17 +3225,17 @@ WHERE id = '${superseded.memory.id}' try { let nowMs = TEST_NOW_MS; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, }); - const created = await store.createMemory({ + const created = await createUserMemory(store, { content: "Prefers release notes with risk callouts.", kind: "preference", idempotencyKey: "memory-test:exact-dedup-original", }); await expect( - store.createMemory({ + createUserMemory(store, { content: " Prefers release notes\nwith risk callouts. ", kind: "preference", expiresAtMs: TEST_NOW_MS + 1, @@ -3218,7 +3247,7 @@ WHERE id = '${superseded.memory.id}' }); nowMs = TEST_NOW_MS + 2; await expect( - store.createMemory({ + createUserMemory(store, { content: "Changed retry content must not create a duplicate.", kind: "preference", idempotencyKey: "memory-test:exact-dedup-repeat", @@ -3227,7 +3256,7 @@ WHERE id = '${superseded.memory.id}' created: false, memory: { id: created.memory.id }, }); - const otherKind = await store.createMemory({ + const otherKind = await createUserMemory(store, { content: "Prefers release notes with risk callouts.", kind: "knowledge", idempotencyKey: "memory-test:exact-dedup-other-kind", @@ -3236,7 +3265,7 @@ WHERE id = '${superseded.memory.id}' created: true, memory: { content: created.memory.content, kind: "knowledge" }, }); - const otherScope = await store.createConversationMemory({ + const otherScope = await createConversationMemory(store, { content: "Prefers release notes with risk callouts.", kind: "preference", idempotencyKey: "memory-test:exact-dedup-other-scope", @@ -3246,14 +3275,16 @@ WHERE id = '${superseded.memory.id}' memory: { content: created.memory.content, kind: "preference" }, }); - await expect(store.listMemories({})).resolves.toEqual( + await expect(readActiveMemoryRows(fixture, nowMs)).resolves.toEqual( expect.arrayContaining([ expect.objectContaining({ id: otherScope.memory.id }), expect.objectContaining({ id: otherKind.memory.id }), expect.objectContaining({ id: created.memory.id }), ]), ); - await expect(store.listMemories({})).resolves.toHaveLength(3); + await expect(readActiveMemoryRows(fixture, nowMs)).resolves.toHaveLength( + 3, + ); } finally { await fixture.close(); } @@ -3269,17 +3300,17 @@ WHERE id = '${superseded.memory.id}' [firstContent]: unitEmbedding(1), [duplicateContent]: unitEmbedding(1), }); - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => TEST_NOW_MS, }); - const first = await store.createMemory({ + const first = await createUserMemory(store, { content: firstContent, kind: "preference", idempotencyKey: "memory-test:vector-dedup-idempotent-original", }); - const second = await store.createMemory({ + const second = await createUserMemory(store, { content: duplicateContent, kind: "preference", idempotencyKey: "memory-test:vector-dedup-idempotent-repeat", @@ -3290,7 +3321,7 @@ WHERE id = '${superseded.memory.id}' }); expect(second.memory.id).not.toBe(first.memory.id); await expect( - store.createMemory({ + createUserMemory(store, { content: "Changed retry content should resolve to its original write.", kind: "preference", @@ -3301,7 +3332,9 @@ WHERE id = '${superseded.memory.id}' memory: { id: second.memory.id, content: duplicateContent }, }); - await expect(store.listMemories({})).resolves.toHaveLength(2); + await expect( + readActiveMemoryRows(fixture, TEST_NOW_MS), + ).resolves.toHaveLength(2); } finally { await fixture.close(); } @@ -3318,17 +3351,17 @@ WHERE id = '${superseded.memory.id}' [firstContent]: unitEmbedding(1), [duplicateContent]: unitEmbedding(1), }); - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => TEST_NOW_MS, }); - const created = await store.createMemory({ + const created = await createUserMemory(store, { content: firstContent, kind: "knowledge", idempotencyKey: "memory-test:vector-dedup-original", }); - const neighbor = await store.createMemory({ + const neighbor = await createUserMemory(store, { content: duplicateContent, kind: "knowledge", idempotencyKey: "memory-test:vector-dedup-repeat", @@ -3339,7 +3372,9 @@ WHERE id = '${superseded.memory.id}' }); expect(neighbor.memory.id).not.toBe(created.memory.id); - await expect(store.listMemories({})).resolves.toHaveLength(2); + await expect( + readActiveMemoryRows(fixture, TEST_NOW_MS), + ).resolves.toHaveLength(2); await expect( memoryDb(fixture).select().from(memorySqlSchema.juniorMemoryEmbeddings), ).resolves.toHaveLength(2); @@ -3353,10 +3388,10 @@ WHERE id = '${superseded.memory.id}' try { const content = "Prefers derived embeddings to be repairable."; - const firstStore = createMemoryStore(memoryDb(fixture), slackContext(), { + const firstStore = memoryFixture(memoryDb(fixture), slackContext(), { now: () => TEST_NOW_MS, }); - const created = await firstStore.createMemory({ + const created = await createUserMemory(firstStore, { content, kind: "preference", idempotencyKey: "memory-test:embedding-retry-backfill", @@ -3366,12 +3401,12 @@ WHERE id = '${superseded.memory.id}' ).resolves.toEqual([]); const embedder = createTestEmbedder(); - const retryStore = createMemoryStore(memoryDb(fixture), slackContext(), { + const retryStore = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => TEST_NOW_MS + 1, }); await expect( - retryStore.createMemory({ + createUserMemory(retryStore, { content: "Changed retry content should not be embedded.", kind: "preference", idempotencyKey: "memory-test:embedding-retry-backfill", @@ -3394,6 +3429,8 @@ WHERE id = '${superseded.memory.id}' it("archives expired visible memories during reads", async () => { const fixture = await createMemoryFixture(); + let nowMs = TEST_NOW_MS; + const now = vi.spyOn(Date, "now").mockImplementation(() => nowMs); try { const expiredContent = "Temporary CLI memory should expire cleanly."; @@ -3404,24 +3441,24 @@ WHERE id = '${superseded.memory.id}' [activeContent]: unitEmbedding(2), [supersededContent]: unitEmbedding(3), }); - let nowMs = TEST_NOW_MS; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const runtime = slackContext(); + const store = memoryFixture(memoryDb(fixture), runtime, { embedder, now: () => nowMs, }); - const expired = await store.createMemory({ + const expired = await createUserMemory(store, { content: expiredContent, kind: "preference", expiresAtMs: TEST_NOW_MS + 10, idempotencyKey: "memory-test:read-expired", }); - const active = await store.createMemory({ + const active = await createUserMemory(store, { content: activeContent, kind: "preference", idempotencyKey: "memory-test:read-active", }); - const superseded = await store.createMemory({ + const superseded = await createUserMemory(store, { content: supersededContent, kind: "preference", expiresAtMs: TEST_NOW_MS + 10, @@ -3438,9 +3475,16 @@ WHERE id = '${superseded.memory.id}' ).resolves.toHaveLength(3); nowMs = TEST_NOW_MS + 11; - await expect(store.listMemories({})).resolves.toEqual([ - expect.objectContaining({ id: active.memory.id }), - ]); + await expect( + createMemoryListTool({ + agent: allowMemory("actor"), + db: memoryDb(fixture), + ...runtime, + users: memoryUsers(runtime.userId), + }).execute({}, {}), + ).resolves.toMatchObject({ + memories: [expect.objectContaining({ id: active.memory.id })], + }); await expect( memoryDb(fixture) .select() @@ -3480,6 +3524,7 @@ WHERE id = '${superseded.memory.id}' }), ]); } finally { + now.mockRestore(); await fixture.close(); } }, 15_000); @@ -3492,12 +3537,12 @@ WHERE id = '${superseded.memory.id}' { "Prefers lexical fallback for vector failures.": [1, 0, 0] }, { dimensions: 3 }, ); - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => TEST_NOW_MS, }); - const created = await store.createMemory({ + const created = await createUserMemory(store, { content: "Prefers lexical fallback for vector failures.", kind: "preference", idempotencyKey: "memory-test:embedding-dimension-mismatch", @@ -3507,7 +3552,7 @@ WHERE id = '${superseded.memory.id}' memoryDb(fixture).select().from(memorySqlSchema.juniorMemoryEmbeddings), ).resolves.toEqual([]); await expect( - store.searchMemories({ query: "lexical fallback" }), + searchMemories(store, { query: "lexical fallback" }), ).resolves.toEqual([expect.objectContaining({ id: created.memory.id })]); } finally { await fixture.close(); @@ -3787,37 +3832,40 @@ WHERE id = '${superseded.memory.id}' try { let nowMs = TEST_NOW_MS; const context = slackContext({ channelId: "D123" }); - const store = createMemoryStore(memoryDb(fixture), context, { + const store = memoryFixture(memoryDb(fixture), context, { now: () => nowMs, }); - const personal = await store.createMemory({ + const personal = await createUserMemory(store, { content: "Prefers PR summaries with risks first.", kind: "preference", idempotencyKey: "memory-test:recall-personal", }); nowMs += 1; - const conversation = await store.createConversationMemory({ + const conversation = await createConversationMemory(store, { content: "Release notes live in Notion.", kind: "knowledge", idempotencyKey: "memory-test:recall-conversation", }); nowMs += 1; - await store.createMemory({ + await createUserMemory(store, { content: "Prefers PR summary obsolete wording.", kind: "preference", expiresAtMs: TEST_NOW_MS + 1, idempotencyKey: "memory-test:recall-expired", }); nowMs += 1; - await createMemoryStore( - memoryDb(fixture), - slackContext({ channelId: "D456", userId: "U456" }), - { now: () => nowMs }, - ).createMemory({ - content: "Prefers PR summary unrelated owner.", - kind: "preference", - idempotencyKey: "memory-test:recall-other-user", - }); + await createUserMemory( + memoryFixture( + memoryDb(fixture), + slackContext({ channelId: "D456", userId: "U456" }), + { now: () => nowMs }, + ), + { + content: "Prefers PR summary unrelated owner.", + kind: "preference", + idempotencyKey: "memory-test:recall-other-user", + }, + ); const emitted: PluginConversationEventValue[] = []; const plugin = memoryPlugin(); @@ -3896,13 +3944,13 @@ WHERE id = '${superseded.memory.id}' try { let nowMs = TEST_NOW_MS; const context = slackContext(); - const store = createMemoryStore(memoryDb(fixture), context, { + const store = memoryFixture(memoryDb(fixture), context, { now: () => nowMs, }); const created = []; for (let index = 0; index < 6; index += 1) { created.push( - await store.createConversationMemory({ + await createConversationMemory(store, { content: `Deploy step ${index + 1} uses checklist item ${index + 1}.`, kind: "procedure", idempotencyKey: `memory-test:recall-budget-${index}`, @@ -3962,13 +4010,16 @@ WHERE id = '${superseded.memory.id}' try { const context = slackContext(); - const conversation = await createMemoryStore(memoryDb(fixture), context, { - now: () => TEST_NOW_MS, - }).createConversationMemory({ - content: "Release notes live in Notion.", - kind: "knowledge", - idempotencyKey: "memory-test:recall-conversation-context", - }); + const conversation = await createConversationMemory( + memoryFixture(memoryDb(fixture), context, { + now: () => TEST_NOW_MS, + }), + { + content: "Release notes live in Notion.", + kind: "knowledge", + idempotencyKey: "memory-test:recall-conversation-context", + }, + ); const plugin = memoryPlugin(); const result = await plugin.hooks?.userPrompt?.({ @@ -4022,20 +4073,20 @@ WHERE id = '${superseded.memory.id}' try { const context = slackContext(); - const store = createMemoryStore(memoryDb(fixture), context, { + const store = memoryFixture(memoryDb(fixture), context, { now: () => TEST_NOW_MS, }); - const relevant = await store.createConversationMemory({ + const relevant = await createConversationMemory(store, { content: "getsentry/junior CI runs package tests with pnpm.", kind: "knowledge", idempotencyKey: "memory-test:recall-gate-relevant", }); - await store.createConversationMemory({ + await createConversationMemory(store, { content: "getsentry/sentry autofix PR tests use a dashboard workflow.", kind: "knowledge", idempotencyKey: "memory-test:recall-gate-vocabulary", }); - await store.createConversationMemory({ + await createConversationMemory(store, { content: "Single-tenant repository access is configured in the admin dashboard.", kind: "knowledge", @@ -4078,13 +4129,17 @@ WHERE id = '${superseded.memory.id}' try { const context = slackContext(); - await createMemoryStore(memoryDb(fixture), context, { - now: () => TEST_NOW_MS, - }).createConversationMemory({ - content: "getsentry/sentry autofix PR tests use a dashboard workflow.", - kind: "knowledge", - idempotencyKey: "memory-test:recall-gate-empty", - }); + await createConversationMemory( + memoryFixture(memoryDb(fixture), context, { + now: () => TEST_NOW_MS, + }), + { + content: + "getsentry/sentry autofix PR tests use a dashboard workflow.", + kind: "knowledge", + idempotencyKey: "memory-test:recall-gate-empty", + }, + ); const emitted: PluginConversationEventValue[] = []; const plugin = memoryPlugin(); @@ -4178,13 +4233,16 @@ WHERE id = '${superseded.memory.id}' try { const context = slackContext(); - await createMemoryStore(memoryDb(fixture), context, { - now: () => TEST_NOW_MS, - }).createConversationMemory({ - content: "Release notes live in Notion.", - kind: "knowledge", - idempotencyKey: "memory-test:recall-gate-failure", - }); + await createConversationMemory( + memoryFixture(memoryDb(fixture), context, { + now: () => TEST_NOW_MS, + }), + { + content: "Release notes live in Notion.", + kind: "knowledge", + idempotencyKey: "memory-test:recall-gate-failure", + }, + ); const plugin = memoryPlugin(); await expect( @@ -4216,7 +4274,7 @@ WHERE id = '${superseded.memory.id}' try { const context = slackContext(); - await createMemoryStore(memoryDb(fixture), context).createMemory({ + await createUserMemory(memoryFixture(memoryDb(fixture), context), { content: "Prefers PR summaries with risks first.", kind: "preference", idempotencyKey: "memory-test:recall-blank", @@ -4253,14 +4311,17 @@ WHERE id = '${superseded.memory.id}' [memory]: unitEmbedding(1), [query]: unitEmbedding(1), }); - await createMemoryStore(memoryDb(fixture), context, { - embedder, - now: () => TEST_NOW_MS, - }).createMemory({ - content: memory, - kind: "preference", - idempotencyKey: "memory-test:recall-semantic", - }); + await createUserMemory( + memoryFixture(memoryDb(fixture), context, { + embedder, + now: () => TEST_NOW_MS, + }), + { + content: memory, + kind: "preference", + idempotencyKey: "memory-test:recall-semantic", + }, + ); const plugin = memoryPlugin(); const result = await plugin.hooks?.userPrompt?.({ @@ -4324,9 +4385,7 @@ WHERE id = '${superseded.memory.id}' ), ).resolves.toMatchObject({ created: true }); - await expect( - createMemoryStore(memoryDb(fixture), slackContext()).listMemories({}), - ).resolves.toEqual([ + await expect(readActiveMemoryRows(fixture)).resolves.toEqual([ expect.objectContaining({ content: "Prefers the second remembered fact.", }), @@ -4339,58 +4398,24 @@ WHERE id = '${superseded.memory.id}' } }, 15_000); - it("keeps private local memory with its User", async () => { + it("searches memory created in a local conversation", async () => { const fixture = await createMemoryFixture(); try { - let nowMs = TEST_NOW_MS; - const store = createMemoryStore(memoryDb(fixture), localContext(), { - now: () => nowMs, - }); - - const personal = await store.createMemory({ - content: "Prefers local CLI memory checks.", - kind: "preference", - idempotencyKey: "memory-test:local-personal", + const store = memoryFixture(memoryDb(fixture), localContext(), { + now: () => TEST_NOW_MS, }); - nowMs = TEST_NOW_MS + 1; - const conversation = await store.createConversationMemory({ + const conversation = await createConversationMemory(store, { content: "Memory plugin validation is tracked in this local session.", kind: "knowledge", idempotencyKey: "memory-test:local-conversation", }); - await expect(store.listMemories({})).resolves.toEqual([ - expect.objectContaining({ id: conversation.memory.id }), - expect.objectContaining({ id: personal.memory.id }), - ]); await expect( - store.searchMemories({ query: "validation" }), + searchMemories(store, { query: "validation" }), ).resolves.toEqual([ expect.objectContaining({ id: conversation.memory.id }), ]); - - const otherConversationStore = createMemoryStore( - memoryDb(fixture), - localContext({ conversationId: "local:junior:other-memory-test" }), - { now: () => nowMs }, - ); - await expect(otherConversationStore.listMemories({})).resolves.toEqual([ - expect.objectContaining({ id: conversation.memory.id }), - expect.objectContaining({ id: personal.memory.id }), - ]); - await expect( - otherConversationStore.archiveMemory({ id: conversation.memory.id }), - ).resolves.toMatchObject({ id: conversation.memory.id }); - - nowMs = TEST_NOW_MS + 2; - const archived = await otherConversationStore.archiveMemory({ - id: personal.memory.id, - }); - expect(archived).toMatchObject({ - archivedAtMs: TEST_NOW_MS + 2, - id: personal.memory.id, - }); } finally { await fixture.close(); } @@ -4402,11 +4427,11 @@ WHERE id = '${superseded.memory.id}' try { let nowMs = TEST_NOW_MS; const context = slackContext({ channelId: "D123" }); - const store = createMemoryStore(memoryDb(fixture), context, { + const store = memoryFixture(memoryDb(fixture), context, { now: () => nowMs, }); - const created = await store.createMemory({ + const created = await createUserMemory(store, { content: "Different content with the same retry key.", kind: "preference", idempotencyKey: "explicit-create-1", @@ -4415,7 +4440,7 @@ WHERE id = '${superseded.memory.id}' nowMs = TEST_NOW_MS + 1; await expect( - store.createMemory({ + createUserMemory(store, { content: "Changed content with the same retry key.", kind: "preference", idempotencyKey: "explicit-create-1", @@ -4482,11 +4507,11 @@ INSERT INTO junior_memory_memories ( }; }, }; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => TEST_NOW_MS, supersessionDecider: createMemoryAgent(model), }); - const existing = await store.createMemory({ + const existing = await createUserMemory(store, { content: "Prefers PR summaries with risks first.", kind: "preference", idempotencyKey: "memory-test:adjudicated-duplicate-original", @@ -4494,7 +4519,7 @@ INSERT INTO junior_memory_memories ( duplicateId = existing.memory.id; await expect( - store.createMemory({ + createUserMemory(store, { content: "Wants danger notes at the beginning of code review recaps.", kind: "preference", idempotencyKey: "memory-test:adjudicated-duplicate-repeat", @@ -4506,9 +4531,9 @@ INSERT INTO junior_memory_memories ( id: existing.memory.id, }, }); - await expect(store.listMemories({})).resolves.toEqual([ - expect.objectContaining({ id: existing.memory.id }), - ]); + await expect(readActiveMemoryRows(fixture, TEST_NOW_MS)).resolves.toEqual( + [expect.objectContaining({ id: existing.memory.id })], + ); } finally { await fixture.close(); } @@ -4528,7 +4553,7 @@ INSERT INTO junior_memory_memories ( (_, index) => `Wants release notes for workflow detail ${index}.`, ); let duplicateId: string | undefined; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, supersessionDecider: { adjudicateSupersession(input) { @@ -4545,14 +4570,14 @@ INSERT INTO junior_memory_memories ( }); for (const [index, content] of distractorContents.entries()) { nowMs = TEST_NOW_MS + index; - await store.createMemory({ + await createUserMemory(store, { content, kind: "preference", idempotencyKey: `memory-test:recent-candidate-distractor-${index}`, }); } nowMs = TEST_NOW_MS + 20; - const existing = await store.createMemory({ + const existing = await createUserMemory(store, { content: existingContent, kind: "preference", idempotencyKey: "memory-test:recent-candidate-existing", @@ -4561,7 +4586,7 @@ INSERT INTO junior_memory_memories ( nowMs = TEST_NOW_MS + 21; await expect( - store.createMemory({ + createUserMemory(store, { content: duplicateContent, kind: "preference", idempotencyKey: "memory-test:recent-candidate-duplicate", @@ -4595,7 +4620,7 @@ INSERT INTO junior_memory_memories ( } const embedder = createTestEmbedder(vectors); const preferenceAdjudicationCalls: MemorySupersessionInput[] = []; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => nowMs, supersessionDecider: { @@ -4612,7 +4637,7 @@ INSERT INTO junior_memory_memories ( }, }); - const oldMemory = await store.createMemory({ + const oldMemory = await createUserMemory(store, { content: oldContent, kind: "preference", idempotencyKey: "memory-test:supersession-old", @@ -4620,7 +4645,7 @@ INSERT INTO junior_memory_memories ( for (const [index, content] of unrelatedContents.entries()) { nowMs = TEST_NOW_MS + index + 1; - await store.createMemory({ + await createUserMemory(store, { content, kind: "preference", idempotencyKey: `memory-test:supersession-unrelated-${index}`, @@ -4629,7 +4654,7 @@ INSERT INTO junior_memory_memories ( preferenceAdjudicationCalls.length = 0; nowMs = TEST_NOW_MS + 20; - const newMemory = await store.createMemory({ + const newMemory = await createUserMemory(store, { content: newContent, kind: "preference", idempotencyKey: "memory-test:supersession-new", @@ -4647,7 +4672,7 @@ INSERT INTO junior_memory_memories ( content: oldContent, id: oldMemory.memory.id, }); - const activeMemories = await store.listMemories({}); + const activeMemories = await readActiveMemoryRows(fixture, nowMs); expect(activeMemories).toContainEqual( expect.objectContaining({ content: newContent, @@ -4682,7 +4707,7 @@ INSERT INTO junior_memory_memories ( nowMs = TEST_NOW_MS + 2; await expect( - store.createMemory({ + createUserMemory(store, { content: "Different content with the superseding retry key.", kind: "preference", idempotencyKey: "memory-test:supersession-new", @@ -4701,7 +4726,7 @@ INSERT INTO junior_memory_memories ( try { let nowMs = TEST_NOW_MS; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, supersessionDecider: { adjudicateSupersession() { @@ -4710,20 +4735,20 @@ INSERT INTO junior_memory_memories ( }, }); - const oldMemory = await store.createMemory({ + const oldMemory = await createUserMemory(store, { content: "Prefers terse PR summaries.", kind: "preference", idempotencyKey: "memory-test:supersession-distinct-old", }); nowMs = TEST_NOW_MS + 1; - const newMemory = await store.createMemory({ + const newMemory = await createUserMemory(store, { content: "Prefers Slack updates in the morning.", kind: "preference", idempotencyKey: "memory-test:supersession-distinct-new", }); - await expect(store.listMemories({})).resolves.toEqual([ + await expect(readActiveMemoryRows(fixture, nowMs)).resolves.toEqual([ expect.objectContaining({ id: newMemory.memory.id }), expect.objectContaining({ id: oldMemory.memory.id }), ]); @@ -4755,26 +4780,26 @@ INSERT INTO junior_memory_memories ( throw new Error("expired replacement should not use supersession"); }, }; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, supersessionDecider: adjudicator, }); - const oldMemory = await store.createMemory({ + const oldMemory = await createUserMemory(store, { content: "Prefers Python for automation scripts.", kind: "preference", idempotencyKey: "memory-test:supersession-expired-old", }); nowMs = TEST_NOW_MS + 1; - await store.createMemory({ + await createUserMemory(store, { content: "Prefers TypeScript for automation scripts.", expiresAtMs: TEST_NOW_MS, kind: "preference", idempotencyKey: "memory-test:supersession-expired-new", }); - await expect(store.listMemories({})).resolves.toEqual([ + await expect(readActiveMemoryRows(fixture, nowMs)).resolves.toEqual([ expect.objectContaining({ id: oldMemory.memory.id }), ]); await expect( @@ -4800,7 +4825,7 @@ INSERT INTO junior_memory_memories ( try { let nowMs = TEST_NOW_MS; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, supersessionDecider: { adjudicateSupersession() { @@ -4811,20 +4836,20 @@ INSERT INTO junior_memory_memories ( }, }); - const oldMemory = await store.createConversationMemory({ + const oldMemory = await createConversationMemory(store, { content: "Prefers Python for automation scripts.", kind: "preference", idempotencyKey: "memory-test:supersession-conversation-old", }); nowMs = TEST_NOW_MS + 1; - const newMemory = await store.createConversationMemory({ + const newMemory = await createConversationMemory(store, { content: "Prefers TypeScript for automation scripts.", kind: "preference", idempotencyKey: "memory-test:supersession-conversation-new", }); - await expect(store.listMemories({})).resolves.toEqual([ + await expect(readActiveMemoryRows(fixture, nowMs)).resolves.toEqual([ expect.objectContaining({ id: newMemory.memory.id }), expect.objectContaining({ id: oldMemory.memory.id }), ]); @@ -4838,7 +4863,7 @@ INSERT INTO junior_memory_memories ( try { let nowMs = TEST_NOW_MS; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, supersessionDecider: { adjudicateSupersession() { @@ -4847,20 +4872,20 @@ INSERT INTO junior_memory_memories ( }, }); - const oldMemory = await store.createMemory({ + const oldMemory = await createUserMemory(store, { content: "Deploy checks use the release runbook.", kind: "knowledge", idempotencyKey: "memory-test:supersession-knowledge-old", }); nowMs = TEST_NOW_MS + 1; - const newMemory = await store.createMemory({ + const newMemory = await createUserMemory(store, { content: "Deploy checks use the release checklist.", kind: "knowledge", idempotencyKey: "memory-test:supersession-knowledge-new", }); - await expect(store.listMemories({})).resolves.toEqual([ + await expect(readActiveMemoryRows(fixture, nowMs)).resolves.toEqual([ expect.objectContaining({ id: newMemory.memory.id }), expect.objectContaining({ id: oldMemory.memory.id }), ]); @@ -4874,23 +4899,34 @@ INSERT INTO junior_memory_memories ( try { let nowMs = TEST_NOW_MS; - const store = createMemoryStore( + const store = memoryFixture( memoryDb(fixture), slackContext({ channelId: "D123" }), { now: () => nowMs }, ); - const archived = await store.createMemory({ + const archived = await createUserMemory(store, { content: "Prefers short deployment summaries.", kind: "preference", idempotencyKey: "explicit-create-archived", }); nowMs = TEST_NOW_MS + 1; - await store.archiveMemory({ id: archived.memory.id }); + await memoryDb(fixture) + .update(memorySqlSchema.juniorMemoryMemories) + .set({ archivedAtMs: nowMs, archiveReason: "user_removed" }) + .where(eq(memorySqlSchema.juniorMemoryMemories.id, archived.memory.id)); + await memoryDb(fixture) + .delete(memorySqlSchema.juniorMemoryEmbeddings) + .where( + eq( + memorySqlSchema.juniorMemoryEmbeddings.memoryId, + archived.memory.id, + ), + ); nowMs = TEST_NOW_MS + 2; - const recreated = await store.createMemory({ + const recreated = await createUserMemory(store, { content: "Prefers short deployment summaries.", kind: "preference", idempotencyKey: "explicit-create-archived", @@ -4903,7 +4939,7 @@ INSERT INTO junior_memory_memories ( nowMs = TEST_NOW_MS + 3; await expect( - store.createMemory({ + createUserMemory(store, { content: "Changed content with the recreated retry key.", kind: "preference", idempotencyKey: "explicit-create-archived", @@ -4915,7 +4951,7 @@ INSERT INTO junior_memory_memories ( content: recreated.memory.content, }, }); - await expect(store.listMemories({})).resolves.toEqual([ + await expect(readActiveMemoryRows(fixture, nowMs)).resolves.toEqual([ expect.objectContaining({ id: recreated.memory.id }), ]); } finally { @@ -4923,32 +4959,27 @@ INSERT INTO junior_memory_memories ( } }, 15_000); - it("treats expired memories as inactive for archive and recreate", async () => { + it("treats expired memories as inactive when recreating", async () => { const fixture = await createMemoryFixture(); try { let nowMs = TEST_NOW_MS; const content = "Temporarily prefers quiet deploy reminders."; const embedder = createTestEmbedder({ [content]: unitEmbedding(1) }); - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { embedder, now: () => nowMs, }); - const expired = await store.createMemory({ + const expired = await createUserMemory(store, { content, kind: "preference", expiresAtMs: TEST_NOW_MS + 10, idempotencyKey: "memory-test:expires", }); - nowMs = TEST_NOW_MS + 11; - await expect( - store.archiveMemory({ id: expired.memory.id }), - ).rejects.toThrow("Memory was not found in the current context."); - nowMs = TEST_NOW_MS + 12; - const recreated = await store.createMemory({ + const recreated = await createUserMemory(store, { content, kind: "preference", idempotencyKey: "memory-test:expires", @@ -4977,7 +5008,7 @@ INSERT INTO junior_memory_memories ( ).resolves.toEqual([ expect.objectContaining({ memoryId: recreated.memory.id }), ]); - await expect(store.searchMemories({ query: "quiet" })).resolves.toEqual([ + await expect(searchMemories(store, { query: "quiet" })).resolves.toEqual([ expect.objectContaining({ id: recreated.memory.id }), ]); } finally { @@ -4990,10 +5021,10 @@ INSERT INTO junior_memory_memories ( try { let nowMs = TEST_NOW_MS; - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, }); - const target = await store.createConversationMemory({ + const target = await createConversationMemory(store, { content: "Release cutover rehearsal is durable.", kind: "knowledge", idempotencyKey: "memory-test:search-target", @@ -5001,7 +5032,7 @@ INSERT INTO junior_memory_memories ( for (let index = 0; index < 205; index += 1) { nowMs = TEST_NOW_MS + index + 1; - await store.createConversationMemory({ + await createConversationMemory(store, { content: `Recent unrelated memory ${index}`, kind: "knowledge", idempotencyKey: `memory-test:search-recent-${index}`, @@ -5010,7 +5041,7 @@ INSERT INTO junior_memory_memories ( nowMs = TEST_NOW_MS + 300; await expect( - store.searchMemories({ query: "cutover rehearsal" }), + searchMemories(store, { query: "cutover rehearsal" }), ).resolves.toEqual([expect.objectContaining({ id: target.memory.id })]); } finally { await fixture.close(); @@ -5023,13 +5054,13 @@ INSERT INTO junior_memory_memories ( try { let nowMs = TEST_NOW_MS; // No embedder: vector leg stays empty so lexical alone must fill limit. - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => nowMs, }); const ids: string[] = []; for (let index = 0; index < 55; index += 1) { nowMs = TEST_NOW_MS + index; - const created = await store.createConversationMemory({ + const created = await createConversationMemory(store, { content: `Deploy freeze checklist item ${index}`, kind: "knowledge", idempotencyKey: `memory-test:search-leg-fill-${index}`, @@ -5037,7 +5068,7 @@ INSERT INTO junior_memory_memories ( ids.push(created.memory.id); } - const results = await store.searchMemories({ + const results = await searchMemories(store, { limit: 50, query: "deploy freeze checklist", }); @@ -5052,27 +5083,23 @@ INSERT INTO junior_memory_memories ( const fixture = await createMemoryFixture(); try { - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => TEST_NOW_MS, }); await expect( - store.createMemory({ + createUserMemory(store, { content: "Prefers short PR summaries.", kind: "preference", idempotencyKey: "memory-test:smuggle", scope: "public", subjectKey: "slack:T123:U999", subjectType: "general", - } as Parameters[0]), + } as Parameters[1]), ).rejects.toThrow(/Invalid input|Unrecognized key/); - await expect( - store.listMemories({ - actor: { platform: "local", userId: "local-user" }, - } as Parameters[0]), - ).rejects.toThrow(/Invalid input|Unrecognized key/); - - await expect(store.listMemories({})).resolves.toEqual([]); + await expect(readActiveMemoryRows(fixture, TEST_NOW_MS)).resolves.toEqual( + [], + ); } finally { await fixture.close(); } @@ -5082,18 +5109,20 @@ INSERT INTO junior_memory_memories ( const fixture = await createMemoryFixture(); try { - const store = createMemoryStore(memoryDb(fixture), slackContext(), { + const store = memoryFixture(memoryDb(fixture), slackContext(), { now: () => TEST_NOW_MS, }); await expect( - store.createMemory({ + createUserMemory(store, { content: " \n\t ", kind: "preference", idempotencyKey: "memory-test:empty-content", }), ).rejects.toThrow("Memory content is required."); - await expect(store.listMemories({})).resolves.toEqual([]); + await expect(readActiveMemoryRows(fixture, TEST_NOW_MS)).resolves.toEqual( + [], + ); } finally { await fixture.close(); } diff --git a/packages/junior-memory/tsup.config.ts b/packages/junior-memory/tsup.config.ts index 42bf368b00..d0caacb4bb 100644 --- a/packages/junior-memory/tsup.config.ts +++ b/packages/junior-memory/tsup.config.ts @@ -3,7 +3,7 @@ import { defineConfig } from "tsup"; export default defineConfig({ clean: true, dts: false, - entry: ["src/index.ts"], + entry: ["src/index.ts", "src/testing.ts"], format: ["esm"], sourcemap: true, target: "node24", diff --git a/packages/junior/tests/component/memory-plugin-storage.test.ts b/packages/junior/tests/component/memory-plugin-storage.test.ts index a0cec0baf8..198cf5c3cf 100644 --- a/packages/junior/tests/component/memory-plugin-storage.test.ts +++ b/packages/junior/tests/component/memory-plugin-storage.test.ts @@ -1,11 +1,8 @@ import path from "node:path"; import { readdirSync } from "node:fs"; import { afterAll, afterEach, describe, expect, it, vi } from "vitest"; -import { - memoryPlugin, - createMemoryStore, - type MemoryDb, -} from "@sentry/junior-memory"; +import { memoryPlugin } from "@sentry/junior-memory"; +import { createMemory, type MemoryDb } from "@sentry/junior-memory/testing"; import { defineJuniorPlugins } from "@/plugins"; import { getPluginTools, setPlugins } from "@/chat/plugins/agent-hooks"; import { migratePluginSchemas } from "@/chat/plugins/migrations"; @@ -394,51 +391,67 @@ WHERE indexname = 'junior_memory_memories_search_idx' nowMs: Date.parse("2026-08-21T12:00:00.000Z"), userId: "U123", }); - const publicMemory = await createMemoryStore(db, { - conversationId: "slack:C123:1718800000.000000", - actor: { platform: "slack", teamId: "T123", userId: "U123" }, - source: createSlackSource({ - teamId: "T123", - channelId: "C123", - messageTs: "1718800000.000000", - visibility: "public", - }), - }).createConversationMemory({ - content: "Public runbooks live in Notion.", - idempotencyKey: "component-public-memory", - kind: "knowledge", + const publicMemory = await createMemory({ + db, + context: { + conversationId: "slack:C123:1718800000.000000", + actor: { platform: "slack", teamId: "T123", userId: "U123" }, + source: createSlackSource({ + teamId: "T123", + channelId: "C123", + messageTs: "1718800000.000000", + visibility: "public", + }), + }, + input: { + content: "Public runbooks live in Notion.", + idempotencyKey: "component-public-memory", + kind: "knowledge", + }, + subjectType: "conversation", }); - const privateMemory = await createMemoryStore(db, { - conversationId: viewerConversationId, - locationId: viewer.locationId, - actor: { platform: "slack", teamId: "T123", userId: "U123" }, - source: createSlackSource({ - teamId: "T123", - channelId: "D123", - messageTs: "1718800001.000000", - visibility: "private", - }), - userId: viewer.user.id, - }).createMemory({ - content: "Prefers terse status updates in this DM.", - idempotencyKey: "component-private-memory", - kind: "preference", + + const privateMemory = await createMemory({ + db, + context: { + conversationId: viewerConversationId, + locationId: viewer.locationId, + actor: { platform: "slack", teamId: "T123", userId: "U123" }, + source: createSlackSource({ + teamId: "T123", + channelId: "D123", + messageTs: "1718800001.000000", + visibility: "private", + }), + userId: viewer.user.id, + }, + input: { + content: "Prefers terse status updates in this DM.", + idempotencyKey: "component-private-memory", + kind: "preference", + }, + subjectType: "user", }); - const otherPrivateMemory = await createMemoryStore(db, { - conversationId: viewerConversationId, - locationId: viewer.locationId, - actor: { platform: "slack", teamId: "T123", userId: "U999" }, - source: createSlackSource({ - teamId: "T123", - channelId: "D123", - messageTs: "1718800002.000000", - visibility: "private", - }), - userId: "other-user", - }).createMemory({ - content: "Only the other User can read this.", - idempotencyKey: "component-other-private-memory", - kind: "knowledge", + const otherPrivateMemory = await createMemory({ + context: { + conversationId: viewerConversationId, + locationId: viewer.locationId, + actor: { platform: "slack", teamId: "T123", userId: "U999" }, + source: createSlackSource({ + teamId: "T123", + channelId: "D123", + messageTs: "1718800002.000000", + visibility: "private", + }), + userId: "other-user", + }, + db, + input: { + content: "Only the other User can read this.", + idempotencyKey: "component-other-private-memory", + kind: "knowledge", + }, + subjectType: "user", }); await expect( fixture.sql.query<{ location_id: string | null }>( @@ -511,26 +524,34 @@ WHERE indexname = 'junior_memory_memories_search_idx' nowMs: Date.parse("2026-08-21T12:00:00.000Z"), userId: actor.userId, }); - const store = createMemoryStore( - // @ts-expect-error non-overlapping boundary cast; rule forbids as-unknown-as chains - fixture.sql.db() as MemoryDb, - { - conversationId, - locationId: userContext.locationId, - actor, - source, - userId: userContext.user.id, + // @ts-expect-error non-overlapping boundary cast; rule forbids as-unknown-as chains + const db = fixture.sql.db() as MemoryDb; + const context = { + conversationId, + locationId: userContext.locationId, + actor, + source, + userId: userContext.user.id, + }; + await createMemory({ + context, + db, + input: { + content: "I prefer host-wired personal recall.", + idempotencyKey: "component-memory-personal", + kind: "preference", }, - ); - await store.createMemory({ - content: "I prefer host-wired personal recall.", - idempotencyKey: "component-memory-personal", - kind: "preference", + subjectType: "user", }); - await store.createConversationMemory({ - content: "This thread tracks host-wired memory context.", - idempotencyKey: "component-memory-conversation", - kind: "knowledge", + await createMemory({ + context, + db, + input: { + content: "This thread tracks host-wired memory context.", + idempotencyKey: "component-memory-conversation", + kind: "knowledge", + }, + subjectType: "conversation", }); const tools = getPluginTools({ diff --git a/packages/junior/vitest.config.ts b/packages/junior/vitest.config.ts index 905123aa8e..4d351ddbff 100644 --- a/packages/junior/vitest.config.ts +++ b/packages/junior/vitest.config.ts @@ -19,6 +19,10 @@ export default defineConfig({ __dirname, "../junior-plugin-api/src/index.ts", ), + "@sentry/junior-memory/testing": path.resolve( + __dirname, + "../junior-memory/src/testing.ts", + ), "@sentry/junior-memory": path.resolve( __dirname, "../junior-memory/src/index.ts", diff --git a/scripts/file-length-exceptions.mjs b/scripts/file-length-exceptions.mjs index 70888fbac5..9f168ef286 100644 --- a/scripts/file-length-exceptions.mjs +++ b/scripts/file-length-exceptions.mjs @@ -16,8 +16,6 @@ export const fileLengthExceptions = { "Existing broad plugin suite; split with plugin modules.", "packages/junior-github/tests/webhook-outcomes.test.ts": "Existing broad webhook outcome suite; split by outcome.", - "packages/junior-memory/src/store.ts": - "Existing memory store; split by storage concern.", "packages/junior-memory/tests/storage.test.ts": "Existing broad memory storage suite; split by storage concern.", "packages/junior/src/chat/agent/index.ts":