Skip to content

Commit a57cb45

Browse files
authored
notebooks: pipe renderer API postmessages to the renderer itself (microsoft#100414)
* notebooks: pipe renderer API postmessages to the renderer itself Previously the postMessage on acquireNotebookRenderer API was just a proxy to the global vscode postmessage. Now, it's linked to the renderer and will cause an optional `onDidReceiveMessage` method on the renderer to be called. The message still _also_ goes to the global webview message handling for advanced use cases, but this change allows the webview<->renderer communication to be more nicely contained and separate for most use cases. * wip * fixup! pr comments
1 parent 5e6729d commit a57cb45

11 files changed

Lines changed: 206 additions & 85 deletions

File tree

src/vs/vscode.proposed.d.ts

Lines changed: 47 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1607,6 +1607,20 @@ declare module 'vscode' {
16071607
*/
16081608
render(document: NotebookDocument, request: NotebookRenderRequest): string;
16091609

1610+
/**
1611+
* Call before HTML from the renderer is executed, and will be called for
1612+
* every editor associated with notebook documents where the renderer
1613+
* is or was used.
1614+
*
1615+
* The communication object will only send and receive messages to the
1616+
* render API, retrieved via `acquireNotebookRendererApi`, acquired with
1617+
* this specific renderer's ID.
1618+
*
1619+
* If you need to keep an association between the communication object
1620+
* and the document for use in the `render()` method, you can use a WeakMap.
1621+
*/
1622+
resolveNotebook?(document: NotebookDocument, communication: NotebookCommunication): void;
1623+
16101624
readonly preloads?: Uri[];
16111625
}
16121626

@@ -1735,27 +1749,41 @@ declare module 'vscode' {
17351749
readonly backupId?: string;
17361750
}
17371751

1752+
/**
1753+
* Communication object passed to the {@link NotebookContentProvider} and
1754+
* {@link NotebookOutputRenderer} to communicate with the webview.
1755+
*/
1756+
export interface NotebookCommunication {
1757+
/**
1758+
* ID of the editor this object communicates with. A single notebook
1759+
* document can have multiple attached webviews and editors, when the
1760+
* notebook is split for instance. The editor ID lets you differentiate
1761+
* between them.
1762+
*/
1763+
readonly editorId: string;
1764+
1765+
/**
1766+
* Fired when the output hosting webview posts a message.
1767+
*/
1768+
readonly onDidReceiveMessage: Event<any>;
1769+
/**
1770+
* Post a message to the output hosting webview.
1771+
*
1772+
* Messages are only delivered if the editor is live.
1773+
*
1774+
* @param message Body of the message. This must be a string or other json serilizable object.
1775+
*/
1776+
postMessage(message: any): Thenable<boolean>;
1777+
1778+
/**
1779+
* Convert a uri for the local file system to one that can be used inside outputs webview.
1780+
*/
1781+
asWebviewUri(localResource: Uri): Uri;
1782+
}
1783+
17381784
export interface NotebookContentProvider {
17391785
openNotebook(uri: Uri, openContext: NotebookDocumentOpenContext): NotebookData | Promise<NotebookData>;
1740-
resolveNotebook(document: NotebookDocument, webview: {
1741-
/**
1742-
* Fired when the output hosting webview posts a message.
1743-
*/
1744-
readonly onDidReceiveMessage: Event<any>;
1745-
/**
1746-
* Post a message to the output hosting webview.
1747-
*
1748-
* Messages are only delivered if the editor is live.
1749-
*
1750-
* @param message Body of the message. This must be a string or other json serilizable object.
1751-
*/
1752-
postMessage(message: any): Thenable<boolean>;
1753-
1754-
/**
1755-
* Convert a uri for the local file system to one that can be used inside outputs webview.
1756-
*/
1757-
asWebviewUri(localResource: Uri): Uri;
1758-
}): Promise<void>;
1786+
resolveNotebook(document: NotebookDocument, webview: NotebookCommunication): Promise<void>;
17591787
saveNotebook(document: NotebookDocument, cancellation: CancellationToken): Promise<void>;
17601788
saveNotebookAs(targetResource: Uri, document: NotebookDocument, cancellation: CancellationToken): Promise<void>;
17611789
readonly onDidChangeNotebook: Event<NotebookDocumentContentChangeEvent>;

src/vs/workbench/api/browser/mainThreadNotebook.ts

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,7 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
253253
return;
254254
}
255255

256-
this._proxy.$acceptDocumentAndEditorsDelta(delta);
256+
return this._proxy.$acceptDocumentAndEditorsDelta(delta);
257257
}
258258

