Commit 2351f4f0 authored by Shekhar Patnaik's avatar Shekhar Patnaik 💬 Committed by Tomas Vik (OOO back on 2026-08-31)
Browse files

feat: Add circuit breaker to code completion

parent a08a7993

src/__mocks__/vscode.js

+4 −0
Original line number Diff line number Diff line
@@ -108,4 +108,8 @@ module.exports = { 
  InlineCompletionTriggerKind: {

    Automatic: true,

  },

  InlineCompletionItem: function InlineCompletionItem(insertText, range) {

    this.insertText = insertText;

    this.range = range;

  },

};

src/completion/circuit_breaker.test.ts

0 → 100644
+46 −0
Original line number Diff line number Diff line 
import { CircuitBreaker } from './circuit_breaker';



describe('CircuitBreaker', () => {

  it('does not break by default', () => {

    const cb = new CircuitBreaker(2, 1000);

    expect(cb.isBreaking()).toBe(false);

  });



  it('breaks after the max number of errors has been reached', () => {

    const cb = new CircuitBreaker(2, 1000);

    cb.error();

    expect(cb.isBreaking()).toBe(false);

    cb.error();

    expect(cb.isBreaking()).toBe(true);

  });



  it('success removes the break', () => {

    const cb = new CircuitBreaker(1, 1000);

    cb.error();

    expect(cb.isBreaking()).toBe(true);

    cb.success();

    expect(cb.isBreaking()).toBe(false);

  });



  it('removes the break after the breakTimeMs', () => {

    const errorDate = new Date();

    jest.useFakeTimers().setSystemTime(errorDate);

    const cb = new CircuitBreaker(1, 1000);

    cb.error();

    expect(cb.isBreaking()).toBe(true);

    // after 1001 ms, the break is removed

    jest.useFakeTimers().setSystemTime(new Date(errorDate.valueOf() + 1001));

    expect(cb.isBreaking()).toBe(false);

  });



  it('instantly breaks if an error comes after expired break', () => {

    const errorDate = new Date();

    jest.useFakeTimers().setSystemTime(errorDate);

    const cb = new CircuitBreaker(1, 1000);

    cb.error();

    // after 1001 ms, the break is removed

    jest.useFakeTimers().setSystemTime(new Date(errorDate.valueOf() + 1001));

    cb.error();

    expect(cb.isBreaking()).toBe(true);

  });

});

src/completion/circuit_breaker.ts

0 → 100644
+30 −0
Original line number Diff line number Diff line 
export class CircuitBreaker {

  readonly #maxErrorsBeforeBreaking: number;



  readonly #breakTimeMs: number;



  #errorCount = 0;



  #tryAgainAfterTimestamp = 0;



  constructor(maxErrorsBeforeBreaking: number, breakTimeMs: number) {

    this.#maxErrorsBeforeBreaking = maxErrorsBeforeBreaking;

    this.#breakTimeMs = breakTimeMs;

  }



  error() {

    this.#errorCount += 1;

    if (this.#errorCount >= this.#maxErrorsBeforeBreaking) {

      this.#tryAgainAfterTimestamp = Date.now() + this.#breakTimeMs;

    }

  }



  isBreaking() {

    return Date.now() < this.#tryAgainAfterTimestamp;

  }



  success() {

    this.#errorCount = 0;

    this.#tryAgainAfterTimestamp = 0;

  }

}

src/completion/gitlab_code_completion_provider.test.ts

+62 −15
Original line number Diff line number Diff line 
import * as vscode from 'vscode';

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



import { GitLabCodeCompletionProvider } from './gitlab_code_completion_provider';

import {

  GitLabCodeCompletionProvider,

  CIRCUIT_BREAK_INTERVAL_MS,

} from './gitlab_code_completion_provider';

import { GitLabPlatform, GitLabPlatformProject } from '../gitlab/gitlab_platform';



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

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



const mockPrompt = 'const areaOfCube = ';

const mockDocumentPartial: Partial<vscode.TextDocument> = {

  fileName: 'test.js',

  getText: () => mockPrompt,

  lineAt: () => ({ text: mockPrompt } as vscode.TextLine),

};

const mockDocument = mockDocumentPartial as unknown as vscode.TextDocument;

const choice = '(side) => ';

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



const mockPosition = {

  line: 0,

  character: mockPrompt.length,

} as vscode.Position;



