feat: Move streaming decision to the LS (99a1a0fa) · Commits · GitLab.org / gitlab-vscode-extension

package-lock.json

+9 −7

Original line number	Diff line number	Diff line
		@@ -11,7 +11,7 @@
		"license": "MIT",
		"dependencies": {
		"@anycable/core": "^0.7.12",
		"@gitlab-org/gitlab-lsp": "^3.30.0",
		"@gitlab-org/gitlab-lsp": "^3.31.0",
		"@snowplow/tracker-core": "3.19.0",
		"cross-fetch": "^4.0.0",
		"dayjs": "^1.11.10",
		@@ -1088,9 +1088,9 @@
		}
		},
		"node_modules/@gitlab-org/gitlab-lsp": {
		"version": "3.30.0",
		"resolved": "https://gitlab.com/api/v4/projects/46519181/packages/npm/@gitlab-org/gitlab-lsp/-/@gitlab-org/gitlab-lsp-3.30.0.tgz",
		"integrity": "sha1-Gs00bQ5/XAG9gQp6jHhMLqSEJBo=",
		"version": "3.31.0",
		"resolved": "https://gitlab.com/api/v4/projects/46519181/packages/npm/@gitlab-org/gitlab-lsp/-/@gitlab-org/gitlab-lsp-3.31.0.tgz",
		"integrity": "sha1-DE9tcP3T6YW8dKh3SMj1+riC/Jw=",
		"dependencies": {
		"@snowplow/tracker-core": "^3.15.0",
		"ajv": "^8.12.0",
		@@ -1099,6 +1099,7 @@
		"dayjs": "^1.11.10",
		"events": "^3.3.0",
		"get-proxy-settings": "^0.1.13",
		"lodash": "^4.17.21",
		"lru-cache": "^10.1.0",
		"proxy-agent": "^6.3.1",
		"semver": "^7.5.4",
		@@ -14889,9 +14890,9 @@
		"dev": true
		},
		"@gitlab-org/gitlab-lsp": {
		"version": "3.30.0",
		"resolved": "https://gitlab.com/api/v4/projects/46519181/packages/npm/@gitlab-org/gitlab-lsp/-/@gitlab-org/gitlab-lsp-3.30.0.tgz",
		"integrity": "sha1-Gs00bQ5/XAG9gQp6jHhMLqSEJBo=",
		"version": "3.31.0",
		"resolved": "https://gitlab.com/api/v4/projects/46519181/packages/npm/@gitlab-org/gitlab-lsp/-/@gitlab-org/gitlab-lsp-3.31.0.tgz",
		"integrity": "sha1-DE9tcP3T6YW8dKh3SMj1+riC/Jw=",
		"requires": {
		"@snowplow/tracker-core": "^3.15.0",
		"ajv": "^8.12.0",
		@@ -14900,6 +14901,7 @@
		"dayjs": "^1.11.10",
		"events": "^3.3.0",
		"get-proxy-settings": "^0.1.13",
		"lodash": "^4.17.21",
		"lru-cache": "^10.1.0",
		"proxy-agent": "^6.3.1",
		"semver": "^7.5.4",

package.json

+1 −1

Original line number	Diff line number	Diff line
		@@ -294,7 +294,7 @@
		},
		"dependencies": {
		"@anycable/core": "^0.7.12",
		"@gitlab-org/gitlab-lsp": "^3.30.0",
		"@gitlab-org/gitlab-lsp": "^3.31.0",
		"@snowplow/tracker-core": "3.19.0",
		"cross-fetch": "^4.0.0",
		"dayjs": "^1.11.10",

src/common/language_server/completion_stream.ts

+2 −2

Original line number	Diff line number	Diff line
		import vscode from 'vscode';
		import { uniqueId } from 'lodash';
		import {
		BaseLanguageClient,
		InlineCompletionTriggerKind,
		@@ -33,8 +32,9 @@ export class CompletionStream {
		client: BaseLanguageClient,
		document: vscode.TextDocument,
		position: vscode.Position,
		streamId: string,
		) {
		this.#id = uniqueId('code-suggestion-stream-');
		this.#id = streamId;
		this.#client = client;
		const protocolDocument = this.#client.code2ProtocolConverter.asTextDocumentPositionParams(
		document,

src/common/language_server/language_client_middleware.test.ts

+87 −111

Original line number	Diff line number	Diff line
		import vscode from 'vscode';
		import { BaseLanguageClient, TextDocumentPositionParams } from 'vscode-languageclient';
		import { CompletionIntentRequest, StreamingCompletionRequest } from '@gitlab-org/gitlab-lsp';
		import { START_STREAMING_COMMAND, StreamingCompletionRequest } 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';
		@@ -34,7 +34,6 @@ 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'),
		});
		@@ -118,11 +117,14 @@ describe('LanguageClientMiddleware', () => {
		});

		it('sets loading to false even if fetching suggestions throws an error', async () => {
		const next = jest.fn().mockRejectedValue(new Error());

		await expect(
		middleware.provideInlineCompletionItems(d, p, ctx, cancellationTokenSource.token, next),
		).rejects.toThrow();
		const next = jest.fn().mockRejectedValue('Failed to fetch');
		await middleware.provideInlineCompletionItems(
		d,
		p,
		ctx,
		cancellationTokenSource.token,
		next,
		);

		expect(jest.mocked(stateManager.setLoading).mock.calls).toEqual([[true], [false]]);
		});
		@@ -146,27 +148,13 @@ describe('LanguageClientMiddleware', () => {
		await expect(result).resolves.toEqual([]);
		expect(jest.mocked(stateManager.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();
		});
		});
		});

		describe('when streaming is enabled', () => {
		let nextSpy: ProvideInlineCompletionItemsSignature;
		let middleware: LanguageClientMiddleware;
		const streamId = 'code-suggestion-stream-LS-id';

		const callMiddleware = async ({
		times = 1,
		@@ -185,7 +173,17 @@ describe('LanguageClientMiddleware', () => {
		);

		beforeEach(() => {
		nextSpy = jest.fn().mockRejectedValue(new Error()); // we shouldn't be calling next for streaming
		nextSpy = jest.fn().mockResolvedValue({
		items: [
		{
		insertText: '',
		command: {
		command: START_STREAMING_COMMAND,
		arguments: [streamId],
		},
		},
		],
		});

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

		middleware.client = client;
		@@ -250,81 +247,41 @@ describe('LanguageClientMiddleware', () => {
		});
		}

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

		beforeEach(() => {
		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: false },
		]);
		});

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

		expect(client.sendRequest).toHaveBeenCalledWith(CompletionIntentRequest, {
		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 callMiddleware({ next: nonStreamingCompletionHandler });
		describe('streaming', () => {
		it('starts with the request to the server', async () => {
		await callMiddleware();

		expect(nonStreamingCompletionHandler).toHaveBeenCalled();
		expect(client.sendNotification).not.toHaveBeenCalled();
		});

		it('should proceed to streaming completion when intent is `generation`', async () => {
		jest.mocked(client.sendRequest).mockResolvedValue({
		intent: 'generation',
		});

		await callMiddleware({ next: nonStreamingCompletionHandler });

		expect(nonStreamingCompletionHandler).not.toHaveBeenCalled();
		expect(client.sendNotification).toHaveBeenCalledWith(
		StreamingCompletionRequest,
		expect.anything(),
		expect.objectContaining({
		method: StreamingCompletionRequest.method,
		}),
		expect.objectContaining({
		id: streamId,
		}),
		);
		});
		});

		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 },
		]);
		const result = (await callMiddleware()) as vscode.InlineCompletionItem[];

		expect(client.sendNotification).toHaveBeenCalledTimes(1);
		expect(result[0].insertText).toEqual('');
		});

		it('getStreamingCompletion keeps receiving notifications until done', async () => {
		it('keeps receiving notifications until done', async () => {
		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test', done: false },
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: true },
		{ id: streamId, completion: 'test', done: false },
		{ id: streamId, completion: '', done: true },
		]);

		const result = (await callMiddleware()) as vscode.InlineCompletionItem[];
		/* Here and further we need to make one additional call of the callMiddleware.
		The first one will imitate the first request to the LS
		that returns the response with the streaming command and starts the stream.
		Next calls imitate streaming of the chunks
		So there is always one additional call which detects and starts the stream
		and the next ones that actually receive the stream */
		const result = (await callMiddleware({ times: 2 })) as vscode.InlineCompletionItem[];

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

		it('if position does not change. Returns the existing response', async () => {
		it('returns the existing response if position does not change', async () => {
		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test 123', done: false },
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: true },
		{ id: streamId, completion: 'test 123', done: false },
		{ id: streamId, completion: '', done: true },
		]);

		const result = (await callMiddleware({ times: 2 })) as vscode.InlineCompletionItem[];
		@@ -332,60 +289,73 @@ describe('LanguageClientMiddleware', () => {
		expect(result[0].insertText).toEqual('test 123');
		});

		it('if position does change. Returns the new stream', async () => {
		it('returns the new stream if position does change', async () => {
		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test', done: false },
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: true },
		{ id: streamId, completion: 'test', done: false },
		{ id: streamId, completion: '', done: true },
		]);

		await callMiddleware({ times: 2 });

		const p2 = {
		line: 1,
		character: 0,
		} as vscode.Position;

		client.onNotification = invokeNotifications([
		{ id: 'code-suggestion-stream-uniqueId', completion: 'test2', done: false },
		{ id: 'code-suggestion-stream-uniqueId', completion: '', done: true },
		{ id: streamId, completion: 'test2', done: false },
		{ id: streamId, completion: '', done: true },
		]);

		const result = (await callMiddleware({
		times: 2,
		position: p2,
		})) as vscode.InlineCompletionItem[];

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

		it('handles setLoading gracefully', async () => {
		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 },
		{ id: streamId, completion: 'test', done: false },
		{ id: streamId, completion: 'test 123', done: false },
		{ id: streamId, completion: 'test 123 abc', done: true },
		]);

		await callMiddleware({ times: 3 });
		await callMiddleware({ times: 4 });

		expect(jest.mocked(stateManager.setLoading).mock.calls).toEqual([[true], [false]]);
		expect(jest.mocked(stateManager.setLoading).mock.calls).toEqual([
		[true], // request to detect streaming started
		[true], // stream started
		[false], // request to detect streaming ended (finally called)
		[false], // stream completed
		]);
		});

		it('when canceled, handles setLoading gracefully', async () => {
		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 },
		{ id: streamId, completion: 'test', done: false },
		{ id: streamId, completion: 'test 123', done: false },
		]);

		const result = callMiddleware();
		const result = callMiddleware({ times: 2 });

		cancellationTokenSource.cancel();

		await result;

		expect(jest.mocked(stateManager.setLoading).mock.calls).toEqual([[true], [false]]);
		expect(jest.mocked(stateManager.setLoading).mock.calls).toEqual([
		[true], // request to detect streaming started
		[true], // stream started
		[false], // request to detect streaming ended (finally called)
		[false], // stream cancelled
		]);
		});

		it('when positioned changed, handles loading stack gracefully', async () => {
		it('when position changes, 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 },
		{ id: streamId, completion: 'test', done: false },
		{ id: streamId, completion: 'test 123', done: true },
		]);

		const firstCancellation = new vscode.CancellationTokenSource();
		@@ -394,19 +364,25 @@ describe('LanguageClientMiddleware', () => {
		character: 0,
		});

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

		firstCancellation.cancel();
		await Promise.all([firstCall, callMiddleware({ times: 2 })]);
		const streamWithNewPosition = callMiddleware({ times: 3 });
		await Promise.all([firstStream, streamWithNewPosition]);

		expect(jest.mocked(stateManager.setLoading).mock.calls).toEqual([
		[true],
		[true],
		[false],
		[false],
		[true], // request to detect streaming started for firstStream
		[true], // request to detect streaming started for streamWithNewPosition
		[true], // started firstStream
		[true], // started streamWithNewPosition
		[false], // request to detect streaming ended for firstStream (finally called)
		[false], // request to detect streaming ended for streamWithNewPosition (finally called)
		[false], // firstStream cancelled
		[false], // streamWithNewPosition completed
		]);
		});
		});