259259
registerListeners() {
@@ -476,16 +476,11 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
476476
return this._proxy.$executeNotebook(viewType, uri, undefined, useAttachedKernel, token);
477477
}
478478

479-
async $postMessage(handle: number, value: any): Promise<boolean> {
480-
481-
const activeEditorPane = this.editorService.activeEditorPane as any | undefined;
482-
if (activeEditorPane?.isNotebookEditor) {
483-
const notebookEditor = (activeEditorPane.getControl() as INotebookEditor);
484-
485-
if (notebookEditor.viewModel?.handle === handle) {
486-
notebookEditor.postMessage(value);
487-
return true;
488-
}
479+
async $postMessage(editorId: string, forRendererId: string | undefined, value: any): Promise<boolean> {
480+
const editor = this._notebookService.getNotebookEditor(editorId) as INotebookEditor | undefined;
481+
if (editor?.isNotebookEditor) {
482+
editor.postMessage(forRendererId, value);
483+
return true;
489484
}
490485

491486
return false;
@@ -645,8 +640,8 @@ export class MainThreadNotebookController implements IMainNotebookController {
645640
return this._mainThreadNotebook.executeNotebook(viewType, uri, useAttachedKernel, token);
646641
}
647642

648-
onDidReceiveMessage(editorId: string, message: any): void {
649-
this._proxy.$onDidReceiveMessage(editorId, message);
643+
onDidReceiveMessage(editorId: string, rendererType: string | undefined, message: unknown): void {
644+
this._proxy.$onDidReceiveMessage(editorId, rendererType, message);
650645
}
651646

652647
async removeNotebookDocument(notebook: INotebookTextModel): Promise<void> {

src/vs/workbench/api/common/extHost.protocol.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -711,7 +711,7 @@ export interface MainThreadNotebookShape extends IDisposable {
711711
$updateNotebookMetadata(viewType: string, resource: UriComponents, metadata: NotebookDocumentMetadata): Promise<void>;
712712
$updateNotebookCellMetadata(viewType: string, resource: UriComponents, handle: number, metadata: NotebookCellMetadata | undefined): Promise<void>;
713713
$spliceNotebookCellOutputs(viewType: string, resource: UriComponents, cellHandle: number, splices: NotebookCellOutputsSplice[], renderers: number[]): Promise<void>;
714-
$postMessage(handle: number, value: any): Promise<boolean>;
714+
$postMessage(editorId: string, forRendererId: string | undefined, value: any): Promise<boolean>;
715715

716716
$onDidEdit(resource: UriComponents, viewType: string, editId: number, label: string | undefined): void;
717717
$onContentChange(resource: UriComponents, viewType: string): void;
@@ -1610,7 +1610,7 @@ export interface ExtHostNotebookShape {
16101610
$acceptDisplayOrder(displayOrder: INotebookDisplayOrder): void;
16111611
$renderOutputs(uriComponents: UriComponents, id: string, request: IOutputRenderRequest<UriComponents>): Promise<IOutputRenderResponse<UriComponents> | undefined>;
16121612
$renderOutputs2<T>(uriComponents: UriComponents, id: string, request: IOutputRenderRequest<T>): Promise<IOutputRenderResponse<T> | undefined>;
1613-
$onDidReceiveMessage(editorId: string, message: any): void;
1613+
$onDidReceiveMessage(editorId: string, rendererId: string | undefined, message: unknown): void;
16141614
$acceptModelChanged(uriComponents: UriComponents, event: NotebookCellsChangedEvent): void;
16151615
$acceptEditorPropertiesChanged(uriComponents: UriComponents, data: INotebookEditorPropertiesChangeData): void;
16161616
$acceptDocumentAndEditorsDelta(delta: INotebookDocumentsAndEditorsDelta): Promise<void>;

src/vs/workbench/api/common/extHostNotebook.ts

Lines changed: 84 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -494,6 +494,10 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
494494
cell.detachTextDocument();
495495
}
496496
}
497+
498+
allEditors() {
499+
return this._documentsAndEditors.allEditors();
500+
}
497501
}
498502

499503
export class NotebookEditorCellEditBuilder implements vscode.NotebookEditorCellEdit {
@@ -553,27 +557,48 @@ export class NotebookEditorCellEditBuilder implements vscode.NotebookEditorCellE
553557
}
554558
}
555559

556-
class ExtHostWebviewComm extends Disposable {
557-
558-
onDidReceiveMessage: vscode.Event<any> = this._onDidReceiveMessage.event;
560+
class ExtHostWebviewCommWrapper extends Disposable {
561+
private readonly _onDidReceiveDocumentMessage = new Emitter<any>();
562+
private readonly _rendererIdToEmitters = new Map<string, Emitter<any>>();
559563

560564
constructor(
561-
readonly id: string,
565+
private _editorId: string,
562566
public uri: URI,
563567
private _proxy: MainThreadNotebookShape,
564-
private _onDidReceiveMessage: Emitter<any>,
565568
private _webviewInitData: WebviewInitData,
566569
public document: ExtHostNotebookDocument,
567570
) {
568571
super();
569572
}
570573

571-
async postMessage(message: any): Promise<boolean> {
572-
return this._proxy.$postMessage(this.document.handle, message);
574+
public onDidReceiveMessage(forRendererId: string | undefined, message: any) {
575+
this._onDidReceiveDocumentMessage.fire(message);
576+
if (forRendererId !== undefined) {
577+
this._rendererIdToEmitters.get(forRendererId)?.fire(message);
578+
}
573579
}
574580

575-
asWebviewUri(localResource: vscode.Uri): vscode.Uri {
576-
return asWebviewUri(this._webviewInitData, this.id, localResource);
581+
public readonly contentProviderComm: vscode.NotebookCommunication = {
582+
editorId: this._editorId,
583+
onDidReceiveMessage: this._onDidReceiveDocumentMessage.event,
584+
postMessage: (message: any) => this._proxy.$postMessage(this._editorId, undefined, message),
585+
asWebviewUri: (uri: vscode.Uri) => this._asWebviewUri(uri),
586+
};
587+
588+
public getRendererComm(rendererId: string): vscode.NotebookCommunication {
589+
const emitter = new Emitter<any>();
590+
this._rendererIdToEmitters.set(rendererId, emitter);
591+
return {
592+
editorId: this._editorId,
593+
onDidReceiveMessage: emitter.event,
594+
postMessage: (message: any) => this._proxy.$postMessage(this._editorId, rendererId, message),
595+
asWebviewUri: (uri: vscode.Uri) => this._asWebviewUri(uri),
596+
};
597+
}
598+
599+
600+
private _asWebviewUri(localResource: vscode.Uri): vscode.Uri {
601+
return asWebviewUri(this._webviewInitData, this._editorId, localResource);
577602
}
578603
}
579604

@@ -618,7 +643,7 @@ export class ExtHostNotebookEditor extends Disposable implements vscode.Notebook
618643
readonly id: string,
619644
public uri: URI,
620645
private _proxy: MainThreadNotebookShape,
621-
private _webComm: ExtHostWebviewComm,
646+
private _webComm: vscode.NotebookCommunication,
622647
public document: ExtHostNotebookDocument,
623648
private _documentsAndEditors: ExtHostDocumentsAndEditors
624649
) {
@@ -721,6 +746,7 @@ export class ExtHostNotebookEditor extends Disposable implements vscode.Notebook
721746

722747
export class ExtHostNotebookOutputRenderer {
723748
private static _handlePool: number = 0;
749+
private resolvedComms = new WeakSet<ExtHostWebviewCommWrapper>();
724750
readonly handle = ExtHostNotebookOutputRenderer._handlePool++;
725751

726752
constructor(
@@ -740,13 +766,19 @@ export class ExtHostNotebookOutputRenderer {
740766
return false;
741767
}
742768

769+
resolveNotebook(document: ExtHostNotebookDocument, comm: ExtHostWebviewCommWrapper) {
770+
if (!this.resolvedComms.has(comm) && this.renderer.resolveNotebook) {
771+
this.renderer.resolveNotebook(document, comm.getRendererComm(this.type));
772+
this.resolvedComms.add(comm);
773+
}
774+
}
775+
743776
render(document: ExtHostNotebookDocument, output: vscode.CellDisplayOutput, outputId: string, mimeType: string): string {
744777
let html = this.renderer.render(document, { output, outputId, mimeType });
745778

746779
return html;
747780
}
748781
}
749-
750782
export interface ExtHostNotebookOutputRenderingHandler {
751783
outputDisplayOrder: INotebookDisplayOrder | undefined;
752784
findBestMatchedRenderer(mimeType: string): ExtHostNotebookOutputRenderer[];
@@ -759,8 +791,9 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
759791
private readonly _documents = new Map<string, ExtHostNotebookDocument>();
760792
private readonly _unInitializedDocuments = new Map<string, ExtHostNotebookDocument>();
761793
private readonly _editors = new Map<string, { editor: ExtHostNotebookEditor }>();
762-
private readonly _webviewComm = new Map<string, { comm: ExtHostWebviewComm, onDidReceiveMessage: Emitter<any> }>();
794+
private readonly _webviewComm = new Map<string, ExtHostWebviewCommWrapper>();
763795
private readonly _notebookOutputRenderers = new Map<string, ExtHostNotebookOutputRenderer>();
796+
private readonly _renderersUsedInNotebooks = new WeakMap<ExtHostNotebookDocument, Set<ExtHostNotebookOutputRenderer>>();
764797
private readonly _onDidChangeNotebookCells = new Emitter<vscode.NotebookCellsChangeEvent>();
765798
readonly onDidChangeNotebookCells = this._onDidChangeNotebookCells.event;
766799
private readonly _onDidChangeCellOutputs = new Emitter<vscode.NotebookCellOutputsChangeEvent>();
@@ -854,6 +887,8 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
854887
}
855888

856889
const renderer = this._notebookOutputRenderers.get(id)!;
890+
this.provideCommToNotebookRenderers(document, renderer);
891+
857892
const cellsResponse: IOutputRenderResponseCellInfo<UriComponents>[] = request.items.map(cellInfo => {
858893
const cell = document.getCell2(cellInfo.key)!;
859894
const outputResponse: IOutputRenderResponseOutputInfo[] = cellInfo.outputs.map(output => {
@@ -890,6 +925,8 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
890925
}
891926

892927
const renderer = this._notebookOutputRenderers.get(id)!;
928+
this.provideCommToNotebookRenderers(document, renderer);
929+
893930
const cellsResponse: IOutputRenderResponseCellInfo<T>[] = request.items.map(cellInfo => {
894931
const outputResponse: IOutputRenderResponseOutputInfo[] = cellInfo.outputs.map(output => {
895932
return {
@@ -1035,19 +1072,36 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
10351072
return;
10361073
}
10371074

1075+
let webComm = this._webviewComm.get(editorId);
1076+
if (!webComm) {
1077+
webComm = new ExtHostWebviewCommWrapper(editorId, revivedUri, this._proxy, this._webviewInitData, document);
1078+
this._webviewComm.set(editorId, webComm);
1079+
}
1080+
10381081
if (!provider.provider.resolveNotebook) {
10391082
return;
10401083
}
10411084

1042-
let webComm = this._webviewComm.get(editorId)?.comm;
1085+
await provider.provider.resolveNotebook(document, webComm.contentProviderComm);
1086+
}
10431087

1044-
if (webComm) {
1045-
await provider.provider.resolveNotebook(document, webComm);
1046-
} else {
1047-
const onDidReceiveMessage = new Emitter<any>();
1048-
webComm = new ExtHostWebviewComm(editorId, revivedUri, this._proxy, onDidReceiveMessage, this._webviewInitData, document);
1049-
this._webviewComm.set(editorId, { comm: webComm, onDidReceiveMessage });
1050-
await provider.provider.resolveNotebook(document, webComm);
1088+
private provideCommToNotebookRenderers(document: ExtHostNotebookDocument, renderer: ExtHostNotebookOutputRenderer) {
1089+
let alreadyRegistered = this._renderersUsedInNotebooks.get(document);
1090+
if (!alreadyRegistered) {
1091+
alreadyRegistered = new Set();
1092+
this._renderersUsedInNotebooks.set(document, alreadyRegistered);
1093+
}
1094+
1095+
if (alreadyRegistered.has(renderer)) {
1096+
return;
1097+
}
1098+
1099+
alreadyRegistered.add(renderer);
1100+
for (const editor of document.allEditors()) {
1101+
const comm = this._webviewComm.get(editor.id);
1102+
if (comm) {
1103+
renderer.resolveNotebook(document, comm);
1104+
}
10511105
}
10521106
}
10531107

@@ -1182,12 +1236,8 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
11821236
return editor;
11831237
}
11841238

1185-
$onDidReceiveMessage(editorId: string, message: any): void {
1186-
let messageEmitter = this._webviewComm.get(editorId)?.onDidReceiveMessage;
1187-
1188-
if (messageEmitter) {
1189-
messageEmitter.fire(message);
1190-
}
1239+
$onDidReceiveMessage(editorId: string, forRendererType: string | undefined, message: any): void {
1240+
this._webviewComm.get(editorId)?.onDidReceiveMessage(forRendererType, message);
11911241
}
11921242

11931243
$acceptModelChanged(uriComponents: UriComponents, event: NotebookCellsChangedEvent): void {
@@ -1226,20 +1276,19 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
12261276

12271277
private _createExtHostEditor(document: ExtHostNotebookDocument, editorId: string, selections: number[]) {
12281278
const revivedUri = document.uri;
1229-
let webComm = this._webviewComm.get(editorId)?.comm;
1279+
let webComm = this._webviewComm.get(editorId);
12301280

12311281
if (!webComm) {
1232-
const onDidReceiveMessage = new Emitter<any>();
1233-
webComm = new ExtHostWebviewComm(editorId, revivedUri, this._proxy, onDidReceiveMessage, this._webviewInitData, document);
1234-
this._webviewComm.set(editorId, { comm: webComm!, onDidReceiveMessage });
1282+
webComm = new ExtHostWebviewCommWrapper(editorId, revivedUri, this._proxy, this._webviewInitData, document);
1283+
this._webviewComm.set(editorId, webComm);
12351284
}
12361285

12371286
let editor = new ExtHostNotebookEditor(
12381287
document.viewType,
12391288
editorId,
12401289
revivedUri,
12411290
this._proxy,
1242-
webComm,
1291+
webComm.contentProviderComm,
12431292
document,
12441293
this._documentsAndEditors
12451294
);
@@ -1255,6 +1304,10 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
12551304

12561305
this._editors.get(editorId)?.editor.dispose();
12571306

1307+
for (const renderer of this._renderersUsedInNotebooks.get(document) ?? []) {
1308+
renderer.resolveNotebook(document, webComm);
1309+
}
1310+
12581311
this._editors.set(editorId, { editor });
12591312
}
12601313

src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -270,7 +270,7 @@ export interface INotebookEditor extends IEditor {
270270
/**
271271
* Send message to the webview for outputs.
272272
*/
273-
postMessage(message: any): void;
273+
postMessage(forRendererId: string | undefined, message: any): void;
274274

275275
/**
276276
* Trigger the editor to scroll from scroll event programmatically

0 commit comments

Comments
 (0)