const platform: GitLabPlatform = {

  getProject: async () => ({ project: 'gitlab' } as GitLabPlatformProject),

  fetchFromApi: async <T>(): Promise<T> =>
@@ -48,15 +66,6 @@ describe('GitLabCodeCompletionProvider', () => { 
  });



  describe('provideInlineCompletionItems', () => {

    const mockPrompt = 'const areaOfCube = ';

    const mockDocument: Partial<vscode.TextDocument> = {

      getText: () => mockPrompt,

      lineAt: () => ({ text: mockPrompt } as vscode.TextLine),

    };

    const mockPosition = {

      line: 0,

      character: mockPrompt.length,

    } as vscode.Position;

    const mockInlineCompletions = [] as vscode.InlineCompletionItem[];

    const mockContext = {

      triggerKind: vscode.InlineCompletionTriggerKind.Automatic,
@@ -68,15 +77,53 @@ describe('GitLabCodeCompletionProvider', () => { 
      glcp.getCompletions = jest.fn().mockResolvedValue(mockInlineCompletions);



      jest.runAllTimers();

      await glcp.provideInlineCompletionItems(

        mockDocument as unknown as vscode.TextDocument,

        mockPosition,

        mockContext,

      );

      await glcp.provideInlineCompletionItems(mockDocument, mockPosition, mockContext);

      jest.runAllTimers();



      expect(glcp.getCompletions).toHaveBeenCalled();

      jest.runAllTimers();

    });

  });



  describe(`circuit breaking`, () => {

    const turnOnCircuitBreaker = async (glcp: GitLabCodeCompletionProvider) => {

      await glcp.getCompletions(mockDocument, mockPosition);

      await glcp.getCompletions(mockDocument, mockPosition);

      await glcp.getCompletions(mockDocument, mockPosition);

      await glcp.getCompletions(mockDocument, mockPosition);

    };



    it(`starts breaking after 4 errors`, async () => {

      const glcp = new GitLabCodeCompletionProvider(platform);



      glcp.fetchCompletions = jest.fn().mockRejectedValue(new Error('test problem'));



      await turnOnCircuitBreaker(glcp);



      glcp.fetchCompletions = jest.fn().mockResolvedValue(mockCompletions);



      const result = await glcp.getCompletions(mockDocument, mockPosition);

      expect(result).toEqual([]);

      expect(glcp.fetchCompletions).not.toHaveBeenCalled();

    });



    describe("after circuit breaker's break time elapses", () => {

      it('fetches completions again', async () => {

        const glcp = new GitLabCodeCompletionProvider(platform);

        glcp.fetchCompletions = jest.fn().mockRejectedValue(new Error('test problem'));



        await turnOnCircuitBreaker(glcp);



        jest

          .useFakeTimers({ advanceTimers: true })

          .setSystemTime(new Date(Date.now() + CIRCUIT_BREAK_INTERVAL_MS));



        glcp.fetchCompletions = jest.fn().mockResolvedValue(mockCompletions);



        await glcp.getCompletions(mockDocument, mockPosition);



        expect(glcp.fetchCompletions).toHaveBeenCalled();

      });

    });

  });

});

src/completion/gitlab_code_completion_provider.ts

+58 −12
Original line number Diff line number Diff line
@@ -14,8 +14,12 @@ import { 
} from '../utils/extension_configuration';

import { getActiveProject } from '../commands/run_with_valid_project';

import { GitLabPlatform } from '../gitlab/gitlab_platform';

import { CodeCompletionTokenManager } from './token_manager';

import { CodeCompletionTokenManager, CompletionToken } from './token_manager';

import { getUserAgentHeader } from '../gitlab/http/get_user_agent_header';

import { CircuitBreaker } from './circuit_breaker';



export const CIRCUIT_BREAK_INTERVAL_MS = 10000;

export const MAX_ERRORS_BEFORE_CIRCUIT_BREAK = 4;



interface Choice {

  text: string;
@@ -57,6 +61,11 @@ export class GitLabCodeCompletionProvider implements vscode.InlineCompletionItem 


  private tokenManager: CodeCompletionTokenManager;



  private circuitBreaker = new CircuitBreaker(

    MAX_ERRORS_BEFORE_CIRCUIT_BREAK,

    CIRCUIT_BREAK_INTERVAL_MS,

  );



  constructor(platform: GitLabPlatform, noDebounce = false) {

    this.model = 'gitlab';

    this.server = GitLabCodeCompletionProvider.#getServer();
@@ -136,6 +145,10 @@ export class GitLabCodeCompletionProvider implements vscode.InlineCompletionItem 
    document: vscode.TextDocument,

    position: vscode.Position,

  ): Promise<vscode.InlineCompletionItem[]> {

    if (this.circuitBreaker.isBreaking()) {

      return [];

    }



    const prompt: CodeSuggestionPrompt = GitLabCodeCompletionProvider.#getPrompt(

      document,

      position,
@@ -156,21 +169,22 @@ export class GitLabCodeCompletionProvider implements vscode.InlineCompletionItem 


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



    const fetchResponse = await fetch(this.server, {

      method: `POST`,

      body: JSON.stringify(prompt),

      headers: {

        ...getUserAgentHeader(),

        'X-Gitlab-Authentication-Type': 'oidc',

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

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

      },

    });

    let response: CodeSuggestionResponse;



    try {

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

      this.circuitBreaker.success();

    } catch (e) {

      log.error(`AI Assist: Error fetching completions: ${e.toString()}`);

      this.circuitBreaker.error();

      return [];

    }



    const response: CodeSuggestionResponse = await fetchResponse.json();

    const choices = response.choices || [];



    log.debug(`AI Assist: got ${choices.length} completions`);

    console.log('vscode', vscode);

    console.log('vscode.InlineCompletionItem', vscode.InlineCompletionItem);



    return choices.map(

      choice => new vscode.InlineCompletionItem(choice.text, new vscode.Range(position, position)),
@@ -199,4 +213,36 @@ export class GitLabCodeCompletionProvider implements vscode.InlineCompletionItem 
      }

    });

  }



  async fetchCompletions(

    token: CompletionToken,

    prompt: CodeSuggestionPrompt,

  ): Promise<CodeSuggestionResponse> {

    const requestOptions = {

      method: 'POST',

      headers: {

        ...getUserAgentHeader(),

        'X-Gitlab-Authentication-Type': 'oidc',

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

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

      },

      body: JSON.stringify(prompt),

    };



    const response = await fetch(this.server, requestOptions);



    await this.handleErrorReponse(response);



    const data = await response.json();

    return data;

  }



  private async handleErrorReponse(response: Response) {

    if (!response.ok) {

      const body = await response.text().catch(() => undefined);

      throw new Error(

        `Fetching code suggestions from ${response.url} failed for server ${this.server}. Body: ${body}`,

      );

    }

  }

}