feat: Decide whether to stream data based on the intent (140cb27f) · Commits · GitLab.org / gitlab-vscode-extension

src/common/language_server/language_client_middleware.test.ts

+81 −3

Original line number	Diff line number	Diff line
		import vscode from 'vscode';
		import { BaseLanguageClient, TextDocumentPositionParams } from 'vscode-languageclient';
		import { GET_COMPLETION_INTENT_REQUEST } from '@gitlab-org/gitlab-lsp';
		import { ProvideInlineCompletionItemsSignature } from 'vscode-languageclient/lib/common/inlineCompletion';
		import { CodeSuggestionsStateManager } from '../code_suggestions/code_suggestions_state_manager';
		import { createFakePartial } from '../test_utils/create_fake_partial';
		@@ -31,6 +32,7 @@ describe('LanguageClientMiddleware', () => {
		});

		describe('provideInlineCompletionItem', () => {
		const documentFilePath = 'file:///home/user/dev/test.md';
		const d = createFakePartial<vscode.TextDocument>({
		uri: vscode.Uri.parse('file:///home/user/dev/test.md'),
		});
		@@ -69,6 +71,9 @@ describe('LanguageClientMiddleware', () => {
		describe('when suggestions are active', () => {
		let stateManager: CodeSuggestionsStateManager;
		let middleware: LanguageClientMiddleware;
		const client = createFakePartial<BaseLanguageClient>({
		sendRequest: jest.fn(),
		});

		let setLoading: jest.Func;

		@@ -79,6 +84,7 @@ describe('LanguageClientMiddleware', () => {
		setLoading,
		});
		middleware = new LanguageClientMiddleware(stateManager);
		middleware.client = client;
		});

		it('calls through to default logic if suggestions are enabled', async () => {
		@@ -145,6 +151,21 @@ describe('LanguageClientMiddleware', () => {
		await expect(result).resolves.toEqual([]);
		expect(jest.mocked(setLoading).mock.calls).toEqual([[true], [false]]);
		});

		it('should not make a call to detect intent', async () => {
		const mockItem = createFakePartial<vscode.InlineCompletionItem>({});
		const next = jest.fn().mockResolvedValue([mockItem]);

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

		expect(client.sendRequest).not.toHaveBeenCalled();
		});
		});
		});

		@@ -172,16 +193,18 @@ describe('LanguageClientMiddleware', () => {
		middleware = new LanguageClientMiddleware(stateManager);
		});

		it('returns [] if client is not set', async () => {
		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,
		next,
		nextReturnsItem,
		);

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

		@@ -197,6 +220,7 @@ describe('LanguageClientMiddleware', () => {
		asTextDocumentPositionParams,
		},
		onNotification: jest.fn().mockImplementation(() => Promise.resolve()),
		sendRequest: jest.fn(),
		});

		middleware.client = client;
		@@ -227,6 +251,59 @@ describe('LanguageClientMiddleware', () => {
		});
		}

		describe('Intent detection', () => {
		const nonStreamingCompletionHandler = jest.fn().mockResolvedValue([]);
		let streamingCompletionHandler: jest.SpyInstance;

		beforeEach(() => {
		streamingCompletionHandler = jest
		.spyOn(middleware, 'provideStreamingInlineCompletionItems')
		.mockImplementationOnce(() => Promise.resolve([]));
		});

		const triggerInlineCompletion = async () => {
		await middleware.provideInlineCompletionItems(
		d,
		p,
		ctx,
		cancellationTokenSource.token,
		nonStreamingCompletionHandler,
		);
		};
		it('should make a call to detect intent', async () => {
		await triggerInlineCompletion();
		expect(client.sendRequest).toHaveBeenCalledWith(GET_COMPLETION_INTENT_REQUEST, {
		documentUri: documentFilePath,
		position: p,
		});
		});

		it('should proceed to non-streaming completion when intent is not `generation`', async () => {
		jest.mocked(client.sendRequest).mockResolvedValue({
		intent: 'completion',
		});
		await triggerInlineCompletion();
		expect(nonStreamingCompletionHandler).toHaveBeenCalled();
		expect(streamingCompletionHandler).not.toHaveBeenCalled();
		});

		it('should proceed to streaming completion when intent is `generation`', async () => {
		jest.mocked(client.sendRequest).mockResolvedValue({
		intent: 'generation',
		});
		await triggerInlineCompletion();
		jest.advanceTimersByTime(10);
		expect(nonStreamingCompletionHandler).not.toHaveBeenCalled();
		expect(streamingCompletionHandler).toHaveBeenCalled();
		});
		});

		describe('generation', () => {
		beforeEach(() => {
		jest.mocked(client.sendRequest).mockResolvedValue({
		intent: 'generation',
		});
		});
		it('calls getStreamingCompletion on language client', async () => {
		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: false },
		@@ -324,3 +401,4 @@ describe('LanguageClientMiddleware', () => {
		});
		});
		});
		});

src/common/language_server/language_client_middleware.ts

+34 −1

Original line number	Diff line number	Diff line
		import vscode from 'vscode';
		import { BaseLanguageClient, Middleware } from 'vscode-languageclient';
		import {
		Intent,
		ICompletionIntentResponse,
		GET_COMPLETION_INTENT_REQUEST,
		} from '@gitlab-org/gitlab-lsp';
		import { CodeSuggestionsStateManager } from '../code_suggestions/code_suggestions_state_manager';
		import { waitForCancellationToken } from '../utils/wait_for_cancellation_token';
		import { waitForMs } from '../utils/wait_for_ms';
		@@ -63,7 +68,7 @@ export class LanguageClientMiddleware implements Middleware {
		return [];
		}

		if (!this.#hasStreamingEnabled) {
		const provideInlineCompletionItems = async () => {
		this.#stateManager.setLoading(true);

		// Short circuit after both cancellation and time have passed
		@@ -76,6 +81,15 @@ export class LanguageClientMiddleware implements Middleware {
		} finally {
		this.#stateManager.setLoading(false);
		}
		};

		if (!this.#hasStreamingEnabled) {
		return provideInlineCompletionItems();
		}

		const intent = await this.#getIntent(document, position);
		if (intent !== 'generation') {
		return provideInlineCompletionItems();
		}

		try {
		@@ -176,4 +190,23 @@ export class LanguageClientMiddleware implements Middleware {
		await editor.edit(edit => edit.replace(replaceRange, value));
		await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger');
		}

		async #getIntent(
		document: vscode.TextDocument,
		position: vscode.Position,
		): Promise<Intent \| undefined> {
		let intentTypeResponse: ICompletionIntentResponse \| undefined;

		try {
		intentTypeResponse = await this.#client?.sendRequest(GET_COMPLETION_INTENT_REQUEST, {
		documentUri: document.uri.toString(),
		position,
		});
		} catch (error) {
		log.error(`Failed to detect completion intent`, error);
		}

		log.info(`INTENT: ${intentTypeResponse?.intent}`);
		return intentTypeResponse?.intent;
		}
		}