@@ -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
499503export 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
722747export 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-
750782export 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
0 commit comments