CoolFace
Datasetpublic

basant307/AI_Governance_Project

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
0likes45downloads
uiTelemetry.test.ts1233 linesDownload Raw Back to telemetry
1/**2 * @license3 * Copyright 2025 Google LLC4 * SPDX-License-Identifier: Apache-2.05 */6 7import { describe, it, expect, vi, beforeEach } from 'vitest';8import { UiTelemetryService, MAIN_SOURCE } from './uiTelemetry.js';9import { ToolCallDecision } from './tool-call-decision.js';10import type { ApiErrorEvent, ApiResponseEvent } from './types.js';11import { ToolCallEvent } from './types.js';12import {13  EVENT_API_ERROR,14  EVENT_API_RESPONSE,15  EVENT_TOOL_CALL,16} from './constants.js';17import type {18  CancelledToolCall,19  CompletedToolCall,20  ErroredToolCall,21  SuccessfulToolCall,22} from '../core/coreToolScheduler.js';23import { ToolErrorType } from '../tools/tool-error.js';24import { ToolConfirmationOutcome } from '../tools/tools.js';25import { MockTool } from '../test-utils/mock-tool.js';26 27const createFakeCompletedToolCall = (28  name: string,29  success: boolean | 'cancelled',30  duration = 100,31  outcome?: ToolConfirmationOutcome,32  error?: Error,33): CompletedToolCall => {34  const request = {35    callId: `call_${name}_${Date.now()}`,36    name,37    args: { foo: 'bar' },38    isClientInitiated: false,39    prompt_id: 'prompt-id-1',40  };41  const tool = new MockTool({ name });42 43  if (success === true) {44    return {45      status: 'success',46      request,47      tool,48      invocation: tool.build({ param: 'test' }),49      response: {50        callId: request.callId,51        responseParts: [52          {53            functionResponse: {54              id: request.callId,55              name,56              response: { output: 'Success!' },57            },58          },59        ],60        error: undefined,61        errorType: undefined,62        resultDisplay: 'Success!',63      },64      durationMs: duration,65      outcome,66    } as SuccessfulToolCall;67  } else if (success === 'cancelled') {68    return {69      status: 'cancelled',70      request,71      tool,72      invocation: tool.build({ param: 'test' }),73      response: {74        callId: request.callId,75        responseParts: [76          {77            functionResponse: {78              id: request.callId,79              name,80              response: { error: 'Tool cancelled' },81            },82          },83        ],84        error: new Error('Tool cancelled'),85        errorType: ToolErrorType.UNKNOWN,86        resultDisplay: 'Cancelled!',87      },88      durationMs: duration,89      outcome,90    } as CancelledToolCall;91  } else {92    return {93      status: 'error',94      request,95      tool,96      response: {97        callId: request.callId,98        responseParts: [99          {100            functionResponse: {101              id: request.callId,102              name,103              response: { error: 'Tool failed' },104            },105          },106        ],107        error: error || new Error('Tool failed'),108        errorType: ToolErrorType.UNKNOWN,109        resultDisplay: 'Failure!',110      },111      durationMs: duration,112      outcome,113    } as ErroredToolCall;114  }115};116 117describe('UiTelemetryService', () => {118  let service: UiTelemetryService;119 120  beforeEach(() => {121    service = new UiTelemetryService();122  });123 124  it('should have correct initial metrics', () => {125    const metrics = service.getMetrics();126    expect(metrics).toEqual({127      models: {},128      tools: {129        totalCalls: 0,130        totalSuccess: 0,131        totalFail: 0,132        totalDurationMs: 0,133        totalDecisions: {134          [ToolCallDecision.ACCEPT]: 0,135          [ToolCallDecision.REJECT]: 0,136          [ToolCallDecision.MODIFY]: 0,137          [ToolCallDecision.AUTO_ACCEPT]: 0,138        },139        byName: {},140      },141      files: {142        totalLinesAdded: 0,143        totalLinesRemoved: 0,144      },145      skills: {146        totalCalls: 0,147        totalSuccess: 0,148        totalFail: 0,149        byName: {},150      },151    });152    expect(service.getLastPromptTokenCount()).toBe(0);153  });154 155  it('should emit an update event when an event is added', () => {156    const spy = vi.fn();157    service.on('update', spy);158 159    const event = {160      'event.name': EVENT_API_RESPONSE,161      model: 'gemini-2.5-pro',162      duration_ms: 500,163      input_token_count: 10,164      output_token_count: 20,165      total_token_count: 30,166      cached_content_token_count: 5,167      thoughts_token_count: 2,168    } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };169 170    service.addEvent(event);171 172    expect(spy).toHaveBeenCalledOnce();173    const { metrics, lastPromptTokenCount } = spy.mock.calls[0][0];174    expect(metrics).toBeDefined();175    expect(lastPromptTokenCount).toBe(0);176  });177 178  describe('API Response Event Processing', () => {179    it('should process a single ApiResponseEvent', () => {180      const event = {181        'event.name': EVENT_API_RESPONSE,182        model: 'gemini-2.5-pro',183        duration_ms: 500,184        input_token_count: 10,185        output_token_count: 20,186        total_token_count: 30,187        cached_content_token_count: 5,188        thoughts_token_count: 2,189      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };190 191      service.addEvent(event);192 193      const metrics = service.getMetrics();194      const modelAggregate = {195        api: {196          totalRequests: 1,197          totalErrors: 0,198          totalLatencyMs: 500,199        },200        tokens: {201          prompt: 10,202          candidates: 20,203          total: 30,204          cached: 5,205          thoughts: 2,206        },207      };208      expect(metrics.models['gemini-2.5-pro']).toEqual({209        ...modelAggregate,210        bySource: {211          [MAIN_SOURCE]: modelAggregate,212        },213      });214      expect(service.getLastPromptTokenCount()).toBe(0);215    });216 217    it('should aggregate multiple ApiResponseEvents for the same model', () => {218      const event1 = {219        'event.name': EVENT_API_RESPONSE,220        model: 'gemini-2.5-pro',221        duration_ms: 500,222        input_token_count: 10,223        output_token_count: 20,224        total_token_count: 30,225        cached_content_token_count: 5,226        thoughts_token_count: 2,227      } as ApiResponseEvent & {228        'event.name': typeof EVENT_API_RESPONSE;229      };230      const event2 = {231        'event.name': EVENT_API_RESPONSE,232        model: 'gemini-2.5-pro',233        duration_ms: 600,234        input_token_count: 15,235        output_token_count: 25,236        total_token_count: 40,237        cached_content_token_count: 10,238        thoughts_token_count: 4,239      } as ApiResponseEvent & {240        'event.name': typeof EVENT_API_RESPONSE;241      };242 243      service.addEvent(event1);244      service.addEvent(event2);245 246      const metrics = service.getMetrics();247      const modelAggregate = {248        api: {249          totalRequests: 2,250          totalErrors: 0,251          totalLatencyMs: 1100,252        },253        tokens: {254          prompt: 25,255          candidates: 45,256          total: 70,257          cached: 15,258          thoughts: 6,259        },260      };261      expect(metrics.models['gemini-2.5-pro']).toEqual({262        ...modelAggregate,263        bySource: {264          [MAIN_SOURCE]: modelAggregate,265        },266      });267      expect(service.getLastPromptTokenCount()).toBe(0);268    });269 270    it('should handle ApiResponseEvents for different models', () => {271      const event1 = {272        'event.name': EVENT_API_RESPONSE,273        model: 'gemini-2.5-pro',274        duration_ms: 500,275        input_token_count: 10,276        output_token_count: 20,277        total_token_count: 30,278        cached_content_token_count: 5,279        thoughts_token_count: 2,280      } as ApiResponseEvent & {281        'event.name': typeof EVENT_API_RESPONSE;282      };283      const event2 = {284        'event.name': EVENT_API_RESPONSE,285        model: 'gemini-2.5-flash',286        duration_ms: 1000,287        input_token_count: 100,288        output_token_count: 200,289        total_token_count: 300,290        cached_content_token_count: 50,291        thoughts_token_count: 20,292      } as ApiResponseEvent & {293        'event.name': typeof EVENT_API_RESPONSE;294      };295 296      service.addEvent(event1);297      service.addEvent(event2);298 299      const metrics = service.getMetrics();300      expect(metrics.models['gemini-2.5-pro']).toBeDefined();301      expect(metrics.models['gemini-2.5-flash']).toBeDefined();302      expect(metrics.models['gemini-2.5-pro'].api.totalRequests).toBe(1);303      expect(metrics.models['gemini-2.5-flash'].api.totalRequests).toBe(1);304      expect(service.getLastPromptTokenCount()).toBe(0);305    });306  });307 308  describe('API Error Event Processing', () => {309    it('should process a single ApiErrorEvent', () => {310      const event = {311        'event.name': EVENT_API_ERROR,312        model: 'gemini-2.5-pro',313        duration_ms: 300,314        error_message: 'Something went wrong',315      } as ApiErrorEvent & { 'event.name': typeof EVENT_API_ERROR };316 317      service.addEvent(event);318 319      const metrics = service.getMetrics();320      const modelAggregate = {321        api: {322          totalRequests: 1,323          totalErrors: 1,324          totalLatencyMs: 300,325        },326        tokens: {327          prompt: 0,328          candidates: 0,329          total: 0,330          cached: 0,331          thoughts: 0,332        },333      };334      expect(metrics.models['gemini-2.5-pro']).toEqual({335        ...modelAggregate,336        bySource: {337          [MAIN_SOURCE]: modelAggregate,338        },339      });340    });341 342    it('should aggregate ApiErrorEvents and ApiResponseEvents', () => {343      const responseEvent = {344        'event.name': EVENT_API_RESPONSE,345        model: 'gemini-2.5-pro',346        duration_ms: 500,347        input_token_count: 10,348        output_token_count: 20,349        total_token_count: 30,350        cached_content_token_count: 5,351        thoughts_token_count: 2,352      } as ApiResponseEvent & {353        'event.name': typeof EVENT_API_RESPONSE;354      };355      const errorEvent = {356        'event.name': EVENT_API_ERROR,357        model: 'gemini-2.5-pro',358        duration_ms: 300,359        error_message: 'Something went wrong',360      } as ApiErrorEvent & { 'event.name': typeof EVENT_API_ERROR };361 362      service.addEvent(responseEvent);363      service.addEvent(errorEvent);364 365      const metrics = service.getMetrics();366      const modelAggregate = {367        api: {368          totalRequests: 2,369          totalErrors: 1,370          totalLatencyMs: 800,371        },372        tokens: {373          prompt: 10,374          candidates: 20,375          total: 30,376          cached: 5,377          thoughts: 2,378        },379      };380      expect(metrics.models['gemini-2.5-pro']).toEqual({381        ...modelAggregate,382        bySource: {383          [MAIN_SOURCE]: modelAggregate,384        },385      });386    });387  });388 389  describe('Subagent Source Attribution', () => {390    it('attributes API calls without subagent_name to MAIN_SOURCE', () => {391      const event = {392        'event.name': EVENT_API_RESPONSE,393        model: 'glm-5',394        duration_ms: 100,395        input_token_count: 10,396        output_token_count: 5,397        total_token_count: 15,398        cached_content_token_count: 0,399        thoughts_token_count: 0,400      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };401 402      service.addEvent(event);403 404      const modelMetrics = service.getMetrics().models['glm-5'];405      expect(Object.keys(modelMetrics.bySource)).toEqual([MAIN_SOURCE]);406      expect(modelMetrics.bySource[MAIN_SOURCE].api.totalRequests).toBe(1);407      expect(modelMetrics.api.totalRequests).toBe(1);408    });409 410    it('splits a single model between main and a subagent', () => {411      const mainEvent = {412        'event.name': EVENT_API_RESPONSE,413        model: 'glm-5',414        duration_ms: 200,415        input_token_count: 100,416        output_token_count: 50,417        total_token_count: 150,418        cached_content_token_count: 20,419        thoughts_token_count: 0,420      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };421      const subagentEvent = {422        'event.name': EVENT_API_RESPONSE,423        model: 'glm-5',424        duration_ms: 80,425        input_token_count: 40,426        output_token_count: 10,427        total_token_count: 50,428        cached_content_token_count: 0,429        thoughts_token_count: 0,430        subagent_name: 'echoer',431      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };432 433      service.addEvent(mainEvent);434      service.addEvent(subagentEvent);435 436      const modelMetrics = service.getMetrics().models['glm-5'];437      // Aggregate spans both main and subagent calls438      expect(modelMetrics.api.totalRequests).toBe(2);439      expect(modelMetrics.api.totalLatencyMs).toBe(280);440      expect(modelMetrics.tokens.prompt).toBe(140);441      expect(modelMetrics.tokens.total).toBe(200);442      // Per-source breakdown isolates each contributor443      expect(new Set(Object.keys(modelMetrics.bySource))).toEqual(444        new Set([MAIN_SOURCE, 'echoer']),445      );446      expect(modelMetrics.bySource[MAIN_SOURCE].api.totalRequests).toBe(1);447      expect(modelMetrics.bySource[MAIN_SOURCE].tokens.prompt).toBe(100);448      expect(modelMetrics.bySource['echoer'].api.totalRequests).toBe(1);449      expect(modelMetrics.bySource['echoer'].tokens.prompt).toBe(40);450    });451 452    it('splits two subagents sharing a model into distinct source buckets', () => {453      const makeEvent = (454        subagentName: string,455        duration: number,456      ): ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE } =>457        ({458          'event.name': EVENT_API_RESPONSE,459          model: 'glm-5',460          duration_ms: duration,461          input_token_count: 10,462          output_token_count: 5,463          total_token_count: 15,464          cached_content_token_count: 0,465          thoughts_token_count: 0,466          subagent_name: subagentName,467        }) as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };468 469      service.addEvent(makeEvent('alpha', 50));470      service.addEvent(makeEvent('bravo', 70));471 472      const modelMetrics = service.getMetrics().models['glm-5'];473      expect(modelMetrics.api.totalRequests).toBe(2);474      expect(Object.keys(modelMetrics.bySource).sort()).toEqual([475        'alpha',476        'bravo',477      ]);478      expect(modelMetrics.bySource['alpha'].api.totalRequests).toBe(1);479      expect(modelMetrics.bySource['bravo'].api.totalRequests).toBe(1);480      // Main bucket should NOT be created when no main-origin event arrived481      expect(modelMetrics.bySource[MAIN_SOURCE]).toBeUndefined();482    });483 484    it('handles a subagent named after an Object.prototype member without crashing', () => {485      // `constructor` is a valid subagent name per the naming regex. A486      // plain-object `bySource` would return `Object.prototype.constructor`487      // from a truthiness check, short-circuiting the bucket creation and488      // crashing the aggregation path. The prototype-free map prevents this.489      const event = {490        'event.name': EVENT_API_RESPONSE,491        model: 'glm-5',492        duration_ms: 100,493        input_token_count: 10,494        output_token_count: 5,495        total_token_count: 15,496        cached_content_token_count: 0,497        thoughts_token_count: 0,498        subagent_name: 'constructor',499      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };500 501      expect(() => service.addEvent(event)).not.toThrow();502 503      const modelMetrics = service.getMetrics().models['glm-5'];504      expect(modelMetrics.bySource['constructor']).toBeDefined();505      expect(modelMetrics.bySource['constructor'].api.totalRequests).toBe(1);506      expect(modelMetrics.bySource['constructor'].tokens.prompt).toBe(10);507      // Sanity: the Object prototype member was not actually mutated.508      expect(typeof modelMetrics.bySource['constructor']).toBe('object');509    });510 511    it('attributes API errors to the subagent source bucket', () => {512      const errorEvent = {513        'event.name': EVENT_API_ERROR,514        model: 'glm-5',515        duration_ms: 150,516        error_message: 'boom',517        subagent_name: 'alpha',518      } as ApiErrorEvent & { 'event.name': typeof EVENT_API_ERROR };519 520      service.addEvent(errorEvent);521 522      const modelMetrics = service.getMetrics().models['glm-5'];523      expect(modelMetrics.api.totalErrors).toBe(1);524      expect(modelMetrics.bySource['alpha'].api.totalErrors).toBe(1);525      expect(modelMetrics.bySource[MAIN_SOURCE]).toBeUndefined();526    });527  });528 529  describe('Tool Call Event Processing', () => {530    it('should process a single successful ToolCallEvent', () => {531      const toolCall = createFakeCompletedToolCall(532        'test_tool',533        true,534        150,535        ToolConfirmationOutcome.ProceedOnce,536      );537      service.addEvent({538        ...structuredClone(new ToolCallEvent(toolCall)),539        'event.name': EVENT_TOOL_CALL,540      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });541 542      const metrics = service.getMetrics();543      const { tools } = metrics;544 545      expect(tools.totalCalls).toBe(1);546      expect(tools.totalSuccess).toBe(1);547      expect(tools.totalFail).toBe(0);548      expect(tools.totalDurationMs).toBe(150);549      expect(tools.totalDecisions[ToolCallDecision.ACCEPT]).toBe(1);550      expect(tools.byName['test_tool']).toEqual({551        count: 1,552        success: 1,553        fail: 0,554        durationMs: 150,555        decisions: {556          [ToolCallDecision.ACCEPT]: 1,557          [ToolCallDecision.REJECT]: 0,558          [ToolCallDecision.MODIFY]: 0,559          [ToolCallDecision.AUTO_ACCEPT]: 0,560        },561      });562    });563 564    it('should process a single failed ToolCallEvent', () => {565      const toolCall = createFakeCompletedToolCall(566        'test_tool',567        false,568        200,569        ToolConfirmationOutcome.Cancel,570      );571      service.addEvent({572        ...structuredClone(new ToolCallEvent(toolCall)),573        'event.name': EVENT_TOOL_CALL,574      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });575 576      const metrics = service.getMetrics();577      const { tools } = metrics;578 579      expect(tools.totalCalls).toBe(1);580      expect(tools.totalSuccess).toBe(0);581      expect(tools.totalFail).toBe(1);582      expect(tools.totalDurationMs).toBe(200);583      expect(tools.totalDecisions[ToolCallDecision.REJECT]).toBe(1);584      expect(tools.byName['test_tool']).toEqual({585        count: 1,586        success: 0,587        fail: 1,588        durationMs: 200,589        decisions: {590          [ToolCallDecision.ACCEPT]: 0,591          [ToolCallDecision.REJECT]: 1,592          [ToolCallDecision.MODIFY]: 0,593          [ToolCallDecision.AUTO_ACCEPT]: 0,594        },595      });596    });597 598    it('should process a single cancelled ToolCallEvent', () => {599      const toolCall = createFakeCompletedToolCall(600        'test_tool',601        'cancelled',602        180,603        ToolConfirmationOutcome.Cancel,604      );605      service.addEvent({606        ...structuredClone(new ToolCallEvent(toolCall)),607        'event.name': EVENT_TOOL_CALL,608      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });609 610      const metrics = service.getMetrics();611      const { tools } = metrics;612 613      expect(tools.totalCalls).toBe(1);614      expect(tools.totalSuccess).toBe(0);615      expect(tools.totalFail).toBe(1);616      expect(tools.totalDurationMs).toBe(180);617      expect(tools.totalDecisions[ToolCallDecision.REJECT]).toBe(1);618      expect(tools.byName['test_tool']).toEqual({619        count: 1,620        success: 0,621        fail: 1,622        durationMs: 180,623        decisions: {624          [ToolCallDecision.ACCEPT]: 0,625          [ToolCallDecision.REJECT]: 1,626          [ToolCallDecision.MODIFY]: 0,627          [ToolCallDecision.AUTO_ACCEPT]: 0,628        },629      });630    });631 632    it('should process a ToolCallEvent with modify decision', () => {633      const toolCall = createFakeCompletedToolCall(634        'test_tool',635        true,636        250,637        ToolConfirmationOutcome.ModifyWithEditor,638      );639      service.addEvent({640        ...structuredClone(new ToolCallEvent(toolCall)),641        'event.name': EVENT_TOOL_CALL,642      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });643 644      const metrics = service.getMetrics();645      const { tools } = metrics;646 647      expect(tools.totalDecisions[ToolCallDecision.MODIFY]).toBe(1);648      expect(tools.byName['test_tool'].decisions[ToolCallDecision.MODIFY]).toBe(649        1,650      );651    });652 653    it('should process a ToolCallEvent without a decision', () => {654      const toolCall = createFakeCompletedToolCall('test_tool', true, 100);655      service.addEvent({656        ...structuredClone(new ToolCallEvent(toolCall)),657        'event.name': EVENT_TOOL_CALL,658      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });659 660      const metrics = service.getMetrics();661      const { tools } = metrics;662 663      expect(tools.totalDecisions).toEqual({664        [ToolCallDecision.ACCEPT]: 0,665        [ToolCallDecision.REJECT]: 0,666        [ToolCallDecision.MODIFY]: 0,667        [ToolCallDecision.AUTO_ACCEPT]: 0,668      });669      expect(tools.byName['test_tool'].decisions).toEqual({670        [ToolCallDecision.ACCEPT]: 0,671        [ToolCallDecision.REJECT]: 0,672        [ToolCallDecision.MODIFY]: 0,673        [ToolCallDecision.AUTO_ACCEPT]: 0,674      });675    });676 677    it('should aggregate multiple ToolCallEvents for the same tool', () => {678      const toolCall1 = createFakeCompletedToolCall(679        'test_tool',680        true,681        100,682        ToolConfirmationOutcome.ProceedOnce,683      );684      const toolCall2 = createFakeCompletedToolCall(685        'test_tool',686        false,687        150,688        ToolConfirmationOutcome.Cancel,689      );690 691      service.addEvent({692        ...structuredClone(new ToolCallEvent(toolCall1)),693        'event.name': EVENT_TOOL_CALL,694      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });695      service.addEvent({696        ...structuredClone(new ToolCallEvent(toolCall2)),697        'event.name': EVENT_TOOL_CALL,698      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });699 700      const metrics = service.getMetrics();701      const { tools } = metrics;702 703      expect(tools.totalCalls).toBe(2);704      expect(tools.totalSuccess).toBe(1);705      expect(tools.totalFail).toBe(1);706      expect(tools.totalDurationMs).toBe(250);707      expect(tools.totalDecisions[ToolCallDecision.ACCEPT]).toBe(1);708      expect(tools.totalDecisions[ToolCallDecision.REJECT]).toBe(1);709      expect(tools.byName['test_tool']).toEqual({710        count: 2,711        success: 1,712        fail: 1,713        durationMs: 250,714        decisions: {715          [ToolCallDecision.ACCEPT]: 1,716          [ToolCallDecision.REJECT]: 1,717          [ToolCallDecision.MODIFY]: 0,718          [ToolCallDecision.AUTO_ACCEPT]: 0,719        },720      });721    });722 723    it('should handle ToolCallEvents for different tools', () => {724      const toolCall1 = createFakeCompletedToolCall('tool_A', true, 100);725      const toolCall2 = createFakeCompletedToolCall('tool_B', false, 200);726      service.addEvent({727        ...structuredClone(new ToolCallEvent(toolCall1)),728        'event.name': EVENT_TOOL_CALL,729      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });730      service.addEvent({731        ...structuredClone(new ToolCallEvent(toolCall2)),732        'event.name': EVENT_TOOL_CALL,733      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });734 735      const metrics = service.getMetrics();736      const { tools } = metrics;737 738      expect(tools.totalCalls).toBe(2);739      expect(tools.totalSuccess).toBe(1);740      expect(tools.totalFail).toBe(1);741      expect(tools.byName['tool_A']).toBeDefined();742      expect(tools.byName['tool_B']).toBeDefined();743      expect(tools.byName['tool_A'].count).toBe(1);744      expect(tools.byName['tool_B'].count).toBe(1);745    });746 747    it('redacts function_args for structured_output calls while preserving metrics', () => {748      const toolCall = createFakeCompletedToolCall(749        'structured_output',750        true,751        250,752        ToolConfirmationOutcome.ProceedOnce,753      );754      // The fake helper hardcodes args to { foo: 'bar' }; in the real755      // structured-output flow this would be the user's extracted payload.756      // ToolCallEvent must not pass that through to telemetry.757      (toolCall.request as { args: Record<string, unknown> }).args = {758        secret: 'extracted private value',759      };760 761      const event = new ToolCallEvent(toolCall);762 763      expect(event.function_name).toBe('structured_output');764      expect(event.function_args).not.toHaveProperty('secret');765      expect(event.function_args).toEqual({766        __redacted: 'structured_output payload (see stdout result)',767      });768 769      // Metrics still flow through normally — duration, success, decision.770      service.addEvent({771        ...structuredClone(event),772        'event.name': EVENT_TOOL_CALL,773      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });774 775      const { tools } = service.getMetrics();776      expect(tools.totalCalls).toBe(1);777      expect(tools.totalSuccess).toBe(1);778      expect(tools.totalDurationMs).toBe(250);779      expect(tools.byName['structured_output']).toMatchObject({780        count: 1,781        success: 1,782        durationMs: 250,783      });784    });785 786    it('does not redact function_args for non-structured_output tools', () => {787      const toolCall = createFakeCompletedToolCall(788        'write_file',789        true,790        100,791        ToolConfirmationOutcome.ProceedOnce,792      );793      (toolCall.request as { args: Record<string, unknown> }).args = {794        path: '/tmp/x',795        content: 'hello',796      };797 798      const event = new ToolCallEvent(toolCall);799 800      expect(event.function_args).toEqual({801        path: '/tmp/x',802        content: 'hello',803      });804    });805  });806 807  describe('Skill Invocation Metrics', () => {808    it('aggregates successful and failed skill invocations by name', () => {809      service.recordSkillInvocation('review', true);810      service.recordSkillInvocation('review', false);811      service.recordSkillInvocation('testing', true);812 813      expect(service.getMetrics().skills).toEqual({814        totalCalls: 3,815        totalSuccess: 2,816        totalFail: 1,817        byName: {818          review: { count: 2, success: 1, fail: 1 },819          testing: { count: 1, success: 1, fail: 0 },820        },821      });822    });823 824    it('handles skill names that collide with object prototype keys', () => {825      service.recordSkillInvocation('constructor', true);826      service.recordSkillInvocation('__proto__', false);827 828      expect(service.getMetrics().skills?.byName['constructor']).toEqual({829        count: 1,830        success: 1,831        fail: 0,832      });833      expect(service.getMetrics().skills?.byName['__proto__']).toEqual({834        count: 1,835        success: 0,836        fail: 1,837      });838    });839  });840 841  describe('resetLastPromptTokenCount', () => {842    it('should reset the last prompt token count to 0', () => {843      // First, set up some initial token count844      const event = {845        'event.name': EVENT_API_RESPONSE,846        model: 'gemini-2.5-pro',847        duration_ms: 500,848        input_token_count: 100,849        output_token_count: 200,850        total_token_count: 300,851        cached_content_token_count: 50,852        thoughts_token_count: 20,853      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };854 855      service.addEvent(event);856      expect(service.getLastPromptTokenCount()).toBe(0);857 858      // Now reset the token count859      service.setLastPromptTokenCount(0);860      expect(service.getLastPromptTokenCount()).toBe(0);861    });862 863    it('should emit an update event when resetLastPromptTokenCount is called', () => {864      const spy = vi.fn();865      service.on('update', spy);866 867      // Set up initial token count868      const event = {869        'event.name': EVENT_API_RESPONSE,870        model: 'gemini-2.5-pro',871        duration_ms: 500,872        input_token_count: 100,873        output_token_count: 200,874        total_token_count: 300,875        cached_content_token_count: 50,876        thoughts_token_count: 20,877      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };878 879      service.addEvent(event);880      spy.mockClear(); // Clear the spy to focus on the reset call881 882      service.setLastPromptTokenCount(0);883 884      expect(spy).toHaveBeenCalledOnce();885      const { metrics, lastPromptTokenCount } = spy.mock.calls[0][0];886      expect(metrics).toBeDefined();887      expect(lastPromptTokenCount).toBe(0);888    });889 890    it('should not affect other metrics when resetLastPromptTokenCount is called', () => {891      // Set up initial state with some metrics892      const event = {893        'event.name': EVENT_API_RESPONSE,894        model: 'gemini-2.5-pro',895        duration_ms: 500,896        input_token_count: 100,897        output_token_count: 200,898        total_token_count: 300,899        cached_content_token_count: 50,900        thoughts_token_count: 20,901      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };902 903      service.addEvent(event);904 905      const metricsBefore = service.getMetrics();906 907      service.setLastPromptTokenCount(0);908 909      const metricsAfter = service.getMetrics();910 911      // Metrics should be unchanged912      expect(metricsAfter).toEqual(metricsBefore);913 914      // Only the last prompt token count should be reset915      expect(service.getLastPromptTokenCount()).toBe(0);916    });917 918    it('should work correctly when called multiple times', () => {919      const spy = vi.fn();920      service.on('update', spy);921 922      // Set up initial token count923      const event = {924        'event.name': EVENT_API_RESPONSE,925        model: 'gemini-2.5-pro',926        duration_ms: 500,927        input_token_count: 100,928        output_token_count: 200,929        total_token_count: 300,930        cached_content_token_count: 50,931        thoughts_token_count: 20,932      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };933 934      service.addEvent(event);935      expect(service.getLastPromptTokenCount()).toBe(0);936 937      // Reset once938      service.setLastPromptTokenCount(0);939      expect(service.getLastPromptTokenCount()).toBe(0);940 941      // Reset again - should still be 0 and still emit event942      spy.mockClear();943      service.setLastPromptTokenCount(0);944      expect(service.getLastPromptTokenCount()).toBe(0);945      expect(spy).toHaveBeenCalledOnce();946    });947 948    it('should correctly set status field for success/error/cancelled calls', () => {949      const successCall = createFakeCompletedToolCall(950        'success_tool',951        true,952        100,953      );954      const errorCall = createFakeCompletedToolCall('error_tool', false, 150);955      const cancelledCall = createFakeCompletedToolCall(956        'cancelled_tool',957        'cancelled',958        200,959      );960 961      const successEvent = new ToolCallEvent(successCall);962      const errorEvent = new ToolCallEvent(errorCall);963      const cancelledEvent = new ToolCallEvent(cancelledCall);964 965      // Verify status field is correctly set966      expect(successEvent.status).toBe('success');967      expect(errorEvent.status).toBe('error');968      expect(cancelledEvent.status).toBe('cancelled');969 970      // Verify backward compatibility with success field971      expect(successEvent.success).toBe(true);972      expect(errorEvent.success).toBe(false);973      expect(cancelledEvent.success).toBe(false);974    });975  });976 977  describe('Tool Call Event with Line Count Metadata', () => {978    it('should aggregate valid line count metadata', () => {979      const toolCall = createFakeCompletedToolCall('test_tool', true, 100);980      const event = {981        ...structuredClone(new ToolCallEvent(toolCall)),982        'event.name': EVENT_TOOL_CALL,983        metadata: {984          model_added_lines: 10,985          model_removed_lines: 5,986        },987      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL };988 989      service.addEvent(event);990 991      const metrics = service.getMetrics();992      expect(metrics.files.totalLinesAdded).toBe(10);993      expect(metrics.files.totalLinesRemoved).toBe(5);994    });995 996    it('should ignore null/undefined values in line count metadata', () => {997      const toolCall = createFakeCompletedToolCall('test_tool', true, 100);998      const event = {999        ...structuredClone(new ToolCallEvent(toolCall)),1000        'event.name': EVENT_TOOL_CALL,1001        metadata: {1002          model_added_lines: null,1003          model_removed_lines: undefined,1004        },1005      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL };1006 1007      service.addEvent(event);1008 1009      const metrics = service.getMetrics();1010      expect(metrics.files.totalLinesAdded).toBe(0);1011      expect(metrics.files.totalLinesRemoved).toBe(0);1012    });1013  });1014 1015  describe('Per-Session Metrics Isolation', () => {1016    const SESSION_A = 'session-aaa';1017    const SESSION_B = 'session-bbb';1018 1019    const makeApiEvent = (model: string, inputTokens: number) =>1020      ({1021        'event.name': EVENT_API_RESPONSE,1022        model,1023        duration_ms: 100,1024        input_token_count: inputTokens,1025        output_token_count: 10,1026        total_token_count: inputTokens + 10,1027        cached_content_token_count: 0,1028        thoughts_token_count: 0,1029      }) as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };1030 1031    const makeToolEvent = (name: string) =>1032      ({1033        'event.name': EVENT_TOOL_CALL,1034        function_name: name,1035        duration_ms: 50,1036        success: true,1037        decision: ToolCallDecision.AUTO_ACCEPT,1038        prompt_id: 'p1',1039      }) as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL };1040 1041    it('should isolate metrics by sessionId', () => {1042      service.addEvent(makeApiEvent('model-a', 100), SESSION_A);1043      service.addEvent(makeApiEvent('model-b', 200), SESSION_B);1044 1045      const metricsA = service.getMetricsForSession(SESSION_A);1046      const metricsB = service.getMetricsForSession(SESSION_B);1047 1048      expect(metricsA.models['model-a']?.tokens.prompt).toBe(100);1049      expect(metricsA.models['model-b']).toBeUndefined();1050 1051      expect(metricsB.models['model-b']?.tokens.prompt).toBe(200);1052      expect(metricsB.models['model-a']).toBeUndefined();1053    });1054 1055    it('should still accumulate to global metrics', () => {1056      service.addEvent(makeApiEvent('model-x', 100), SESSION_A);1057      service.addEvent(makeApiEvent('model-x', 200), SESSION_B);1058 1059      const global = service.getMetrics();1060      expect(global.models['model-x']?.tokens.prompt).toBe(300);1061    });1062 1063    it('should return empty metrics for unknown session', () => {1064      const metrics = service.getMetricsForSession('unknown');1065      expect(metrics.models).toEqual({});1066      expect(metrics.tools.totalCalls).toBe(0);1067    });1068 1069    it('should handle events without sessionId (global only)', () => {1070      service.addEvent(makeApiEvent('model-z', 50));1071 1072      const global = service.getMetrics();1073      expect(global.models['model-z']?.tokens.prompt).toBe(50);1074 1075      const sessionMetrics = service.getMetricsForSession('any-session');1076      expect(sessionMetrics.models).toEqual({});1077    });1078 1079    it('resetSession should clear only that session', () => {1080      service.addEvent(makeApiEvent('m', 100), SESSION_A);1081      service.addEvent(makeApiEvent('m', 200), SESSION_B);1082 1083      service.resetSession(SESSION_A);1084 1085      const metricsA = service.getMetricsForSession(SESSION_A);1086      const metricsB = service.getMetricsForSession(SESSION_B);1087 1088      expect(metricsA.models).toEqual({});1089      expect(metricsB.models['m']?.tokens.prompt).toBe(200);1090 1091      // Global should not be affected1092      const global = service.getMetrics();1093      expect(global.models['m']?.tokens.prompt).toBe(300);1094    });1095 1096    it('removeSession should prevent late events from recreating bucket', () => {1097      service.addEvent(makeApiEvent('m', 100), SESSION_A);1098      service.removeSession(SESSION_A);1099 1100      // Late event after removal1101      service.addEvent(makeApiEvent('m', 50), SESSION_A);1102 1103      // Session bucket should not be recreated1104      const metricsA = service.getMetricsForSession(SESSION_A);1105      expect(metricsA.models).toEqual({});1106 1107      // But global should still accumulate1108      const global = service.getMetrics();1109      expect(global.models['m']?.tokens.prompt).toBe(150);1110    });1111 1112    it('resetSession should re-enable a closed session', () => {1113      service.addEvent(makeApiEvent('m', 100), SESSION_A);1114      service.removeSession(SESSION_A);1115 1116      // Re-open the session1117      service.resetSession(SESSION_A);1118      service.addEvent(makeApiEvent('m', 50), SESSION_A);1119 1120      const metricsA = service.getMetricsForSession(SESSION_A);1121      expect(metricsA.models['m']?.tokens.prompt).toBe(50);1122    });1123 1124    it('should isolate tool call metrics by session', () => {1125      service.addEvent(makeToolEvent('Read'), SESSION_A);1126      service.addEvent(makeToolEvent('Write'), SESSION_B);1127      service.addEvent(makeToolEvent('Read'), SESSION_B);1128 1129      const metricsA = service.getMetricsForSession(SESSION_A);1130      const metricsB = service.getMetricsForSession(SESSION_B);1131 1132      expect(metricsA.tools.totalCalls).toBe(1);1133      expect(metricsA.tools.byName['Read']?.count).toBe(1);1134      expect(metricsA.tools.byName['Write']).toBeUndefined();1135 1136      expect(metricsB.tools.totalCalls).toBe(2);1137      expect(metricsB.tools.byName['Write']?.count).toBe(1);1138      expect(metricsB.tools.byName['Read']?.count).toBe(1);1139    });1140 1141    it('should isolate skill invocation metrics by session', () => {1142      service.recordSkillInvocation('review', true, SESSION_A);1143      service.recordSkillInvocation('review', false, SESSION_B);1144      service.recordSkillInvocation('testing', true, SESSION_B);1145 1146      const metricsA = service.getMetricsForSession(SESSION_A);1147      const metricsB = service.getMetricsForSession(SESSION_B);1148 1149      expect(metricsA.skills).toEqual({1150        totalCalls: 1,1151        totalSuccess: 1,1152        totalFail: 0,1153        byName: {1154          review: { count: 1, success: 1, fail: 0 },1155        },1156      });1157      expect(metricsB.skills).toEqual({1158        totalCalls: 2,1159        totalSuccess: 1,1160        totalFail: 1,1161        byName: {1162          review: { count: 1, success: 0, fail: 1 },1163          testing: { count: 1, success: 1, fail: 0 },1164        },1165      });1166    });1167 1168    it('removeSession should prevent late skill metrics from recreating bucket', () => {1169      service.recordSkillInvocation('review', true, SESSION_A);1170      service.removeSession(SESSION_A);1171 1172      service.recordSkillInvocation('review', false, SESSION_A);1173 1174      expect(service.getMetricsForSession(SESSION_A).skills).toEqual({1175        totalCalls: 0,1176        totalSuccess: 0,1177        totalFail: 0,1178        byName: {},1179      });1180      expect(service.getMetrics().skills?.byName['review']).toEqual({1181        count: 2,1182        success: 1,1183        fail: 1,1184      });1185    });1186 1187    it('resetSession should not clear global metrics (replay scenario)', () => {1188      // Simulate: session A active, session B being resumed1189      service.addEvent(makeApiEvent('m', 100), SESSION_A);1190      service.addEvent(makeApiEvent('m', 200), SESSION_B);1191 1192      // Resume session B: resetSession only clears B's bucket1193      service.resetSession(SESSION_B);1194 1195      // Session A untouched1196      const metricsA = service.getMetricsForSession(SESSION_A);1197      expect(metricsA.models['m']?.tokens.prompt).toBe(100);1198 1199      // Session B cleared1200      const metricsB = service.getMetricsForSession(SESSION_B);

Showing the first 1,200 of 1233 lines. Download the file for the rest.

basant307/AI_Governance_Project · CoolFace