basant307/AI_Governance_Project

1/**2 * @license3 * Copyright 2025 Google LLC4 * SPDX-License-Identifier: Apache-2.05 */6 7import { beforeEach, describe, expect, it, vi } from 'vitest';8import {9  buildRelevantAutoMemoryPrompt,10  resolveRelevantAutoMemoryPromptForQuery,11  selectRelevantAutoMemoryDocuments,12} from './recall.js';13import type { ScannedAutoMemoryDocument } from './scan.js';14import type { Config } from '../config/config.js';15import { scanAutoMemoryTopicDocuments } from './scan.js';16import { selectRelevantAutoMemoryDocumentsByModel } from './relevanceSelector.js';17 18vi.mock('./scan.js', async (importOriginal) => {19  const actual = await importOriginal<typeof import('./scan.js')>();20  return {21    ...actual,22    scanAutoMemoryTopicDocuments: vi.fn(),23    // Explicit mock — recall now unions user-level docs into the pool, so24    // leaving this on the real implementation would silently fall through25    // to the filesystem (only "works" because the path doesn't exist and26    // listMarkdownFiles swallows ENOENT). Defaults to an empty pool.27    scanUserAutoMemoryTopicDocuments: vi.fn().mockResolvedValue([]),28  };29});30 31vi.mock('./relevanceSelector.js', () => ({32  selectRelevantAutoMemoryDocumentsByModel: vi.fn(),33}));34 35const docs: ScannedAutoMemoryDocument[] = [36  {37    type: 'reference',38    filePath: '/tmp/reference.md',39    relativePath: 'reference.md',40    filename: 'reference.md',41    title: 'Reference Memory',42    description: 'Dashboards and external docs',43    body: '# Reference Memory\n\n- Grafana dashboard: grafana.internal/d/api-latency',44    mtimeMs: 3,45  },46  {47    type: 'project',48    filePath: '/tmp/project.md',49    relativePath: 'project.md',50    filename: 'project.md',51    title: 'Project Memory',52    description: 'Project constraints and release context',53    body: '# Project Memory\n\n- Release freeze starts Friday.',54    mtimeMs: 2,55  },56  {57    type: 'user',58    filePath: '/tmp/user.md',59    relativePath: 'user.md',60    filename: 'user.md',61    title: 'User Memory',62    description: 'User preferences',63    body: '# User Memory\n\n- User prefers terse responses.',64    mtimeMs: 1,65  },66];67 68const activeToolDocs: ScannedAutoMemoryDocument[] = [69  {70    type: 'reference',71    filePath: '/tmp/ata-tool.md',72    relativePath: 'ata-tool.md',73    filename: 'ata-tool.md',74    title: 'ATA tool schema notes',75    description:76      'article-list-query parameter schema and failed tool-call attempts',77    body: '# ATA tool schema notes\n\n- ata::article-list-query failed with guessed field mappings.',78    mtimeMs: 4,79  },80  {81    type: 'reference',82    filePath: '/tmp/ata-gotcha.md',83    relativePath: 'ata-gotcha.md',84    filename: 'ata-gotcha.md',85    title: 'ATA tool gotcha',86    description: 'article-list-query known workaround for transient failures',87    body: '# ATA tool gotcha\n\n- mcp__ata__article-list-query can return systemError during index rotation; retry after checking the ATA oncall note.',88    mtimeMs: 6,89  },90  {91    type: 'reference',92    filePath: '/tmp/ata-owner.md',93    relativePath: 'ata-owner.md',94    filename: 'ata-owner.md',95    title: 'ATA escalation',96    description: 'ATA service owner and escalation path',97    body: '# ATA escalation\n\n- Ask the ATA oncall when the service returns systemError.',98    mtimeMs: 5,99  },100];101 102describe('auto-memory relevant recall', () => {103  beforeEach(() => {104    vi.clearAllMocks();105  });106 107  it('selects the most relevant documents for a query', () => {108    const selected = selectRelevantAutoMemoryDocuments(109      'check the dashboard reference for latency',110      docs,111    );112 113    expect(selected[0]?.type).toBe('reference');114    expect(selected.map((doc) => doc.type)).toContain('reference');115  });116 117  it('returns an empty list for an empty query', () => {118    expect(selectRelevantAutoMemoryDocuments('   ', docs)).toEqual([]);119  });120 121  it('formats selected documents as a prompt block', () => {122    const prompt = buildRelevantAutoMemoryPrompt([docs[0], docs[2]]);123 124    expect(prompt).toContain('## Relevant memory');125    expect(prompt).toContain('Reference Memory (reference.md)');126    expect(prompt).toContain('User Memory (user.md)');127  });128 129  it('uses model-driven selection when config is provided', async () => {130    vi.mocked(scanAutoMemoryTopicDocuments).mockResolvedValue(docs);131    vi.mocked(selectRelevantAutoMemoryDocumentsByModel).mockResolvedValue([132      docs[0],133    ]);134 135    const result = await resolveRelevantAutoMemoryPromptForQuery(136      '/tmp/project',137      'check the dashboard reference for latency',138      {139        config: {} as Config,140      },141    );142 143    expect(result.strategy).toBe('model');144    expect(result.selectedDocs).toEqual([docs[0]]);145    expect(result.prompt).toContain('Reference Memory (reference.md)');146  });147 148  it('falls back to heuristic selection when model-driven selection fails', async () => {149    vi.mocked(scanAutoMemoryTopicDocuments).mockResolvedValue(docs);150    vi.mocked(selectRelevantAutoMemoryDocumentsByModel).mockRejectedValue(151      new Error('selector failed'),152    );153 154    const result = await resolveRelevantAutoMemoryPromptForQuery(155      '/tmp/project',156      'check the dashboard reference for latency',157      {158        config: {} as Config,159        excludedFilePaths: ['/tmp/user.md'],160      },161    );162 163    expect(result.strategy).toBe('heuristic');164    expect(result.selectedDocs.map((doc) => doc.filePath)).toContain(165      '/tmp/reference.md',166    );167    expect(result.selectedDocs.map((doc) => doc.filePath)).not.toContain(168      '/tmp/user.md',169    );170  });171 172  it('keeps active tool schemas out of heuristic fallback', async () => {173    vi.mocked(scanAutoMemoryTopicDocuments).mockResolvedValue(activeToolDocs);174    vi.mocked(selectRelevantAutoMemoryDocumentsByModel).mockRejectedValue(175      new Error('selector failed'),176    );177 178    const result = await resolveRelevantAutoMemoryPromptForQuery(179      '/tmp/project',180      'read the ATA article with article-list-query',181      {182        config: {} as Config,183        recentTools: ['mcp__ata__article-list-query'],184      },185    );186 187    expect(result.strategy).toBe('heuristic');188    expect(result.selectedDocs.map((doc) => doc.filePath)).not.toContain(189      '/tmp/ata-tool.md',190    );191    expect(result.selectedDocs.map((doc) => doc.filePath)).toContain(192      '/tmp/ata-gotcha.md',193    );194    expect(result.selectedDocs.map((doc) => doc.filePath)).toContain(195      '/tmp/ata-owner.md',196    );197  });198});199