basant307/AI_Governance_Project

1/**2 * @license3 * Copyright 2025 Google LLC4 * SPDX-License-Identifier: Apache-2.05 */6 7import { describe, it, expect, vi } from 'vitest';8import type { ToolInvocation, ToolResult } from './tools.js';9import type { PermissionDecision } from '../permissions/types.js';10import { DeclarativeTool, hasCycleInSchema, Kind } from './tools.js';11import { ToolErrorType } from './tool-error.js';12 13class TestToolInvocation implements ToolInvocation<object, ToolResult> {14  constructor(15    readonly params: object,16    private readonly executeFn: () => Promise<ToolResult>,17  ) {}18 19  getDescription(): string {20    return 'A test invocation';21  }22 23  toolLocations() {24    return [];25  }26 27  getDefaultPermission(): Promise<PermissionDecision> {28    return Promise.resolve('allow');29  }30 31  getConfirmationDetails(): Promise<never> {32    throw new Error('Not implemented');33  }34 35  execute(): Promise<ToolResult> {36    return this.executeFn();37  }38}39 40class TestTool extends DeclarativeTool<object, ToolResult> {41  private readonly buildFn: (params: object) => TestToolInvocation;42 43  constructor(buildFn: (params: object) => TestToolInvocation) {44    super('test-tool', 'Test Tool', 'A tool for testing', Kind.Other, {});45    this.buildFn = buildFn;46  }47 48  build(params: object): ToolInvocation<object, ToolResult> {49    return this.buildFn(params);50  }51}52 53describe('DeclarativeTool', () => {54  describe('validateBuildAndExecute', () => {55    const abortSignal = new AbortController().signal;56 57    it('should return INVALID_TOOL_PARAMS error if build fails', async () => {58      const buildError = new Error('Invalid build parameters');59      const buildFn = vi.fn().mockImplementation(() => {60        throw buildError;61      });62      const tool = new TestTool(buildFn);63      const params = { foo: 'bar' };64 65      const result = await tool.validateBuildAndExecute(params, abortSignal);66 67      expect(buildFn).toHaveBeenCalledWith(params);68      expect(result).toEqual({69        llmContent: `Error: Invalid parameters provided. Reason: ${buildError.message}`,70        returnDisplay: buildError.message,71        error: {72          message: buildError.message,73          type: ToolErrorType.INVALID_TOOL_PARAMS,74        },75      });76    });77 78    it('should return EXECUTION_FAILED error if execute fails', async () => {79      const executeError = new Error('Execution failed');80      const executeFn = vi.fn().mockRejectedValue(executeError);81      const invocation = new TestToolInvocation({}, executeFn);82      const buildFn = vi.fn().mockReturnValue(invocation);83      const tool = new TestTool(buildFn);84      const params = { foo: 'bar' };85 86      const result = await tool.validateBuildAndExecute(params, abortSignal);87 88      expect(buildFn).toHaveBeenCalledWith(params);89      expect(executeFn).toHaveBeenCalled();90      expect(result).toEqual({91        llmContent: `Error: Tool call execution failed. Reason: ${executeError.message}`,92        returnDisplay: executeError.message,93        error: {94          message: executeError.message,95          type: ToolErrorType.EXECUTION_FAILED,96        },97      });98    });99 100    it('should return the result of execute on success', async () => {101      const successResult: ToolResult = {102        llmContent: 'Success!',103        returnDisplay: 'Success!',104      };105      const executeFn = vi.fn().mockResolvedValue(successResult);106      const invocation = new TestToolInvocation({}, executeFn);107      const buildFn = vi.fn().mockReturnValue(invocation);108      const tool = new TestTool(buildFn);109      const params = { foo: 'bar' };110 111      const result = await tool.validateBuildAndExecute(params, abortSignal);112 113      expect(buildFn).toHaveBeenCalledWith(params);114      expect(executeFn).toHaveBeenCalled();115      expect(result).toEqual(successResult);116    });117  });118});119 120describe('hasCycleInSchema', () => {121  it('should detect a simple direct cycle', () => {122    const schema = {123      properties: {124        data: {125          $ref: '#/properties/data',126        },127      },128    };129    expect(hasCycleInSchema(schema)).toBe(true);130  });131 132  it('should detect a cycle from object properties referencing parent properties', () => {133    const schema = {134      type: 'object',135      properties: {136        data: {137          type: 'object',138          properties: {139            child: { $ref: '#/properties/data' },140          },141        },142      },143    };144    expect(hasCycleInSchema(schema)).toBe(true);145  });146 147  it('should detect a cycle from array items referencing parent properties', () => {148    const schema = {149      type: 'object',150      properties: {151        data: {152          type: 'array',153          items: {154            type: 'object',155            properties: {156              child: { $ref: '#/properties/data/items' },157            },158          },159        },160      },161    };162    expect(hasCycleInSchema(schema)).toBe(true);163  });164 165  it('should detect a cycle between sibling properties', () => {166    const schema = {167      type: 'object',168      properties: {169        a: {170          type: 'object',171          properties: {172            child: { $ref: '#/properties/b' },173          },174        },175        b: {176          type: 'object',177          properties: {178            child: { $ref: '#/properties/a' },179          },180        },181      },182    };183    expect(hasCycleInSchema(schema)).toBe(true);184  });185 186  it('should not detect a cycle in a valid schema', () => {187    const schema = {188      type: 'object',189      properties: {190        name: { type: 'string' },191        address: { $ref: '#/definitions/address' },192      },193      definitions: {194        address: {195          type: 'object',196          properties: {197            street: { type: 'string' },198            city: { type: 'string' },199          },200        },201      },202    };203    expect(hasCycleInSchema(schema)).toBe(false);204  });205 206  it('should handle non-cyclic sibling refs', () => {207    const schema = {208      properties: {209        a: { $ref: '#/definitions/stringDef' },210        b: { $ref: '#/definitions/stringDef' },211      },212      definitions: {213        stringDef: { type: 'string' },214      },215    };216    expect(hasCycleInSchema(schema)).toBe(false);217  });218 219  it('should handle nested but not cyclic refs', () => {220    const schema = {221      properties: {222        a: { $ref: '#/definitions/defA' },223      },224      definitions: {225        defA: { properties: { b: { $ref: '#/definitions/defB' } } },226        defB: { type: 'string' },227      },228    };229    expect(hasCycleInSchema(schema)).toBe(false);230  });231 232  it('should return false for an empty schema', () => {233    expect(hasCycleInSchema({})).toBe(false);234  });235});236