basant307/AI_Governance_Project

1/**2 * @license3 * Copyright 2025 Qwen Team4 * SPDX-License-Identifier: Apache-2.05 */6 7import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';8 9const runSideQueryMock = vi.fn();10const debugLoggerMock = vi.hoisted(() => ({11  debug: vi.fn(),12  warn: vi.fn(),13}));14 15vi.mock('../utils/sideQuery.js', () => ({16  runSideQuery: (...args: unknown[]) => runSideQueryMock(...args),17}));18 19vi.mock('../utils/debugLogger.js', () => ({20  createDebugLogger: () => debugLoggerMock,21}));22 23import {24  classifyAction,25  sanitizeClassifierReason,26  STAGE1_TIMEOUT_MS,27  STAGE2_TIMEOUT_MS,28  type ClassifierInput,29} from './classifier.js';30import type { Config } from '../config/config.js';31import type { ToolRegistry } from '../tools/tool-registry.js';32 33function makeConfig(34  autoModeSettings: ReturnType<Config['getAutoModeSettings']> = {},35): Config {36  return {37    getFastModel: () => 'qwen-turbo-test',38    getModel: () => 'qwen-max-test',39    getAutoModeSettings: () => autoModeSettings,40    getToolRegistry: () =>41      ({ getTool: () => undefined }) as unknown as ToolRegistry,42  } as unknown as Config;43}44 45function makeInput(over: Partial<ClassifierInput> = {}): ClassifierInput {46  return {47    toolName: 'run_shell_command',48    toolParams: { command: 'ls' },49    messages: [],50    config: makeConfig(),51    signal: new AbortController().signal,52    ...over,53  };54}55 56beforeEach(() => {57  runSideQueryMock.mockReset();58  debugLoggerMock.debug.mockReset();59  debugLoggerMock.warn.mockReset();60});61 62afterEach(() => {63  vi.restoreAllMocks();64});65 66describe('classifyAction — stage 1 happy path', () => {67  it('returns allow without calling stage 2 when stage 1 says shouldBlock=false', async () => {68    runSideQueryMock.mockResolvedValueOnce({ shouldBlock: false });69 70    const result = await classifyAction(makeInput());71 72    expect(result.shouldBlock).toBe(false);73    expect(result.reason).toBe('');74    expect(result.unavailable).toBeUndefined();75    expect(result.stage).toBe('fast');76    expect(runSideQueryMock).toHaveBeenCalledTimes(1);77  });78 79  it('passes the fast-stage purpose to sideQuery', async () => {80    runSideQueryMock.mockResolvedValueOnce({ shouldBlock: false });81    await classifyAction(makeInput());82    const call = runSideQueryMock.mock.calls[0]?.[1] as { purpose?: string };83    expect(call?.purpose).toBe('permission_classifier_stage1');84  });85});86 87describe('classifyAction — stage 1 escalates to stage 2', () => {88  it('returns stage 2 verdict (block + reason) when stage 2 confirms block', async () => {89    runSideQueryMock90      .mockResolvedValueOnce({ shouldBlock: true })91      .mockResolvedValueOnce({92        thinking: 'rm -rf / destroys the root filesystem',93        shouldBlock: true,94        reason: 'Irreversible filesystem destruction',95      });96 97    const result = await classifyAction(makeInput());98 99    expect(result.shouldBlock).toBe(true);100    expect(result.reason).toBe('Irreversible filesystem destruction');101    expect(result.thinking).toContain('rm -rf');102    expect(result.unavailable).toBeUndefined();103    expect(result.stage).toBe('thinking');104    expect(runSideQueryMock).toHaveBeenCalledTimes(2);105  });106 107  it('downgrades stage 1 block to allow when stage 2 says shouldBlock=false', async () => {108    runSideQueryMock109      .mockResolvedValueOnce({ shouldBlock: true })110      .mockResolvedValueOnce({111        thinking: 'cleanup of node_modules is consistent with user intent',112        shouldBlock: false,113        reason: 'safe cleanup',114      });115 116    const result = await classifyAction(makeInput());117 118    expect(result.shouldBlock).toBe(false);119    // Allow path discards the reason field.120    expect(result.reason).toBe('');121    expect(result.stage).toBe('thinking');122  });123 124  it('passes the thinking-stage purpose for the second call', async () => {125    runSideQueryMock126      .mockResolvedValueOnce({ shouldBlock: true })127      .mockResolvedValueOnce({ thinking: 't', shouldBlock: false, reason: '' });128    await classifyAction(makeInput());129    const call = runSideQueryMock.mock.calls[1]?.[1] as { purpose?: string };130    expect(call?.purpose).toBe('permission_classifier_stage2');131  });132});133 134describe('classifyAction — fail-closed on stage 1 failure', () => {135  it('returns unavailable=true when stage 1 throws an API error', async () => {136    runSideQueryMock.mockRejectedValueOnce(new Error('API 500'));137    const result = await classifyAction(makeInput());138    expect(result.shouldBlock).toBe(true);139    expect(result.unavailable).toBe(true);140    expect(result.stage).toBe('fast');141    expect(result.reason).toMatch(/blocked for safety/);142  });143 144  it('surfaces a context-overflow reason when stage 1 fails with that error', async () => {145    runSideQueryMock.mockRejectedValueOnce(146      new Error('Prompt is too long: 200000 tokens > 128000 maximum'),147    );148    const result = await classifyAction(makeInput());149    expect(result.shouldBlock).toBe(true);150    expect(result.unavailable).toBe(true);151    expect(result.reason).toMatch(/context window/i);152  });153 154  it('re-throws when the user signal is aborted (not converted to block)', async () => {155    const controller = new AbortController();156    runSideQueryMock.mockImplementationOnce(async () => {157      controller.abort();158      throw new Error('aborted');159    });160    await expect(161      classifyAction(makeInput({ signal: controller.signal })),162    ).rejects.toThrow();163  });164});165 166describe('classifyAction — fail-closed on stage 2 failure', () => {167  it('honors stage 1 block when stage 2 fails (unavailable=true)', async () => {168    runSideQueryMock169      .mockResolvedValueOnce({ shouldBlock: true })170      .mockRejectedValueOnce(new Error('API 500'));171 172    const result = await classifyAction(makeInput());173 174    expect(result.shouldBlock).toBe(true);175    expect(result.unavailable).toBe(true);176    expect(result.stage).toBe('thinking');177    expect(result.reason).toMatch(/Stage 1 flagged/i);178  });179 180  it('re-throws when the user signal aborts during stage 2', async () => {181    const controller = new AbortController();182    runSideQueryMock183      .mockResolvedValueOnce({ shouldBlock: true })184      .mockImplementationOnce(async () => {185        controller.abort();186        throw new Error('aborted');187      });188    await expect(189      classifyAction(makeInput({ signal: controller.signal })),190    ).rejects.toThrow();191  });192});193 194describe('classifier configuration', () => {195  it('uses configured stage timeouts when provided', async () => {196    const timeoutSpy = vi197      .spyOn(AbortSignal, 'timeout')198      .mockImplementation(() => new AbortController().signal);199    runSideQueryMock200      .mockResolvedValueOnce({ shouldBlock: true })201      .mockResolvedValueOnce({ thinking: 't', shouldBlock: false, reason: '' });202 203    await classifyAction(204      makeInput({205        config: makeConfig({206          classifier: {207            timeouts: {208              stage1Ms: 12_345,209              stage2Ms: 67_890,210            },211          },212        }),213      }),214    );215 216    expect(timeoutSpy).toHaveBeenNthCalledWith(1, 12_345);217    expect(timeoutSpy).toHaveBeenNthCalledWith(2, 67_890);218  });219 220  it('falls back when configured stage timeouts are too low', async () => {221    const timeoutSpy = vi222      .spyOn(AbortSignal, 'timeout')223      .mockImplementation(() => new AbortController().signal);224    runSideQueryMock225      .mockResolvedValueOnce({ shouldBlock: true })226      .mockResolvedValueOnce({ thinking: 't', shouldBlock: false, reason: '' });227 228    await classifyAction(229      makeInput({230        config: makeConfig({231          classifier: {232            timeouts: {233              stage1Ms: 1,234              stage2Ms: 999,235            },236          },237        }),238      }),239    );240 241    expect(timeoutSpy).toHaveBeenNthCalledWith(1, STAGE1_TIMEOUT_MS);242    expect(timeoutSpy).toHaveBeenNthCalledWith(2, STAGE2_TIMEOUT_MS);243    expect(debugLoggerMock.warn).toHaveBeenCalledWith(244      `Classifier timeout 1ms below 1000ms floor, using default ${STAGE1_TIMEOUT_MS}ms`,245    );246    expect(debugLoggerMock.warn).toHaveBeenCalledWith(247      `Classifier timeout 999ms below 1000ms floor, using default ${STAGE2_TIMEOUT_MS}ms`,248    );249  });250 251  it('uses temperature 0 and max_output_tokens=32 with thinking disabled for stage 1', async () => {252    runSideQueryMock.mockResolvedValueOnce({ shouldBlock: false });253    await classifyAction(makeInput());254    const opts = runSideQueryMock.mock.calls[0]?.[1] as {255      config?: {256        temperature?: number;257        maxOutputTokens?: number;258        thinkingConfig?: { includeThoughts?: boolean };259      };260    };261    expect(opts.config?.temperature).toBe(0);262    expect(opts.config?.maxOutputTokens).toBe(32);263    expect(opts.config?.thinkingConfig?.includeThoughts).toBe(false);264  });265 266  it('uses max_output_tokens=4096 with thinking disabled for stage 2', async () => {267    runSideQueryMock268      .mockResolvedValueOnce({ shouldBlock: true })269      .mockResolvedValueOnce({ thinking: 't', shouldBlock: false, reason: '' });270    await classifyAction(makeInput());271    const opts = runSideQueryMock.mock.calls[1]?.[1] as {272      config?: {273        maxOutputTokens?: number;274        thinkingConfig?: { includeThoughts?: boolean };275      };276    };277    expect(opts.config?.maxOutputTokens).toBe(4096);278    // Thinking is disabled in every stage (latency-sensitive permission gate).279    expect(opts.config?.thinkingConfig?.includeThoughts).toBe(false);280  });281 282  it('enables API thinking only for stage 2 when configured', async () => {283    runSideQueryMock284      .mockResolvedValueOnce({ shouldBlock: true })285      .mockResolvedValueOnce({ thinking: 't', shouldBlock: false, reason: '' });286 287    await classifyAction(288      makeInput({289        config: makeConfig({290          classifier: {291            thinking: {292              stage2Enabled: true,293            },294          },295        }),296      }),297    );298 299    const stage1 = runSideQueryMock.mock.calls[0]?.[1] as {300      config?: { thinkingConfig?: { includeThoughts?: boolean } };301    };302    const stage2 = runSideQueryMock.mock.calls[1]?.[1] as {303      config?: { thinkingConfig?: { includeThoughts?: boolean } };304    };305 306    expect(stage1.config?.thinkingConfig?.includeThoughts).toBe(false);307    expect(stage2.config?.thinkingConfig?.includeThoughts).toBe(true);308  });309 310  it('does not pin a model — defaults to the fast model via sideQuery', async () => {311    runSideQueryMock.mockResolvedValueOnce({ shouldBlock: false });312    await classifyAction(makeInput());313    const opts = runSideQueryMock.mock.calls[0]?.[1] as { model?: string };314    expect(opts.model).toBeUndefined();315  });316});317 318// Context-overflow detection now delegated to the shared319// `isContextLengthExceededError` utility; tests covering its behavior live320// alongside that module (utils/contextLengthError.test.ts).321 322describe('sanitizeClassifierReason', () => {323  // Security-critical: the classifier reason is LLM-generated and gets324  // interpolated into the main model's tool-error message. A hostile325  // reason can stage a prompt injection if not sanitized.326 327  it('passes empty / falsy through unchanged', () => {328    expect(sanitizeClassifierReason('')).toBe('');329  });330 331  it('strips simple pseudo-tags like <system>...</system>', () => {332    expect(sanitizeClassifierReason('safe <system>danger</system> tail')).toBe(333      'safe danger tail',334    );335  });336 337  it('iterates strip until stable — no complete <...> tag can survive', () => {338    // The threat is a pseudo-tag like `<system>...` confusing the339    // downstream model. A single /<[^>]*>/g pass on a nested input340    // like `<scr<script>extra>` leaves `>` orphaned tokens which is341    // fine — what must NOT survive is any complete `<...>` pair.342    const result = sanitizeClassifierReason('<scr<script>extra>payload');343    expect(result).not.toMatch(/<[^>]*>/);344  });345 346  it('bounds iteration so adversarial inputs cannot create unbounded work', () => {347    // 8-iteration cap means the function is O(n) regardless of how the348    // attacker structures the input. Even a degenerate string with many349    // overlapping tags terminates promptly.350    const adversarial = '<a'.repeat(2000) + '>'.repeat(2000);351    const t0 = Date.now();352    sanitizeClassifierReason(adversarial);353    expect(Date.now() - t0).toBeLessThan(1000);354  });355 356  it('collapses whitespace and newlines to single spaces', () => {357    expect(sanitizeClassifierReason('line1\nline2\n\n\nline3')).toBe(358      'line1 line2 line3',359    );360  });361 362  it('hard-caps length at 200 characters', () => {363    expect(sanitizeClassifierReason('a'.repeat(500)).length).toBe(200);364  });365 366  it('trims surrounding whitespace after collapse', () => {367    expect(sanitizeClassifierReason('   leading  trailing   ')).toBe(368      'leading trailing',369    );370  });371});372