diff --git a/src/services/memory/cache.js b/src/services/memory/cache.js index 7d84d8c..58aabdd 100644 --- a/src/services/memory/cache.js +++ b/src/services/memory/cache.js @@ -152,6 +152,14 @@ export function createMemoryCache(options = {}) { runtimeNamespaces.clear(); totalRuntimeEvents = 0; }, + prune(predicate) { + if (typeof predicate !== 'function') return; + for (const key of store.keys()) { + if (predicate(key)) { + store.delete(key); + } + } + }, setRuntimeEvent(scope, event, ttlMs = MINUTE_MS) { const now = Date.now(); removeExpired(now); diff --git a/src/services/memory/index.js b/src/services/memory/index.js index 1c9ec82..d109b85 100644 --- a/src/services/memory/index.js +++ b/src/services/memory/index.js @@ -227,6 +227,24 @@ function recordThroughput(samples, now) { } } +function invalidateAgentCache(agentId) { + if (!agentId) return; + // If the cache supports prune (it should), use it to remove entries for this agent. + if (typeof cache.prune === 'function') { + cache.prune((key) => { + try { + const parsed = JSON.parse(key); + return parsed.agent_id === agentId; + } catch { + return false; + } + }); + } else { + // Fallback if cache implementation doesn't support prune + cache.clear(); + } +} + /** * Ingests agent events and stores them as memory records. * @@ -255,7 +273,22 @@ export async function ingestEvents(events, options = {}) { } }, }); - cache.clear(); + + const uniqueAgents = new Set(); + if (options.agent_id) uniqueAgents.add(options.agent_id); + for (const event of events) { + if (event.agent_id) uniqueAgents.add(event.agent_id); + } + for (const record of records) { + if (record.agent_id) uniqueAgents.add(record.agent_id); + } + + if (uniqueAgents.size > 0) { + for (const agentId of uniqueAgents) invalidateAgentCache(agentId); + } else { + cache.clear(); + } + return records; } @@ -372,7 +405,16 @@ export async function runPruneCycle(options = {}) { }); emitMetric(options, 'memory_pruned_total', { count: prunable.length, retention_ms: retentionMs }); - cache.clear(); + const uniqueAgents = new Set(); + for (const memory of prunable) { + if (memory.agent_id) uniqueAgents.add(memory.agent_id); + } + + if (uniqueAgents.size > 0) { + for (const agentId of uniqueAgents) invalidateAgentCache(agentId); + } else if (prunable.length > 0) { + cache.clear(); + } const result = { pruned: prunable }; if (Number.isFinite(options.scheduleEveryMs) && options.scheduleEveryMs > 0) { @@ -397,7 +439,11 @@ export async function runPruneCycle(options = {}) { export async function pinMemory(id, options = {}) { const repository = options.repository ?? memoryRepository; const updated = await repository.setPinned(id, true); - cache.clear(); + if (updated && updated.agent_id) { + invalidateAgentCache(updated.agent_id); + } else { + cache.clear(); + } return updated; } @@ -414,7 +460,11 @@ export async function pinMemory(id, options = {}) { export async function unpinMemory(id, options = {}) { const repository = options.repository ?? memoryRepository; const updated = await repository.setPinned(id, false); - cache.clear(); + if (updated && updated.agent_id) { + invalidateAgentCache(updated.agent_id); + } else { + cache.clear(); + } return updated; } @@ -431,8 +481,14 @@ export async function unpinMemory(id, options = {}) { */ export async function markMemoryMerged(id, mergedInto, options = {}) { const repository = options.repository ?? memoryRepository; + const memory = await inspectMemory(id, { repository }); const merged = await repository.markMerged(id, mergedInto); - cache.clear(); + + if (memory && memory.agent_id) { + invalidateAgentCache(memory.agent_id); + } else { + cache.clear(); + } return merged; } diff --git a/test/memory-cache.test.mjs b/test/memory-cache.test.mjs index 702fbc6..3fa252b 100644 --- a/test/memory-cache.test.mjs +++ b/test/memory-cache.test.mjs @@ -77,3 +77,16 @@ test('blocks decrypted/session-sensitive data from durable promotion without san assert.equal(events[1].durable_eligible, true); assert.equal(cache.getMetrics().runtime_sensitive_events_blocked, 1); }); + +test('prunes cache entries based on predicate', () => { + const cache = createMemoryCache(); + cache.set('key1', 'value1'); + cache.set('key2', 'value2'); + cache.set('other', 'value3'); + + cache.prune((key) => key.startsWith('key')); + + assert.equal(cache.get('key1'), undefined); + assert.equal(cache.get('key2'), undefined); + assert.equal(cache.get('other'), 'value3'); +});