Commit 1688f537 authored by Olena Horal-Koretska's avatar Olena Horal-Koretska 2️⃣
Browse files

feat: Add Language Server startup monitoring

parent 9ee960f4

src/common/language_server/language_client_wrapper.test.ts

+56 −2
Original line number Diff line number Diff line
@@ -11,6 +11,7 @@ import { 
  DidChangeConfigurationNotification,

  GenericNotificationHandler,

  NotificationType,

  State,

} from 'vscode-languageclient';

import { createFakePartial } from '../test_utils/create_fake_partial';

import { GitLabPlatformManagerForCodeSuggestions } from '../code_suggestions/gitlab_platform_manager_for_code_suggestions';
@@ -113,13 +114,17 @@ describe('LanguageClientWrapper', () => { 
    const gitLabPlatform: GitLabPlatformForAccount = gitlabPlatformForAccount;

    getGitLabPlatformMock.mockResolvedValue(gitLabPlatform);

    client = createFakePartial<BaseLanguageClient>({

      start: jest.fn(),

      start: jest.fn().mockResolvedValue(undefined),

      stop: jest.fn(),

      registerProposedFeatures: jest.fn(),

      onNotification: jest.fn(),

      sendNotification: jest.fn(),

      sendRequest: jest.fn(),

      onRequest: jest.fn(),

      onDidChangeState: jest.fn(handler => {

        handler({ oldState: State.Stopped, newState: State.Running });

        return { dispose: jest.fn() };

      }),

    });

    lsGitProvider = createFakePartial<LSGitProvider>({

      getDiffWithHead: jest.fn(),
@@ -158,6 +163,51 @@ describe('LanguageClientWrapper', () => { 
  });



  describe('initAndStart', () => {

    describe('startup monitor', () => {

      it('resolves when the language server reaches Running state', async () => {

        // default mock fires State.Running immediately - just verify it resolves

        wrapper = createWrapper();



        await expect(wrapper.initAndStart()).resolves.toBeUndefined();

      });



      it('rejects when the language server does not spawn within the spawn timeout', async () => {

        jest.useFakeTimers();

        client = createFakePartial<BaseLanguageClient>({

          ...client,

          onDidChangeState: jest.fn(() => ({ dispose: jest.fn() })), // never fires

        });

        wrapper = createWrapper();



        const promise = wrapper.initAndStart();

        jest.advanceTimersByTime(10000);



        await expect(promise).rejects.toThrow('Language Server process did not start within 10s');

        jest.useRealTimers();

      });



      it('rejects when the handshake does not complete within the handshake timeout', async () => {

        jest.useFakeTimers();

        client = createFakePartial<BaseLanguageClient>({

          ...client,

          onDidChangeState: jest.fn(handler => {

            // process spawns but never reaches Running

            handler({ oldState: State.Stopped, newState: State.Starting });

            return { dispose: jest.fn() };

          }),

        });

        wrapper = createWrapper();



        const promise = wrapper.initAndStart();

        jest.advanceTimersByTime(30000);



        await expect(promise).rejects.toThrow(

          'Language Server process started but LSP initialize handshake did not complete within 30s',

        );

        jest.useRealTimers();

      });

    });



    it('starts the client and synchronizes the configuration', async () => {

      wrapper = createWrapper();
@@ -411,12 +461,16 @@ describe('LanguageClientWrapper', () => { 
        getGitLabPlatformMock.mockResolvedValue(gitlabPlatformForProject);



        client = createFakePartial<BaseLanguageClient>({

          start: jest.fn(),

          start: jest.fn().mockResolvedValue(undefined),

          stop: jest.fn(),

          registerProposedFeatures: jest.fn(),

          onNotification: jest.fn(),

          onRequest: jest.fn(),

          sendNotification: jest.fn(),

          onDidChangeState: jest.fn(handler => {

            handler({ oldState: State.Stopped, newState: State.Running });

            return { dispose: jest.fn() };

          }),

        });

      });

src/common/language_server/language_client_wrapper.ts

+14 −11
Original line number Diff line number Diff line
@@ -52,6 +52,7 @@ import { extensionConfigurationService } from '../utils/extension_configuration_ 
import { TerminalManager } from '../duo_workflow/terminal_manager';

import { getActiveFileContext } from '../chat/gitlab_chat_file_context';

import { LanguageServerFeatureStateProvider } from './language_server_feature_state_provider';

import { LanguageServerStartupMonitor } from './language_server_startup_monitor';

import { ApplyEditClientWrapper } from './apply_edit_client_wrapper';

import { DocumentQualityHandler, GET_DIAGNOSTICS_REQUEST_METHOD } from './document_quality_handler';

import { DiagnosticsDelayMiddleware } from './diagnostics_delay_middleware';
@@ -64,8 +65,6 @@ import { RepositoryClient, RepositoryRequestFunction } from './repository_client 
// Interface IDs for dependency injection

export const BaseLanguageClientId = createInterfaceId<BaseLanguageClient>('BaseLanguageClient');



const LS_STARTUP_TIMEOUT_MS = 40000;



const createNotifyFn =

  <T>(client: BaseLanguageClient, method: string) =>

  (param: T) =>
@@ -77,6 +76,7 @@ const createRequestFn = 
    client.sendRequest(method, param);



export interface LanguageClientWrapper {

  startupMonitor: LanguageServerStartupMonitor;

  setCustomPlatformConfig(clientConfig: Partial<ClientConfig>): void;

  initAndStart(): Promise<void>;

  syncConfig(): Promise<void>;
@@ -128,6 +128,13 @@ export class LanguageClientWrapperImpl implements LanguageClientWrapper { 


  #platformClientConfig?: Partial<ClientConfig>;



  #startupMonitor = new LanguageServerStartupMonitor();



  set startupMonitor(monitor: LanguageServerStartupMonitor) {

    this.#startupMonitor = monitor;

    this.#subscriptions.push(monitor);

  }



  constructor(

    client: BaseLanguageClient,

    suggestionsManager: GitLabPlatformManagerForCodeSuggestions,
@@ -249,15 +256,11 @@ export class LanguageClientWrapperImpl implements LanguageClientWrapper { 
      this.#documentQualityHandler.getDiagnostics,

    );



    const timeoutError = new Error(

      `The GitLab Language Server failed to start in ${LS_STARTUP_TIMEOUT_MS / 1000} seconds. Try to restart the GitLab extension.`,

    );

    await Promise.race([

      this.#client.start(),

      new Promise((_, reject) => {

        setTimeout(() => reject(timeoutError), LS_STARTUP_TIMEOUT_MS);

      }),

    ]);

    const startupComplete = this.#startupMonitor.observe(this.#client);

    this.#client.start().catch(() => {

      // Rejection here is expected during crash/restart cycles - the monitor tracks the outcome

    });

    await startupComplete;

    this.#subscriptions.push({ dispose: () => this.#client.stop() });

    await this.syncConfig();

    await this.#sendOpenTabs();

src/common/language_server/language_server_manager.test.ts

+29 −6
Original line number Diff line number Diff line
@@ -64,6 +64,9 @@ describe('LanguageServerManager', () => { 
      sendQuickChatMessageEvent: jest.fn(),

      syncConfig: jest.fn(),

      dispose: jest.fn(),

      startupMonitor: createFakePartial({

        notifyHandshakeStarted: jest.fn(),

      }),

    });

    clientContext = {

      ide: {
@@ -159,9 +162,16 @@ describe('LanguageServerManager', () => { 
    expect(languageClientFactory.createLanguageClient).toHaveBeenCalledWith(

      context,

      expect.objectContaining({

        initializationOptions: expect.objectContaining({

          baseAssetsUrl: expect.stringContaining('/assets/language-server/'),

        initializationOptions: expect.any(Function),

      }),

    );



    const { initializationOptions } = jest.mocked(languageClientFactory.createLanguageClient).mock

      .calls[0][1];

    const options = (initializationOptions as () => unknown)();

    expect(options).toEqual(

      expect.objectContaining({

        baseAssetsUrl: expect.stringContaining('/assets/language-server/'),

      }),

    );

  });
@@ -191,14 +201,14 @@ describe('LanguageServerManager', () => { 
    customLanguageServerManager.setLanguageClientWrapper(clientWrapper);

    await customLanguageServerManager.startLanguageServer();



    expect(languageClientFactory.createLanguageClient).toHaveBeenCalledWith(

      context,

    const { initializationOptions } = jest.mocked(languageClientFactory.createLanguageClient).mock

      .calls[1][1];

    const options = (initializationOptions as () => unknown)();

    expect(options).toEqual(

      expect.objectContaining({

        initializationOptions: expect.objectContaining({

        ...customClientContext,

        baseAssetsUrl: expect.stringContaining('/assets/language-server/'),

      }),

      }),

    );

  });
@@ -206,6 +216,19 @@ describe('LanguageServerManager', () => { 
    expect(clientWrapper.initAndStart).toHaveBeenCalled();

  });



  it('calls notifyHandshakeStarted on the startupMonitor when initializationOptions is invoked', () => {

    const { initializationOptions } = jest.mocked(languageClientFactory.createLanguageClient).mock

      .calls[0][1];

    const notifyHandshakeStarted = jest.spyOn(

      clientWrapper.startupMonitor,

      'notifyHandshakeStarted',

    );



    (initializationOptions as () => unknown)();



    expect(notifyHandshakeStarted).toHaveBeenCalled();

  });



  it('initializes state manager', async () => {

    expect(stateManager.init).toHaveBeenCalled();

  });

src/common/language_server/language_server_manager.ts

+8 −1
Original line number Diff line number Diff line
@@ -49,6 +49,7 @@ import { LanguageClientFactory } from './client_factory'; 
import { LanguageClientMiddleware } from './language_client_middleware';

import { LanguageClientWrapper } from './language_client_wrapper';

import { LanguageServerFeatureStateProvider } from './language_server_feature_state_provider';

import { LanguageServerStartupMonitor } from './language_server_startup_monitor';



export class LanguageServerManager implements WebviewManager, VersionProvider {

  #client: BaseLanguageClient | undefined;
@@ -110,6 +111,7 @@ export class LanguageServerManager implements WebviewManager, VersionProvider { 
    const statusBarItem = new CodeSuggestionsStatusBarItem(stateManager);

    const gutterIcon = new CodeSuggestionsGutterIcon(this.#context, stateManager);

    const middleware = new LanguageClientMiddleware(stateManager);

    const startupMonitor = new LanguageServerStartupMonitor();

    const baseAssetsUrl = vscode.Uri.joinPath(

      this.#context.extensionUri,

      './assets/language-server/',
@@ -125,9 +127,12 @@ export class LanguageServerManager implements WebviewManager, VersionProvider { 
        { scheme: 'gitlab-web-ide' },

        { scheme: 'untitled' },

      ],

      initializationOptions: {

      initializationOptions: () => {

        startupMonitor.notifyHandshakeStarted();

        return {

          ...this.#clientContext,

          baseAssetsUrl,

        };

      },

      middleware,

      outputChannel,
@@ -145,6 +150,8 @@ export class LanguageServerManager implements WebviewManager, VersionProvider { 
      );

    }



    this.#wrapper.startupMonitor = startupMonitor;



    await this.#wrapper.initAndStart();

    const subscriptions = [

      this.#wrapper,

src/common/language_server/language_server_startup_monitor.test.ts

0 → 100644
+228 −0
Original line number Diff line number Diff line 
import { BaseLanguageClient, State, StateChangeEvent } from 'vscode-languageclient';

import { createFakePartial } from '../test_utils/create_fake_partial';

import { log } from '../log';

import { LanguageServerStartupMonitor } from './language_server_startup_monitor';



jest.mock('../log');



describe('LanguageServerStartupMonitor', () => {

  let monitor: LanguageServerStartupMonitor;

  let onDidChangeStateHandler: (event: StateChangeEvent) => void;

  let client: BaseLanguageClient;



  beforeEach(() => {

    jest.useFakeTimers();

    jest.clearAllMocks();

    monitor = new LanguageServerStartupMonitor();

    client = createFakePartial<BaseLanguageClient>({

      onDidChangeState: jest.fn(handler => {

        onDidChangeStateHandler = handler;

        return { dispose: jest.fn() };

      }),

    });

  });



  afterEach(() => {

    monitor.dispose();

    jest.useRealTimers();

  });



  it('starts in idle phase', () => {

    expect(monitor.phase).toBe('idle');

  });



  describe('notifyHandshakeStarted()', () => {

    it('is a no-op before State.Starting fires', () => {

      monitor.observe(client).catch(() => {});



      monitor.notifyHandshakeStarted();



      expect(monitor.phase).toBe('idle');

    });



    it('is a no-op on subsequent calls within the same spawn attempt', () => {

      monitor.observe(client).catch(() => {});

      onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });



      monitor.notifyHandshakeStarted();

      monitor.notifyHandshakeStarted();



      expect(monitor.phase).toBe('spawning');

    });



    it('resets per spawn attempt so each new attempt records its own handshake start time', () => {

      monitor.observe(client).catch(() => {});



      onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

      monitor.notifyHandshakeStarted(); // attempt 1



      onDidChangeStateHandler({ oldState: State.Starting, newState: State.Stopped });

      onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

      monitor.notifyHandshakeStarted(); // attempt 2 - should not be a no-op



      expect(monitor.phase).toBe('spawning');

    });

  });



  describe('dispose()', () => {

    it('disposes the onDidChangeState listener', () => {

      const dispose = jest.fn();

      client = createFakePartial<BaseLanguageClient>({

        onDidChangeState: jest.fn(handler => {

          onDidChangeStateHandler = handler;

          return { dispose };

        }),

      });



      monitor.observe(client).catch(() => {});

      monitor.dispose();



      expect(dispose).toHaveBeenCalledTimes(1);

    });



    it('is safe to call multiple times', () => {

      const dispose = jest.fn();

      client = createFakePartial<BaseLanguageClient>({

        onDidChangeState: jest.fn(handler => {

          onDidChangeStateHandler = handler;

          return { dispose };

        }),

      });



      monitor.observe(client).catch(() => {});

      monitor.dispose();

      monitor.dispose();



      expect(dispose).toHaveBeenCalledTimes(1);

    });

  });



  describe('observe()', () => {

    it('registers onDidChangeState listener', () => {

      monitor.observe(client).catch(() => {});



      expect(client.onDidChangeState).toHaveBeenCalledWith(expect.any(Function));

    });



    it('transitions to spawning when State.Starting fires', () => {

      monitor.observe(client).catch(() => {});

      onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });



      expect(monitor.phase).toBe('spawning');

    });



    it('resolves and transitions to running when State.Running fires', async () => {

      const promise = monitor.observe(client);

      onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

      onDidChangeStateHandler({ oldState: State.Starting, newState: State.Running });



      await expect(promise).resolves.toBeUndefined();

      expect(monitor.phase).toBe('running');

    });



    it('does not reject when State.Stopped fires - lets the library handle restarts', async () => {

      const promise = monitor.observe(client);

      onDidChangeStateHandler({ oldState: State.Starting, newState: State.Stopped });



      // resolve via Running after restart - promise should still work

      onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

      onDidChangeStateHandler({ oldState: State.Starting, newState: State.Running });



      await expect(promise).resolves.toBeUndefined();

    });



    describe('spawn timeout', () => {

      it('rejects if State.Starting never fires within 10s', async () => {

        const promise = monitor.observe(client);

        jest.advanceTimersByTime(10000);



        await expect(promise).rejects.toThrow('Language Server process did not start within 10s');

        expect(monitor.phase).toBe('failed');

      });



      it('clears spawn timeout when State.Starting fires', async () => {

        const promise = monitor.observe(client);

        jest.advanceTimersByTime(9999);

        onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

        onDidChangeStateHandler({ oldState: State.Starting, newState: State.Running });



        await expect(promise).resolves.toBeUndefined();

      });

    });



    describe('handshake timeout', () => {

      it('rejects with handshake timeout message when process never completes handshake', async () => {

        const promise = monitor.observe(client);

        onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

        jest.advanceTimersByTime(30000);



        await expect(promise).rejects.toThrow(

          'Language Server process started but LSP initialize handshake did not complete within 30s',

        );

        expect(monitor.phase).toBe('failed');

      });



      it('clears handshake timeout when State.Running fires', async () => {

        const promise = monitor.observe(client);

        onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

        onDidChangeStateHandler({ oldState: State.Starting, newState: State.Running });

        jest.advanceTimersByTime(30000);



        await expect(promise).resolves.toBeUndefined();

      });



      it('resets handshake timeout on each State.Starting so the full 30s is allowed per attempt', async () => {

        const promise = monitor.observe(client);



        // first attempt - advance 29s without completing

        onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

        jest.advanceTimersByTime(29000);



        // crash and restart - timer should reset

        onDidChangeStateHandler({ oldState: State.Starting, newState: State.Stopped });

        onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });



        // advance another 29s - should not have fired yet (29s into the new timer)

        jest.advanceTimersByTime(29000);

        expect(monitor.phase).toBe('spawning');



        // now let it expire

        jest.advanceTimersByTime(1000);

        await expect(promise).rejects.toThrow(

          'Language Server process started but LSP initialize handshake did not complete within 30s',

        );

      });

    });



    describe('restart grace period', () => {

      it('transitions to failed and logs when State.Starting does not follow State.Stopped within grace period', async () => {

        const promise = monitor.observe(client);

        onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

        onDidChangeStateHandler({ oldState: State.Starting, newState: State.Running });

        await promise;



        onDidChangeStateHandler({ oldState: State.Running, newState: State.Stopped });

        jest.advanceTimersByTime(5000);



        expect(monitor.phase).toBe('failed');

        expect(log.debug).toHaveBeenCalledWith(

          expect.stringContaining('failed to restart after 1 attempt(s)'),

        );

      });



      it('does not reject when State.Starting follows State.Stopped within grace period', async () => {

        const promise = monitor.observe(client);

        onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

        onDidChangeStateHandler({ oldState: State.Starting, newState: State.Running });

        await promise;



        onDidChangeStateHandler({ oldState: State.Running, newState: State.Stopped });

        jest.advanceTimersByTime(4999);

        onDidChangeStateHandler({ oldState: State.Stopped, newState: State.Starting });

        jest.advanceTimersByTime(1); // grace period would have fired - but was cancelled



        expect(monitor.phase).toBe('spawning');

      });

    });

  });

});