Commit ce78d2c4 authored by Paul Slaughter's avatar Paul Slaughter 2️⃣
Browse files

fix: End LS request when canceled 

- This should also be fixed in the LS, but there
   are cache concerns to be considered.
- For a quick fix lets handle this in the extension
parent 0452c1ac

src/__mocks__/vscode.js

+16 −2
Original line number Diff line number Diff line
@@ -98,10 +98,24 @@ module.exports = { 
    return { start: args[0], end: args[1] };

  },

  CancellationTokenSource: function CancellationTokenSource() {

    const controller = new AbortController();



    return {

      token: { isCancellationRequested: false },

      token: {

        get isCancellationRequested() {

          return controller.signal.aborted;

        },

        set isCancellationRequested(val) {

          throw new Error(

            'Cannot set isCancellationRequested. Try using the CancellationTokenSource.',

          );

        },

        onCancellationRequested(callback) {

          return controller.signal.addEventListener('abort', callback);

        },

      },

      cancel() {

        this.token.isCancellationRequested = true;

        controller.abort();

      },

    };

  },

src/common/code_suggestions/code_suggestions_provider.test.ts

+4 −3
Original line number Diff line number Diff line
@@ -731,13 +731,14 @@ describe('CodeSuggestionsProvider', () => { 
    });



    it('should send a state of cancelled when suggestion is cancelled', async () => {

      const cancelToken = new vscode.CancellationTokenSource().token;

      cancelToken.isCancellationRequested = true;

      const cancelTokenSource = new vscode.CancellationTokenSource();



      cancelTokenSource.cancel();



      await glcp.getCompletions({

        document: mockDocument,

        position: mockPosition,

        cancellationToken: cancelToken,

        cancellationToken: cancelTokenSource.token,

      });



      expect(updateSuggestionState.mock.calls[0]).toEqual([

src/common/language_server/language_client_middleware.test.ts

+55 −5
Original line number Diff line number Diff line
@@ -3,8 +3,13 @@ import { CodeSuggestionsStateManager } from '../code_suggestions/code_suggestion 
import { createFakePartial } from '../test_utils/create_fake_partial';

import { LanguageClientMiddleware } from './language_client_middleware';

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

import { waitForCancellationToken } from '../utils/wait_for_cancellation_token';



describe('LanguageClientMiddleware', () => {

  beforeAll(() => {

    jest.useFakeTimers();

  });



  it('disables standard completion - provideCompletionItem always returns empty array', () => {

    const stateManager = new CodeSuggestionsStateManager(

      createFakePartial<GitLabPlatformManager>({}),
@@ -17,7 +22,12 @@ describe('LanguageClientMiddleware', () => { 
    const d = createFakePartial<vscode.TextDocument>({});

    const p = createFakePartial<vscode.Position>({});

    const ctx = createFakePartial<vscode.InlineCompletionContext>({});

    const tkn = createFakePartial<vscode.CancellationToken>({});



    let cancellationTokenSource: vscode.CancellationTokenSource;



    beforeEach(() => {

      cancellationTokenSource = new vscode.CancellationTokenSource();

    });



    it('returns empty array if suggestions are not active', async () => {

      const stateManager = createFakePartial<CodeSuggestionsStateManager>({
@@ -26,7 +36,13 @@ describe('LanguageClientMiddleware', () => { 
      const middleware = new LanguageClientMiddleware(stateManager);

      const next = jest.fn();



      const result = await middleware.provideInlineCompletionItems(d, p, ctx, tkn, next);

      const result = await middleware.provideInlineCompletionItems(

        d,

        p,

        ctx,

        cancellationTokenSource.token,

        next,

      );



      expect(result).toEqual([]);

      expect(next).not.toHaveBeenCalled();
@@ -51,7 +67,13 @@ describe('LanguageClientMiddleware', () => { 
        const mockItem = createFakePartial<vscode.InlineCompletionItem>({});

        const next = jest.fn().mockResolvedValue([mockItem]);



        const result = await middleware.provideInlineCompletionItems(d, p, ctx, tkn, next);

        const result = await middleware.provideInlineCompletionItems(

          d,

          p,

          ctx,

          cancellationTokenSource.token,

          next,

        );



        expect(result).toEqual([mockItem]);

      });
@@ -59,7 +81,13 @@ describe('LanguageClientMiddleware', () => { 
      it('sets suggestions to loading state', async () => {

        const next = jest.fn().mockResolvedValue([]);



        await middleware.provideInlineCompletionItems(d, p, ctx, tkn, next);

        await middleware.provideInlineCompletionItems(

          d,

          p,

          ctx,

          cancellationTokenSource.token,

          next,

        );



        expect(setLoading).toHaveBeenCalledTimes(2);

        expect(setLoading).toHaveBeenCalledWith(true);
@@ -70,13 +98,35 @@ describe('LanguageClientMiddleware', () => { 
        const next = jest.fn().mockRejectedValue(new Error());



        await expect(

          middleware.provideInlineCompletionItems(d, p, ctx, tkn, next),

          middleware.provideInlineCompletionItems(d, p, ctx, cancellationTokenSource.token, next),

        ).rejects.toThrow();



        expect(setLoading).toHaveBeenCalledTimes(2);

        expect(setLoading).toHaveBeenCalledWith(true);

        expect(setLoading).toHaveBeenLastCalledWith(false);

      });



      it('sets loading to false if token is canceled and next never resolves', async () => {

        const next = jest.fn().mockReturnValue(new Promise(() => {}));



        const result = middleware.provideInlineCompletionItems(

          d,

          p,

          ctx,

          cancellationTokenSource.token,

          next,

        );



        // We need to flush promises after canceling, so we use our waitFor helper here

        const waitForCanceled = waitForCancellationToken(cancellationTokenSource.token);



        cancellationTokenSource.cancel();

        await waitForCanceled;

        jest.advanceTimersByTime(150);



        await expect(result).resolves.toEqual([]);

        expect(jest.mocked(setLoading).mock.calls).toEqual([[true], [false]]);

      });

    });

  });

});

src/common/language_server/language_client_middleware.ts

+12 −1
Original line number Diff line number Diff line
@@ -5,6 +5,11 @@ import { 
  ProvideInlineCompletionItemsSignature,

} from 'vscode-languageclient/lib/common/inlineCompletion';

import { CodeSuggestionsStateManager } from '../code_suggestions/code_suggestions_state_manager';

import { waitForCancellationToken } from '../utils/wait_for_cancellation_token';

import { waitForMs } from '../utils/wait_for_ms';



// We need to wait just a bit after cancellation, otherwise the loading icon flickers while someone types

const CANCELLATION_DELAY = 150;



export class LanguageClientMiddleware implements InlineCompletionMiddleware, CompletionMiddleware {

  #stateManager: CodeSuggestionsStateManager;
@@ -29,8 +34,14 @@ export class LanguageClientMiddleware implements InlineCompletionMiddleware, Com 
    }



    this.#stateManager.setLoading(true);



    // Short circuit after both cancellation and time have passed

    const shortCircuit = waitForCancellationToken(token)

      .then(() => waitForMs(CANCELLATION_DELAY))

      .then(() => []);



    try {

      return await next(document, position, context, token);

      return await Promise.race([shortCircuit, next(document, position, context, token)]);

    } finally {

      this.#stateManager.setLoading(false);

    }

src/common/utils/wait_for_cancellation_token.test.ts

0 → 100644
+14 −0
Original line number Diff line number Diff line 
import * as vscode from 'vscode';

import { waitForCancellationToken } from './wait_for_cancellation_token';



describe('common/utils/wait_for_cancellation_token', () => {

  it('resolves when cancellation token is canceled', async () => {

    const tokenSource = new vscode.CancellationTokenSource();



    const result = waitForCancellationToken(tokenSource.token);



    tokenSource.cancel();



    await expect(result).resolves.toEqual(expect.anything());

  });

});