Skip to content

Commit 3494341

Browse files
authored
feat(agent-memory): Phase 7 — discovery marker (#255)
* feat(agent-memory): Phase 7 — discovery marker - MemoryDiscovery registers an agent_memory marker on the shared __betterdb:caches registry: type 'agent_memory', prefix, version, capabilities ['recall','consolidate','reinforce'], stats_key - Reuses the agent-cache discovery protocol constants (additively re-exported) so Monitor reads memory markers identically; writes via MemoryStore's .call client (no method-style client needed) - Heartbeat on an unref'd interval with a TTL'd heartbeat key; best-effort writes; name-collision detection against a different cache type - MemoryStore gains an opt-in discovery option and close() to stop the heartbeat and remove the marker - Re-export discovery constants/MarkerMetadata from @betterdb/agent-cache * fix(agent-memory): namespace the memory discovery marker under {name}:mem Register the agent_memory marker under the {name}:mem registry field and heartbeat key so a memory store and an agent-cache sharing the same name no longer collide on the same __betterdb:caches field / heartbeat key (reported on the AgentMemory facade, which discovers both tiers). * fix(agent-memory): make discovery collisions visible and trim heartbeat work - a cross-type marker collision now emits a visible console.warn and proceeds last-writer-wins, instead of throwing into a swallowed registration promise (the memory marker lives under {name}:mem, so it never collides with agent-cache's {name} field — the check is purely defensive) - read package.json lazily inside createDiscovery so non-discovery importers don't pay a disk read (and avoid the bundler-emit hazard) at module load - drop the redundant SET protocol NX from every heartbeat tick (no-op after register); note the best-effort HGET->HSET TOCTOU * fix(agent-memory): await an in-flight heartbeat tick before stop deletes it stop() cleared the interval and DEL'd the heartbeat but didn't wait for a tick already running, so a tick that fired just before close() could re-write the heartbeat/marker after the DEL and make the store look alive post-shutdown. Track the in-flight tick and await it before deleting.
1 parent af6542f commit 3494341

6 files changed

Lines changed: 473 additions & 6 deletions

File tree

