Verified Commit 867641ee authored by Olena Horal-Koretska's avatar Olena Horal-Koretska 2️⃣ Committed by GitLab
Browse files

feat: Interrupt running command when LS notifies to do so

parent ba83692e

src/common/duo_workflow/terminal_manager.test.ts

+173 −42
Original line number Diff line number Diff line
@@ -24,9 +24,56 @@ describe('Terminal Manager', () => { 
  let mockExecution: TerminalShellExecution;

  let mockTerminal: Terminal;



  // Shared handler references

  let eventFn: RunCommandEventHandler;

  let cancelNotificationFn: (params: { workflowId: string }) => void;

  let shellIntegrationFn: (e: TerminalShellIntegrationChangeEvent) => void;

  let shellExecutionFn: (e: TerminalShellExecutionEndEvent) => void;

  let onDidCloseTerminalListener: (terminal: Terminal) => void;



  // Helper to setup standard event handlers

  const setupEventHandlers = (

    options: {

      shellExecDispose?: jest.Mock;

      shellIntDispose?: jest.Mock;

      requestDispose?: jest.Mock;

      closeTerminalDispose?: jest.Mock;

    } = {},

  ) => {

    jest.mocked(mockClient.onRequest).mockImplementation((event: string, fn) => {

      if (event === '$/gitlab/runCommand') {

        eventFn = fn as RunCommandEventHandler;

      }

      return { dispose: options.requestDispose || jest.fn() };

    });



    jest.mocked(mockClient.onNotification).mockImplementation((event: string, fn) => {

      if (event === '$/gitlab/cancelRunningCommand') {

        cancelNotificationFn = fn as (params: { workflowId: string }) => void;

      }

      return { dispose: jest.fn() };

    });



    jest.mocked(window.onDidChangeTerminalShellIntegration).mockImplementation(fn => {

      shellIntegrationFn = fn;

      return { dispose: options.shellIntDispose || jest.fn() };

    });



    jest.mocked(window.onDidEndTerminalShellExecution).mockImplementation(fn => {

      shellExecutionFn = fn;

      return { dispose: options.shellExecDispose || jest.fn() };

    });



    jest.mocked(window.onDidCloseTerminal).mockImplementation(listener => {

      onDidCloseTerminalListener = listener;

      return { dispose: options.closeTerminalDispose || jest.fn() };

    });

  };



  beforeEach(() => {

    mockClient = createFakePartial<BaseLanguageClient>({

      onRequest: jest.fn(),

      onNotification: jest.fn(),

    });

    mockShellIntegrationGetter = jest.fn();

    mockExecution = createFakePartial<TerminalShellExecution>({
@@ -41,9 +88,11 @@ describe('Terminal Manager', () => { 
      },

      show: jest.fn(),

      dispose: jest.fn(),

      sendText: jest.fn(),

    });



    jest.mocked(window.createTerminal).mockReturnValue(mockTerminal);

    jest.mocked(mockShellIntegration.executeCommand).mockReturnValue(mockExecution);



    terminalManager = new TerminalManagerImpl();

    jest.useFakeTimers();
@@ -71,28 +120,14 @@ describe('Terminal Manager', () => { 
  });



  describe('dispose', () => {

    let eventFn: RunCommandEventHandler;

    let shellExecutionFn: (e: TerminalShellExecutionEndEvent) => Disposable;

    let listenerDispose: jest.Mock;

    let closeTerminalDispose: jest.Mock;



    beforeEach(async () => {

      listenerDispose = jest.fn();

      closeTerminalDispose = jest.fn();

      jest.mocked(mockClient.onRequest).mockImplementation((event: string, fn) => {

        if (event === '$/gitlab/runCommand') {

          eventFn = fn as RunCommandEventHandler;

        }

        return { dispose: listenerDispose };

      });

      jest.mocked(window.onDidEndTerminalShellExecution).mockImplementation(fn => {

        shellExecutionFn = fn;

        return { dispose() {} };

      });

      jest.mocked(window.onDidCloseTerminal).mockImplementationOnce(() => {

        return { dispose: closeTerminalDispose };

      });

      jest.mocked(mockShellIntegration.executeCommand).mockReturnValue(mockExecution);



      setupEventHandlers({ requestDispose: listenerDispose, closeTerminalDispose });



      terminalManager.setupRequests(mockClient);
@@ -107,7 +142,7 @@ describe('Terminal Manager', () => { 


      const result = eventFn({ workflowId, command, args, silent: false });



      await jest.advanceTimersByTimeAsync(3000); // Shell integration timer

      await jest.advanceTimersByTimeAsync(3000);



      shellExecutionFn({

        execution: mockExecution,
@@ -134,38 +169,14 @@ describe('Terminal Manager', () => { 
    const command = 'npm';

    const silent = false;

    const args = ['run', 'test:unit'];

    let eventFn: RunCommandEventHandler;

    let shellIntegrationFn: (e: TerminalShellIntegrationChangeEvent) => Disposable;

    let shellExecutionFn: (e: TerminalShellExecutionEndEvent) => Disposable;

    let onDidCloseTerminalListener: (terminal: Terminal) => void;

    let shellExecDispose: jest.Mock;

    let shellIntDispose: jest.Mock;



    beforeEach(() => {

      shellExecDispose = jest.fn();

      shellIntDispose = jest.fn();

      onDidCloseTerminalListener = jest.fn();



      jest.mocked(mockClient.onRequest).mockImplementation((event: string, fn) => {

        if (event === '$/gitlab/runCommand') {

          eventFn = fn as RunCommandEventHandler;

        }

        return { dispose() {} };

      });

      jest.mocked(window.onDidChangeTerminalShellIntegration).mockImplementation(fn => {

        shellIntegrationFn = fn;

        return { dispose: shellIntDispose };

      });

      jest.mocked(window.onDidEndTerminalShellExecution).mockImplementation(fn => {

        shellExecutionFn = fn;

        return { dispose: shellExecDispose };

      });

      jest.mocked(mockShellIntegration.executeCommand).mockReturnValue(mockExecution);

      jest.mocked(window.onDidCloseTerminal).mockImplementationOnce(listener => {

        onDidCloseTerminalListener = listener;



        return { dispose() {} };

      });

      setupEventHandlers({ shellExecDispose, shellIntDispose });



      terminalManager.setupRequests(mockClient);

    });
@@ -576,4 +587,124 @@ describe('Terminal Manager', () => { 
      });

    });

  });



  describe('$/gitlab/cancelRunningCommand', () => {

    const workflowId = '1234';

    const command = 'npm';

    const args = ['run', 'test:unit'];

    const silent = false;



    beforeEach(() => {

      setupEventHandlers();

      terminalManager.setupRequests(mockClient);

    });



    it('listens to `$/gitlab/cancelRunningCommand` notification', () => {

      expect(mockClient.onNotification).toHaveBeenCalledWith(

        '$/gitlab/cancelRunningCommand',

        expect.any(Function),

      );

    });



    it('cancels a running command and returns partial output', async () => {

      mockShellIntegrationGetter.mockReturnValue(mockShellIntegration);



      // Simulate streaming output that gets interrupted

      const streamedOutput = ['line 1\n', 'line 2\n'];



      jest.mocked(mockExecution.read).mockImplementation(async function* read() {

        yield streamedOutput[0];

        yield streamedOutput[1];

      });



      const resultPromise = eventFn({ workflowId, command, args, silent });



      await jest.advanceTimersByTimeAsync(3000);



      // Let some output be collected

      await Promise.resolve();

      await Promise.resolve();



      // Cancel the command before it completes

      cancelNotificationFn({ workflowId });



      const result = (await resultPromise) as RunCommandResult;



      expect(mockTerminal.sendText).toHaveBeenCalledWith('\x03', false); // Ctrl+C without newline

      expect(result.output).toContain('line 1');

      expect(result.exitCode).toBe(0);

    });



    it('handles cancel for non-existent workflow gracefully', () => {

      const nonExistentWorkflowId = 'does-not-exist';



      expect(() => {

        cancelNotificationFn({ workflowId: nonExistentWorkflowId });

      }).not.toThrow();

    });



    it('does not affect command if cancel arrives after completion', async () => {

      mockShellIntegrationGetter.mockReturnValue(mockShellIntegration);



      const fullOutput = 'command completed successfully\n';

      jest.mocked(mockExecution.read).mockImplementation(async function* read() {

        yield fullOutput;

      });



      const resultPromise = eventFn({ workflowId, command, args, silent });



      await jest.advanceTimersByTimeAsync(3000);



      // Complete the execution normally

      shellExecutionFn({

        execution: mockExecution,

        exitCode: 0,

        terminal: mockTerminal,

        shellIntegration: mockShellIntegration,

      });



      const result = (await resultPromise) as RunCommandResult;



      // Now try to cancel after completion

      cancelNotificationFn({ workflowId });



      // Should not send Ctrl+C since command already completed

      expect(mockTerminal.sendText).not.toHaveBeenCalled();

      expect(result.output).toBe(fullOutput);

      expect(result.exitCode).toBe(0);

    });

  });



  describe('dispose with running commands', () => {

    const workflowId = '1234';

    const command = 'npm';

    const args = ['run', 'test:unit'];

    const silent = false;



    beforeEach(() => {

      setupEventHandlers();

      terminalManager.setupRequests(mockClient);

    });



    it('cancels running commands when disposed', async () => {

      mockShellIntegrationGetter.mockReturnValue(mockShellIntegration);



      jest.mocked(mockExecution.read).mockImplementation(async function* read() {

        yield 'some output\n';

      });



      const resultPromise = eventFn({ workflowId, command, args, silent });



      await jest.advanceTimersByTimeAsync(3000);

      await Promise.resolve();



      // Dispose should cancel the running command

      terminalManager.dispose();



      const result = (await resultPromise) as RunCommandResult;



      expect(mockTerminal.sendText).toHaveBeenCalledWith('\x03', false);

      expect(result.output).toBe('some output\n');

    });

  });

});

src/common/duo_workflow/terminal_manager.ts

+78 −8
Original line number Diff line number Diff line 
import { Terminal, window, Disposable } from 'vscode';

import { Terminal, window, Disposable, TerminalShellExecution } from 'vscode';

import { BaseLanguageClient } from 'vscode-languageclient';

import { createInterfaceId } from '@gitlab/needle';

import { log } from '../log';
@@ -9,6 +9,13 @@ const lookupKey = (workflowId: string, silent: boolean) => { 
  return `${workflowId}-${silent}`;

};



type RunningExecution = {

  execution: TerminalShellExecution;

  terminal: Terminal;

  outputCollector: string[];

  resolver: (value: { exitCode: number; output: string }) => void;

};



export interface TerminalManager extends Disposable {

  setupRequests(client: BaseLanguageClient): void;

}
@@ -18,10 +25,13 @@ export const TerminalManager = createInterfaceId<TerminalManager>('TerminalManag 
export class TerminalManagerImpl implements TerminalManager {

  #terminals: Map<string, Terminal>;



  #runningExecutions: Map<string, RunningExecution>;



  #disposables: Disposable[] = [];



  constructor() {

    this.#terminals = new Map();

    this.#runningExecutions = new Map();

  }



  setupRequests(client: BaseLanguageClient) {
@@ -33,6 +43,13 @@ export class TerminalManagerImpl implements TerminalManager { 
      }),

    );



    this.#disposables.push(

      client.onNotification('$/gitlab/cancelRunningCommand', ({ workflowId }) => {

        log.debug(`Received cancel request for workflow: ${workflowId}`);

        this.#cancelRunningCommand(workflowId);

      }),

    );



    this.#disposables.push(

      window.onDidCloseTerminal(closedTerminal => {

        for (const [workflowId, terminal] of this.#terminals.entries()) {
@@ -68,17 +85,32 @@ export class TerminalManagerImpl implements TerminalManager { 
      ? terminal.shellIntegration.executeCommand(command, args)

      : terminal.shellIntegration.executeCommand(command);



    const executionEndPromise = new Promise<number>(resolve => {

      const disposable = window.onDidEndTerminalShellExecution(({ execution, exitCode }) => {

    const output: string[] = [];

    let executionEndDisposable: Disposable | undefined;



    // Create promises for both normal completion and cancellation

    const executionEndPromise = new Promise<{ exitCode: number; cancelled: false }>(resolve => {

      executionEndDisposable = window.onDidEndTerminalShellExecution(({ execution, exitCode }) => {

        if (exec === execution) {

          disposable.dispose();

          resolve(exitCode ?? 0);

          executionEndDisposable?.dispose();

          resolve({ exitCode: exitCode ?? 0, cancelled: false });

        }

      });

    });



    const output: string[] = [];

    const cancellationPromise = new Promise<{ exitCode: number; cancelled: true }>(resolve => {

      // Store the running execution so it can be cancelled

      this.#runningExecutions.set(workflowId, {

        execution: exec,

        terminal,

        outputCollector: output,

        resolver: (value: { exitCode: number; output: string }) => {

          resolve({ exitCode: value.exitCode, cancelled: true });

        },

      });

    });



    const streamPromise = (async () => {

      try {

        const stream = exec.read();

        for await (const data of stream) {
@@ -87,10 +119,22 @@ export class TerminalManagerImpl implements TerminalManager { 
      } catch (error) {

        log.error('Error reading stream:', error);

      }

    })();



    const exitCode = await executionEndPromise;

    // Wait for either normal completion or cancellation

    const result = await Promise.race([executionEndPromise, cancellationPromise]);



    return { exitCode, output: output.join('') };

    executionEndDisposable?.dispose();

    this.#runningExecutions.delete(workflowId);



    // If cancelled, return immediately with partial output

    if (result.cancelled) {

      return { exitCode: result.exitCode, output: output.join('') };

    }



    // Otherwise wait for stream to finish and return full output

    await streamPromise;

    return { exitCode: result.exitCode, output: output.join('') };

  }



  async #getOrCreateTerminal(workflowId: string, silent: boolean) {
@@ -116,6 +160,27 @@ export class TerminalManagerImpl implements TerminalManager { 
    }

  }



  #cancelRunningCommand(workflowId: string) {

    const runningExec = this.#runningExecutions.get(workflowId);



    if (!runningExec) {

      log.warn(`No running command found for workflow ${workflowId}, may have already completed`);

      return;

    }



    log.info(`Cancelling running command for workflow ${workflowId}`);



    // Remove entry first to prevent race with #executeCommand cleanup

    this.#runningExecutions.delete(workflowId);



    // Send Ctrl+C (SIGINT) to the terminal to interrupt the running process

    runningExec.terminal.sendText('\x03', false);



    // Trigger the cancellation promise with partial output

    const partialOutput = runningExec.outputCollector.join('');

    runningExec.resolver({ exitCode: 0, output: partialOutput });

  }



  #listenForShellIntegration(term: Terminal) {

    return new Promise<Terminal>((resolve, reject) => {

      if (term.shellIntegration) {
@@ -146,6 +211,11 @@ export class TerminalManagerImpl implements TerminalManager { 
  }



  dispose() {

    this.#runningExecutions.forEach((_, workflowId) => {

      this.#cancelRunningCommand(workflowId);

    });

    this.#runningExecutions.clear();



    this.#terminals.forEach(term => {

      term.dispose();

    });