Commit ac2ff72b authored by Alejandro Rodríguez's avatar Alejandro Rodríguez 🌴 Committed by Tomas Vik (OOO back on 2026-08-31)
Browse files

feat: Include model data on telemetry

parent 4b24d6ab

src/code_suggestions/code_suggestions_provider.test.ts

+48 −35
Original line number Diff line number Diff line 
import * as vscode from 'vscode';

import fetch from '../desktop/gitlab/fetch_logged';



import {

  TELEMETRY_HEADER_REQUESTS,

  TELEMETRY_HEADER_ERRORS,

  TELEMETRY_HEADER_ACCEPTS,

  codeSuggestionsTelemetry,

} from './code_suggestions_telemetry';

import { codeSuggestionsTelemetry } from './code_suggestions_telemetry';



import { GitLabPlatformManager } from '../common/platform/gitlab_platform';

import { project } from '../common/test_utils/entities';

import { log } from '../common/log';

import { asMock } from '../desktop/test_utils/as_mock';

import { PROGRAMMATIC_COMMANDS } from '../desktop/command_names';

import { CIRCUIT_BREAK_INTERVAL_MS, CodeSuggestionsProvider } from './code_suggestions_provider';
@@ -17,10 +13,11 @@ import { CodeSuggestionsState, CodeSuggestionsStateManager } from './code_sugges 
import { GitLabProject } from '../common/platform/gitlab_project';



jest.mock('../desktop/gitlab/fetch_logged');

jest.mock('../common/log');



const crossFetchCallArgument = () => JSON.parse(asMock(fetch).mock.calls[0][1].body);

const lastFetchCallHeaders = () =>

  asMock(fetch).mock.calls[asMock(fetch).mock.calls.length - 1][1].headers;

const lastFetchCallBody = () =>

  JSON.parse(asMock(fetch).mock.calls[asMock(fetch).mock.calls.length - 1][1].body);



const mockPrompt = 'const areaOfCube = ';

