diff --git a/src/client/datascience/ipywidgets/ipyWidgetMessageDispatcher.ts b/src/client/datascience/ipywidgets/ipyWidgetMessageDispatcher.ts index 7c594f59cce..f7d86df41aa 100644 --- a/src/client/datascience/ipywidgets/ipyWidgetMessageDispatcher.ts +++ b/src/client/datascience/ipywidgets/ipyWidgetMessageDispatcher.ts @@ -54,13 +54,23 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher { private totalWaitTime: number = 0; private totalWaitedMessages: number = 0; private hookCount: number = 0; - private fullHandleMessage?: { id: string; promise: Deferred }; /** - * This will be true if user has executed something that has resulted in the use of ipywidgets. - * We make this determinination based on whether we see messages coming from backend kernel of a specific shape. - * E.g. if it contains ipywidget mime type, then widgets are in use. + * The Output widget's model can set up or tear down a kernel message hook on state change. + * We need to wait until the kernel message hook has been connected before it's safe to send + * more messages to the UI kernel. + * + * To do this we: + * - Keep track of the id of all the Output widget models in the outputWidgetIds instance variable. + * We add/remove these ids by inspecting messages in onKernelSocketMessage. + * - When a state update message is sent to one of these widgets, we synchronize with the UI and + * stop sending messages until we receive a reply indicating that the state change has been fully handled. + * We keep track of the message we're waiting for in the fullHandleMessage instance variable. + * We start waiting for the state change to finish processing in onKernelSocketMessage, + * and we stop waiting in iopubMessageHandled. */ - private isUsingIPyWidgets?: boolean; + private outputWidgetIds = new Set(); + private fullHandleMessage?: { id: string; promise: Deferred }; + private isUsingIPyWidgets = false; private readonly deserialize: (data: string | ArrayBuffer) => KernelMessage.IMessage; constructor(private readonly kernelProvider: IKernelProvider, public readonly document: NotebookDocument) { @@ -251,14 +261,12 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher { // fully handled on both the UI and extension side before we process the next message incoming private messageNeedsFullHandle(message: any) { // We only get a handled callback for iopub messages, so this channel must be iopub - if (message.channel === 'iopub') { - if (message.header?.msg_type === 'comm_msg') { - // IOPub comm messages need to be fully handled - return true; - } - } - - return false; + return ( + message.channel === 'iopub' && + message.header?.msg_type === 'comm_msg' && + message.content?.data?.method === 'update' && + this.outputWidgetIds.has(message.content?.comm_id) + ); } // Callback from the UI kernel when an iopubMessage has been fully handled @@ -272,46 +280,11 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher { } private async onKernelSocketMessage(data: WebSocketData): Promise { // Hooks expect serialized data as this normally comes from a WebSocket - let message: undefined | KernelMessage.ICommOpenMsg; // = this.deserialize(data as any) as any; - if (!this.isUsingIPyWidgets) { - // Lets deserialize only if we know we have a potential case - // where this message contains some data we're interested in. - let mustDeserialize = false; - if (typeof data === 'string') { - mustDeserialize = data.includes(WIDGET_MIMETYPE) || data.includes(Identifiers.DefaultCommTarget); - } else { - // Array buffers (non-plain text data) must be deserialized. - mustDeserialize = true; - } - if (!message && mustDeserialize) { - message = this.deserialize(data as any) as any; - } - - // Check for hints that would indicate whether ipywidgest are used in outputs. - if ( - message && - message.content && - message.content.data && - (message.content.data[WIDGET_MIMETYPE] || message.content.target_name === Identifiers.DefaultCommTarget) - ) { - this.isUsingIPyWidgets = true; - } - } const msgUuid = uuid(); const promise = createDeferred(); this.waitingMessageIds.set(msgUuid, { startTime: Date.now(), resultPromise: promise }); - // Check if we need to fully handle this message on UI and Extension side before we move to the next - if (this.isUsingIPyWidgets) { - if (!message) { - message = this.deserialize(data as any) as any; - } - if (this.messageNeedsFullHandle(message)) { - this.fullHandleMessage = { id: message!.header.msg_id, promise: createDeferred() }; - } - } - if (typeof data === 'string') { this.raisePostMessage(IPyWidgetMessages.IPyWidgets_msg, { id: msgUuid, data }); } else { @@ -321,21 +294,42 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher { }); } - // There are three handling states that we have for messages here - // 1. If we have not detected ipywidget usage at all, we just forward messages to the kernel - // 2. If we have detected ipywidget usage. We wait on our message to be received, but not - // possibly processed yet by the UI kernel. This make sure our ordering is in sync - // 3. For iopub comm messages we wait for them to be fully handled by the UI kernel - // and the Extension kernel as they may be required to do things like - // register message hooks on both sides before we process the nextExtension message - - // If there are no ipywidgets thusfar in the notebook, then no need to synchronize messages. - if (this.isUsingIPyWidgets) { - await promise.promise; - - // Comm specific iopub messages we need to wait until they are full handled - // by both the UI and extension side before we move forward - if (this.fullHandleMessage) { + // Lets deserialize only if we know we have a potential case + // where this message contains some data we're interested in. + const mustDeserialize = + typeof data !== 'string' || + data.includes(WIDGET_MIMETYPE) || + data.includes(Identifiers.DefaultCommTarget) || + data.includes('comm_open') || + data.includes('comm_close') || + data.includes('comm_msg'); + if (mustDeserialize) { + const message = this.deserialize(data as any) as any; + + // Check for hints that would indicate whether ipywidgest are used in outputs. + if ( + message && + message.content && + message.content.data && + (message.content.data[WIDGET_MIMETYPE] || message.content.target_name === Identifiers.DefaultCommTarget) + ) { + this.isUsingIPyWidgets = true; + } + + const isIPYWidgetOutputModelOpen = + message.header?.msg_type === 'comm_open' && + message.content?.data?.state?._model_module === '@jupyter-widgets/output' && + message.content?.data?.state?._model_name === 'OutputModel'; + const isIPYWidgetOutputModelClose = + message.header?.msg_type === 'comm_close' && this.outputWidgetIds.has(message.content?.comm_id); + + if (isIPYWidgetOutputModelOpen) { + this.outputWidgetIds.add(message.content.comm_id); + } else if (isIPYWidgetOutputModelClose) { + this.outputWidgetIds.delete(message.content.comm_id); + } else if (this.messageNeedsFullHandle(message)) { + this.fullHandleMessage = { id: message.header.msg_id, promise: createDeferred() }; + await promise.promise; await this.fullHandleMessage.promise.promise; this.fullHandleMessage = undefined; }