src/common/language_server/language_client_middleware.ts

+45 −88

Original line number	Diff line number	Diff line
		import vscode from 'vscode';
		import { BaseLanguageClient, Middleware } from 'vscode-languageclient';
		import { Intent, ICompletionIntentResponse, CompletionIntentRequest } from '@gitlab-org/gitlab-lsp';
		import { START_STREAMING_COMMAND } 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';
		import { log } from '../log';
		import { COMMAND_CODE_SUGGESTION_STREAM_ACCEPTED } from '../code_suggestions/commands/code_suggestion_stream_accepted';
		import { CompletionStream, getStreamContextId } from './completion_stream';
		import { FeatureFlag, isEnabled } from '../feature_flags';
		import { CONFIG_NAMESPACE } from '../constants';
		import { serializeInlineCompletionContext, serializePosition } from './serialization_utils';
		import { serializeInlineCompletionContext } from './serialization_utils';
		import { isInlineCompletionList } from '../utils/code_suggestions';

		// We need to wait just a bit after cancellation, otherwise the loading icon flickers while someone types
		const CANCELLATION_DELAY = 150;
		@@ -19,22 +18,12 @@ export class LanguageClientMiddleware implements Middleware {

		#subscriptions: vscode.Disposable[] = [];

		#hasStreamingEnabled: boolean \| undefined;

		#client?: BaseLanguageClient;

		#activeStreams: Record<string, CompletionStream> = {};

		constructor(stateManager: CodeSuggestionsStateManager) {
		this.#stateManager = stateManager;
		this.#hasStreamingEnabled = isEnabled(FeatureFlag.StreamCodeGenerations);
		this.#subscriptions.push(
		vscode.workspace.onDidChangeConfiguration(e => {
		if (e.affectsConfiguration(CONFIG_NAMESPACE)) {
		this.#hasStreamingEnabled = isEnabled(FeatureFlag.StreamCodeGenerations);
		}
		}),
		);
		}

		dispose() {
		@@ -45,7 +34,7 @@ export class LanguageClientMiddleware implements Middleware {
		this.#client = client;
		}

		// we disable standard completion form LS and only use inline completion
		// we disable standard completion from LS and only use inline completion
		// eslint-disable-next-line class-methods-use-this
		provideCompletionItem = () => [];

		@@ -65,75 +54,60 @@ export class LanguageClientMiddleware implements Middleware {
		return [];
		}

		const shouldStream = await this.#shouldStream(document, position);

		if (!shouldStream) {
		if (this.#hasStreamingEnabled) {
		this.#cancelAllStreams();
		}

		return this.#provideDefaultInlineCompletionItems(document, position, context, token, next);
		const activeStream = this.#getStreamAt(document, position);
		if (activeStream) {
		return this.#provideStreamingInlineCompletionItems(document, position, token, activeStream);
		}

		try {
		return await this.#provideStreamingInlineCompletionItems(document, position, token);
		} catch (e) {
		log.error(e);
		return [];
		}
		}

		async #provideDefaultInlineCompletionItems(
		document: vscode.TextDocument,
		position: vscode.Position,
		context: vscode.InlineCompletionContext,
		token: vscode.CancellationToken,
		next: (
		document: vscode.TextDocument,
		position: vscode.Position,
		context: vscode.InlineCompletionContext,
		token: vscode.CancellationToken,
		) => vscode.ProviderResult<vscode.InlineCompletionItem[] \| vscode.InlineCompletionList>,
		) {
		this.#stateManager.setLoading(true);
		this.#cancelAllStreams();

		// Short circuit after both cancellation and time have passed
		const shortCircuit = waitForCancellationToken(token)
		.then(() => waitForMs(CANCELLATION_DELAY))
		.then(() => []);

		try {
		return await Promise.race([
		const response = await Promise.race([
		shortCircuit,
		next(document, position, serializeInlineCompletionContext(context), token),
		]);

		if (isInlineCompletionList(response)) {
		const command = response.items?.[0]?.command?.command;
		if (command === START_STREAMING_COMMAND) {
		const streamId = response.items?.[0]?.command?.arguments?.[0];

		await this.#startStreaming(document, position, streamId);
		return [];
		}
		}

		return response;
		} catch (e) {
		log.error(e);
		return [];
		} finally {
		this.#stateManager.setLoading(false);
		}
		}

		async #provideStreamingInlineCompletionItems(
		async #startStreaming(
		document: vscode.TextDocument,
		position: vscode.Position,
		token: vscode.CancellationToken,
		): Promise<vscode.InlineCompletionItem[] \| vscode.InlineCompletionList \| undefined> {
		streamId: string,
		) {
		if (!this.#client) {
		log.error(
		'Invoking LanguageServer client without initializing the inline completion middleware',
		);
		return [];
		return;
		}

		let stream = this.#getStreamAt(document, position);

		if (!stream) {
		stream = new CompletionStream(this.#client, document, position);

		// Add new loading state __before__ canceling old streams so that icon doesn't flicker
		this.#stateManager.setLoadingResource(stream, true);
		this.#cancelAllStreams();
		const stream = new CompletionStream(this.#client, document, position, streamId);

		try {
		this.#stateManager.setLoadingResource(stream, true);
		await stream.start();
		} catch (e) {
		this.#stateManager.setLoadingResource(stream, false);
		@@ -141,6 +115,21 @@ export class LanguageClientMiddleware implements Middleware {
		}

		this.#saveStreamAt(document, position, stream);

		LanguageClientMiddleware.forceTriggerInlineCompletion().catch(e => log.error(e));
		}

		async #provideStreamingInlineCompletionItems(
		document: vscode.TextDocument,
		position: vscode.Position,
		token: vscode.CancellationToken,
		stream: CompletionStream,
		): Promise<vscode.InlineCompletionItem[] \| vscode.InlineCompletionList \| undefined> {
		if (!this.#client) {
		log.error(
		'Invoking LanguageServer client without initializing the inline completion middleware',
		);
		return [];
		}

		if (token.isCancellationRequested) {
		@@ -200,7 +189,6 @@ export class LanguageClientMiddleware implements Middleware {
		const value = document.getText(replaceRange);

		await editor.edit(edit => edit.replace(replaceRange, value));
		await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger');
		}

		#getStreamAt(
		@@ -229,35 +217,4 @@ export class LanguageClientMiddleware implements Middleware {
		delete this.#activeStreams[contextId];
		});
		}

		async #shouldStream(document: vscode.TextDocument, position: vscode.Position): Promise<boolean> {
		if (!this.#hasStreamingEnabled) {
		return false;
		}
		if (this.#getStreamAt(document, position)) {
		return true;
		}

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

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

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

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