CoolFace
Datasetpublic

basant307/AI_Governance_Project

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
0likes48downloads
dreamAgentPlanner.test.ts210 linesDownload Raw Back to memory
1/**2 * @license3 * Copyright 2026 Qwen Team4 * SPDX-License-Identifier: Apache-2.05 */6 7import * as fs from 'node:fs/promises';8import * as os from 'node:os';9import * as path from 'node:path';10import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';11import type { Config } from '../config/config.js';12import type { PermissionManager } from '../permissions/permission-manager.js';13import { ToolNames } from '../tools/tool-names.js';14import { Storage } from '../config/storage.js';15import type { ForkedAgentResult } from '../utils/forkedAgent.js';16import { runForkedAgent } from '../utils/forkedAgent.js';17import { escapeShellArg, getShellConfiguration } from '../utils/shell-utils.js';18import {19  getAutoMemoryRoot,20  getUserAutoMemoryRoot,21  clearAutoMemoryRootCache,22} from './paths.js';23import {24  buildConsolidationTaskPrompt,25  getTranscriptDir,26  planManagedAutoMemoryDreamByAgent,27} from './dreamAgentPlanner.js';28import { ensureAutoMemoryScaffold } from './store.js';29 30vi.mock('../utils/forkedAgent.js', () => ({31  runForkedAgent: vi.fn(),32}));33 34describe('dreamAgentPlanner', () => {35  const originalMemoryBase = process.env['QWEN_CODE_MEMORY_BASE_DIR'];36  let tempDir: string;37  let projectRoot: string;38  let config: Config;39 40  beforeEach(async () => {41    tempDir = await fs.mkdtemp(42      path.join(os.tmpdir(), 'auto-memory-dream-agent-'),43    );44    projectRoot = path.join(tempDir, 'project');45    await fs.mkdir(projectRoot, { recursive: true });46    process.env['QWEN_CODE_MEMORY_BASE_DIR'] = path.join(tempDir, 'memory');47    clearAutoMemoryRootCache();48    await ensureAutoMemoryScaffold(projectRoot);49    config = {50      getSessionId: vi.fn().mockReturnValue('session-1'),51      getModel: vi.fn().mockReturnValue('qwen-test'),52      getApprovalMode: vi.fn(),53    } as unknown as Config;54    vi.mocked(runForkedAgent).mockReset();55  });56 57  afterEach(async () => {58    Storage.setRuntimeBaseDir(null);59    if (originalMemoryBase === undefined) {60      delete process.env['QWEN_CODE_MEMORY_BASE_DIR'];61    } else {62      process.env['QWEN_CODE_MEMORY_BASE_DIR'] = originalMemoryBase;63    }64    clearAutoMemoryRootCache();65    await fs.rm(tempDir, {66      recursive: true,67      force: true,68      maxRetries: 3,69      retryDelay: 10,70    });71  });72 73  it('returns project-scoped session transcript directory', () => {74    const runtimeDir = path.join(tempDir, 'runtime');75    Storage.setRuntimeBaseDir(runtimeDir);76 77    expect(getTranscriptDir(projectRoot)).toBe(78      path.join(new Storage(projectRoot).getProjectDir(), 'chats'),79    );80    expect(getTranscriptDir(projectRoot)).toContain(81      path.join(runtimeDir, 'projects'),82    );83    expect(getTranscriptDir(projectRoot)).not.toContain(84      `${path.sep}.qwen${path.sep}tmp${path.sep}`,85    );86  });87 88  it('shell-quotes the transcript directory in the grep example', () => {89    const transcriptDir = path.join(90      tempDir,91      'runtime dir; touch BAD',92      'projects',93      '-tmp-project',94      'chats',95    );96    const quotedTranscriptDir = escapeShellArg(97      `${transcriptDir}${path.sep}`,98      getShellConfiguration().shell,99    );100    const prompt = buildConsolidationTaskPrompt(101      path.join(tempDir, 'memory'),102      transcriptDir,103    );104 105    expect(prompt).toContain(106      `grep -rn "<narrow term>" ${quotedTranscriptDir} --include="*.jsonl" | tail -50`,107    );108    expect(prompt).not.toContain(109      `grep -rn "<narrow term>" ${transcriptDir}${path.sep} --include="*.jsonl" | tail -50`,110    );111  });112 113  it('returns the forked agent result', async () => {114    const mockResult: ForkedAgentResult = {115      status: 'completed',116      finalText: 'Merged 2 duplicate Vim entries into prefers-vim.md.',117      filesTouched: [118        path.join(projectRoot, '.qwen', 'memory', 'user', 'prefers-vim.md'),119      ],120    };121 122    vi.mocked(runForkedAgent).mockResolvedValue(mockResult);123 124    const result = await planManagedAutoMemoryDreamByAgent(config, projectRoot);125 126    expect(result).toBe(mockResult);127    expect(runForkedAgent).toHaveBeenCalledWith(128      expect.objectContaining({129        maxTurns: 8,130        maxTimeMinutes: 5,131        tools: [132          'read_file',133          'grep_search',134          'glob',135          'list_directory',136          'run_shell_command',137          'write_file',138          'edit',139        ],140      }),141    );142  });143 144  it('can read transcripts while keeping writes project-memory-only', async () => {145    vi.mocked(runForkedAgent).mockResolvedValue({146      status: 'completed',147      filesTouched: [],148    } satisfies ForkedAgentResult);149 150    await planManagedAutoMemoryDreamByAgent(config, projectRoot);151    const params = vi.mocked(runForkedAgent).mock.calls[0]?.[0] as {152      config: Config;153    };154    const pm = params.config.getPermissionManager?.() as PermissionManager;155 156    await expect(157      pm.evaluate({158        toolName: ToolNames.GREP,159        filePath: getTranscriptDir(projectRoot),160      }),161    ).resolves.toBe('default');162    await expect(163      pm.evaluate({164        toolName: ToolNames.WRITE_FILE,165        filePath: path.join(getAutoMemoryRoot(projectRoot), 'project.md'),166      }),167    ).resolves.toBe('allow');168    await expect(169      pm.evaluate({170        toolName: ToolNames.WRITE_FILE,171        filePath: path.join(getUserAutoMemoryRoot(), 'user', 'a.md'),172      }),173    ).resolves.toBe('deny');174  });175 176  it('throws when the agent fails', async () => {177    vi.mocked(runForkedAgent).mockResolvedValue({178      status: 'failed',179      terminateReason: 'Model timed out',180      filesTouched: [],181    } satisfies ForkedAgentResult);182 183    await expect(184      planManagedAutoMemoryDreamByAgent(config, projectRoot),185    ).rejects.toThrow('Model timed out');186  });187 188  it('throws when the agent terminates as cancelled', async () => {189    // runForkedAgent maps AgentTerminateMode.CANCELLED to a resolved190    // `{status: 'cancelled'}` rather than a rejection. Without191    // re-throwing here, `runDreamByAgent` and downstream callers would192    // treat an aborted run as a normal completion — bumping193    // `lastDreamAt` metadata and overwriting a user-cancelled task194    // record with `'completed'`. The throw lets the manager's existing195    // catch path (which checks `signal.aborted && status === 'cancelled'`)196    // do the right thing.197    const mockResult: ForkedAgentResult = {198      status: 'cancelled',199      terminateReason: 'CANCELLED',200      filesTouched: [],201    };202 203    vi.mocked(runForkedAgent).mockResolvedValue(mockResult);204 205    await expect(206      planManagedAutoMemoryDreamByAgent(config, projectRoot),207    ).rejects.toThrow(/cancelled/i);208  });209});210 
basant307/AI_Governance_Project · CoolFace