feat: add toggle on/off for code suggestions (efee9a13) · Commits · GitLab.org / gitlab-vscode-extension

src/mocks/vscode.js

+1 −1

Original line number	Diff line number	Diff line
		@@ -52,7 +52,7 @@ module.exports = {
		registerCommand: jest.fn(),
		},
		languages: {
		registerInlineCompletionItemProvider: jest.fn(),
		registerInlineCompletionItemProvider: jest.fn().mockReturnValue({ dispose: jest.fn() }),
		},
		workspace: {
		openTextDocument: jest.fn(),

src/common/code_suggestions/code_suggestions.test.ts

+40 −1

Original line number	Diff line number	Diff line
		@@ -4,7 +4,19 @@ import { project } from '../test_utils/entities';
		import { CodeSuggestions } from './code_suggestions';
		import { VisibleCodeSuggestionsState } from './code_suggestions_state';

		jest.mock('./code_suggestions_state');
		jest.mock('./code_suggestions_state', () => ({
		...jest.requireActual('./code_suggestions_state'),
		CodeSuggestionsStateManager: function CodeSuggestionsStateManager() {
		return {
		isEnabled: jest.fn().mockResolvedValue(true),
		setGlobalState: jest.fn(),
		onDidChangeEnabledState: jest.fn(),
		setUnsupportedLanguageDocument: jest.fn(),
		};
		},
		}));

		jest.mock('./code_suggestions_provider');
		jest.mock('./code_suggestions_status_bar_item');

		const manager: GitLabPlatformManager = {
		@@ -31,6 +43,33 @@ describe('CodeSuggestions', () => {
		codeSuggestions.dispose();
		});

		describe('code suggestions registration', () => {
		let changeListener: (enabled: boolean) => void;

		beforeEach(() => {
		[[changeListener]] = jest.mocked(
		codeSuggestions.stateManager.onDidChangeEnabledState,
		).mock.calls;
		});

		it('registers code suggestion provider when state becomes enabled', () => {
		jest.mocked(codeSuggestions.stateManager.isEnabled).mockReturnValue(true);
		jest.mocked(vscode.languages.registerInlineCompletionItemProvider).mockClear();

		changeListener(true);

		expect(vscode.languages.registerInlineCompletionItemProvider).toHaveBeenCalled();
		});

		it('unregisters code suggestion provider when state becomes disabled', () => {
		jest.mocked(codeSuggestions.stateManager.isEnabled).mockReturnValue(false);

		changeListener(false);

		expect(codeSuggestions.providerDisposable?.dispose).toHaveBeenCalled();
		});
		});

		describe('state updates', () => {
		it.each`
		isEnabledInSettings \| codeSuggestionState

src/common/code_suggestions/code_suggestions.ts

+22 −5

Original line number	Diff line number	Diff line
		@@ -9,17 +9,25 @@ import {
		AI_ASSISTED_CODE_SUGGESTIONS_MODE,
		getAiAssistedCodeSuggestionsConfiguration,
		} from '../utils/extension_configuration';
		import { COMMAND_TOGGLE_CODE_SUGGESTIONS, toggleCodeSuggestions } from './commands/toggle';

		export class CodeSuggestions {
		stateManager = new CodeSuggestionsStateManager();

		statusBarItem: CodeSuggestionsStatusBarItem;

		toggleSuggestionsCommand?: vscode.Disposable;

		providerDisposable?: vscode.Disposable;

		activeTextEditorChangeDisposable?: vscode.Disposable;

		constructor(manager: GitLabPlatformManager) {
		this.toggleSuggestionsCommand = vscode.commands.registerCommand(
		COMMAND_TOGGLE_CODE_SUGGESTIONS,
		() => toggleCodeSuggestions({ stateManager: this.stateManager }),
		);

		this.statusBarItem = new CodeSuggestionsStatusBarItem(this.stateManager);

		const updateCodeSuggestionsStateForEditor = (editor?: vscode.TextEditor) => {
		@@ -42,29 +50,38 @@ export class CodeSuggestions {
		};

		const enableOrDisableSuggestions = () => {
		if (getAiAssistedCodeSuggestionsConfiguration().enabled) {
		const enabled = this.stateManager.isEnabled();
		if (enabled) {
		log.debug('Enabling code completion');
		this.stateManager.setGlobalState(GlobalCodeSuggestionsState.READY);
		register();
		} else {
		log.debug('Disabling code completion');
		this.stateManager.setGlobalState(GlobalCodeSuggestionsState.DISABLED_VIA_SETTINGS);
		this.providerDisposable?.dispose();
		this.activeTextEditorChangeDisposable?.dispose();
		}
		};
		const syncExtensionGlobalState = () => {
		const newState = getAiAssistedCodeSuggestionsConfiguration().enabled
		? GlobalCodeSuggestionsState.READY
		: GlobalCodeSuggestionsState.DISABLED_VIA_SETTINGS;
		this.stateManager.setGlobalState(newState);
		};

		enableOrDisableSuggestions();
		vscode.workspace.onDidChangeConfiguration(e => {
		if (e.affectsConfiguration(AI_ASSISTED_CODE_SUGGESTIONS_MODE)) {
		enableOrDisableSuggestions();
		syncExtensionGlobalState();
		}
		});
		syncExtensionGlobalState();

		this.stateManager.onDidChangeEnabledState(enableOrDisableSuggestions);
		enableOrDisableSuggestions();
		}

		dispose() {
		this.statusBarItem?.dispose();
		this.providerDisposable?.dispose();
		this.activeTextEditorChangeDisposable?.dispose();
		this.toggleSuggestionsCommand?.dispose();
		}
		}

src/common/code_suggestions/code_suggestions_state.test.ts

+28 −13

Original line number	Diff line number	Diff line
		@@ -7,6 +7,17 @@ import {

		describe('Code suggestions state manager', () => {
		let stateManager: CodeSuggestionsStateManager;

		function itShouldIgnoreAnyTemporaryStateAndReturn(state: VisibleCodeSuggestionsState) {
		it('should ignore any temporary state', () => {
		Object.values(TemporaryState).forEach(temporaryState => {
		stateManager.setTemporaryState(temporaryState);

		expect(stateManager.getVisibleState()).toBe(state);
		});
		});
		}

		beforeEach(() => {
		stateManager = new CodeSuggestionsStateManager();
		});
		@@ -22,13 +33,13 @@ describe('Code suggestions state manager', () => {
		expect(stateManager.getVisibleState()).toBe(state);
		});

		it('should ignore any temporary state', () => {
		Object.values(TemporaryState).forEach(temporaryState => {
		stateManager.setTemporaryState(temporaryState);
		it('should ignore temporary disabled flag', () => {
		stateManager.setTemporaryDisabled(true);

		expect(stateManager.getVisibleState()).toBe(state);
		});
		});

		itShouldIgnoreAnyTemporaryStateAndReturn(state);
		});

		describe('when global state is ready', () => {
		@@ -36,6 +47,18 @@ describe('Code suggestions state manager', () => {
		stateManager.setGlobalState(GlobalState.READY);
		});

		describe('when is temporary disabled', () => {
		beforeEach(() => {
		stateManager.setTemporaryDisabled(true);
		});

		it('should report as temporary disabled', () => {
		expect(stateManager.getVisibleState()).toBe(VisibleCodeSuggestionsState.DISABLED_BY_USER);
		});

		itShouldIgnoreAnyTemporaryStateAndReturn(VisibleCodeSuggestionsState.DISABLED_BY_USER);
		});

		describe('when in document with unsupported language', () => {
		beforeEach(() => {
		stateManager.setUnsupportedLanguageDocument(true);
		@@ -47,15 +70,7 @@ describe('Code suggestions state manager', () => {
		);
		});

		it('should ignore any temporary state', () => {
		Object.values(TemporaryState).forEach(temporaryState => {
		stateManager.setTemporaryState(temporaryState);

		expect(stateManager.getVisibleState()).toBe(
		VisibleCodeSuggestionsState.UNSUPPORTED_LANGUAGE,
		);
		});
		});
		itShouldIgnoreAnyTemporaryStateAndReturn(VisibleCodeSuggestionsState.UNSUPPORTED_LANGUAGE);
		});

		it('should report state as ready', () => {

src/common/code_suggestions/code_suggestions_state.ts

+34 −3

Original line number	Diff line number	Diff line
		@@ -17,6 +17,7 @@ export const TemporaryCodeSuggestionsState = {
		export type VisibleCodeSuggestionsState = ValueOf<typeof VisibleCodeSuggestionsState>;
		export const VisibleCodeSuggestionsState = {
		...GlobalCodeSuggestionsState,
		DISABLED_BY_USER: 'code-suggestions-disabled-by-user',
		...TemporaryCodeSuggestionsState,
		UNSUPPORTED_LANGUAGE: 'code-suggestions-document-unsupported-language',
		} as const;
		@@ -24,20 +25,37 @@ export const VisibleCodeSuggestionsState = {
		export class CodeSuggestionsStateManager {
		#globalState: GlobalCodeSuggestionsState = GlobalCodeSuggestionsState.DISABLED_VIA_SETTINGS;

		#isDisabledByUserForSession = false;

		#temporaryState: TemporaryCodeSuggestionsState \| null = null;

		isUnsupportedLanguageDocument = false;
		#isUnsupportedLanguageDocument = false;

		#changeVisibleStateEmitter = new vscode.EventEmitter<VisibleCodeSuggestionsState>();

		onDidChangeVisibleState = this.#changeVisibleStateEmitter.event;

		#changeEnabledStateEmitter = new vscode.EventEmitter<boolean>();

		onDidChangeEnabledState = this.#changeEnabledStateEmitter.event;

		isEnabled() {
		return !(
		this.#globalState === GlobalCodeSuggestionsState.DISABLED_VIA_SETTINGS \|\|
		this.#isDisabledByUserForSession
		);
		}

		getVisibleState(): VisibleCodeSuggestionsState {
		if (this.#globalState === GlobalCodeSuggestionsState.DISABLED_VIA_SETTINGS) {
		return this.#globalState;
		}

		if (this.isUnsupportedLanguageDocument) {
		if (this.#isDisabledByUserForSession) {
		return VisibleCodeSuggestionsState.DISABLED_BY_USER;
		}

		if (this.#isUnsupportedLanguageDocument) {
		return VisibleCodeSuggestionsState.UNSUPPORTED_LANGUAGE;
		}

		@@ -46,17 +64,24 @@ export class CodeSuggestionsStateManager {

		updateState(handler: () => void) {
		const previousVisibleState = this.getVisibleState();
		const previousEnabled = this.isEnabled();
		handler();
		const newVisibleState = this.getVisibleState();
		const newEnabled = this.isEnabled();

		if (previousVisibleState !== newVisibleState) {
		this.#changeVisibleStateEmitter.fire(newVisibleState);
		}

		if (previousEnabled !== newEnabled) {
		this.#changeEnabledStateEmitter.fire(newEnabled);
		}
		}

		setGlobalState(newState: GlobalCodeSuggestionsState) {
		this.updateState(() => {
		this.#globalState = newState;
		this.#isDisabledByUserForSession = false;
		});
		}

		@@ -66,9 +91,15 @@ export class CodeSuggestionsStateManager {
		});
		}

		setTemporaryDisabled(newState: boolean) {
		this.updateState(() => {
		this.#isDisabledByUserForSession = newState;
		});
		}

		setUnsupportedLanguageDocument(newState: boolean) {
		this.updateState(() => {
		this.isUnsupportedLanguageDocument = newState;
		this.#isUnsupportedLanguageDocument = newState;
		});
		}
		}