feat(memory): inject retrieval context into final prompt rendering

src/services/memory/formatter.js

@@ -1,6 +1,7 @@
 const DEFAULT_TOKEN_BUDGET = 1000;
 const DEFAULT_MAX_SUMMARY_CHARS = 240;
 const DEFAULT_MAX_CONTEXT_CHARS = 320;
+const MEMORY_CONTEXT_HEADER = '[CONTEXT — Relevant memories]';
 
 function clampString(value, maxChars) {
   if (typeof value !== 'string') return '';
@@ -104,3 +105,114 @@ export function formatMemoriesForPrompt(memories, options = {}) {
     truncated: blocks.length < Math.min(k, memories.length),
   };
 }
+
+function toContextString(agentContext) {
+  if (typeof agentContext === 'string') {
+    return agentContext.trim();
+  }
+
+  if (agentContext == null) {
+    return '';
+  }
+
+  if (Array.isArray(agentContext)) {
+    return agentContext.map((value) => String(value)).join('\n').trim();
+  }
+
+  if (typeof agentContext === 'object') {
+    try {
+      return JSON.stringify(agentContext);
+    } catch {
+      return String(agentContext).trim();
+    }
+  }
+
+  return String(agentContext).trim();
+}
+
+/**
+ * Supports safe rollout and canarying with one env var:
+ * - true/1/on/all => enabled for all agents
+ * - false/0/off/no/empty => disabled
+ * - comma-separated agent ids => enabled only for matching ids
+ *
+ * @param {string} agentId
+ * @param {Record<string, string | undefined>} [env]
+ */
+export function isMemoryRetrievalEnabled(agentId, env = process.env) {
+  const raw = String(env.TORCH_MEMORY_RETRIEVAL_ENABLED ?? '').trim();
+  if (!raw) return false;
+
+  const normalized = raw.toLowerCase();
+  if (['0', 'false', 'off', 'no', 'disabled'].includes(normalized)) return false;
+  if (['1', 'true', 'on', 'yes', 'enabled', 'all'].includes(normalized)) return true;
+
+  const allowList = raw
+    .split(',')
+    .map((value) => value.trim())
+    .filter(Boolean);
+
+  if (allowList.length === 0) return false;
+  return allowList.includes(agentId);
+}
+
+/**
+ * Builds a final prompt with retrieved memory context inserted ahead of the base prompt.
+ *
+ * @param {{
+ *  basePrompt: string,
+ *  userRequest: string,
+ *  agentContext?: unknown,
+ *  agent_id: string,
+ *  memoryService: { getRelevantMemories: (params: { agent_id: string, query: string, k?: number }) => Promise<import('./schema.js').MemoryRecord[]> | import('./schema.js').MemoryRecord[] },
+ *  retrievalK?: number,
+ *  tokenBudget?: number,
+ *  env?: Record<string, string | undefined>
+ * }} params
+ */
+export async function renderPromptWithMemoryContext(params) {
+  const {
+    basePrompt,
+    userRequest,
+    agentContext,
+    agent_id,
+    memoryService,
+    retrievalK = 5,
+    tokenBudget = 320,
+    env = process.env,
+  } = params;
+
+  if (!isMemoryRetrievalEnabled(agent_id, env)) {
+    return String(basePrompt ?? '');
+  }
+
+  const requestText = String(userRequest ?? '').trim();
+  const contextText = toContextString(agentContext);
+  const query = [requestText, contextText].filter(Boolean).join('\n\n');
+
+  if (!query || !memoryService?.getRelevantMemories) {
+    return String(basePrompt ?? '');
+  }
+
+  const memories = await memoryService.getRelevantMemories({
+    agent_id,
+    query,
+    k: retrievalK,
+  });
+
+  if (memoryService?.updateMemoryUsage) {
+    const now = Date.now();
+    await Promise.all((memories ?? []).map((memory) => memoryService.updateMemoryUsage(memory.id, now)));
+  }
+
+  const rendered = formatMemoriesForPrompt(memories ?? [], {
+    k: retrievalK,
+    tokenBudget,
+  });
+
+  if (!rendered.text) {
+    return String(basePrompt ?? '');
+  }
+
+  return `${MEMORY_CONTEXT_HEADER}\n${rendered.text}\n\n${String(basePrompt ?? '')}`;
+}