const mockDocumentPartial: Partial<vscode.TextDocument> = {
@@ -30,7 +27,10 @@ const mockDocumentPartial: Partial<vscode.TextDocument> = { 
};

const mockDocument = mockDocumentPartial as unknown as vscode.TextDocument;

const choice = '(side) => ';

const mockCompletions = { choices: [{ text: choice }] };

const mockCompletions = {

  choices: [{ text: choice }],

  model: { name: 'ensemble', engine: 'codegen' },

};



const mockPosition = {

  line: 0,
@@ -221,30 +221,30 @@ describe('CodeSuggestionsProvider', () => { 
    it('increases requests count for success request', async () => {

      await glcp.getCompletions(mockDocument, mockPosition);

      await glcp.getCompletions(mockDocument, mockPosition);

      const headers = lastFetchCallHeaders();

      const body = lastFetchCallBody();



      // We are always sending previous amount of requests, so it is off-by-one

      expect(headers[TELEMETRY_HEADER_REQUESTS]).toBe('1');

      expect(body.telemetry[0].requests).toBe(1);

    });



    it('includes telemetry headers', async () => {

    it('sends model information with the telemetry', async () => {

      asMock(fetch).mockResolvedValue({ ok: true, json: async () => mockCompletions });

      await glcp.getCompletions(mockDocument, mockPosition);

      await glcp.getCompletions(mockDocument, mockPosition);

      const headers = lastFetchCallHeaders();



      [TELEMETRY_HEADER_ACCEPTS, TELEMETRY_HEADER_ERRORS, TELEMETRY_HEADER_REQUESTS].forEach(

        header => {

          expect(headers).toHaveProperty(header);

        },

      );

    });

      const body = lastFetchCallBody();



      expect(body.telemetry[0].model_engine).toBe('codegen');

      expect(body.telemetry[0].model_name).toBe('ensemble');

    });

    it('increases requests count for success request', async () => {

      await glcp.getCompletions(mockDocument, mockPosition);

      await glcp.getCompletions(mockDocument, mockPosition);

      const headers = lastFetchCallHeaders();



      const body = lastFetchCallBody();

      // We are always sending previous amount of requests, so it is off-by-one

      expect(headers[TELEMETRY_HEADER_REQUESTS]).toBe('1');

      expect(body.telemetry[0].requests).toBe(1);

      expect(body.telemetry[0].model_engine).toBe('codegen');

      expect(body.telemetry[0].model_name).toBe('ensemble');

    });



    it('increases request count and request errors for failed requests', async () => {
@@ -252,10 +252,10 @@ describe('CodeSuggestionsProvider', () => { 
      await glcp.getCompletions(mockDocument, mockPosition);



      await glcp.getCompletions(mockDocument, mockPosition);

      const headers = lastFetchCallHeaders();

      const body = lastFetchCallBody();



      expect(headers[TELEMETRY_HEADER_REQUESTS]).toBe('1');

      expect(headers[TELEMETRY_HEADER_ERRORS]).toBe('1');

      expect(body.telemetry[0].requests).toBe(1);

      expect(body.telemetry[0].errors).toBe(1);

    });



    it('does not reset request count and request errors for failed requests', async () => {
@@ -266,10 +266,10 @@ describe('CodeSuggestionsProvider', () => { 
      await glcp.getCompletions(mockDocument, mockPosition);



      await glcp.getCompletions(mockDocument, mockPosition);

      const headers = lastFetchCallHeaders();

      const body = lastFetchCallBody();



      expect(headers[TELEMETRY_HEADER_REQUESTS]).toBe('2');

      expect(headers[TELEMETRY_HEADER_ERRORS]).toBe('2');

      expect(body.telemetry[0].requests).toBe(2);

      expect(body.telemetry[0].errors).toBe(2);

    });



    it('resets counters on successful requests', async () => {
@@ -278,10 +278,10 @@ describe('CodeSuggestionsProvider', () => { 


      await glcp.getCompletions(mockDocument, mockPosition);

      await glcp.getCompletions(mockDocument, mockPosition);

      const headers = lastFetchCallHeaders();

      const body = lastFetchCallBody();



      expect(headers[TELEMETRY_HEADER_REQUESTS]).toBe('1');

      expect(headers[TELEMETRY_HEADER_ERRORS]).toBe('0');

      expect(body.telemetry[0].requests).toBe(1);

      expect(body.telemetry[0].errors).toBe(0);

    });



    it('includes correct command when completion is accepted', async () => {
@@ -292,12 +292,25 @@ describe('CodeSuggestionsProvider', () => { 
    });



    it('sends correct accepted value when it is increased in telemetry', async () => {

      codeSuggestionsTelemetry.incAcceptCount();

      codeSuggestionsTelemetry.incAcceptCount({ name: 'testModel', engine: 'testEngine' });

      await glcp.getCompletions(mockDocument, mockPosition);

      const headers = lastFetchCallHeaders();

      const body = lastFetchCallBody();



      expect(headers[TELEMETRY_HEADER_REQUESTS]).toBe('0');

      expect(headers[TELEMETRY_HEADER_ACCEPTS]).toBe('1');

      expect(body.telemetry[0].requests).toBe(0);

      expect(body.telemetry[0].accepts).toBe(1);

    });



    describe('logging', () => {

      const getLoggedMessage = () => jest.mocked(log.debug).mock.calls[1][0];



      it('logs the telemetry details', async () => {

        codeSuggestionsTelemetry.incAcceptCount({ name: 'testModel', engine: 'testEngine' });

        await glcp.getCompletions(mockDocument, mockPosition);



        expect(getLoggedMessage()).toContain(

          'AI Assist: fetching completions ... (telemetry: [{"model_engine":"testEngine","model_name":"testModel","requests":0,"accepts":1,"errors":0}])',

        );

      });

    });

  });

});

src/code_suggestions/code_suggestions_provider.ts

+20 −16
Original line number Diff line number Diff line
@@ -11,14 +11,10 @@ import { CodeSuggestionsTokenManager, CompletionToken } from './code_suggestions 
import { getUserAgentHeader } from '../desktop/gitlab/http/get_user_agent_header';

import { CircuitBreaker } from './circuit_breaker';

import { PROGRAMMATIC_COMMANDS } from '../desktop/command_names';

import {

  TELEMETRY_HEADER_ACCEPTS,

  TELEMETRY_HEADER_ERRORS,

  TELEMETRY_HEADER_REQUESTS,

  codeSuggestionsTelemetry,

} from './code_suggestions_telemetry';

import { Model, Telemetry, codeSuggestionsTelemetry } from './code_suggestions_telemetry';

import { GitLabProject } from '../common/platform/gitlab_project';

import { CodeSuggestionsState, CodeSuggestionsStateManager } from './code_suggestions_state';

import { prettyJson } from '../common/utils/json';



export const CIRCUIT_BREAK_INTERVAL_MS = 10000;

export const MAX_ERRORS_BEFORE_CIRCUIT_BREAK = 4;
@@ -31,7 +27,7 @@ interface Choice { 


interface CodeSuggestionsResponse {

  id: string;

  model: string;

  model: Model;

  object: string;

  created: number;

  choices: Choice[];
@@ -43,11 +39,13 @@ interface CurrentFile { 
  content_below_cursor: string;

  file_name: string;

}



interface CodeSuggestionPrompt {

  current_file: CurrentFile;

  prompt_version: number;

  project_id?: number;

  project_path?: string;

  telemetry: Telemetry[];

}



export class CodeSuggestionsProvider implements vscode.InlineCompletionItemProvider {
@@ -125,6 +123,7 @@ export class CodeSuggestionsProvider implements vscode.InlineCompletionItemProvi 
        content_below_cursor: contentBelowCursor,

      },

      ...projectInfo,

      telemetry: codeSuggestionsTelemetry.toArray(),

    };



    return payload;
@@ -156,14 +155,23 @@ export class CodeSuggestionsProvider implements vscode.InlineCompletionItemProvi 
      return [];

    }



    log.debug(`AI Assist: fetching completions ...`);

    log.debug(

      `AI Assist: fetching completions ... (telemetry: ${prettyJson(

        codeSuggestionsTelemetry.toArray(),

        0,

      )})`,

    );



    let response: CodeSuggestionsResponse;

    let model: Model = { engine: '', name: '' }; // Defaults for telemetry of failed requests



    try {

      this.stateManager.setState(CodeSuggestionsState.LOADING);

      response = await this.fetchCompletions(token, prompt);

      this.stateManager.setState(CodeSuggestionsState.OK);



      if (response.model !== undefined) model = response.model;



      this.circuitBreaker.success();



      // The previous counts were successfully sent...
@@ -172,11 +180,11 @@ export class CodeSuggestionsProvider implements vscode.InlineCompletionItemProvi 
      log.error(`AI Assist: Error fetching completions: ${e.toString()}`);

      this.circuitBreaker.error();

      this.stateManager.setState(CodeSuggestionsState.ERROR);

      codeSuggestionsTelemetry.incErrorCount();

      codeSuggestionsTelemetry.incErrorCount(model);

      return [];

    } finally {

      // Keep track of this request for next send..

      codeSuggestionsTelemetry.incRequestCount();

      codeSuggestionsTelemetry.incRequestCount(model);

    }



    const choices = response.choices || [];
@@ -187,7 +195,7 @@ export class CodeSuggestionsProvider implements vscode.InlineCompletionItemProvi 
    const acceptedCommand: vscode.Command = {

      title: 'Code Suggestion Accepted',

      command: PROGRAMMATIC_COMMANDS.CODE_SUGGESTION_ACCEPTED,

      arguments: [],

      arguments: [model],

    };



    return choices.map(
@@ -227,8 +235,7 @@ export class CodeSuggestionsProvider implements vscode.InlineCompletionItemProvi 
    token: CompletionToken,

    prompt: CodeSuggestionPrompt,

  ): Promise<CodeSuggestionsResponse> {

    const telemetryString = `requests:${codeSuggestionsTelemetry.requestCount}/errors:${codeSuggestionsTelemetry.errorCount}/accepts:${codeSuggestionsTelemetry.acceptCount}`;

    log.debug(`AI Assist: fetching completions... (telemetry: ${telemetryString})`);

    log.debug(`AI Assist: fetching completions...`);



    const requestOptions = {

      method: 'POST',
@@ -237,9 +244,6 @@ export class CodeSuggestionsProvider implements vscode.InlineCompletionItemProvi 
        'X-Gitlab-Authentication-Type': 'oidc',

        Authorization: `Bearer ${token.access_token}`,

        'Content-Type': 'application/json',

        [TELEMETRY_HEADER_REQUESTS]: codeSuggestionsTelemetry.requestCount.toString(),

        [TELEMETRY_HEADER_ERRORS]: codeSuggestionsTelemetry.errorCount.toString(),

        [TELEMETRY_HEADER_ACCEPTS]: codeSuggestionsTelemetry.acceptCount.toString(),

      },

      body: JSON.stringify(prompt),

    };

src/code_suggestions/code_suggestions_telemetry.test.ts

0 → 100644
+65 −0
Original line number Diff line number Diff line 
import { compareBy } from '../desktop/utils/compare_by';

import { CodeSuggestionsTelemetry, Model, Telemetry } from './code_suggestions_telemetry';



describe('CodeSuggestionsTelemetry', () => {

  let telemetry: CodeSuggestionsTelemetry;



  const codegenModel: Model = {

    engine: 'gitlab-native',

    name: 'codegen-v2-1.0.0',

  };



  const codegenTelemetry: Telemetry = {

    model_name: codegenModel.name,

    model_engine: codegenModel.engine,

    accepts: 0,

    errors: 0,

    requests: 0,

  };



  beforeEach(() => {

    telemetry = new CodeSuggestionsTelemetry();

  });



  it('increases request count', () => {

    telemetry.incRequestCount(codegenModel);



    expect(telemetry.toArray()).toEqual([{ ...codegenTelemetry, requests: 1 }]);

  });



  it('increases accept count', () => {

    telemetry.incAcceptCount(codegenModel);



    expect(telemetry.toArray()).toEqual([{ ...codegenTelemetry, accepts: 1 }]);

  });



  it('increases error count', () => {

    telemetry.incErrorCount(codegenModel);



    expect(telemetry.toArray()).toEqual([{ ...codegenTelemetry, errors: 1 }]);

  });



  it('handles multiple models', () => {

    const bisonModel: Model = {

      engine: 'code-bison',

      name: 'v2-1.0.0',

    };



    telemetry.incRequestCount(codegenModel);

    telemetry.incAcceptCount(codegenModel);



    telemetry.incRequestCount(bisonModel);

    telemetry.incErrorCount(bisonModel);



    expect(telemetry.toArray().sort(compareBy('model_engine'))).toEqual([

      {

        model_engine: bisonModel.engine,

        model_name: bisonModel.name,

        requests: 1,

        errors: 1,

        accepts: 0,

      },

      { ...codegenTelemetry, requests: 1, accepts: 1 },

    ]);

  });

});

src/code_suggestions/code_suggestions_telemetry.ts

+44 −16
Original line number Diff line number Diff line 
export const TELEMETRY_HEADER_REQUESTS = 'X-GitLab-CS-Requests';

export const TELEMETRY_HEADER_ERRORS = 'X-GitLab-CS-Errors';

export const TELEMETRY_HEADER_ACCEPTS = 'X-GitLab-CS-Accepts';

export interface Model {

  engine: string;

  name: string;

}



export class CodeSuggestionsTelemetry {

  requestCount = 0;

export interface Telemetry {

  model_engine: string;

  model_name: string;

  requests: number;

  accepts: number;

  errors: number;

}



type ModelKey = string;



const DELIMITER = '\u200B'; // zero width space minimises chance of conflict with model name or engine characters



const modelToKey = (m: Model): ModelKey => `${m.engine}${DELIMITER}${m.name}`;



  errorCount = 0;

const createEmptyTelemetry = (m: Model): Telemetry => ({

  model_engine: m.engine,

  model_name: m.name,

  requests: 0,

  accepts: 0,

  errors: 0,

});



  acceptCount = 0;

export class CodeSuggestionsTelemetry {

  #telemetryByModel: Record<ModelKey, Telemetry> = {};



  incRequestCount() {

    this.requestCount += 1;

  incRequestCount(model: Model) {

    this.#getTelemetry(model).requests += 1;

  }



  incErrorCount() {

    this.errorCount += 1;

  incErrorCount(model: Model) {

    this.#getTelemetry(model).errors += 1;

  }



  incAcceptCount() {

    this.acceptCount += 1;

  incAcceptCount(model: Model) {

    this.#getTelemetry(model).accepts += 1;

  }



  resetCounts() {

    this.requestCount = 0;

    this.errorCount = 0;

    this.acceptCount = 0;

    this.#telemetryByModel = {};

  }



  #getTelemetry(model: Model) {

    if (this.#telemetryByModel[modelToKey(model)] === undefined) {

      this.#telemetryByModel[modelToKey(model)] = createEmptyTelemetry(model);

    }

    return this.#telemetryByModel[modelToKey(model)];

  }



  toArray() {

    return Object.values(this.#telemetryByModel);

  }

}

src/code_suggestions/commands/code_suggestion_accepted.test.ts

+2 −2
Original line number Diff line number Diff line
@@ -3,8 +3,8 @@ import { codeSuggestionsTelemetry } from '../code_suggestions_telemetry'; 


describe('code suggestion accepted command', () => {

  it('updates codeSuggestionsTelemetry with the correct value', async () => {

    await codeSuggestionAccepted();

    await codeSuggestionAccepted({ name: 'ensemble', engine: 'codegen' });



    expect(codeSuggestionsTelemetry.acceptCount).toBe(1);

    expect(codeSuggestionsTelemetry.toArray()[0].accepts).toBe(1);

  });

});