diff --git a/src/kernels/execution/inputFlushStartupCodeProvider.ts b/src/kernels/execution/inputFlushStartupCodeProvider.ts new file mode 100644 index 00000000000..b50939c76c7 --- /dev/null +++ b/src/kernels/execution/inputFlushStartupCodeProvider.ts @@ -0,0 +1,60 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { inject, injectable } from 'inversify'; +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { IKernel, IStartupCodeProvider, IStartupCodeProviders, StartupCodePriority } from '../types'; +import { isPythonKernelConnection } from '../helpers'; +import { InteractiveWindowView, JupyterNotebookView } from '../../platform/common/constants'; + +/** + * Startup code that monkey-patches Python's input() function to flush stdout before requesting input. + * This ensures that any pending output (like print statements) is displayed before the input prompt. + */ +const inputFlushStartupCode = ` +# Monkey patch input() to flush stdout before requesting input +import builtins +import sys + +def __vscode_input_with_flush(*args, **kwargs): + """ + Wrapper around input() that flushes stdout before requesting input. + This ensures that any pending output is displayed before the input prompt. + """ + try: + # Flush stdout to ensure all output is displayed before the input prompt + sys.stdout.flush() + except: + # If flushing fails for any reason, continue without error + pass + + # Call the original input function + return __vscode_original_input(*args, **kwargs) + +# Store reference to original input and replace with our wrapper +__vscode_original_input = builtins.input +builtins.input = __vscode_input_with_flush + +# Clean up temporary variables +del __vscode_input_with_flush +`.trim(); + +@injectable() +export class InputFlushStartupCodeProvider implements IStartupCodeProvider, IExtensionSyncActivationService { + public priority = StartupCodePriority.Base; + + constructor(@inject(IStartupCodeProviders) private readonly registry: IStartupCodeProviders) {} + + activate(): void { + this.registry.register(this, JupyterNotebookView); + this.registry.register(this, InteractiveWindowView); + } + + async getCode(kernel: IKernel): Promise { + // Only apply this monkey patch to Python kernels + if (!isPythonKernelConnection(kernel.kernelConnectionMetadata)) { + return []; + } + return [inputFlushStartupCode]; + } +} diff --git a/src/kernels/execution/inputFlushStartupCodeProvider.unit.test.ts b/src/kernels/execution/inputFlushStartupCodeProvider.unit.test.ts new file mode 100644 index 00000000000..844cda07250 --- /dev/null +++ b/src/kernels/execution/inputFlushStartupCodeProvider.unit.test.ts @@ -0,0 +1,85 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { expect } from 'chai'; +import { mock, instance, when } from 'ts-mockito'; +import { InputFlushStartupCodeProvider } from './inputFlushStartupCodeProvider'; +import { IKernel, IStartupCodeProviders, KernelConnectionMetadata } from '../types'; + +suite('InputFlushStartupCodeProvider', () => { + let provider: InputFlushStartupCodeProvider; + let mockRegistry: IStartupCodeProviders; + let mockKernel: IKernel; + + setup(() => { + mockRegistry = mock(); + mockKernel = mock(); + provider = new InputFlushStartupCodeProvider(instance(mockRegistry)); + }); + + test('Should return startup code for Python kernels', async () => { + // Arrange - Create a Python kernel connection metadata + const pythonConnection: KernelConnectionMetadata = { + kind: 'startUsingPythonInterpreter', + id: 'test-python-kernel' + } as any; + when(mockKernel.kernelConnectionMetadata).thenReturn(pythonConnection); + + // Act + const code = await provider.getCode(instance(mockKernel)); + + // Assert + expect(code).to.have.length(1); + expect(code[0]).to.contain('builtins.input'); + expect(code[0]).to.contain('sys.stdout.flush()'); + expect(code[0]).to.contain('__vscode_input_with_flush'); + }); + + test('Should return empty array for non-Python kernels', async () => { + // Arrange - Create a non-Python kernel connection metadata + const nonPythonConnection: KernelConnectionMetadata = { + kind: 'startUsingLocalKernelSpec', + id: 'test-non-python-kernel' + } as any; + when(mockKernel.kernelConnectionMetadata).thenReturn(nonPythonConnection); + + // Act + const code = await provider.getCode(instance(mockKernel)); + + // Assert + expect(code).to.be.empty; + }); + + test('Startup code should monkey patch input correctly', async () => { + // Arrange + const pythonConnection: KernelConnectionMetadata = { + kind: 'startUsingPythonInterpreter', + id: 'test-python-kernel' + } as any; + when(mockKernel.kernelConnectionMetadata).thenReturn(pythonConnection); + + // Act + const code = await provider.getCode(instance(mockKernel)); + + // Assert + const startupCode = code[0]; + + // Should import required modules + expect(startupCode).to.contain('import builtins'); + expect(startupCode).to.contain('import sys'); + + // Should define wrapper function + expect(startupCode).to.contain('def __vscode_input_with_flush'); + + // Should flush stdout before calling original input + expect(startupCode).to.contain('sys.stdout.flush()'); + expect(startupCode).to.contain('__vscode_original_input(*args, **kwargs)'); + + // Should replace builtins.input with wrapper + expect(startupCode).to.contain('__vscode_original_input = builtins.input'); + expect(startupCode).to.contain('builtins.input = __vscode_input_with_flush'); + + // Should clean up temporary variables + expect(startupCode).to.contain('del __vscode_input_with_flush'); + }); +}); diff --git a/src/kernels/serviceRegistry.node.ts b/src/kernels/serviceRegistry.node.ts index 74a8c5594f8..fa9843989de 100644 --- a/src/kernels/serviceRegistry.node.ts +++ b/src/kernels/serviceRegistry.node.ts @@ -47,6 +47,7 @@ import { IJupyterVariables } from './variables/types'; import { LastCellExecutionTracker } from './execution/lastCellExecutionTracker'; import { ClearJupyterServersCommand } from './jupyter/clearJupyterServersCommand'; import { KernelChatStartupCodeProvider } from './chat/kernelStartupCodeProvider'; +import { InputFlushStartupCodeProvider } from './execution/inputFlushStartupCodeProvider'; import { KernelWorkingDirectory } from './raw/session/kernelWorkingDirectory.node'; export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) { @@ -143,4 +144,8 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, KernelChatStartupCodeProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + InputFlushStartupCodeProvider + ); } diff --git a/src/kernels/serviceRegistry.web.ts b/src/kernels/serviceRegistry.web.ts index 7779e8c213b..e950ec92acd 100644 --- a/src/kernels/serviceRegistry.web.ts +++ b/src/kernels/serviceRegistry.web.ts @@ -35,6 +35,7 @@ import { KernelStartupCodeProviders } from './kernelStartupCodeProviders.web'; import { LastCellExecutionTracker } from './execution/lastCellExecutionTracker'; import { ClearJupyterServersCommand } from './jupyter/clearJupyterServersCommand'; import { KernelChatStartupCodeProvider } from './chat/kernelStartupCodeProvider'; +import { InputFlushStartupCodeProvider } from './execution/inputFlushStartupCodeProvider'; @injectable() class RawNotebookSupportedService implements IRawNotebookSupportedService { @@ -103,4 +104,8 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, KernelChatStartupCodeProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + InputFlushStartupCodeProvider + ); }