test/memory-retriever.test.mjs

@@ -2,7 +2,11 @@ import test from 'node:test';
 import assert from 'node:assert/strict';
 
 import { filterAndRankMemories } from '../src/services/memory/retriever.js';
-import { formatMemoriesForPrompt } from '../src/services/memory/formatter.js';
+import {
+  formatMemoriesForPrompt,
+  isMemoryRetrievalEnabled,
+  renderPromptWithMemoryContext,
+} from '../src/services/memory/formatter.js';
 
 function createMemory(overrides = {}) {
   return {
@@ -127,3 +131,68 @@ test('formatMemoriesForPrompt returns prompt blocks and enforces token budget tr
     assert.equal(result.truncated, true);
   }
 });
+
+test('isMemoryRetrievalEnabled supports global toggle and agent canary list', () => {
+  assert.equal(isMemoryRetrievalEnabled('agent-a', { TORCH_MEMORY_RETRIEVAL_ENABLED: 'true' }), true);
+  assert.equal(isMemoryRetrievalEnabled('agent-a', { TORCH_MEMORY_RETRIEVAL_ENABLED: 'false' }), false);
+  assert.equal(
+    isMemoryRetrievalEnabled('agent-a', { TORCH_MEMORY_RETRIEVAL_ENABLED: 'agent-a,agent-b' }),
+    true,
+  );
+  assert.equal(
+    isMemoryRetrievalEnabled('agent-c', { TORCH_MEMORY_RETRIEVAL_ENABLED: 'agent-a,agent-b' }),
+    false,
+  );
+});
+
+test('renderPromptWithMemoryContext injects bounded memory block before base prompt', async () => {
+  const memories = [
+    createMemory({
+      id: 'm-ctx',
+      summary: 'Known issue in deploy pipeline',
+      content: 'Retry build after rotating temporary token',
+      created_at: 1_700_000_000_000,
+      importance: 0.91,
+      tags: ['deploy', 'incident'],
+    }),
+  ];
+
+  const calls = [];
+  const usageUpdates = [];
+  const memoryService = {
+    async getRelevantMemories(params) {
+      calls.push(params);
+      return memories;
+    },
+    async updateMemoryUsage(id, lastSeen) {
+      usageUpdates.push({ id, lastSeen });
+    },
+  };
+
+  const output = await renderPromptWithMemoryContext({
+    basePrompt: 'Base final prompt body',
+    userRequest: 'Please debug deployment failure',
+    agentContext: { branch: 'main', tool: 'torch' },
+    agent_id: 'agent-a',
+    memoryService,
+    retrievalK: 1,
+    tokenBudget: 60,
+    env: { TORCH_MEMORY_RETRIEVAL_ENABLED: 'agent-a,agent-b' },
+  });
+
+  assert.equal(calls.length, 1);
+  assert.equal(calls[0].agent_id, 'agent-a');
+  assert.equal(calls[0].k, 1);
+  assert.match(calls[0].query, /Please debug deployment failure/);
+  assert.match(calls[0].query, /"branch":"main"/);
+
+  const headerIndex = output.indexOf('[CONTEXT — Relevant memories]');
+  const baseIndex = output.indexOf('Base final prompt body');
+  assert.ok(headerIndex >= 0);
+  assert.ok(baseIndex > headerIndex);
+  assert.ok(output.includes('[m-ctx,'));
+
+  assert.equal(usageUpdates.length, 1);
+  assert.equal(usageUpdates[0].id, 'm-ctx');
+  assert.equal(typeof usageUpdates[0].lastSeen, 'number');
+});