CoolFace
Datasetpublic

basant307/AI_Governance_Project

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
0likes48downloads
notebook-edit.test.ts988 linesDownload Raw Back to tools
1/**2 * @license3 * Copyright 2026 Qwen Team4 * SPDX-License-Identifier: Apache-2.05 */6 7import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';8import fs from 'node:fs';9import os from 'node:os';10import path from 'node:path';11import { execFileSync } from 'node:child_process';12import type { Config } from '../config/config.js';13import { ApprovalMode } from '../config/config.js';14import { FileDiscoveryService } from '../services/fileDiscoveryService.js';15import { FileReadCache } from '../services/fileReadCache.js';16import { StandardFileSystemService } from '../services/fileSystemService.js';17import { CommitAttributionService } from '../services/commitAttribution.js';18import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js';19import { ToolErrorType } from './tool-error.js';20import type { ToolInvocation, ToolResult } from './tools.js';21import { applyNotebookEdit, NotebookEditTool } from './notebook-edit.js';22 23vi.mock('../telemetry/loggers.js', () => ({24  logFileOperation: vi.fn(),25}));26 27describe('NotebookEditTool', () => {28  let tempDir: string;29  let fileReadCache: FileReadCache;30  let config: Config;31  let tool: NotebookEditTool;32  let mockFileHistoryService: { trackEdit: ReturnType<typeof vi.fn> };33  const abortSignal = new AbortController().signal;34 35  beforeEach(() => {36    CommitAttributionService.resetInstance();37    tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'notebook-edit-test-'));38    fileReadCache = new FileReadCache();39    mockFileHistoryService = { trackEdit: vi.fn() };40    config = {41      getTargetDir: () => tempDir,42      getProjectRoot: () => tempDir,43      getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),44      setApprovalMode: vi.fn(),45      getWorkspaceContext: () => createMockWorkspaceContext(tempDir),46      getFileService: () => new FileDiscoveryService(tempDir),47      getFileSystemService: () => new StandardFileSystemService(),48      getDefaultFileEncoding: () => 'utf-8',49      getFileReadCache: () => fileReadCache,50      getFileHistoryService: () => mockFileHistoryService,51      getFileReadCacheDisabled: () => false,52      getGeminiClient: vi.fn(),53      getBaseLlmClient: vi.fn(),54      getIdeMode: () => false,55      getApiKey: () => 'test-api-key',56      getModel: () => 'test-model',57      getSandbox: () => false,58      getDebugMode: () => false,59      getQuestion: () => undefined,60      getFullContext: () => false,61      getToolDiscoveryCommand: () => undefined,62      getToolCallCommand: () => undefined,63      getMcpServerCommand: () => undefined,64      getMcpServers: () => undefined,65      getUserAgent: () => 'test-agent',66      getUserMemory: () => '',67      setUserMemory: vi.fn(),68      getGeminiMdFileCount: () => 0,69      setGeminiMdFileCount: vi.fn(),70      getToolRegistry: () => ({}) as never,71    } as unknown as Config;72    tool = new NotebookEditTool(config);73  });74 75  afterEach(() => {76    CommitAttributionService.resetInstance();77    fs.rmSync(tempDir, { recursive: true, force: true });78  });79 80  function writeNotebook(name: string, notebook: Record<string, unknown>) {81    const filePath = path.join(tempDir, name);82    fs.writeFileSync(filePath, JSON.stringify(notebook, null, 1), 'utf-8');83    return filePath;84  }85 86  function seedNotebookRead(filePath: string) {87    fileReadCache.recordRead(filePath, fs.statSync(filePath), {88      full: true,89      cacheable: false,90    });91  }92 93  function buildInvocation(params: Parameters<NotebookEditTool['build']>[0]) {94    return tool.build(params) as ToolInvocation<95      Parameters<NotebookEditTool['build']>[0],96      ToolResult97    >;98  }99 100  it('replaces a code cell by real ID and clears stale outputs', async () => {101    const filePath = writeNotebook('analysis.ipynb', {102      nbformat: 4,103      nbformat_minor: 5,104      cells: [105        {106          cell_type: 'code',107          id: 'load-data',108          source: ['x = 1\n'],109          execution_count: 7,110          outputs: [{ output_type: 'stream', text: ['old\n'] }],111          metadata: {},112        },113      ],114      metadata: { language_info: { name: 'python' } },115    });116    seedNotebookRead(filePath);117 118    const result = await buildInvocation({119      notebook_path: filePath,120      cell_id: 'load-data',121      new_source: 'x = 2\nprint(x)',122    }).execute(abortSignal);123 124    expect(result.error).toBeUndefined();125    const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8'));126    expect(updated.cells[0].source).toEqual(['x = 2\n', 'print(x)']);127    expect(updated.cells[0].execution_count).toBeNull();128    expect(updated.cells[0].outputs).toEqual([]);129    expect(result.llmContent).toContain('replace cell load-data');130 131    const cacheState = fileReadCache.check(fs.statSync(filePath));132    expect(cacheState.state).toBe('fresh');133    if (cacheState.state === 'fresh') {134      expect(cacheState.entry.lastReadWasFull).toBe(true);135      expect(cacheState.entry.lastReadCacheable).toBe(false);136    }137  });138 139  it('blocks writing a secret into a team-memory notebook', async () => {140    const teamDir = path.join(tempDir, '.qwen', 'team-memory');141    fs.mkdirSync(teamDir, { recursive: true });142    const filePath = path.join(teamDir, 'analysis.ipynb');143    fs.writeFileSync(144      filePath,145      JSON.stringify(146        {147          nbformat: 4,148          nbformat_minor: 5,149          cells: [150            {151              cell_type: 'code',152              id: 'load-data',153              source: ['x = 1\n'],154              execution_count: null,155              outputs: [],156              metadata: {},157            },158          ],159          metadata: { language_info: { name: 'python' } },160        },161        null,162        1,163      ),164      'utf-8',165    );166    seedNotebookRead(filePath);167    const originalContent = fs.readFileSync(filePath, 'utf-8');168 169    // Rejected at validate/build time (parity with edit/write-file), before any170    // invocation is created — so the serialized notebook never reaches disk.171    expect(() =>172      buildInvocation({173        notebook_path: filePath,174        cell_id: 'load-data',175        new_source: `token = "ghp_${'a'.repeat(36)}"`,176      }),177    ).toThrow(/shared with all repository collaborators/i);178    expect(fs.readFileSync(filePath, 'utf-8')).toBe(originalContent);179  });180 181  it('blocks a secret in a sibling cell at execute time (full-notebook backstop)', async () => {182    // A team-memory notebook that already carries a secret in one cell. Editing183    // a DIFFERENT, clean cell passes the validate-time single-cell scan (its184    // new_source has no secret), so only execute()'s scan of the whole185    // serialized notebook — the backstop edit/write-file can't run on an186    // .ipynb — catches it. Exercises the execute-time path, not the build-time one.187    const teamDir = path.join(tempDir, '.qwen', 'team-memory');188    fs.mkdirSync(teamDir, { recursive: true });189    const filePath = path.join(teamDir, 'analysis.ipynb');190    fs.writeFileSync(191      filePath,192      JSON.stringify(193        {194          nbformat: 4,195          nbformat_minor: 5,196          cells: [197            {198              cell_type: 'code',199              id: 'creds',200              source: [`token = "ghp_${'a'.repeat(36)}"`],201              execution_count: null,202              outputs: [],203              metadata: {},204            },205            {206              cell_type: 'code',207              id: 'clean',208              source: ['x = 1\n'],209              execution_count: null,210              outputs: [],211              metadata: {},212            },213          ],214          metadata: { language_info: { name: 'python' } },215        },216        null,217        1,218      ),219      'utf-8',220    );221    seedNotebookRead(filePath);222    const originalContent = fs.readFileSync(filePath, 'utf-8');223 224    // new_source is clean, so build() succeeds; the rejection comes from225    // execute()'s full-notebook scan, returned as an error result (not a throw).226    const result = await buildInvocation({227      notebook_path: filePath,228      cell_id: 'clean',229      new_source: 'x = 2\n',230    }).execute(abortSignal);231 232    expect(result.error?.type).toBe(ToolErrorType.INVALID_TOOL_PARAMS);233    expect(result.llmContent).toMatch(234      /shared with all repository collaborators/i,235    );236    // Blocked before any disk write — the notebook is untouched.237    expect(fs.readFileSync(filePath, 'utf-8')).toBe(originalContent);238  });239 240  it('replaces a code cell in a UTF-8 BOM notebook and preserves the BOM', async () => {241    const filePath = path.join(tempDir, 'bom-replace.ipynb');242    fs.writeFileSync(243      filePath,244      `\ufeff${JSON.stringify(245        {246          nbformat: 4,247          nbformat_minor: 5,248          cells: [249            {250              cell_type: 'code',251              id: 'load-data',252              source: ['x = 1\n'],253              execution_count: 7,254              outputs: [{ output_type: 'stream', text: ['old\n'] }],255              metadata: {},256            },257          ],258          metadata: { language_info: { name: 'python' } },259        },260        null,261        1,262      )}`,263      'utf-8',264    );265    seedNotebookRead(filePath);266 267    const result = await buildInvocation({268      notebook_path: filePath,269      cell_id: 'load-data',270      new_source: 'x = 2\nprint(x)',271    }).execute(abortSignal);272 273    expect(result.error).toBeUndefined();274    const updatedBuffer = fs.readFileSync(filePath);275    expect([...updatedBuffer.subarray(0, 3)]).toEqual([0xef, 0xbb, 0xbf]);276    const updated = JSON.parse(updatedBuffer.toString('utf-8').slice(1));277    expect(updated.cells[0].source).toEqual(['x = 2\n', 'print(x)']);278    expect(updated.cells[0].execution_count).toBeNull();279    expect(updated.cells[0].outputs).toEqual([]);280  });281 282  it('replaces by cell-N fallback and converts code to markdown cleanly', async () => {283    const filePath = writeNotebook('convert.ipynb', {284      nbformat: 4,285      nbformat_minor: 5,286      cells: [287        {288          cell_type: 'code',289          source: 'print("old")',290          execution_count: 1,291          outputs: [{ output_type: 'stream', text: 'old\n' }],292          metadata: {},293        },294      ],295      metadata: {},296    });297    seedNotebookRead(filePath);298 299    const result = await buildInvocation({300      notebook_path: filePath,301      cell_id: 'cell-0',302      cell_type: 'markdown',303      new_source: '# Notes',304    }).execute(abortSignal);305 306    expect(result.error).toBeUndefined();307    const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8'));308    expect(updated.cells[0].cell_type).toBe('markdown');309    expect(updated.cells[0].source).toBe('# Notes');310    expect(updated.cells[0]).not.toHaveProperty('outputs');311    expect(updated.cells[0]).not.toHaveProperty('execution_count');312  });313 314  it('converts markdown to code with code-only fields', async () => {315    const filePath = writeNotebook('convert-to-code.ipynb', {316      nbformat: 4,317      nbformat_minor: 5,318      cells: [319        {320          cell_type: 'markdown',321          id: 'intro',322          source: ['# Intro'],323          metadata: {},324        },325      ],326      metadata: {},327    });328    seedNotebookRead(filePath);329 330    const result = await buildInvocation({331      notebook_path: filePath,332      cell_id: 'intro',333      cell_type: 'code',334      new_source: 'print("hi")',335    }).execute(abortSignal);336 337    expect(result.error).toBeUndefined();338    const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8'));339    expect(updated.cells[0].cell_type).toBe('code');340    expect(updated.cells[0].source).toEqual(['print("hi")']);341    expect(updated.cells[0].execution_count).toBeNull();342    expect(updated.cells[0].outputs).toEqual([]);343  });344 345  it('inserts after a target cell and generates an nbformat 4.5 cell ID', async () => {346    const raw = JSON.stringify({347      nbformat: 4,348      nbformat_minor: 5,349      cells: [350        { cell_type: 'markdown', id: 'cell-1', source: ['# A'], metadata: {} },351        { cell_type: 'code', id: 'cell-2', source: ['a = 1'], metadata: {} },352      ],353      metadata: {},354    });355 356    const result = applyNotebookEdit(raw, {357      notebook_path: '/tmp/insert.ipynb',358      edit_mode: 'insert',359      cell_id: 'cell-1',360      cell_type: 'markdown',361      new_source: '## Inserted',362    });363 364    const updated = JSON.parse(result.updatedContent);365    expect(updated.cells).toHaveLength(3);366    expect(updated.cells[1].cell_type).toBe('markdown');367    expect(updated.cells[1].source).toEqual(['## Inserted']);368    expect(updated.cells[1].id).toBe('qwen-cell-1');369    expect(result.editedCellId).toBe('qwen-cell-1');370  });371 372  it('preserves adjacent source style for inserted cells in mixed-format notebooks', async () => {373    const raw = JSON.stringify({374      nbformat: 4,375      nbformat_minor: 5,376      cells: [377        { cell_type: 'markdown', id: 'intro', source: '# Intro', metadata: {} },378        {379          cell_type: 'code',380          id: 'code',381          source: ['value = 1\n'],382          metadata: {},383        },384      ],385      metadata: {},386    });387 388    const result = applyNotebookEdit(raw, {389      notebook_path: '/tmp/insert.ipynb',390      edit_mode: 'insert',391      cell_id: 'intro',392      cell_type: 'markdown',393      new_source: '## Inserted',394    });395 396    const updated = JSON.parse(result.updatedContent);397    expect(updated.cells[1].source).toBe('## Inserted');398  });399 400  it('preserves notebook JSON indentation and trailing newline style on edit', () => {401    const raw = JSON.stringify(402      {403        nbformat: 4,404        nbformat_minor: 5,405        cells: [406          {407            cell_type: 'markdown',408            id: 'intro',409            source: '# Intro',410            metadata: {},411          },412        ],413        metadata: {},414      },415      null,416      2,417    );418 419    const result = applyNotebookEdit(raw, {420      notebook_path: '/tmp/format.ipynb',421      cell_id: 'intro',422      new_source: '# Updated',423    });424 425    expect(result.updatedContent).toContain('\n  "cells"');426    expect(result.updatedContent.endsWith('\n')).toBe(false);427  });428 429  it('rejects ambiguous fallback-like cell IDs', async () => {430    const filePath = writeNotebook('ambiguous.ipynb', {431      nbformat: 4,432      nbformat_minor: 5,433      cells: [434        {435          cell_type: 'markdown',436          id: 'cell-1',437          source: ['real id'],438          metadata: {},439        },440        {441          cell_type: 'markdown',442          source: ['fallback id'],443          metadata: {},444        },445      ],446      metadata: {},447    });448    seedNotebookRead(filePath);449 450    const result = await buildInvocation({451      notebook_path: filePath,452      cell_id: 'cell-1',453      new_source: 'updated',454    }).execute(abortSignal);455 456    expect(result.error?.type).toBe(ToolErrorType.INVALID_TOOL_PARAMS);457    expect(result.llmContent).toContain('ambiguous');458  });459 460  it('inserts at the beginning when no cell_id is provided', async () => {461    const filePath = writeNotebook('insert-start.ipynb', {462      nbformat: 4,463      nbformat_minor: 4,464      cells: [{ cell_type: 'code', source: ['x = 1'], metadata: {} }],465      metadata: {},466    });467    seedNotebookRead(filePath);468 469    const result = await buildInvocation({470      notebook_path: filePath,471      edit_mode: 'insert',472      cell_type: 'code',473      new_source: 'print("first")',474    }).execute(abortSignal);475 476    expect(result.error).toBeUndefined();477    const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8'));478    expect(updated.cells[0].source).toEqual(['print("first")']);479    expect(updated.cells[0]).not.toHaveProperty('id');480  });481 482  it('deletes a cell without requiring new_source', async () => {483    const filePath = writeNotebook('delete.ipynb', {484      nbformat: 4,485      nbformat_minor: 5,486      cells: [487        { cell_type: 'markdown', id: 'keep', source: ['keep'], metadata: {} },488        { cell_type: 'markdown', id: 'drop', source: ['drop'], metadata: {} },489      ],490      metadata: {},491    });492    seedNotebookRead(filePath);493 494    const result = await buildInvocation({495      notebook_path: filePath,496      edit_mode: 'delete',497      cell_id: 'drop',498    }).execute(abortSignal);499 500    expect(result.error).toBeUndefined();501    const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8'));502    expect(updated.cells.map((cell: { id: string }) => cell.id)).toEqual([503      'keep',504    ]);505  });506 507  it('requires a fresh read after structural edits when fallback IDs can shift', async () => {508    const filePath = writeNotebook('fallback-shift.ipynb', {509      nbformat: 4,510      nbformat_minor: 5,511      cells: [512        { cell_type: 'markdown', source: ['A'], metadata: {} },513        { cell_type: 'markdown', source: ['B'], metadata: {} },514      ],515      metadata: {},516    });517    seedNotebookRead(filePath);518 519    const result = await buildInvocation({520      notebook_path: filePath,521      edit_mode: 'insert',522      cell_type: 'markdown',523      new_source: 'inserted',524    }).execute(abortSignal);525 526    expect(result.error).toBeUndefined();527    expect(fileReadCache.check(fs.statSync(filePath)).state).toBe('unknown');528  });529 530  it('preserves fresh read state after structural edits when all IDs are stable', async () => {531    const filePath = writeNotebook('stable-ids.ipynb', {532      nbformat: 4,533      nbformat_minor: 5,534      cells: [535        { cell_type: 'markdown', id: 'a', source: ['A'], metadata: {} },536        { cell_type: 'markdown', id: 'b', source: ['B'], metadata: {} },537      ],538      metadata: {},539    });540    seedNotebookRead(filePath);541 542    const result = await buildInvocation({543      notebook_path: filePath,544      edit_mode: 'insert',545      cell_id: 'a',546      cell_type: 'markdown',547      new_source: 'inserted',548    }).execute(abortSignal);549 550    expect(result.error).toBeUndefined();551    const cacheState = fileReadCache.check(fs.statSync(filePath));552    expect(cacheState.state).toBe('fresh');553    if (cacheState.state === 'fresh') {554      expect(cacheState.entry.lastReadWasFull).toBe(true);555    }556  });557 558  it('requires a fresh full notebook read before editing', async () => {559    const filePath = writeNotebook('unread.ipynb', {560      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],561      metadata: {},562    });563 564    const result = await buildInvocation({565      notebook_path: filePath,566      cell_id: 'a',567      new_source: 'x = 2',568    }).execute(abortSignal);569 570    expect(result.error?.type).toBe(ToolErrorType.EDIT_REQUIRES_PRIOR_READ);571    expect(result.llmContent).toContain('has not been fully read');572  });573 574  it('rejects edits after a truncated notebook read', async () => {575    const filePath = writeNotebook('truncated-read.ipynb', {576      cells: [577        { cell_type: 'code', id: 'visible', source: ['x = 1'], metadata: {} },578        { cell_type: 'code', id: 'tail', source: ['x = 2'], metadata: {} },579      ],580      metadata: {},581    });582    fileReadCache.recordRead(filePath, fs.statSync(filePath), {583      full: false,584      cacheable: false,585    });586 587    const result = await buildInvocation({588      notebook_path: filePath,589      cell_id: 'tail',590      new_source: 'x = 3',591    }).execute(abortSignal);592 593    expect(result.error?.type).toBe(ToolErrorType.EDIT_REQUIRES_PRIOR_READ);594    expect(result.llmContent).toContain('too large for cell-level editing');595    expect(result.llmContent).not.toContain('without offset or limit');596  });597 598  it('rejects notebook directory targets with TARGET_IS_DIRECTORY', async () => {599    const dirPath = path.join(tempDir, 'directory.ipynb');600    fs.mkdirSync(dirPath);601 602    const result = await buildInvocation({603      notebook_path: dirPath,604      cell_id: 'a',605      new_source: 'x = 2',606    }).execute(abortSignal);607 608    expect(result.error?.type).toBe(ToolErrorType.TARGET_IS_DIRECTORY);609    expect(result.llmContent).toContain('is a directory');610  });611 612  it('returns FILE_CHANGED_SINCE_READ when a notebook disappears after content read', async () => {613    const filePath = writeNotebook('disappears-after-read.ipynb', {614      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],615      metadata: {},616    });617    seedNotebookRead(filePath);618    const realFileSystemService = new StandardFileSystemService();619    const fileSystemService = new StandardFileSystemService();620    vi.spyOn(fileSystemService, 'readTextFile').mockImplementation(621      async (args) => {622        const result = await realFileSystemService.readTextFile(args);623        fs.unlinkSync(filePath);624        return result;625      },626    );627    vi.spyOn(config, 'getFileSystemService').mockReturnValue(fileSystemService);628 629    const result = await buildInvocation({630      notebook_path: filePath,631      cell_id: 'a',632      new_source: 'x = 2',633    }).execute(abortSignal);634 635    expect(result.error?.type).toBe(ToolErrorType.FILE_CHANGED_SINCE_READ);636    expect(result.llmContent).toContain('disappeared after it was read');637  });638 639  it('returns PRIOR_READ_VERIFICATION_FAILED when notebook stat verification fails', async () => {640    const filePath = writeNotebook('stat-fails.ipynb', {641      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],642      metadata: {},643    });644    seedNotebookRead(filePath);645    const statSpy = vi646      .spyOn(fs.promises, 'stat')647      .mockRejectedValueOnce(648        Object.assign(new Error('permission denied'), { code: 'EACCES' }),649      );650    let result: ToolResult | undefined;651 652    try {653      result = await buildInvocation({654        notebook_path: filePath,655        cell_id: 'a',656        new_source: 'x = 2',657      }).execute(abortSignal);658    } finally {659      statSpy.mockRestore();660    }661 662    expect(result?.error?.type).toBe(663      ToolErrorType.PRIOR_READ_VERIFICATION_FAILED,664    );665    expect(result?.llmContent).toContain('Could not stat');666  });667 668  it.skipIf(process.platform === 'win32')(669    'rejects non-regular notebook paths with a dedicated error type',670    async () => {671      const fifoPath = path.join(tempDir, 'notebook-fifo.ipynb');672      execFileSync('mkfifo', [fifoPath]);673 674      const result = await buildInvocation({675        notebook_path: fifoPath,676        cell_id: 'a',677        new_source: 'x = 2',678      }).execute(abortSignal);679 680      expect(result.error?.type).toBe(ToolErrorType.TARGET_NOT_REGULAR_FILE);681      expect(result.llmContent).toContain('not a regular file');682    },683  );684 685  it('rejects stale notebook edits after an external change', async () => {686    const filePath = writeNotebook('stale.ipynb', {687      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],688      metadata: {},689    });690    seedNotebookRead(filePath);691    fs.writeFileSync(692      filePath,693      JSON.stringify({694        cells: [695          { cell_type: 'code', id: 'a', source: ['x = 100'], metadata: {} },696        ],697        metadata: {},698      }),699      'utf-8',700    );701 702    const result = await buildInvocation({703      notebook_path: filePath,704      cell_id: 'a',705      new_source: 'x = 2',706    }).execute(abortSignal);707 708    expect(result.error?.type).toBe(ToolErrorType.FILE_CHANGED_SINCE_READ);709  });710 711  it('returns structured errors for missing cells and invalid JSON', async () => {712    const missingCellPath = writeNotebook('missing-cell.ipynb', {713      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],714      metadata: {},715    });716    seedNotebookRead(missingCellPath);717 718    const missingCellResult = await buildInvocation({719      notebook_path: missingCellPath,720      cell_id: 'missing',721      new_source: 'x = 2',722    }).execute(abortSignal);723 724    expect(missingCellResult.error?.type).toBe(725      ToolErrorType.NOTEBOOK_CELL_NOT_FOUND,726    );727 728    const invalidPath = path.join(tempDir, 'bad.ipynb');729    fs.writeFileSync(invalidPath, 'not json', 'utf-8');730    seedNotebookRead(invalidPath);731 732    const invalidResult = await buildInvocation({733      notebook_path: invalidPath,734      edit_mode: 'insert',735      new_source: 'x = 1',736    }).execute(abortSignal);737 738    expect(invalidResult.error?.type).toBe(ToolErrorType.NOTEBOOK_INVALID_JSON);739  });740 741  it('keeps invalid original notebook errors structured for user-modified content', async () => {742    const invalidPath = writeNotebook('bad-original.ipynb', {743      nbformat: 4,744      nbformat_minor: 5,745      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],746      metadata: {},747    });748    seedNotebookRead(invalidPath);749    const originalParams = {750      notebook_path: invalidPath,751      cell_id: 'a',752      new_source: 'x = 2',753    };754    const modifyContext = tool.getModifyContext(abortSignal);755    const currentContent =756      await modifyContext.getCurrentContent(originalParams);757    const proposedContent =758      await modifyContext.getProposedContent(originalParams);759    const updatedParams = modifyContext.createUpdatedParams(760      currentContent,761      proposedContent,762      originalParams,763    );764    fs.writeFileSync(invalidPath, 'not json', 'utf-8');765    seedNotebookRead(invalidPath);766 767    const result = await buildInvocation(768      structuredClone(updatedParams),769    ).execute(abortSignal);770 771    expect(result.error?.type).toBe(ToolErrorType.NOTEBOOK_INVALID_JSON);772  });773 774  it('rejects direct attempts to set internal modified notebook content params', () => {775    const filePath = writeNotebook('injected-modified-content.ipynb', {776      nbformat: 4,777      nbformat_minor: 5,778      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],779      metadata: {},780    });781 782    expect(() =>783      tool.build({784        notebook_path: filePath,785        cell_id: 'a',786        new_source: 'x = 2',787        modified_notebook_content: JSON.stringify({788          cells: [],789          metadata: {},790        }),791      } as Parameters<NotebookEditTool['build']>[0]),792    ).toThrow(/additional properties|modified_notebook_content/i);793  });794 795  it('rejects qwenignored notebooks during validation', () => {796    fs.writeFileSync(path.join(tempDir, '.qwenignore'), '*.ipynb\n', 'utf-8');797    const filePath = writeNotebook('ignored.ipynb', {798      cells: [],799      metadata: {},800    });801 802    expect(() =>803      tool.build({804        notebook_path: filePath,805        edit_mode: 'insert',806        new_source: 'x = 1',807      }),808    ).toThrow(/ignored by \.qwenignore/);809  });810 811  it('rejects notebooks ignored by .agentignore during validation', () => {812    fs.writeFileSync(path.join(tempDir, '.agentignore'), '*.ipynb\n', 'utf-8');813    const filePath = writeNotebook('agent-ignored.ipynb', {814      cells: [],815      metadata: {},816    });817 818    expect(() =>819      tool.build({820        notebook_path: filePath,821        edit_mode: 'insert',822        new_source: 'x = 1',823      }),824    ).toThrow(/ignored by \.agentignore/);825  });826 827  it('returns a notebook diff for confirmation', async () => {828    const filePath = writeNotebook('confirm.ipynb', {829      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],830      metadata: {},831    });832    seedNotebookRead(filePath);833 834    const details = await buildInvocation({835      notebook_path: filePath,836      cell_id: 'a',837      new_source: 'x = 2',838    }).getConfirmationDetails(abortSignal);839 840    const editDetails = details as Extract<typeof details, { type: 'edit' }>;841    expect(editDetails.fileDiff).toContain('-    "x = 1"');842    expect(editDetails.fileDiff).toContain('+    "x = 2"');843    expect((editDetails as { originalContent: string }).originalContent).toBe(844      fs.readFileSync(filePath, 'utf-8'),845    );846  });847 848  it('applies IDE or inline modified full-notebook content instead of the original cell proposal', async () => {849    const filePath = writeNotebook('modified-content.ipynb', {850      nbformat: 4,851      nbformat_minor: 5,852      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],853      metadata: {},854    });855    seedNotebookRead(filePath);856 857    const originalParams = {858      notebook_path: filePath,859      cell_id: 'a',860      new_source: 'x = 2',861    };862    const modifyContext = tool.getModifyContext(abortSignal);863    const currentContent =864      await modifyContext.getCurrentContent(originalParams);865    const proposedContent =866      await modifyContext.getProposedContent(originalParams);867    const modifiedContent = proposedContent.replace('x = 2', 'x = 99');868    const updatedParams = modifyContext.createUpdatedParams(869      currentContent,870      modifiedContent,871      originalParams,872    );873 874    const result = await buildInvocation(875      structuredClone(updatedParams),876    ).execute(abortSignal);877 878    expect(result.error).toBeUndefined();879    const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8'));880    expect(updated.cells[0].source).toEqual(['x = 99']);881    expect(result.llmContent).toContain('modified by the user');882    expect(883      CommitAttributionService.getInstance().getFileAttribution(filePath),884    ).toBeUndefined();885  });886 887  it('uses one current-content snapshot for notebook modify previews', async () => {888    const filePath = writeNotebook('modify-snapshot.ipynb', {889      nbformat: 4,890      nbformat_minor: 5,891      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],892      metadata: {},893    });894 895    const params = {896      notebook_path: filePath,897      cell_id: 'a',898      new_source: 'x = 2',899    };900    const modifyContext = tool.getModifyContext(abortSignal);901    const currentContent = await modifyContext.getCurrentContent(params);902    fs.writeFileSync(903      filePath,904      JSON.stringify(905        {906          nbformat: 4,907          nbformat_minor: 5,908          cells: [909            { cell_type: 'code', id: 'a', source: ['x = 999'], metadata: {} },910          ],911          metadata: {},912        },913        null,914        1,915      ),916      'utf-8',917    );918 919    const proposedContent = await modifyContext.getProposedContent(params);920 921    expect(currentContent).toContain('x = 1');922    expect(proposedContent).toContain('x = 2');923    expect(proposedContent).not.toContain('x = 999');924  });925 926  it('records AI-originated notebook writes for commit attribution', async () => {927    const filePath = writeNotebook('attribution.ipynb', {928      nbformat: 4,929      nbformat_minor: 5,930      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],931      metadata: {},932    });933    seedNotebookRead(filePath);934 935    const result = await buildInvocation({936      notebook_path: filePath,937      cell_id: 'a',938      new_source: 'x = 2',939    }).execute(abortSignal);940 941    expect(result.error).toBeUndefined();942    const attribution =943      CommitAttributionService.getInstance().getFileAttribution(filePath);944    expect(attribution).toBeDefined();945    expect(attribution!.aiContribution).toBeGreaterThan(0);946    expect(mockFileHistoryService.trackEdit).toHaveBeenCalledWith(filePath);947  });948 949  it('tracks file history before the final freshness check', async () => {950    const filePath = writeNotebook('history-before-check.ipynb', {951      nbformat: 4,952      nbformat_minor: 5,953      cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }],954      metadata: {},955    });956    seedNotebookRead(filePath);957    mockFileHistoryService.trackEdit.mockImplementation(async () => {958      fs.writeFileSync(959        filePath,960        JSON.stringify(961          {962            nbformat: 4,963            nbformat_minor: 5,964            cells: [965              { cell_type: 'code', id: 'a', source: ['x = 100'], metadata: {} },966            ],967            metadata: {},968          },969          null,970          1,971        ),972        'utf-8',973      );974    });975 976    const result = await buildInvocation({977      notebook_path: filePath,978      cell_id: 'a',979      new_source: 'x = 2',980    }).execute(abortSignal);981 982    expect(mockFileHistoryService.trackEdit).toHaveBeenCalledWith(filePath);983    expect(result.error?.type).toBe(ToolErrorType.FILE_CHANGED_SINCE_READ);984    const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8'));985    expect(updated.cells[0].source).toEqual(['x = 100']);986  });987});988 
basant307/AI_Governance_Project · CoolFace