Fix stdout flushing for input() calls in notebook cells

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<string[]> {
+        // Only apply this monkey patch to Python kernels
+        if (!isPythonKernelConnection(kernel.kernelConnectionMetadata)) {
+            return [];
+        }
+        return [inputFlushStartupCode];
+    }
+}

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<IStartupCodeProviders>();
+        mockKernel = mock<IKernel>();
+        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');
+    });
+});

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>(
+        IExtensionSyncActivationService,
+        InputFlushStartupCodeProvider
+    );
 }

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>(
+        IExtensionSyncActivationService,
+        InputFlushStartupCodeProvider
+    );
 }