packages/agent-cache/src/index.ts

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,16 @@ export type {
2020
TierDefaults,
2121
ConfigRefreshOptions,
2222
} from './types';
23+
export { AgentCacheError, AgentCacheUsageError, ValkeyCommandError } from './errors';
24+
export type { DiscoveryOptions, MarkerMetadata } from './discovery';
2325
export {
24-
AgentCacheError,
25-
AgentCacheUsageError,
26-
ValkeyCommandError,
27-
} from './errors';
28-
export type { DiscoveryOptions } from './discovery';
26+
PROTOCOL_VERSION,
27+
REGISTRY_KEY,
28+
PROTOCOL_KEY,
29+
HEARTBEAT_KEY_PREFIX,
30+
DEFAULT_HEARTBEAT_INTERVAL_MS,
31+
HEARTBEAT_TTL_SECONDS,
32+
} from './discovery';
2933
export type { Analytics } from './analytics';
3034
export type {
3135
ContentBlock,

packages/agent-memory/src/MemoryStore.ts

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import {
1010
import { parseMemoryItem } from './parseMemoryItem';
1111
import { compositeScore, similarityFromDistance, type RecallWeights } from './compositeScore';
1212
import { selectEvictions, type EvictionCandidate } from './selectEvictions';
13+
import { MemoryDiscovery } from './discovery';
1314
import type {
1415
ConsolidateOptions,
1516
ConsolidateResult,
@@ -33,6 +34,17 @@ const CONSOLIDATE_SCAN_LIMIT = 10000;
3334
const DEFAULT_SUMMARY_IMPORTANCE = 0.7;
3435
const SUMMARY_SOURCE = 'summary';
3536

37+
// Read lazily so only discovery users pay the disk read on import (and avoid a
38+
// bundler hazard, since package.json is not always emitted).
39+
function packageVersion(): string {
40+
return (require('../package.json') as { version: string }).version;
41+
}
42+
43+
export interface MemoryDiscoveryConfig {
44+
version?: string;
45+
heartbeatIntervalMs?: number;
46+
}
47+
3648
export interface MemoryStoreOptions {
3749
client: MemoryStoreClient;
3850
name: string;
@@ -41,6 +53,7 @@ export interface MemoryStoreOptions {
4153
weights?: RecallWeights;
4254
halfLifeSeconds?: number;
4355
maxItemsPerScope?: number;
56+
discovery?: boolean | MemoryDiscoveryConfig;
4457
}
4558

4659
export class MemoryStore {
@@ -51,6 +64,8 @@ export class MemoryStore {
5164
private readonly weights: RecallWeights;
5265
private readonly halfLifeSeconds: number;
5366
private readonly maxItemsPerScope?: number;
67+
private readonly discovery: MemoryDiscovery | null;
68+
private discoveryReady: Promise<void> | null = null;
5469
private dims?: number;
5570

5671
constructor(options: MemoryStoreOptions) {
@@ -61,6 +76,38 @@ export class MemoryStore {
6176
this.weights = options.weights ?? DEFAULT_WEIGHTS;
6277
this.halfLifeSeconds = options.halfLifeSeconds ?? DEFAULT_HALF_LIFE_SECONDS;
6378
this.maxItemsPerScope = options.maxItemsPerScope;
79+
this.discovery = this.createDiscovery(options.discovery);
80+
}
81+
82+
private createDiscovery(config?: boolean | MemoryDiscoveryConfig): MemoryDiscovery | null {
83+
if (!config) {
84+
return null;
85+
}
86+
const settings = config === true ? {} : config;
87+
const discovery = new MemoryDiscovery({
88+
client: this.client,
89+
name: this.name,
90+
version: settings.version ?? packageVersion(),
91+
statsKey: `${this.name}:__mem_stats`,
92+
heartbeatIntervalMs: settings.heartbeatIntervalMs,
93+
});
94+
// Registration is fire-and-forget so construction stays synchronous;
95+
// close() awaits it before tearing the marker down. The floating catch
96+
// keeps any rejected registration from surfacing as an unhandled rejection
97+
// when close() is never called.
98+
const ready = discovery.register();
99+
ready.catch(() => undefined);
100+
this.discoveryReady = ready;
101+
return discovery;
102+
}
103+
104+
async close(): Promise<void> {
105+
if (this.discoveryReady) {
106+
await this.discoveryReady.catch(() => undefined);
107+
}
108+
if (this.discovery) {
109+
await this.discovery.stop({ deleteHeartbeat: true });
110+
}
64111
}
65112

66113
async recall(query: string, options: RecallOptions = {}): Promise<MemoryHit[]> {
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import { describe, it, expect } from 'vitest';
2+
import { REGISTRY_KEY, HEARTBEAT_KEY_PREFIX } from '@betterdb/agent-cache';
3+
import { MemoryStore } from '../MemoryStore';
4+
import { fakeEmbed } from './helpers/fakeEmbed';
5+
import { mockClient } from './helpers/mockClient';
6+
7+
const HEARTBEAT_KEY = `${HEARTBEAT_KEY_PREFIX}mem:mem`;
8+
9+
describe('MemoryStore discovery wiring', () => {
10+
it('registers a discovery marker on construct when discovery is enabled', async () => {
11+
const client = mockClient((command) => (command === 'HGET' ? null : 'OK'));
12+
const store = new MemoryStore({
13+
client,
14+
name: 'mem',
15+
embedFn: fakeEmbed(8),
16+
discovery: { version: '1.0.0', heartbeatIntervalMs: 999_999 },
17+
});
18+
19+
await store.close();
20+
21+
const hset = client.call.mock.calls.find((c) => c[0] === 'HSET' && c[1] === REGISTRY_KEY);
22+
expect(hset?.[2]).toBe('mem:mem');
23+
const marker = JSON.parse(hset?.[3] as string);
24+
expect(marker.type).toBe('agent_memory');
25+
expect(marker.stats_key).toBe('mem:__mem_stats');
26+
const del = client.call.mock.calls.find((c) => c[0] === 'DEL' && c[1] === HEARTBEAT_KEY);
27+
expect(del).toBeDefined();
28+
});
29+
30+
it('does not touch the registry when discovery is not enabled', async () => {
31+
const client = mockClient(() => 'OK');
32+
const store = new MemoryStore({ client, name: 'mem', embedFn: fakeEmbed(8) });
33+
34+
await store.close();
35+
36+
expect(client.call.mock.calls.some((c) => c[0] === 'HSET' && c[1] === REGISTRY_KEY)).toBe(
37+
false,
38+
);
39+
});
40+
});
Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
import { describe, it, expect, vi } from 'vitest';
2+
import {
3+
REGISTRY_KEY,
4+
PROTOCOL_KEY,
5+
HEARTBEAT_KEY_PREFIX,
6+
PROTOCOL_VERSION,
7+
} from '@betterdb/agent-cache';
8+
import { MemoryDiscovery } from '../discovery';
9+
import { mockClient } from './helpers/mockClient';
10+
11+
const HEARTBEAT_KEY = `${HEARTBEAT_KEY_PREFIX}mem:mem`;
12+
13+
function freshClient() {
14+
return mockClient((command) => (command === 'HGET' ? null : 'OK'));
15+
}
16+
17+
function makeDiscovery(client: ReturnType<typeof mockClient>, overrides = {}) {
18+
return new MemoryDiscovery({
19+
client,
20+
name: 'mem',
21+
version: '1.2.3',
22+
statsKey: 'mem:__mem_stats',
23+
heartbeatIntervalMs: 999_999,
24+
...overrides,
25+
});
26+
}
27+
28+
describe('MemoryDiscovery', () => {
29+
it('registers an agent_memory marker with the memory capabilities and stats key', async () => {
30+
const client = freshClient();
31+
const disco = makeDiscovery(client);
32+
33+
await disco.register();
34+
await disco.stop({ deleteHeartbeat: false });
35+
36+
const hset = client.call.mock.calls.find((c) => c[0] === 'HSET' && c[1] === REGISTRY_KEY);
37+
expect(hset?.[2]).toBe('mem:mem');
38+
const marker = JSON.parse(hset?.[3] as string);
39+
expect(marker.type).toBe('agent_memory');
40+
expect(marker.prefix).toBe('mem');
41+
expect(marker.version).toBe('1.2.3');
42+
expect(marker.protocol_version).toBe(PROTOCOL_VERSION);
43+
expect(marker.capabilities).toEqual(['recall', 'consolidate', 'reinforce']);
44+
expect(marker.stats_key).toBe('mem:__mem_stats');
45+
});
46+
47+
it('sets the protocol key with NX and writes a heartbeat with a TTL', async () => {
48+
const client = freshClient();
49+
const disco = makeDiscovery(client);
50+
51+
await disco.register();
52+
await disco.stop({ deleteHeartbeat: false });
53+
54+
const sets = client.call.mock.calls.filter((c) => c[0] === 'SET');
55+
expect(sets.some((c) => c[1] === PROTOCOL_KEY && c[3] === 'NX')).toBe(true);
56+
const heartbeat = sets.find((c) => c[1] === HEARTBEAT_KEY);
57+
expect(heartbeat?.[3]).toBe('EX');
58+
expect(heartbeat?.[4]).toBe('60');
59+
});
60+
61+
it('warns (visibly) and overwrites on a collision with a different cache type', async () => {
62+
const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined);
63+
const client = mockClient((command) =>
64+
command === 'HGET' ? JSON.stringify({ type: 'agent_cache' }) : 'OK',
65+
);
66+
const disco = makeDiscovery(client);
67+
68+
// Registration must not reject into the swallowed promise; the collision is
69+
// surfaced via a visible warning and registration proceeds last-writer-wins.
70+
await expect(disco.register()).resolves.toBeUndefined();
71+
expect(warn).toHaveBeenCalledTimes(1);
72+
expect(warn.mock.calls[0][0]).toMatch(/marker/i);
73+
expect(client.call.mock.calls.some((c) => c[0] === 'HSET' && c[1] === REGISTRY_KEY)).toBe(true);
74+
warn.mockRestore();
75+
});
76+
77+
it('overwrites an existing marker of the same type without throwing', async () => {
78+
const client = mockClient((command) =>
79+
command === 'HGET' ? JSON.stringify({ type: 'agent_memory', version: '0.0.1' }) : 'OK',
80+
);
81+
const disco = makeDiscovery(client);
82+
83+
await disco.register();
84+
await disco.stop({ deleteHeartbeat: false });
85+
86+
expect(client.call.mock.calls.some((c) => c[0] === 'HSET' && c[1] === REGISTRY_KEY)).toBe(true);
87+
});
88+
89+
it('deletes the heartbeat key on stop when asked', async () => {
90+
const client = freshClient();
91+
const disco = makeDiscovery(client);
92+
93+
await disco.register();
94+
await disco.stop({ deleteHeartbeat: true });
95+
96+
expect(client.call.mock.calls.some((c) => c[0] === 'DEL' && c[1] === HEARTBEAT_KEY)).toBe(true);
97+
});
98+
99+
it('re-writes the heartbeat and marker on tickHeartbeat', async () => {
100+
const client = freshClient();
101+
const disco = makeDiscovery(client);
102+
await disco.register();
103+
const before = client.call.mock.calls.length;
104+
105+
await disco.tickHeartbeat();
106+
await disco.stop({ deleteHeartbeat: false });
107+
108+
const after = client.call.mock.calls.slice(before);
109+
expect(after.some((c) => c[0] === 'SET' && c[1] === HEARTBEAT_KEY)).toBe(true);
110+
expect(after.some((c) => c[0] === 'HSET' && c[1] === REGISTRY_KEY)).toBe(true);
111+
});
112+
113+
it('heartbeats on the configured interval', async () => {
114+
vi.useFakeTimers();
115+
try {
116+
const client = freshClient();
117+
const disco = makeDiscovery(client, { heartbeatIntervalMs: 1000 });
118+
await disco.register();
119+
const before = client.call.mock.calls.filter((c) => c[0] === 'HSET').length;
120+
121+
await vi.advanceTimersByTimeAsync(1000);
122+
123+
const after = client.call.mock.calls.filter((c) => c[0] === 'HSET').length;
124+
expect(after).toBeGreaterThan(before);
125+
await disco.stop({ deleteHeartbeat: false });
126+
} finally {
127+
vi.useRealTimers();
128+
}
129+
});
130+
131+
it('waits for an in-flight heartbeat tick before deleting, so no write lands after stop', async () => {
132+
vi.useFakeTimers();
133+
try {
134+
let releaseTick: (value: unknown) => void = () => {};
135+
let gateArmed = false;
136+
const order: string[] = [];
137+
const client = mockClient((command, ...args) => {
138+
order.push(command);
139+
if (gateArmed && command === 'SET' && args[0] === HEARTBEAT_KEY) {
140+
return new Promise((resolve) => {
141+
releaseTick = resolve;
142+
});
143+
}
144+
return command === 'HGET' ? null : 'OK';
145+
});
146+
const disco = makeDiscovery(client, { heartbeatIntervalMs: 1000 });
147+
await disco.register();
148+
149+
gateArmed = true;
150+
await vi.advanceTimersByTimeAsync(1000); // fire one tick; its heartbeat SET blocks
151+
152+
const stopP = disco.stop({ deleteHeartbeat: true });
153+
expect(order.includes('DEL')).toBe(false); // DEL waits behind the in-flight tick
154+
155+
releaseTick('OK');
156+
await stopP;
157+
expect(order[order.length - 1]).toBe('DEL'); // DEL is the final write
158+
} finally {
159+
vi.useRealTimers();
160+
}
161+
});
162+
163+
it('never throws when a registry write fails (best-effort)', async () => {
164+
const onWriteFailed = vi.fn();
165+
const client = mockClient((command) => {
166+
if (command === 'HGET') {
167+
return null;
168+
}
169+
if (command === 'HSET') {
170+
throw new Error('registry boom');
171+
}
172+
return 'OK';
173+
});
174+
const disco = makeDiscovery(client, { onWriteFailed });
175+
176+
await expect(disco.register()).resolves.toBeUndefined();
177+
await disco.stop({ deleteHeartbeat: false });
178+
expect(onWriteFailed).toHaveBeenCalled();
179+
});
180+
});

0 commit comments

Comments
 (0)