basant307/AI_Governance_Project

1/**2 * @license3 * Copyright 2025 Google LLC4 * SPDX-License-Identifier: Apache-2.05 */6 7import { EventEmitter } from 'node:events';8import {9  EVENT_API_ERROR,10  EVENT_API_RESPONSE,11  EVENT_TOOL_CALL,12} from './constants.js';13 14import { ToolCallDecision } from './tool-call-decision.js';15import type {16  ApiErrorEvent,17  ApiResponseEvent,18  ToolCallEvent,19} from './types.js';20import { MAIN_SOURCE } from '../utils/subagentNameContext.js';21 22export { MAIN_SOURCE } from '../utils/subagentNameContext.js';23 24export type UiEvent =25  | (ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE })26  | (ApiErrorEvent & { 'event.name': typeof EVENT_API_ERROR })27  | (ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });28 29export {30  EVENT_API_ERROR,31  EVENT_API_RESPONSE,32  EVENT_TOOL_CALL,33} from './constants.js';34 35export interface ToolCallStats {36  count: number;37  success: number;38  fail: number;39  durationMs: number;40  decisions: {41    [ToolCallDecision.ACCEPT]: number;42    [ToolCallDecision.REJECT]: number;43    [ToolCallDecision.MODIFY]: number;44    [ToolCallDecision.AUTO_ACCEPT]: number;45  };46}47 48export interface SkillCallStats {49  count: number;50  success: number;51  fail: number;52}53 54export interface SkillMetrics {55  totalCalls: number;56  totalSuccess: number;57  totalFail: number;58  byName: Record<string, SkillCallStats>;59}60 61/**62 * Per-model counters without the nested source breakdown. Used both as the63 * aggregate `ModelMetrics` shape (via extension) and as the value type of the64 * `bySource` map — keeping the type non-recursive.65 */66export interface ModelMetricsCore {67  api: {68    totalRequests: number;69    totalErrors: number;70    totalLatencyMs: number;71  };72  tokens: {73    prompt: number;74    candidates: number;75    total: number;76    cached: number;77    thoughts: number;78  };79}80 81export interface ModelMetrics extends ModelMetricsCore {82  /**83   * Per-source breakdown. Keys are subagent names, or `MAIN_SOURCE` ("main")84   * for calls originating from the main conversation. Every API call that85   * increments an aggregate counter also increments the matching per-source86   * record so the two views stay consistent.87   */88  bySource: Record<string, ModelMetricsCore>;89}90 91export interface SessionMetrics {92  models: Record<string, ModelMetrics>;93  tools: {94    totalCalls: number;95    totalSuccess: number;96    totalFail: number;97    totalDurationMs: number;98    totalDecisions: {99      [ToolCallDecision.ACCEPT]: number;100      [ToolCallDecision.REJECT]: number;101      [ToolCallDecision.MODIFY]: number;102      [ToolCallDecision.AUTO_ACCEPT]: number;103    };104    byName: Record<string, ToolCallStats>;105  };106  files: {107    totalLinesAdded: number;108    totalLinesRemoved: number;109  };110  skills?: SkillMetrics;111}112 113const createInitialModelMetricsCore = (): ModelMetricsCore => ({114  api: {115    totalRequests: 0,116    totalErrors: 0,117    totalLatencyMs: 0,118  },119  tokens: {120    prompt: 0,121    candidates: 0,122    total: 0,123    cached: 0,124    thoughts: 0,125  },126});127 128// `bySource` keys are user-controlled subagent names. Using a prototype-free129// map avoids crashes when a subagent is named after an inherited Object130// member (e.g. `constructor`, `toString`, `hasOwnProperty`), which would131// otherwise short-circuit `!bySource[name]` checks and return the inherited132// prototype member as the "bucket".133const createInitialModelMetrics = (): ModelMetrics => ({134  ...createInitialModelMetricsCore(),135  bySource: Object.create(null) as Record<string, ModelMetricsCore>,136});137 138const createInitialSkillMetrics = (): SkillMetrics => ({139  totalCalls: 0,140  totalSuccess: 0,141  totalFail: 0,142  byName: {},143});144 145const createInitialMetrics = (): SessionMetrics => ({146  models: {},147  tools: {148    totalCalls: 0,149    totalSuccess: 0,150    totalFail: 0,151    totalDurationMs: 0,152    totalDecisions: {153      [ToolCallDecision.ACCEPT]: 0,154      [ToolCallDecision.REJECT]: 0,155      [ToolCallDecision.MODIFY]: 0,156      [ToolCallDecision.AUTO_ACCEPT]: 0,157    },158    byName: {},159  },160  files: {161    totalLinesAdded: 0,162    totalLinesRemoved: 0,163  },164  skills: createInitialSkillMetrics(),165});166 167export class UiTelemetryService extends EventEmitter {168  static readonly #MAX_CLOSED_SESSIONS = 1000;169  #metrics: SessionMetrics = createInitialMetrics();170  #sessionMetrics: Map<string, SessionMetrics> = new Map();171  #closedSessions: Set<string> = new Set();172  #lastPromptTokenCount = 0;173  #lastCachedContentTokenCount = 0;174  #sessionStartTime: Date = new Date();175 176  addEvent(event: UiEvent, sessionId?: string) {177    if (!this.#accumulateEvent(this.#metrics, event)) return;178 179    if (sessionId && !this.#closedSessions.has(sessionId)) {180      if (!this.#sessionMetrics.has(sessionId)) {181        this.#sessionMetrics.set(sessionId, createInitialMetrics());182      }183      this.#accumulateEvent(this.#sessionMetrics.get(sessionId)!, event);184    }185 186    this.emit('update', {187      metrics: this.#metrics,188      lastPromptTokenCount: this.#lastPromptTokenCount,189    });190  }191 192  getMetrics(): SessionMetrics {193    return this.#metrics;194  }195 196  getMetricsForSession(sessionId: string): SessionMetrics {197    return this.#sessionMetrics.get(sessionId) ?? createInitialMetrics();198  }199 200  recordSkillInvocation(201    skillName: string,202    success: boolean,203    sessionId?: string,204  ): void {205    this.#accumulateSkillInvocation(this.#metrics, skillName, success);206 207    if (sessionId && !this.#closedSessions.has(sessionId)) {208      if (!this.#sessionMetrics.has(sessionId)) {209        this.#sessionMetrics.set(sessionId, createInitialMetrics());210      }211      this.#accumulateSkillInvocation(212        this.#sessionMetrics.get(sessionId)!,213        skillName,214        success,215      );216    }217 218    this.emit('update', {219      metrics: this.#metrics,220      lastPromptTokenCount: this.#lastPromptTokenCount,221    });222  }223 224  getLastPromptTokenCount(): number {225    return this.#lastPromptTokenCount;226  }227 228  setLastPromptTokenCount(lastPromptTokenCount: number): void {229    this.#lastPromptTokenCount = lastPromptTokenCount;230    this.emit('update', {231      metrics: this.#metrics,232      lastPromptTokenCount: this.#lastPromptTokenCount,233    });234  }235 236  getSessionStartTime(): Date {237    return this.#sessionStartTime;238  }239 240  getLastCachedContentTokenCount(): number {241    return this.#lastCachedContentTokenCount;242  }243 244  setLastCachedContentTokenCount(count: number): void {245    this.#lastCachedContentTokenCount = count;246  }247 248  /**249   * Resets metrics to the initial state (used when resuming a session).250   */251  reset(): void {252    this.#metrics = createInitialMetrics();253    this.#sessionMetrics.clear();254    this.#closedSessions.clear();255    this.#lastPromptTokenCount = 0;256    this.#lastCachedContentTokenCount = 0;257    this.#sessionStartTime = new Date();258    this.emit('update', {259      metrics: this.#metrics,260      lastPromptTokenCount: this.#lastPromptTokenCount,261    });262  }263 264  resetSession(sessionId: string): void {265    this.#sessionMetrics.set(sessionId, createInitialMetrics());266    this.#closedSessions.delete(sessionId);267  }268 269  removeSession(sessionId: string): void {270    this.#sessionMetrics.delete(sessionId);271    this.#closedSessions.add(sessionId);272    if (this.#closedSessions.size > UiTelemetryService.#MAX_CLOSED_SESSIONS) {273      const oldest = this.#closedSessions.values().next().value;274      if (oldest) this.#closedSessions.delete(oldest);275    }276  }277 278  #accumulateEvent(metrics: SessionMetrics, event: UiEvent): boolean {279    switch (event['event.name']) {280      case EVENT_API_RESPONSE:281        this.#accumulateApiResponse(metrics, event);282        return true;283      case EVENT_API_ERROR:284        this.#accumulateApiError(metrics, event);285        return true;286      case EVENT_TOOL_CALL:287        this.#accumulateToolCall(metrics, event);288        return true;289      default:290        return false;291    }292  }293 294  #accumulateApiResponse(295    metrics: SessionMetrics,296    event: ApiResponseEvent,297  ): void {298    const modelMetrics = this.#getOrCreateModelMetrics(metrics, event.model);299    const sourceMetrics = this.#getOrCreateSourceMetrics(300      modelMetrics,301      event.subagent_name ?? MAIN_SOURCE,302    );303 304    for (const bucket of [modelMetrics, sourceMetrics]) {305      bucket.api.totalRequests++;306      bucket.api.totalLatencyMs += event.duration_ms;307 308      bucket.tokens.prompt += event.input_token_count;309      bucket.tokens.candidates += event.output_token_count;310      bucket.tokens.total += event.total_token_count;311      bucket.tokens.cached += event.cached_content_token_count;312      bucket.tokens.thoughts += event.thoughts_token_count;313    }314  }315 316  #accumulateApiError(metrics: SessionMetrics, event: ApiErrorEvent): void {317    const modelMetrics = this.#getOrCreateModelMetrics(metrics, event.model);318    const sourceMetrics = this.#getOrCreateSourceMetrics(319      modelMetrics,320      event.subagent_name ?? MAIN_SOURCE,321    );322 323    for (const bucket of [modelMetrics, sourceMetrics]) {324      bucket.api.totalRequests++;325      bucket.api.totalErrors++;326      bucket.api.totalLatencyMs += event.duration_ms;327    }328  }329 330  #accumulateToolCall(metrics: SessionMetrics, event: ToolCallEvent): void {331    const { tools, files } = metrics;332    tools.totalCalls++;333    tools.totalDurationMs += event.duration_ms;334 335    if (event.success) {336      tools.totalSuccess++;337    } else {338      tools.totalFail++;339    }340 341    if (!tools.byName[event.function_name]) {342      tools.byName[event.function_name] = {343        count: 0,344        success: 0,345        fail: 0,346        durationMs: 0,347        decisions: {348          [ToolCallDecision.ACCEPT]: 0,349          [ToolCallDecision.REJECT]: 0,350          [ToolCallDecision.MODIFY]: 0,351          [ToolCallDecision.AUTO_ACCEPT]: 0,352        },353      };354    }355 356    const toolStats = tools.byName[event.function_name];357    toolStats.count++;358    toolStats.durationMs += event.duration_ms;359    if (event.success) {360      toolStats.success++;361    } else {362      toolStats.fail++;363    }364 365    if (event.decision) {366      tools.totalDecisions[event.decision]++;367      toolStats.decisions[event.decision]++;368    }369 370    if (event.metadata) {371      if (event.metadata['model_added_lines'] !== undefined) {372        files.totalLinesAdded += event.metadata['model_added_lines'];373      }374      if (event.metadata['model_removed_lines'] !== undefined) {375        files.totalLinesRemoved += event.metadata['model_removed_lines'];376      }377    }378  }379 380  #accumulateSkillInvocation(381    metrics: SessionMetrics,382    skillName: string,383    success: boolean,384  ): void {385    const skills = metrics.skills ?? createInitialSkillMetrics();386    metrics.skills = skills;387 388    skills.totalCalls++;389    if (success) {390      skills.totalSuccess++;391    } else {392      skills.totalFail++;393    }394 395    if (!Object.prototype.hasOwnProperty.call(skills.byName, skillName)) {396      Object.defineProperty(skills.byName, skillName, {397        value: {398          count: 0,399          success: 0,400          fail: 0,401        },402        enumerable: true,403        configurable: true,404        writable: true,405      });406    }407 408    const skillStats = skills.byName[skillName];409    if (!skillStats) {410      return;411    }412    skillStats.count++;413    if (success) {414      skillStats.success++;415    } else {416      skillStats.fail++;417    }418  }419 420  #getOrCreateModelMetrics(421    metrics: SessionMetrics,422    modelName: string,423  ): ModelMetrics {424    if (!metrics.models[modelName]) {425      metrics.models[modelName] = createInitialModelMetrics();426    }427    return metrics.models[modelName];428  }429 430  #getOrCreateSourceMetrics(431    modelMetrics: ModelMetrics,432    source: string,433  ): ModelMetricsCore {434    if (!modelMetrics.bySource[source]) {435      modelMetrics.bySource[source] = createInitialModelMetricsCore();436    }437    return modelMetrics.bySource[source];438  }439}440 441export const uiTelemetryService = new UiTelemetryService();442