Skip to content

Commit c07b550

Browse files
author
Tehan
committed
fix(memory): make external-recall model-guard test exercise real dedup
1 parent a518920 commit c07b550

3 files changed

Lines changed: 98 additions & 63 deletions

File tree

packages/pi-plugin/src/inject-compartments-pi.test.ts

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1562,7 +1562,13 @@ describe("Pi external m[1] delta pressure-refold exclusion (cache parity)", () =
15621562
// Cache-busting pass: recomputeM1ThisPass=true. The external delta
15631563
// will be rendered into m[1]. The pressure backstop must NOT fire.
15641564
const secondPass = [userMessage("hello", 11)];
1565-
const r1 = injectM0M1Pi(state, db, secondPass as never, ["entry-0"], true);
1565+
const r1 = injectM0M1Pi(
1566+
state,
1567+
db,
1568+
secondPass as never,
1569+
["entry-0"],
1570+
true,
1571+
);
15661572

15671573
// (a) m[0] NOT re-materialized — external delta must not trigger refold.
15681574
expect(r1.m0Materialized).toBe(false);

packages/pi-plugin/src/inject-compartments-pi.ts

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2067,20 +2067,20 @@ function softRefreshCachedM1Pi(args: {
20672067
args.db,
20682068
args.state.sessionId,
20692069
);
2070-
return {
2071-
...applyCachedPiRow({
2072-
row: sibling,
2073-
state: args.state,
2074-
compartmentsForNormalization: siblingCompartments,
2075-
}),
2076-
memoryUpdateCount: 0,
2077-
recomputed: false,
2078-
// Sibling-adoption replay: the bytes are persisted, not freshly
2079-
// rendered. The external delta is unknown from the persisted row;
2080-
// use "" so the pressure backstop (which only fires on recomputed
2081-
// bytes) is never triggered by a replayed sibling m[1].
2082-
externalDeltaText: "",
2083-
};
2070+
return {
2071+
...applyCachedPiRow({
2072+
row: sibling,
2073+
state: args.state,
2074+
compartmentsForNormalization: siblingCompartments,
2075+
}),
2076+
memoryUpdateCount: 0,
2077+
recomputed: false,
2078+
// Sibling-adoption replay: the bytes are persisted, not freshly
2079+
// rendered. The external delta is unknown from the persisted row;
2080+
// use "" so the pressure backstop (which only fires on recomputed
2081+
// bytes) is never triggered by a replayed sibling m[1].
2082+
externalDeltaText: "",
2083+
};
20842084
}
20852085

20862086
const markers = markersFromCachedPiRow(

packages/plugin/src/features/magic-context/memory/external-recall.test.ts

Lines changed: 77 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ import { computeRecallSnapshotHash, readExternalRecallSnapshot } from "./externa
2525
const mockEmbedBatch = mock(async () => null);
2626
const mockLog = mock(() => {});
2727

28-
mock.module("../../project-embedding-registry", () => ({
28+
mock.module("../project-embedding-registry", () => ({
2929
embedBatchForProject: mockEmbedBatch,
3030
getProjectEmbeddingSnapshot: () => null,
3131
}));
@@ -37,6 +37,7 @@ mock.module("../../../shared/logger", () => ({
3737
}));
3838

3939
const { insertMemory } = await import("./storage-memory");
40+
const { saveEmbedding } = await import("./storage-memory-embeddings");
4041

4142
let db: Database | null = null;
4243

@@ -428,57 +429,85 @@ describe("embedding model-guard (regression: cross-model cosine dedup)", () => {
428429
expect(snapshot?.project).toEqual([{ content: "unique recalled fact" }]);
429430
});
430431

431-
test("model-guard: unknown/off query model → localVectors empty → no cosine dedup", () => {
432-
// Unit-level test of the fixed filter expression.
433-
// Simulate stored embeddings from two different model IDs.
434-
const storedEmbeddings = new Map([
435-
[1, { embedding: new Float32Array([1, 0]), modelId: "model-A" }],
436-
[2, { embedding: new Float32Array([0, 1]), modelId: "model-B" }],
437-
]);
432+
test("model-guard: unknown/off query model → localVectors empty → no cosine dedup (real dedupAndTrim)", async () => {
433+
// RED-GREEN regression: the old filter
434+
// !queryModelId || queryModelId === "off" || e.modelId === queryModelId
435+
// included ALL stored vectors when queryModelId was "off", producing
436+
// meaningless cosine scores. The fix uses an empty localVectors when the
437+
// query model is unknown/off.
438+
//
439+
// Setup: a local memory with a stored embedding (model-A, vector [1, 0]).
440+
// The recalled item has different text (not a hash-dup) but the same
441+
// direction vector. mockEmbedBatch returns modelId="off" for the recalled
442+
// items, so the model guard must suppress cosine dedup entirely.
443+
//
444+
// OLD BUG: localVectors = [model-A vector] → cosine sim = 1.0 ≥ 0.85 →
445+
// recalled item dropped → snapshot.project = [] → test FAILS.
446+
// FIX: localVectors = [] → no cosine dedup → item survives → PASSES.
447+
const localMemory = insertMemory(db!, {
448+
projectPath: ARGS.projectIdentity,
449+
category: "ARCHITECTURE",
450+
content: "local memory content",
451+
sourceType: "historian",
452+
});
453+
// Store a model-A embedding for the local memory.
454+
saveEmbedding(db!, localMemory.id, new Float32Array([1, 0]), "model-A");
455+
resetEmbeddingCacheForTests();
456+
457+
// mockEmbedBatch returns modelId="off" — unknown model, cosine dedup must
458+
// be suppressed regardless of vector similarity.
459+
mockEmbedBatch.mockImplementation(async () => ({
460+
vectors: [new Float32Array([1, 0])], // same direction as local — would be a cosine dup
461+
modelId: "off",
462+
}));
438463

439-
// When queryModelId is "off", the fixed filter must produce an empty array.
440-
const queryModelIdOff = "off";
441-
const localVectorsOff =
442-
queryModelIdOff && queryModelIdOff !== "off"
443-
? [...storedEmbeddings.values()]
444-
.filter((e) => e.modelId === queryModelIdOff)
445-
.map((e) => e.embedding)
446-
: [];
447-
expect(localVectorsOff).toHaveLength(0);
448-
449-
// When queryModelId is falsy (empty string), same result.
450-
const queryModelIdEmpty = "";
451-
const localVectorsEmpty =
452-
queryModelIdEmpty && queryModelIdEmpty !== "off"
453-
? [...storedEmbeddings.values()]
454-
.filter((e) => e.modelId === queryModelIdEmpty)
455-
.map((e) => e.embedding)
456-
: [];
457-
expect(localVectorsEmpty).toHaveLength(0);
464+
_setTestExternalBackendFactory(() =>
465+
recallBackend({
466+
project: [{ content: "recalled item with different text" }],
467+
}),
468+
);
469+
initializeExternalMemory(HINDSIGHT_TEST_CONFIG);
470+
startSessionRecall({ db: db!, ...ARGS });
471+
await waitForSessionRecall(ARGS.sessionId, 5000);
472+
const snapshot = readExternalRecallSnapshot(db!, ARGS.sessionId).snapshot;
473+
// Item must survive: modelId="off" → localVectors empty → no cosine dedup.
474+
expect(snapshot?.project).toEqual([{ content: "recalled item with different text" }]);
458475
});
459476

460-
test("model-guard: known query model → only same-model stored vectors participate; different-model excluded", () => {
461-
// Unit-level test of the fixed filter expression.
462-
// Simulate stored embeddings from two different model IDs.
463-
const storedEmbeddings = new Map([
464-
[1, { embedding: new Float32Array([1, 0]), modelId: "model-A" }],
465-
[2, { embedding: new Float32Array([0, 1]), modelId: "model-B" }],
466-
]);
477+
test("model-guard: known query model → same-model stored vectors participate; near-dup dropped (real dedupAndTrim)", async () => {
478+
// Complement of the test above: when the query model IS known and matches
479+
// the stored embedding model, cosine dedup fires and drops near-duplicates.
480+
//
481+
// Setup: same local memory + model-A embedding. mockEmbedBatch returns
482+
// modelId="model-A" for the recalled item (same model as stored).
483+
//
484+
// FIX: localVectors = [model-A vector] → cosine sim = 1.0 ≥ 0.85 →
485+
// recalled item dropped → snapshot.project = [] → PASSES.
486+
const localMemory = insertMemory(db!, {
487+
projectPath: ARGS.projectIdentity,
488+
category: "ARCHITECTURE",
489+
content: "local memory content",
490+
sourceType: "historian",
491+
});
492+
saveEmbedding(db!, localMemory.id, new Float32Array([1, 0]), "model-A");
493+
resetEmbeddingCacheForTests();
494+
495+
mockEmbedBatch.mockImplementation(async () => ({
496+
vectors: [new Float32Array([1, 0])], // cosine sim = 1.0 with local
497+
modelId: "model-A",
498+
}));
467499

468-
// When queryModelId is "model-A", only model-A vectors are included.
469-
const queryModelId = "model-A";
470-
const localVectors =
471-
queryModelId && queryModelId !== "off"
472-
? [...storedEmbeddings.values()]
473-
.filter((e) => e.modelId === queryModelId)
474-
.map((e) => e.embedding)
475-
: [];
476-
expect(localVectors).toHaveLength(1);
477-
expect(localVectors[0]).toEqual(new Float32Array([1, 0]));
478-
479-
// model-B vector is excluded — different embedding space.
480-
const modelBPresent = localVectors.some((v) => v[0] === 0 && v[1] === 1);
481-
expect(modelBPresent).toBe(false);
500+
_setTestExternalBackendFactory(() =>
501+
recallBackend({
502+
project: [{ content: "recalled item with different text" }],
503+
}),
504+
);
505+
initializeExternalMemory(HINDSIGHT_TEST_CONFIG);
506+
startSessionRecall({ db: db!, ...ARGS });
507+
await waitForSessionRecall(ARGS.sessionId, 5000);
508+
const snapshot = readExternalRecallSnapshot(db!, ARGS.sessionId).snapshot;
509+
// Item must be dropped: model-A matches → cosine sim = 1.0 ≥ 0.85 → dedup.
510+
expect(snapshot?.project).toEqual([]);
482511
});
483512
});
484513

0 commit comments

Comments
 (0)