Wait for IOPub status before switching to Ready

extensions/jupyter-adapter/src/JupyterKernel.ts

@@ -114,6 +114,18 @@ export class JupyterKernel extends EventEmitter implements vscode.Disposable {
 	 */
 	private readonly _idleMessageIds: Set<string> = new Set();
 
+	/**
+	 * Initialisation state. Used to detect on startup that we (a) get a reply to a
+         * dummy kernel-info request on the Shell socket and (b) get a status message on
+         * the IOPub socket. When both of these happen, it means the kernel has made
+         * sufficient progress in its startup that it's able to service requests, and that
+         * the IOPub socket is correctly connected and that the user won't lose data when
+         * we send user requests to the kernel.
+	 */
+	private _initial_request_id = '';
+	private _receivedInitialStatus = false;
+	private _receivedInitialReply = false;
+
 	constructor(
 		private readonly _context: vscode.ExtensionContext,
 		spec: JupyterKernelSpec,
@@ -399,20 +411,32 @@ export class JupyterKernel extends EventEmitter implements vscode.Disposable {
 					}
 				});
 
-				let got_initial_reply = false;
-				let initial_request_id = uuidv4();
-
 				// Connect the Shell socket
 				this._shell?.onMessage(async (args: any[]) => {
 					const msg = deserializeJupyterMessage(args, this._session!.key, this._channel);
 					if (msg !== null) {
 						// Wait for reply to kernel-info request. When it comes in, we
 						// know the kernel is started up and ready to service
 						// requests. That's our cue to switch to the Ready state.
-						if (!got_initial_reply && msg.parent_header.msg_id === initial_request_id) {
-							got_initial_reply = true;
-							this.setStatus(positron.RuntimeState.Ready);
-							resolve();
+						if (!this._receivedInitialReply && msg.parent_header.msg_id === this._initial_request_id) {
+							// It's possible for the Shell socket to be connected
+							// and ready before the IOPub socket. This is a bad
+							// situation because the server might serve requests
+							// but we won't get the output back, resulting in data
+							// loss. To avoid this, we keep resending dummy
+							// requests until we get an IOPub status message. See
+							// discussion in
+							// https://github.com/jupyter/enhancement-proposals/blob/master/65-jupyter-xpub/jupyter-xpub.md
+							if (this._receivedInitialStatus) {
+								// Got both a reply on Shell and a status on IOPub, we're ready
+								this._receivedInitialReply = true;
+								this.setStatus(positron.RuntimeState.Ready);
+								resolve();
+							} else {
+								// Send the request again until we've received an IOPub status
+								this.sendInitialRequest()
+							}
+							return;
 						}
 						this.emitMessage(JupyterSockets.shell, msg);
 					}
@@ -459,7 +483,7 @@ export class JupyterKernel extends EventEmitter implements vscode.Disposable {
 					this.log('Received initial heartbeat: ' + msg);
 
 					// Send a dummy request. When the reply comes in we'll switch to the Ready state.
-					this.send(initial_request_id, 'kernel_info_request', this._shell!, {});
+					this.sendInitialRequest()
 
 					const seconds = vscode.workspace.getConfiguration('positron.jupyterAdapter').get('heartbeat', 30) as number;
 					this.log(`Starting heartbeat check at ${seconds} second intervals...`);
@@ -482,6 +506,11 @@ export class JupyterKernel extends EventEmitter implements vscode.Disposable {
 		});
 	}
 
+	private sendInitialRequest() {
+		this._initial_request_id = uuidv4();
+		this.send(this._initial_request_id, 'kernel_info_request', this._shell!, {});
+	}
+
 	/**
 	 * Starts the Jupyter kernel. Resolves when the kernel is ready to receive
 	 * messages; rejects with a StartupFailure error if the kernel fails to
@@ -557,6 +586,10 @@ export class JupyterKernel extends EventEmitter implements vscode.Disposable {
 		// during bootup
 		this.setStatus(positron.RuntimeState.Initializing);
 
+		// Reset the initialisation state that we use as cues to switch to Ready
+		this._receivedInitialStatus = false;
+		this._receivedInitialReply = false;
+
 		// Before starting a new process, look for metadata about a running
 		// kernel in the current workspace by checking the value stored for this
 		// runtime ID (we support running exactly one kernel per runtime ID). If
@@ -1293,6 +1326,7 @@ export class JupyterKernel extends EventEmitter implements vscode.Disposable {
 		// during startup (e.g. during the kernel-info request
 		// that we emit to detect kernel readiness)
 		if (this.isStarting() && this.isRuntimeStatus(status)) {
+			this._receivedInitialStatus = true;
 			return;
 		}