From cbd1fdec9d4c848121263659aef0a5966d118d16 Mon Sep 17 00:00:00 2001 From: PatrickSys Date: Fri, 17 Apr 2026 16:34:22 +0200 Subject: [PATCH] feat(memory): add scoped memory retrieval --- src/cli-memory.ts | 70 +++++++++++++++- src/memory/store.ts | 61 +++++++++++++- src/tools/remember.ts | 28 ++++++- src/tools/search-codebase.ts | 84 ++++++++++++++++--- src/types/index.ts | 7 ++ tests/memory-store-scope.test.ts | 65 +++++++++++++++ tests/search-scoped-memory.test.ts | 128 +++++++++++++++++++++++++++++ 7 files changed, 419 insertions(+), 24 deletions(-) create mode 100644 tests/memory-store-scope.test.ts create mode 100644 tests/search-scoped-memory.test.ts diff --git a/src/cli-memory.ts b/src/cli-memory.ts index 6921338..ef66dc5 100644 --- a/src/cli-memory.ts +++ b/src/cli-memory.ts @@ -3,13 +3,15 @@ */ import path from 'path'; -import type { Memory } from './types/index.js'; +import type { Memory, MemoryScope } from './types/index.js'; import { CODEBASE_CONTEXT_DIRNAME, MEMORY_FILENAME } from './constants/codebase-context.js'; import { appendMemoryFile, + buildMemoryIdentityParts, readMemoriesFile, removeMemory, filterMemories, + normalizeMemoryScope, withConfidence } from './memory/store.js'; @@ -45,7 +47,7 @@ export async function handleMemoryCli(args: string[]): Promise { const listUsage = 'Usage: codebase-context memory list [--category ] [--type ] [--query ] [--json]'; const addUsage = - 'Usage: codebase-context memory add --type --category --memory --reason [--json]'; + 'Usage: codebase-context memory add --type --category --memory --reason [--scope-kind global|file|symbol] [--scope-file ] [--scope-symbol ] [--json]'; const removeUsage = 'Usage: codebase-context memory remove [--json]'; const exitWithUsageError = (message: string, usage?: string): never => { @@ -134,6 +136,13 @@ export async function handleMemoryCli(args: string[]): Promise { const staleTag = m.stale ? ' [STALE]' : ''; console.log(`[${m.id}] ${m.type}/${m.category}: ${m.memory}${staleTag}`); console.log(` Reason: ${m.reason}`); + if (m.scope && m.scope.kind !== 'global') { + if (m.scope.kind === 'file') { + console.log(` Scope: file ${m.scope.file}`); + } else { + console.log(` Scope: symbol ${m.scope.file}#${m.scope.symbol}`); + } + } console.log(` Date: ${m.date} | Confidence: ${m.effectiveConfidence}`); console.log(''); } @@ -145,6 +154,9 @@ export async function handleMemoryCli(args: string[]): Promise { let category: CliMemoryCategory | undefined; let memory: string | undefined; let reason: string | undefined; + let scopeKind: MemoryScope['kind'] | undefined; + let scopeFile: string | undefined; + let scopeSymbol: string | undefined; for (let i = 1; i < args.length; i++) { if (args[i] === '--type') { @@ -197,6 +209,34 @@ export async function handleMemoryCli(args: string[]): Promise { } reason = value; i++; + } else if (args[i] === '--scope-kind') { + const value = args[i + 1]; + if (!value || value.startsWith('--')) { + exitWithUsageError('Error: --scope-kind requires a value.', addUsage); + } + if (value === 'global' || value === 'file' || value === 'symbol') { + scopeKind = value; + } else { + exitWithUsageError( + 'Error: invalid --scope-kind. Allowed: global, file, symbol.', + addUsage + ); + } + i++; + } else if (args[i] === '--scope-file') { + const value = args[i + 1]; + if (!value || value.startsWith('--')) { + exitWithUsageError('Error: --scope-file requires a value.', addUsage); + } + scopeFile = value; + i++; + } else if (args[i] === '--scope-symbol') { + const value = args[i + 1]; + if (!value || value.startsWith('--')) { + exitWithUsageError('Error: --scope-symbol requires a value.', addUsage); + } + scopeSymbol = value; + i++; } else if (args[i] === '--json') { // handled above } @@ -210,9 +250,30 @@ export async function handleMemoryCli(args: string[]): Promise { const requiredCategory = category; const requiredMemory = memory; const requiredReason = reason; + const scope = normalizeMemoryScope({ + kind: scopeKind, + file: scopeFile, + symbol: scopeSymbol + }); + + if (scopeKind === 'file' && !scope) { + exitWithUsageError('Error: --scope-kind file requires --scope-file.', addUsage); + } + if (scopeKind === 'symbol' && !scope) { + exitWithUsageError( + 'Error: --scope-kind symbol requires --scope-file and --scope-symbol.', + addUsage + ); + } const crypto = await import('crypto'); - const hashContent = `${type}:${requiredCategory}:${requiredMemory}:${requiredReason}`; + const hashContent = buildMemoryIdentityParts({ + type, + category: requiredCategory, + memory: requiredMemory, + reason: requiredReason, + scope + }); const hash = crypto.createHash('sha256').update(hashContent).digest('hex'); const id = hash.substring(0, 12); @@ -222,7 +283,8 @@ export async function handleMemoryCli(args: string[]): Promise { category: requiredCategory, memory: requiredMemory, reason: requiredReason, - date: new Date().toISOString() + date: new Date().toISOString(), + ...(scope && { scope }) }; const result = await appendMemoryFile(memoryPath, newMemory); diff --git a/src/memory/store.ts b/src/memory/store.ts index e3c2422..003ac8a 100644 --- a/src/memory/store.ts +++ b/src/memory/store.ts @@ -1,6 +1,6 @@ import { promises as fs } from 'fs'; import path from 'path'; -import type { Memory, MemoryCategory, MemoryType } from '../types/index.js'; +import type { Memory, MemoryCategory, MemoryScope, MemoryType } from '../types/index.js'; type RawMemory = Partial<{ id: unknown; @@ -11,6 +11,7 @@ type RawMemory = Partial<{ reason: unknown; date: unknown; source: unknown; + scope: unknown; }>; export type MemoryFilters = { @@ -23,6 +24,35 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null; } +function normalizePathLike(value: string): string { + return value.replace(/\\/g, '/').replace(/^\.\//, ''); +} + +export function normalizeMemoryScope(raw: unknown): MemoryScope | undefined { + if (!isRecord(raw)) return undefined; + const kind = raw.kind; + if (kind === 'global') { + return { kind }; + } + if (kind === 'file' && typeof raw.file === 'string' && raw.file.trim()) { + return { kind, file: normalizePathLike(raw.file.trim()) }; + } + if ( + kind === 'symbol' && + typeof raw.file === 'string' && + raw.file.trim() && + typeof raw.symbol === 'string' && + raw.symbol.trim() + ) { + return { + kind, + file: normalizePathLike(raw.file.trim()), + symbol: raw.symbol.trim() + }; + } + return undefined; +} + export function normalizeMemory(raw: unknown): Memory | null { if (!isRecord(raw)) return null; const m = raw as RawMemory; @@ -42,7 +72,8 @@ export function normalizeMemory(raw: unknown): Memory | null { if (!id || !category || !memory || !reason || !date) return null; const source = m.source === 'git' ? ('git' as const) : undefined; - return { id, type, category, memory, reason, date, ...(source && { source }) }; + const scope = normalizeMemoryScope(m.scope); + return { id, type, category, memory, reason, date, ...(source && { source }), ...(scope && { scope }) }; } export function normalizeMemories(raw: unknown): Memory[] { @@ -104,7 +135,7 @@ export function filterMemories(memories: Memory[], filters: MemoryFilters): Memo const terms = query.toLowerCase().split(/\s+/).filter(Boolean); if (terms.length > 0) { filtered = filtered.filter((m) => { - const haystack = `${m.memory} ${m.reason}`.toLowerCase(); + const haystack = `${m.memory} ${m.reason} ${formatMemoryScopeText(m.scope)}`.toLowerCase(); return terms.some((t) => haystack.includes(t)); }); } @@ -175,6 +206,30 @@ export function withConfidence(memories: Memory[], now?: Date): MemoryWithConfid })); } +export function formatMemoryScopeText(scope?: MemoryScope): string { + if (!scope || scope.kind === 'global') return ''; + if (scope.kind === 'file') { + return scope.file; + } + return `${scope.file} ${scope.symbol}`; +} + +export function buildMemoryIdentityParts(memory: { + type: MemoryType; + category: MemoryCategory; + memory: string; + reason: string; + scope?: MemoryScope; +}): string { + const scopePart = + !memory.scope || memory.scope.kind === 'global' + ? 'global' + : memory.scope.kind === 'file' + ? `file:${normalizePathLike(memory.scope.file)}` + : `symbol:${normalizePathLike(memory.scope.file)}:${memory.scope.symbol}`; + return `${memory.type}:${memory.category}:${memory.memory}:${memory.reason}:${scopePart}`; +} + export function applyUnfilteredLimit( memories: Memory[], filters: MemoryFilters, diff --git a/src/tools/remember.ts b/src/tools/remember.ts index 030e997..ca686d9 100644 --- a/src/tools/remember.ts +++ b/src/tools/remember.ts @@ -1,7 +1,7 @@ import type { Tool } from '@modelcontextprotocol/sdk/types.js'; import type { ToolContext, ToolResponse } from './types.js'; -import type { Memory, MemoryCategory, MemoryType } from '../types/index.js'; -import { appendMemoryFile } from '../memory/store.js'; +import type { Memory, MemoryCategory, MemoryScope, MemoryType } from '../types/index.js'; +import { appendMemoryFile, buildMemoryIdentityParts, normalizeMemoryScope } from '../memory/store.js'; export const definition: Tool = { name: 'remember', @@ -39,6 +39,23 @@ export const definition: Tool = { reason: { type: 'string', description: 'Why this matters or what breaks otherwise' + }, + scope: { + type: 'object', + description: + 'Optional scope for this memory. Use { kind: "file", file } or { kind: "symbol", file, symbol }.', + properties: { + kind: { + type: 'string', + enum: ['global', 'file', 'symbol'] + }, + file: { + type: 'string' + }, + symbol: { + type: 'string' + } + } } }, required: ['type', 'category', 'memory', 'reason'] @@ -54,15 +71,17 @@ export async function handle( category: MemoryCategory; memory: string; reason: string; + scope?: MemoryScope; }; const { type = 'decision', category, memory, reason } = args_typed; + const scope = normalizeMemoryScope(args_typed.scope); try { const crypto = await import('crypto'); const memoryPath = ctx.paths.memory; - const hashContent = `${type}:${category}:${memory}:${reason}`; + const hashContent = buildMemoryIdentityParts({ type, category, memory, reason, scope }); const hash = crypto.createHash('sha256').update(hashContent).digest('hex'); const id = hash.substring(0, 12); @@ -72,7 +91,8 @@ export async function handle( category, memory, reason, - date: new Date().toISOString() + date: new Date().toISOString(), + ...(scope && { scope }) }; const result = await appendMemoryFile(memoryPath, newMemory); diff --git a/src/tools/search-codebase.ts b/src/tools/search-codebase.ts index 7128e45..eb7c607 100644 --- a/src/tools/search-codebase.ts +++ b/src/tools/search-codebase.ts @@ -21,7 +21,11 @@ import { import { assessSearchQuality } from '../core/search-quality.js'; import { getRerankerStatus } from '../core/reranker.js'; import { IndexCorruptedError } from '../errors/index.js'; -import { readMemoriesFile, withConfidence } from '../memory/store.js'; +import { + formatMemoryScopeText, + readMemoriesFile, + withConfidence +} from '../memory/store.js'; import type { MemoryWithConfidence } from '../memory/store.js'; import { InternalFileGraph } from '../utils/usage-tracker.js'; import type { FileExport } from '../utils/usage-tracker.js'; @@ -249,12 +253,7 @@ export async function handle( const allMemoriesWithConf = withConfidence(allMemories); const queryTerms = queryStr.toLowerCase().split(/\s+/).filter(Boolean); - const relatedMemories = allMemoriesWithConf - .filter((m) => { - const searchText = `${m.memory} ${m.reason}`.toLowerCase(); - return queryTerms.some((term: string) => searchText.includes(term)); - }) - .sort((a, b) => b.effectiveConfidence - a.effectiveConfidence); + const queryTermSet = new Set(queryTerms); // Load intelligence data for enrichment (all intents, not just preflight) let intelligence: IntelligenceData | null = null; @@ -324,6 +323,67 @@ export async function handle( return a === b || a.endsWith(b) || b.endsWith(a); } + function normalizeSymbolName(value: string): string { + return value.trim().toLowerCase(); + } + + const resultPathSet = new Set(results.map((result) => normalizeGraphPath(result.filePath))); + const resultSymbolSet = new Set( + results + .map((result) => { + const symbolName = result.metadata?.symbolName; + return typeof symbolName === 'string' ? normalizeSymbolName(symbolName) : null; + }) + .filter((value): value is string => value !== null) + ); + + function getMemoryScopeBoost(memory: MemoryWithConfidence): number { + if (!memory.scope || memory.scope.kind === 'global') return 0; + + const normalizedFile = normalizeGraphPath(memory.scope.file); + if (memory.scope.kind === 'file') { + return resultPathSet.has(normalizedFile) ? 3 : 0; + } + + const symbolMatch = + resultSymbolSet.has(normalizeSymbolName(memory.scope.symbol)) || + queryTermSet.has(normalizeSymbolName(memory.scope.symbol)); + + if (resultPathSet.has(normalizedFile) && symbolMatch) return 4; + if (resultPathSet.has(normalizedFile)) return 2; + if (symbolMatch) return 1; + return 0; + } + + function getMemoryTextMatchCount(memory: MemoryWithConfidence): number { + const haystack = `${memory.memory} ${memory.reason} ${formatMemoryScopeText(memory.scope)}`.toLowerCase(); + return queryTerms.filter((term) => haystack.includes(term)).length; + } + + function formatMemoryForOutput(memory: MemoryWithConfidence): string { + const scopeText = + !memory.scope || memory.scope.kind === 'global' + ? '' + : memory.scope.kind === 'file' + ? ` [${memory.scope.file}]` + : ` [${memory.scope.file}#${memory.scope.symbol}]`; + return `${memory.memory}${scopeText} (${memory.effectiveConfidence})`; + } + + const relatedMemories = allMemoriesWithConf + .map((memory) => ({ + memory, + textMatches: getMemoryTextMatchCount(memory), + scopeBoost: getMemoryScopeBoost(memory) + })) + .filter((entry) => entry.textMatches > 0 || entry.scopeBoost > 0) + .sort((a, b) => { + if (b.scopeBoost !== a.scopeBoost) return b.scopeBoost - a.scopeBoost; + if (b.textMatches !== a.textMatches) return b.textMatches - a.textMatches; + return b.memory.effectiveConfidence - a.memory.effectiveConfidence; + }) + .map((entry) => entry.memory); + function computeIndexConfidence(): 'fresh' | 'aging' | 'stale' { let confidence: 'fresh' | 'aging' | 'stale' = 'stale'; if (intelligence?.generatedAt) { @@ -568,9 +628,9 @@ export async function handle( if (terms.length === 0) return []; return memories .filter((m) => { - const text = `${m.memory} ${m.reason}`.toLowerCase(); + const text = `${m.memory} ${m.reason} ${formatMemoryScopeText(m.scope)}`.toLowerCase(); const matchCount = terms.filter((t) => text.includes(t)).length; - return matchCount >= 2 && m.effectiveConfidence >= 0.5; + return (matchCount >= 2 || getMemoryScopeBoost(m) >= 2) && m.effectiveConfidence >= 0.5; }) .slice(0, 2); } @@ -1114,7 +1174,7 @@ export async function handle( }; }), ...(strongMemories.length > 0 && { - relatedMemories: strongMemories.map((m) => `${m.memory} (${m.effectiveConfidence})`) + relatedMemories: strongMemories.map((m) => formatMemoryForOutput(m)) }) }, { mode: 'compact', pretty: true, transportAware: true } @@ -1175,9 +1235,7 @@ export async function handle( }), totalResults: results.length, ...(relatedMemories.length > 0 && { - relatedMemories: relatedMemories - .slice(0, 3) - .map((m) => `${m.memory} (${m.effectiveConfidence})`) + relatedMemories: relatedMemories.slice(0, 3).map((m) => formatMemoryForOutput(m)) }) }, { mode: 'full', pretty: true, transportAware: true } diff --git a/src/types/index.ts b/src/types/index.ts index 0576996..a292c19 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -595,8 +595,15 @@ export interface Memory { date: string; /** Source of the memory: 'user' (default) or 'git' (auto-extracted from commits) */ source?: 'user' | 'git'; + /** Optional scope for file-specific or symbol-specific guidance */ + scope?: MemoryScope; } +export type MemoryScope = + | { kind: 'global' } + | { kind: 'file'; file: string } + | { kind: 'symbol'; file: string; symbol: string }; + // ============================================================================ // SHARED PRIMITIVES // ============================================================================ diff --git a/tests/memory-store-scope.test.ts b/tests/memory-store-scope.test.ts new file mode 100644 index 0000000..e5108f6 --- /dev/null +++ b/tests/memory-store-scope.test.ts @@ -0,0 +1,65 @@ +import { describe, expect, it } from 'vitest'; +import { + buildMemoryIdentityParts, + normalizeMemory, + normalizeMemoryScope +} from '../src/memory/store.js'; + +describe('memory scope normalization', () => { + it('normalizes file and symbol scopes to forward-slash paths', () => { + expect(normalizeMemoryScope({ kind: 'file', file: '.\\src\\auth\\auth.service.ts' })).toEqual({ + kind: 'file', + file: 'src/auth/auth.service.ts' + }); + + expect( + normalizeMemoryScope({ + kind: 'symbol', + file: '.\\src\\auth\\auth.service.ts', + symbol: 'AuthService' + }) + ).toEqual({ + kind: 'symbol', + file: 'src/auth/auth.service.ts', + symbol: 'AuthService' + }); + }); + + it('keeps scoped and global memories distinct in identity hashing inputs', () => { + const base = { + type: 'decision' as const, + category: 'architecture' as const, + memory: 'Use AuthService for token reads', + reason: 'Direct token reads bypass refresh behavior.' + }; + + expect(buildMemoryIdentityParts(base)).not.toBe( + buildMemoryIdentityParts({ + ...base, + scope: { kind: 'file', file: 'src/auth/auth.service.ts' } + }) + ); + }); + + it('parses scoped memories from raw JSON payloads', () => { + const normalized = normalizeMemory({ + id: 'abc123def456', + type: 'gotcha', + category: 'architecture', + memory: 'Avoid direct token reads', + reason: 'They skip refresh logic.', + date: '2026-04-17T00:00:00.000Z', + scope: { + kind: 'symbol', + file: 'src/auth/auth.service.ts', + symbol: 'AuthService' + } + }); + + expect(normalized?.scope).toEqual({ + kind: 'symbol', + file: 'src/auth/auth.service.ts', + symbol: 'AuthService' + }); + }); +}); diff --git a/tests/search-scoped-memory.test.ts b/tests/search-scoped-memory.test.ts new file mode 100644 index 0000000..37e4b62 --- /dev/null +++ b/tests/search-scoped-memory.test.ts @@ -0,0 +1,128 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { promises as fs } from 'fs'; +import os from 'os'; +import path from 'path'; +import { handle } from '../src/tools/search-codebase.js'; +import type { ToolContext } from '../src/tools/types.js'; + +const searchMocks = vi.hoisted(() => ({ + search: vi.fn() +})); + +vi.mock('../src/core/search.js', async () => { + class CodebaseSearcher { + constructor(_rootPath: string) {} + + async search(query: string, limit: number, filters?: unknown) { + return searchMocks.search(query, limit, filters); + } + } + + return { CodebaseSearcher }; +}); + +describe('search_codebase scoped memories', () => { + let tempRoot: string; + let ctx: ToolContext; + + beforeEach(async () => { + searchMocks.search.mockReset(); + tempRoot = await fs.mkdtemp(path.join(os.tmpdir(), 'search-scoped-memory-')); + + const contextDir = path.join(tempRoot, '.codebase-context'); + await fs.mkdir(contextDir, { recursive: true }); + + ctx = { + indexState: { status: 'ready' }, + paths: { + baseDir: contextDir, + memory: path.join(contextDir, 'memory.json'), + intelligence: path.join(contextDir, 'intelligence.json'), + keywordIndex: path.join(contextDir, 'index.json'), + vectorDb: path.join(contextDir, 'index') + }, + rootPath: tempRoot, + performIndexing: () => undefined + }; + + await fs.writeFile( + ctx.paths.memory, + JSON.stringify( + [ + { + id: 'scoped-memory', + type: 'gotcha', + category: 'architecture', + memory: 'Avoid direct token reads', + reason: 'They bypass AuthService refresh logic.', + date: '2026-04-17T00:00:00.000Z', + scope: { + kind: 'symbol', + file: 'src/auth/auth.service.ts', + symbol: 'AuthService' + } + }, + { + id: 'global-memory', + type: 'decision', + category: 'architecture', + memory: 'Use auth interceptors', + reason: 'They keep HTTP token injection consistent.', + date: '2026-04-17T00:00:00.000Z' + } + ], + null, + 2 + ) + ); + + await fs.writeFile( + ctx.paths.intelligence, + JSON.stringify( + { + header: { buildId: 'build-1', formatVersion: 1 }, + generatedAt: '2026-04-17T00:00:00.000Z', + patterns: { + stateManagement: { + primary: { + name: 'Signals', + frequency: '78%', + trend: 'Stable' + } + } + }, + goldenFiles: [{ file: 'src/auth/auth.service.ts', score: 0.97 }] + }, + null, + 2 + ) + ); + }); + + afterEach(async () => { + await fs.rm(tempRoot, { recursive: true, force: true }); + }); + + it('prioritizes scoped memories in search output', async () => { + searchMocks.search.mockResolvedValueOnce([ + { + summary: 'Auth service token management', + snippet: 'export class AuthService { getToken() { return token; } }', + filePath: 'src/auth/auth.service.ts', + startLine: 1, + endLine: 20, + score: 0.91, + language: 'ts', + metadata: { symbolName: 'AuthService', symbolKind: 'class', symbolPath: ['AuthService'] }, + relevanceReason: 'Matches auth service token query' + } + ]); + + const response = await handle({ query: 'auth service token', intent: 'edit' }, ctx); + const payload = JSON.parse(response.content?.[0]?.text ?? '{}') as { + relatedMemories?: string[]; + }; + + expect(payload.relatedMemories?.[0]).toContain('src/auth/auth.service.ts#AuthService'); + }); +});