fix(cs-stream): Clean loading icon flicker (e3d45bba) · Commits · GitLab.org / gitlab-vscode-extension

src/common/code_suggestions/code_suggestions_state_manager.test.ts

+20 −0

Original line number	Diff line number	Diff line
		@@ -181,6 +181,26 @@ describe('Code suggestions state manager', () => {

		expect(stateManager.getVisibleState()).toBe(VisibleCodeSuggestionsState.LOADING);
		});

		it('can associate loading with a resource', () => {
		const resource1 = {};
		const resource2 = {};

		stateManager.setLoadingResource(resource1, true);
		stateManager.setLoadingResource(resource1, true);
		stateManager.setLoadingResource(resource2, false);
		stateManager.setLoadingResource(resource2, false);

		expect(stateManager.getVisibleState()).toBe(VisibleCodeSuggestionsState.LOADING);

		stateManager.setLoadingResource(resource1, false);

		expect(stateManager.getVisibleState()).toBe(VisibleCodeSuggestionsState.READY);

		stateManager.setLoadingResource(resource2, true);

		expect(stateManager.getVisibleState()).toBe(VisibleCodeSuggestionsState.LOADING);
		});
		});

		describe('disabled & active', () => {

src/common/code_suggestions/code_suggestions_state_manager.ts

+16 −0

Original line number	Diff line number	Diff line
		@@ -39,6 +39,8 @@ export class CodeSuggestionsStateManager {
		// boolean flags and counters indicating temporary states
		#isInErrorState = false;

		#loadingResources = new WeakSet();

		// this can't be a boolean flag because it's possible that response from first
		// request comes after we send second request (which would incorrectly set loading to false)
		#loadingCounter = 0;
		@@ -127,6 +129,20 @@ export class CodeSuggestionsStateManager {
		return VisibleCodeSuggestionsState.READY;
		}

		setLoadingResource(resource: object, isLoading: boolean) {
		const isResourceLoading = this.#loadingResources.has(resource);

		if (isResourceLoading !== isLoading) {
		this.setLoading(isLoading);
		}

		if (isLoading) {
		this.#loadingResources.add(resource);
		} else {
		this.#loadingResources.delete(resource);
		}
		}

		#updateStateWrapper = (handler: (v: boolean) => void) => (value: boolean) => {
		const previousVisibleState = this.getVisibleState();
		const previousDisabledByUser = this.isDisabledByUser();

src/common/language_server/completion_stream.ts

+11 −29

Original line number	Diff line number	Diff line
		@@ -10,6 +10,8 @@ import { STREAMING_COMPLETION_REQUEST_NOTIFICATION } from '@gitlab-org/gitlab-ls
		import { createStreamIterator } from './create_stream_iterator';
		import { log } from '../log';

		type IteratorType = ReturnType<typeof createStreamIterator>;

		/** identifies the document and position for this stream */
		export const getStreamContextId = (document: vscode.TextDocument, position: vscode.Position) =>
		`${document.uri.toString()}\|${position.line}\|${position.character}`;
		@@ -25,13 +27,7 @@ export class CompletionStream {

		#position: Position;

		#iterator?: AsyncIterator<string>;

		// FIXME: this is terrible
		// we use this flag from the clients of this stream to identify if we should set the loading icon or not
		loading = false;

		#lastValue = '';
		#iterator: IteratorType;

		constructor(
		client: BaseLanguageClient,
		@@ -47,9 +43,16 @@ export class CompletionStream {
		this.#textDocument = protocolDocument.textDocument;
		this.#position = protocolDocument.position;
		this.#cancellationTokenSource = new vscode.CancellationTokenSource();
		this.#iterator = createStreamIterator(
		this.#client,
		this.#id,
		this.#cancellationTokenSource.token,
		);
		}

		async start() {
		log.debug(`Start of stream ${this.#id}`);

		// TODO: use StreamingCompletionRequest (which contains the expected notification type)
		// unfortunately the const doesn't get through from LS for some reason
		await this.#client.sendNotification(STREAMING_COMPLETION_REQUEST_NOTIFICATION, {
		@@ -58,35 +61,14 @@ export class CompletionStream {
		context: { triggerKind: InlineCompletionTriggerKind.Automatic },
		id: this.#id,
		});

		this.#iterator = createStreamIterator(
		this.#client,
		this.#id,
		this.#cancellationTokenSource.token,
		);
		}

		get iterator() {
		if (!this.#iterator) {
		throw new Error(`Stream ${this.#id}, has not been started.`);
		}
		return {
		next: async () => {
		const { value, done } = await this.#iterator!.next();
		if (!done) {
		this.#lastValue = value;
		}
		return { done, value };
		},
		};
		return this.#iterator;
		}

		cancel() {
		log.debug(`Cancellation requested for stream ${this.#id}`);
		return this.#cancellationTokenSource.cancel();
		}

		get lastSuggestion() {
		return this.#lastValue;
		}
		}

src/common/language_server/create_stream_iterator.ts

+30 −8

Original line number	Diff line number	Diff line
		@@ -8,6 +8,7 @@ import { log } from '../log';

		interface CompletionPart {
		completion: string;
		canceled?: boolean;
		done: boolean;
		}

		@@ -15,8 +16,10 @@ export const createStreamIterator = (
		client: BaseLanguageClient,
		streamId: string,
		cancellationToken?: CancellationToken,
		): AsyncIterator<string> => {
		): AsyncIterator<{ completion: string }, { completion: string; canceled?: boolean }> => {
		let resolveWait: ((val: CompletionPart) => void) \| null = null;
		let lastCompletion = '';

		const queue: CompletionPart[] = [];
		const subscriptions: Disposable[] = [];

		@@ -35,6 +38,8 @@ export const createStreamIterator = (
		};

		const onNext = (part: CompletionPart): void => {
		lastCompletion = part.completion;

		if (resolveWait) {
		resolveWait(part);
		resolveWait = null;
		@@ -49,10 +54,7 @@ export const createStreamIterator = (
		// TODO use CancelStreaming as the notification type (it is type safe)
		// ATM we can't do that because something is wrong with the LSP exporting types/constants
		await client.sendNotification(CANCEL_STREAMING_COMPLETION_NOTIFICATION, { id: streamId });
		onNext({
		completion: '',
		done: true,
		});
		onNext({ completion: lastCompletion, done: true, canceled: true });
		}),
		);
		}
		@@ -67,8 +69,14 @@ export const createStreamIterator = (
		return;
		}

		let completionValue = completion \|\| '';

		if (!completionValue && done) {
		completionValue = lastCompletion;
		}

		onNext({
		completion: completion \|\| '',
		completion: completionValue,
		done,
		});
		},
		@@ -77,7 +85,18 @@ export const createStreamIterator = (

		return {
		async next() {
		const { completion, done } = await waitForNext();
		const { completion, done, canceled = false } = await waitForNext();

		if (canceled) {
		onComplete();
		return {
		value: {
		completion,
		canceled,
		},
		done: true,
		};
		}

		log.debug(`Streaming Suggestion: ${completion}, Done is ${done}`);

		@@ -86,7 +105,10 @@ export const createStreamIterator = (
		}

		return {
		value: completion,
		value: {
		completion,
		canceled,
		},
		done,
		};
		},

src/common/language_server/language_client_middleware.test.ts

+107 −103

Original line number	Diff line number	Diff line
		@@ -18,14 +18,17 @@ jest.mock('lodash', () => {
		});

		describe('LanguageClientMiddleware', () => {
		beforeAll(() => {
		let stateManager: CodeSuggestionsStateManager;

		beforeEach(() => {
		jest.useFakeTimers();

		stateManager = new CodeSuggestionsStateManager(createFakePartial<GitLabPlatformManager>({}));
		jest.spyOn(stateManager, 'setLoading');
		jest.spyOn(stateManager, 'isActive').mockReturnValue(true);
		});

		it('disables standard completion - provideCompletionItem always returns empty array', () => {
		const stateManager = new CodeSuggestionsStateManager(
		createFakePartial<GitLabPlatformManager>({}),
		);
		const middleware = new LanguageClientMiddleware(stateManager);
		expect(middleware.provideCompletionItem()).toEqual([]);
		});
		@@ -49,9 +52,7 @@ describe('LanguageClientMiddleware', () => {

		describe('when streaming is disabled', () => {
		it('returns empty array if suggestions are not active', async () => {
		const stateManager = createFakePartial<CodeSuggestionsStateManager>({
		isActive: () => false,
		});
		jest.spyOn(stateManager, 'isActive').mockReturnValue(false);
		const middleware = new LanguageClientMiddleware(stateManager);
		const next = jest.fn();

		@@ -68,20 +69,13 @@ describe('LanguageClientMiddleware', () => {
		});

		describe('when suggestions are active', () => {
		let stateManager: CodeSuggestionsStateManager;
		let middleware: LanguageClientMiddleware;
		const client = createFakePartial<BaseLanguageClient>({
		sendRequest: jest.fn(),
		});

		let setLoading: jest.Func;

		beforeEach(() => {
		setLoading = jest.fn();
		stateManager = createFakePartial<CodeSuggestionsStateManager>({
		isActive: () => true,
		setLoading,
		});
		jest.spyOn(stateManager, 'isActive').mockReturnValue(true);
		middleware = new LanguageClientMiddleware(stateManager);
		middleware.client = client;
		});
		@@ -112,9 +106,7 @@ describe('LanguageClientMiddleware', () => {
		next,
		);

		expect(setLoading).toHaveBeenCalledTimes(2);
		expect(setLoading).toHaveBeenCalledWith(true);
		expect(setLoading).toHaveBeenLastCalledWith(false);
		expect(jest.mocked(stateManager.setLoading).mock.calls).toEqual([[true], [false]]);
		});

		it('sets loading to false even if fetching suggestions throws an error', async () => {
		@@ -124,9 +116,7 @@ describe('LanguageClientMiddleware', () => {
		middleware.provideInlineCompletionItems(d, p, ctx, cancellationTokenSource.token, next),
		).rejects.toThrow();

		expect(setLoading).toHaveBeenCalledTimes(2);
		expect(setLoading).toHaveBeenCalledWith(true);
		expect(setLoading).toHaveBeenLastCalledWith(false);
		expect(jest.mocked(stateManager.setLoading).mock.calls).toEqual([[true], [false]]);
		});

		it('sets loading to false if token is canceled and next never resolves', async () => {
		@@ -146,7 +136,7 @@ describe('LanguageClientMiddleware', () => {
		await jest.advanceTimersByTimeAsync(150);

		await expect(result).resolves.toEqual([]);
		expect(jest.mocked(setLoading).mock.calls).toEqual([[true], [false]]);
		expect(jest.mocked(stateManager.setLoading).mock.calls).toEqual([[true], [false]]);
		});

		it('should not make a call to detect intent', async () => {
		@@ -167,14 +157,27 @@ describe('LanguageClientMiddleware', () => {
		});

		describe('when streaming is enabled', () => {
		let next: ProvideInlineCompletionItemsSignature;
		let stateManager: CodeSuggestionsStateManager;
		let nextSpy: ProvideInlineCompletionItemsSignature;
		let middleware: LanguageClientMiddleware;
		let setLoading;

		const callMiddleware = async ({
		times = 1,
		position = p,
		cancellationToken = cancellationTokenSource.token,
		next = nextSpy,
		} = {}) =>
		Array.from({ length: times }).reduce<
		ReturnType<LanguageClientMiddleware['provideInlineCompletionItems']>
		>(
		acc =>
		acc.then(() =>
		middleware.provideInlineCompletionItems(d, position, ctx, cancellationToken, next),
		),
		Promise.resolve(undefined),
		);

		beforeEach(() => {
		jest.useRealTimers();
		next = jest.fn().mockRejectedValue(new Error()); // we shouldn't be calling next for streaming
		nextSpy = jest.fn().mockRejectedValue(new Error()); // we shouldn't be calling next for streaming

		setFakeWorkspaceConfiguration({
		featureFlags: {
		@@ -182,33 +185,24 @@ describe('LanguageClientMiddleware', () => {
		},
		});

		setLoading = jest.fn();
		stateManager = createFakePartial<CodeSuggestionsStateManager>({
		isActive: () => true,
		setLoading,
		});
		middleware = new LanguageClientMiddleware(stateManager);
		});

		it('calls the inlineCompletion (next) if client is not set', async () => {
		const mockItem = createFakePartial<vscode.InlineCompletionItem>({});
		const nextReturnsItem = jest.fn().mockResolvedValue([mockItem]);
		const result = await middleware.provideInlineCompletionItems(
		d,
		p,
		ctx,
		cancellationTokenSource.token,
		nextReturnsItem,
		);
		const result = await callMiddleware({ next: nextReturnsItem });

		expect(result).toEqual([mockItem]);
		expect(next).not.toHaveBeenCalled();
		expect(nextSpy).not.toHaveBeenCalled();
		});

		describe('when the language client is set', () => {
		let client: BaseLanguageClient;

		beforeEach(() => {
		jest.useRealTimers();

		const asTextDocumentPositionParams = jest.fn();

		client = createFakePartial<BaseLanguageClient>({
		@@ -251,16 +245,6 @@ describe('LanguageClientMiddleware', () => {
		describe('Intent detection', () => {
		const nonStreamingCompletionHandler = jest.fn().mockResolvedValue([]);

		const triggerInlineCompletion = async () => {
		await middleware.provideInlineCompletionItems(
		d,
		p,
		ctx,
		cancellationTokenSource.token,
		nonStreamingCompletionHandler,
		);
		};

		beforeEach(() => {
		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: false },
		@@ -268,7 +252,8 @@ describe('LanguageClientMiddleware', () => {
		});

		it('should make a call to detect intent', async () => {
		await triggerInlineCompletion();
		await callMiddleware({ next: nonStreamingCompletionHandler });

		expect(client.sendRequest).toHaveBeenCalledWith(GET_COMPLETION_INTENT_REQUEST, {
		documentUri: documentFilePath,
		position: p,
		@@ -279,7 +264,9 @@ describe('LanguageClientMiddleware', () => {
		jest.mocked(client.sendRequest).mockResolvedValue({
		intent: 'completion',
		});
		await triggerInlineCompletion();

		await callMiddleware({ next: nonStreamingCompletionHandler });

		expect(nonStreamingCompletionHandler).toHaveBeenCalled();
		expect(client.sendNotification).not.toHaveBeenCalled();
		});
		@@ -288,8 +275,9 @@ describe('LanguageClientMiddleware', () => {
		jest.mocked(client.sendRequest).mockResolvedValue({
		intent: 'generation',
		});
		await triggerInlineCompletion();
		jest.advanceTimersByTime(10);

		await callMiddleware({ next: nonStreamingCompletionHandler });

		expect(nonStreamingCompletionHandler).not.toHaveBeenCalled();
		expect(client.sendNotification).toHaveBeenCalledWith(
		'streamingCompletionRequest',
		@@ -308,13 +296,7 @@ describe('LanguageClientMiddleware', () => {
		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: false },
		]);
		const result = (await middleware.provideInlineCompletionItems(
		d,
		p,
		ctx,
		cancellationTokenSource.token,
		next,
		)) as vscode.InlineCompletionItem[];
		const result = (await callMiddleware()) as vscode.InlineCompletionItem[];

		expect(client.sendNotification).toHaveBeenCalledTimes(1);
		expect(result[0].insertText).toEqual('');
		@@ -326,13 +308,8 @@ describe('LanguageClientMiddleware', () => {
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: true },
		]);

		const result = (await middleware.provideInlineCompletionItems(
		d,
		p,
		ctx,
		cancellationTokenSource.token,
		next,
		)) as vscode.InlineCompletionItem[];
		const result = (await callMiddleware()) as vscode.InlineCompletionItem[];

		expect(result[0].insertText).toEqual('test');
		});

		@@ -342,20 +319,8 @@ describe('LanguageClientMiddleware', () => {
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: true },
		]);

		await middleware.provideInlineCompletionItems(
		d,
		p,
		ctx,
		cancellationTokenSource.token,
		next,
		);
		const result = (await middleware.provideInlineCompletionItems(
		d,
		p,
		ctx,
		cancellationTokenSource.token,
		next,
		)) as vscode.InlineCompletionItem[];
		const result = (await callMiddleware({ times: 2 })) as vscode.InlineCompletionItem[];

		expect(result[0].insertText).toEqual('test 123');
		});

		@@ -365,20 +330,7 @@ describe('LanguageClientMiddleware', () => {
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: true },
		]);

		await middleware.provideInlineCompletionItems(
		d,
		p,
		ctx,
		cancellationTokenSource.token,
		next,
		);
		await middleware.provideInlineCompletionItems(
		d,
		p,
		ctx,
		cancellationTokenSource.token,
		next,
		);
		await callMiddleware({ times: 2 });
		const p2 = {
		line: 1,
		character: 0,
		@@ -388,15 +340,67 @@ describe('LanguageClientMiddleware', () => {
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test2', done: false },
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: true },
		]);
		const result = (await middleware.provideInlineCompletionItems(
		d,
		p2,
		ctx,
		cancellationTokenSource.token,
		next,
		)) as vscode.InlineCompletionItem[];
		const result = (await callMiddleware({
		position: p2,
		})) as vscode.InlineCompletionItem[];

		expect(result[0].insertText).toEqual('test2');
		});

		it('handles setLoading gracefully', async () => {
		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test', done: false },
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test 123', done: false },
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test 123 abc', done: true },
		]);

		await callMiddleware({ times: 3 });

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

		it('when canceled, handles setLoading gracefully', async () => {
		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test', done: false },
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test 123', done: false },
		]);

		const result = callMiddleware();

		cancellationTokenSource.cancel();

		await result;

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

		it('when positioned changed, handles loading stack gracefully', async () => {
		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test', done: false },
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test 123', done: true },
		]);

		const firstCancellation = new vscode.CancellationTokenSource();
		const firstPosition = createFakePartial<vscode.Position>({
		line: 1,
		character: 0,
		});

		const firstCall = callMiddleware({
		position: firstPosition,
		cancellationToken: firstCancellation.token,
		});

		firstCancellation.cancel();
		await Promise.all([firstCall, callMiddleware({ times: 2 })]);

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