chore(webkit): remove WKPageProxySession, separate connection from browser session (#447)

src/webkit/wkBrowser.ts

@@ -24,7 +24,7 @@ import { ConnectionTransport, SlowMoTransport } from '../transport';
 import * as types from '../types';
 import { Events } from '../events';
 import { Protocol } from './protocol';
-import { WKConnection, WKConnectionEvents, WKPageProxySession } from './wkConnection';
+import { WKConnection, WKSession, kPageProxyMessageReceived, PageProxyMessageReceivedPayload } from './wkConnection';
 import { WKPageProxy } from './wkPageProxy';
 import * as platform from '../platform';
 
@@ -34,7 +34,8 @@ export type WKConnectOptions = {
 };
 
 export class WKBrowser extends platform.EventEmitter implements Browser {
-  readonly _connection: WKConnection;
+  private readonly _connection: WKConnection;
+  private readonly _browserSession: WKSession;
   private readonly _defaultContext: BrowserContext;
   private readonly _contexts = new Map<string, BrowserContext>();
   private readonly _pageProxies = new Map<string, WKPageProxy>();
@@ -53,24 +54,32 @@ export class WKBrowser extends platform.EventEmitter implements Browser {
 
   constructor(transport: ConnectionTransport) {
     super();
-    this._connection = new WKConnection(transport);
+    this._connection = new WKConnection(transport, this._onDisconnect.bind(this));
-    this._connection.on(WKConnectionEvents.Disconnected, () => this.emit(Events.Browser.Disconnected));
+    this._browserSession = this._connection.browserSession;
 
     this._defaultContext = this._createBrowserContext(undefined, {});
 
     this._eventListeners = [
-      helper.addEventListener(this._connection, WKConnectionEvents.PageProxyCreated, this._onPageProxyCreated.bind(this)),
+      helper.addEventListener(this._browserSession, 'Browser.pageProxyCreated', this._onPageProxyCreated.bind(this)),
-      helper.addEventListener(this._connection, WKConnectionEvents.PageProxyDestroyed, this._onPageProxyDestroyed.bind(this))
+      helper.addEventListener(this._browserSession, 'Browser.pageProxyDestroyed', this._onPageProxyDestroyed.bind(this)),
+      helper.addEventListener(this._browserSession, kPageProxyMessageReceived, this._onPageProxyMessageReceived.bind(this)),
     ];
 
     this._firstPageProxyPromise = new Promise<void>(resolve => this._firstPageProxyCallback = resolve);
   }
 
+  _onDisconnect() {
+    for (const pageProxy of this._pageProxies.values())
+      pageProxy.dispose();
+    this._pageProxies.clear();
+    this.emit(Events.Browser.Disconnected);
+  }
+
   async newContext(options: BrowserContextOptions = {}): Promise<BrowserContext> {
-    const { browserContextId } = await this._connection.send('Browser.createContext');
+    const { browserContextId } = await this._browserSession.send('Browser.createContext');
     const context = this._createBrowserContext(browserContextId, options);
     if (options.ignoreHTTPSErrors)
-      await this._connection.send('Browser.setIgnoreCertificateErrors', { browserContextId, ignore: true });
+      await this._browserSession.send('Browser.setIgnoreCertificateErrors', { browserContextId, ignore: true });
     this._contexts.set(browserContextId, context);
     return context;
   }
@@ -88,7 +97,9 @@ export class WKBrowser extends platform.EventEmitter implements Browser {
     await helper.waitWithTimeout(this._firstPageProxyPromise, 'firstPageProxy', timeout);
   }
 
-  _onPageProxyCreated(session: WKPageProxySession, pageProxyInfo: Protocol.Browser.PageProxyInfo) {
+  _onPageProxyCreated(event: Protocol.Browser.pageProxyCreatedPayload) {
+    const { pageProxyInfo } = event;
+    const pageProxyId = pageProxyInfo.pageProxyId;
     let context = null;
     if (pageProxyInfo.browserContextId) {
       // FIXME: we don't know about the default context id, so assume that all targets from
@@ -99,8 +110,11 @@ export class WKBrowser extends platform.EventEmitter implements Browser {
     }
     if (!context)
       context =  this._defaultContext;
-    const pageProxy = new WKPageProxy(session, context);
+    const pageProxySession = new WKSession(this._connection, pageProxyId, `The page has been closed.`, (message: any) => {
-    this._pageProxies.set(pageProxyInfo.pageProxyId, pageProxy);
+      this._connection.rawSend({ ...message, pageProxyId });
+    });
+    const pageProxy = new WKPageProxy(pageProxySession, context);
+    this._pageProxies.set(pageProxyId, pageProxy);
 
     if (pageProxyInfo.openerId) {
       const opener = this._pageProxies.get(pageProxyInfo.openerId);
@@ -114,12 +128,19 @@ export class WKBrowser extends platform.EventEmitter implements Browser {
     }
   }
 
-  _onPageProxyDestroyed(pageProxyId: Protocol.Browser.PageProxyID) {
+  _onPageProxyDestroyed(event: Protocol.Browser.pageProxyDestroyedPayload) {
+    const pageProxyId = event.pageProxyId;
     const pageProxy = this._pageProxies.get(pageProxyId);
+    pageProxy.didClose();
     pageProxy.dispose();
     this._pageProxies.delete(pageProxyId);
   }
 
+  _onPageProxyMessageReceived(event: PageProxyMessageReceivedPayload) {
+    const pageProxy = this._pageProxies.get(event.pageProxyId);
+    pageProxy.dispatchMessageToSession(event.message);
+  }
+
   disconnect() {
     throw new Error('Unsupported operation');
   }
@@ -130,8 +151,8 @@ export class WKBrowser extends platform.EventEmitter implements Browser {
 
   async close() {
     helper.removeEventListeners(this._eventListeners);
-    const disconnected = new Promise(f => this._connection.once(WKConnectionEvents.Disconnected, f));
+    const disconnected = new Promise(f => this.once(Events.Browser.Disconnected, f));
-    await this._connection.send('Browser.close');
+    await this._browserSession.send('Browser.close');
     await disconnected;
   }
 
@@ -143,19 +164,19 @@ export class WKBrowser extends platform.EventEmitter implements Browser {
       },
 
       newPage: async (): Promise<Page> => {
-        const { pageProxyId } = await this._connection.send('Browser.createPage', { browserContextId });
+        const { pageProxyId } = await this._browserSession.send('Browser.createPage', { browserContextId });
         const pageProxy = this._pageProxies.get(pageProxyId);
         return await pageProxy.page();
       },
 
       close: async (): Promise<void> => {
         assert(browserContextId, 'Non-incognito profiles cannot be closed!');
-        await this._connection.send('Browser.deleteContext', { browserContextId });
+        await this._browserSession.send('Browser.deleteContext', { browserContextId });
         this._contexts.delete(browserContextId);
       },
 
       cookies: async (): Promise<network.NetworkCookie[]> => {
-        const { cookies } = await this._connection.send('Browser.getAllCookies', { browserContextId });
+        const { cookies } = await this._browserSession.send('Browser.getAllCookies', { browserContextId });
         return cookies.map((c: network.NetworkCookie) => ({
           ...c,
           expires: c.expires === 0 ? -1 : c.expires
@@ -163,15 +184,14 @@ export class WKBrowser extends platform.EventEmitter implements Browser {
       },
 
       clearCookies: async (): Promise<void> => {
-        await this._connection.send('Browser.deleteAllCookies', { browserContextId });
+        await this._browserSession.send('Browser.deleteAllCookies', { browserContextId });
       },
 
       setCookies: async (cookies: network.SetNetworkCookieParam[]): Promise<void> => {
         const cc = cookies.map(c => ({ ...c, session: c.expires === -1 || c.expires === undefined })) as Protocol.Browser.SetCookieParam[];
-        await this._connection.send('Browser.setCookies', { cookies: cc, browserContextId });
+        await this._browserSession.send('Browser.setCookies', { cookies: cc, browserContextId });
       },
 
-
       setPermissions: async (origin: string, permissions: string[]): Promise<void> => {
         const webPermissionToProtocol = new Map<string, string>([
           ['geolocation', 'geolocation'],
@@ -182,16 +202,16 @@ export class WKBrowser extends platform.EventEmitter implements Browser {
             throw new Error('Unknown permission: ' + permission);
           return protocolPermission;
         });
-        await this._connection.send('Browser.grantPermissions', { origin, browserContextId, permissions: filtered });
+        await this._browserSession.send('Browser.grantPermissions', { origin, browserContextId, permissions: filtered });
       },
 
       clearPermissions: async () => {
-        await this._connection.send('Browser.resetPermissions', { browserContextId });
+        await this._browserSession.send('Browser.resetPermissions', { browserContextId });
       },
 
       setGeolocation: async (geolocation: types.Geolocation | null): Promise<void> => {
         const payload: any = geolocation ? { ...geolocation, timestamp: Date.now() } : undefined;
-        await this._connection.send('Browser.setGeolocationOverride', { browserContextId, geolocation: payload });
+        await this._browserSession.send('Browser.setGeolocationOverride', { browserContextId, geolocation: payload });
       }
     }, options);
     return context;

src/webkit/wkConnection.ts

@@ -23,89 +23,54 @@ import { Protocol } from './protocol';
 const debugProtocol = platform.debug('playwright:protocol');
 const debugWrappedMessage = platform.debug('wrapped');
 
-export const WKConnectionEvents = {
+// WKBrowserServer uses this special id to issue Browser.close command which we
-  Disconnected: Symbol('Disconnected'),
+// should ignore.
-  PageProxyCreated: Symbol('ConnectionEvents.PageProxyCreated'),
-  PageProxyDestroyed: Symbol('Connection.PageProxyDestroyed')
-};
-
 export const kBrowserCloseMessageId = -9999;
 
-export class WKConnection extends platform.EventEmitter {
+// We emulate kPageProxyMessageReceived message to unify it with Browser.pageProxyCreated
+// and Browser.pageProxyDestroyed for easier management.
+export const kPageProxyMessageReceived = 'kPageProxyMessageReceived';
+export type PageProxyMessageReceivedPayload = { pageProxyId: string, message: any };
+
+export class WKConnection {
   private _lastId = 0;
-  private readonly _callbacks = new Map<number, {resolve:(o: any) => void, reject:  (e: Error) => void, error: Error, method: string}>();
   private readonly _transport: ConnectionTransport;
-  private readonly _pageProxySessions = new Map<string, WKPageProxySession>();
-
   private _closed = false;
+  private _onDisconnect: () => void;
 
-  constructor(transport: ConnectionTransport) {
+  readonly browserSession: WKSession;
-    super();
+
+  constructor(transport: ConnectionTransport, onDisconnect: () => void) {
     this._transport = transport;
     this._transport.onmessage = this._dispatchMessage.bind(this);
     this._transport.onclose = this._onClose.bind(this);
+    this._onDisconnect = onDisconnect;
+    this.browserSession = new WKSession(this, '', 'Browser has been closed.', (message: any) => {
+      this.rawSend(message);
+    });
   }
 
   nextMessageId(): number {
     return ++this._lastId;
   }
 
-  send<T extends keyof Protocol.CommandParameters>(
+  rawSend(message: any) {
-    method: T,
+    message = JSON.stringify(message);
-    params?: Protocol.CommandParameters[T],
-    pageProxyId?: string
-  ): Promise<Protocol.CommandReturnValues[T]> {
-    const id = this._rawSend({pageProxyId, method, params});
-    return new Promise((resolve, reject) => {
-      this._callbacks.set(id, {resolve, reject, error: new Error(), method});
-    });
-  }
-
-  _rawSend(message: any): number {
-    const id = this.nextMessageId();
-    message = JSON.stringify(Object.assign({}, message, {id}));
     debugProtocol('SEND ► ' + message);
     this._transport.send(message);
-    return id;
   }
 
   private _dispatchMessage(message: string) {
     debugProtocol('◀ RECV ' + message);
     const object = JSON.parse(message);
-    this._dispatchPageProxyMessage(object, message);
+    if (object.id === kBrowserCloseMessageId)
-    if (object.id) {
+      return;
-      const callback = this._callbacks.get(object.id);
+    if (object.pageProxyId) {
-      // Callbacks could be all rejected if someone has called `.dispose()`.
+      const payload: PageProxyMessageReceivedPayload = { message: object, pageProxyId: object.pageProxyId };
-      if (callback) {
+      this.browserSession.dispatchMessage({ method: kPageProxyMessageReceived, params: payload });
-        this._callbacks.delete(object.id);
+      return;
-        if (object.error)
-          callback.reject(createProtocolError(callback.error, callback.method, object));
-        else
-          callback.resolve(object.result);
-      } else if (object.id !== kBrowserCloseMessageId) {
-        assert(this._closed, 'Received response for unknown callback: ' + object.id);
-      }
-    } else {
-      Promise.resolve().then(() => this.emit(object.method, object.params));
-    }
-  }
-
-  _dispatchPageProxyMessage(object: {method: string, params: any, id?: string, pageProxyId?: string}, message: string) {
-    if (object.method === 'Browser.pageProxyCreated') {
-      const pageProxyId = object.params.pageProxyInfo.pageProxyId;
-      const pageProxySession = new WKPageProxySession(this, pageProxyId);
-      this._pageProxySessions.set(pageProxyId, pageProxySession);
-      Promise.resolve().then(() => this.emit(WKConnectionEvents.PageProxyCreated, pageProxySession, object.params.pageProxyInfo));
-    } else if (object.method === 'Browser.pageProxyDestroyed') {
-      const pageProxyId = object.params.pageProxyId as string;
-      const pageProxySession = this._pageProxySessions.get(pageProxyId);
-      this._pageProxySessions.delete(pageProxyId);
-      pageProxySession.dispose();
-      Promise.resolve().then(() => this.emit(WKConnectionEvents.PageProxyDestroyed, pageProxyId));
-    } else if (!object.id && object.pageProxyId) {
-      const pageProxySession = this._pageProxySessions.get(object.pageProxyId);
-      Promise.resolve().then(() => pageProxySession.emit(object.method, object.params));
     }
+    this.browserSession.dispatchMessage(object);
   }
 
   _onClose() {
@@ -114,14 +79,8 @@ export class WKConnection extends platform.EventEmitter {
     this._closed = true;
     this._transport.onmessage = null;
     this._transport.onclose = null;
-    for (const callback of this._callbacks.values())
+    this.browserSession.dispose();
-      callback.reject(rewriteError(callback.error, `Protocol error (${callback.method}): Target closed.`));
+    this._onDisconnect();
-    this._callbacks.clear();
-
-    for (const pageProxySession of this._pageProxySessions.values())
-      pageProxySession.dispose();
-    this._pageProxySessions.clear();
-    this.emit(WKConnectionEvents.Disconnected);
   }
 
   dispose() {
@@ -130,50 +89,6 @@ export class WKConnection extends platform.EventEmitter {
   }
 }
 
-export const WKSessionEvents = {
-  Disconnected: Symbol('WKSessionEvents.Disconnected')
-};
-
-export class WKPageProxySession extends platform.EventEmitter {
-  _connection: WKConnection;
-  readonly _pageProxyId: string;
-  private readonly _closePromise: Promise<void>;
-  private _closePromiseCallback: () => void;
-  on: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
-  addListener: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
-  off: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
-  removeListener: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
-  once: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
-
-  constructor(connection: WKConnection, pageProxyId: string) {
-    super();
-    this._connection = connection;
-    this._pageProxyId = pageProxyId;
-    this._closePromise = new Promise(r => this._closePromiseCallback = r);
-  }
-
-  send<T extends keyof Protocol.CommandParameters>(
-    method: T,
-    params?: Protocol.CommandParameters[T]
-  ): Promise<Protocol.CommandReturnValues[T]> {
-    if (!this._connection)
-      return Promise.reject(new Error(`Protocol error (${method}): Session closed. Most likely the pageProxy has been closed.`));
-    return Promise.race([
-      this._closePromise.then(() => { throw new Error('Page proxy closed'); }),
-      this._connection.send(method, params, this._pageProxyId)
-    ]);
-  }
-
-  isClosed() {
-    return !this._connection;
-  }
-
-  dispose() {
-    this._closePromiseCallback();
-    this._connection = null;
-  }
-}
-
 export class WKSession extends platform.EventEmitter {
   connection?: WKConnection;
   errorText: string;
@@ -221,7 +136,6 @@ export class WKSession extends platform.EventEmitter {
       callback.reject(rewriteError(callback.error, `Protocol error (${callback.method}): ${this.errorText}`));
     this._callbacks.clear();
     this.connection = undefined;
-    this.emit(WKSessionEvents.Disconnected);
   }
 
   dispatchMessage(object: any) {

src/webkit/wkInput.ts

@@ -18,7 +18,7 @@
 import * as input from '../input';
 import { helper } from '../helper';
 import { macEditingCommands } from '../usKeyboardLayout';
-import { WKPageProxySession, WKSession } from './wkConnection';
+import { WKSession } from './wkConnection';
 
 function toModifiersMask(modifiers: Set<input.Modifier>): number {
   // From Source/WebKit/Shared/WebEvent.h
@@ -35,10 +35,10 @@ function toModifiersMask(modifiers: Set<input.Modifier>): number {
 }
 
 export class RawKeyboardImpl implements input.RawKeyboard {
-  private readonly _pageProxySession: WKPageProxySession;
+  private readonly _pageProxySession: WKSession;
   private _session: WKSession;
 
-  constructor(session: WKPageProxySession) {
+  constructor(session: WKSession) {
     this._pageProxySession = session;
   }
 
@@ -88,9 +88,9 @@ export class RawKeyboardImpl implements input.RawKeyboard {
 }
 
 export class RawMouseImpl implements input.RawMouse {
-  private readonly _pageProxySession: WKPageProxySession;
+  private readonly _pageProxySession: WKSession;
 
-  constructor(session: WKPageProxySession) {
+  constructor(session: WKSession) {
     this._pageProxySession = session;
   }
 

src/webkit/wkNetworkManager.ts

@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-import { WKSession, WKPageProxySession } from './wkConnection';
+import { WKSession } from './wkConnection';
 import { Page } from '../page';
 import { helper, RegisteredListener, assert } from '../helper';
 import { Protocol } from './protocol';
@@ -26,13 +26,13 @@ import * as platform from '../platform';
 
 export class WKNetworkManager {
   private readonly _page: Page;
-  private readonly _pageProxySession: WKPageProxySession;
+  private readonly _pageProxySession: WKSession;
   private _session: WKSession;
   private readonly _requestIdToRequest = new Map<string, InterceptableRequest>();
   private _userCacheDisabled = false;
   private _sessionListeners: RegisteredListener[] = [];
 
-  constructor(page: Page, pageProxySession: WKPageProxySession) {
+  constructor(page: Page, pageProxySession: WKSession) {
     this._page = page;
     this._pageProxySession = pageProxySession;
   }

src/webkit/wkPage.ts

@@ -19,7 +19,7 @@ import * as frames from '../frames';
 import { debugError, helper, RegisteredListener } from '../helper';
 import * as dom from '../dom';
 import * as network from '../network';
-import { WKSession, WKSessionEvents, WKPageProxySession } from './wkConnection';
+import { WKSession } from './wkConnection';
 import { Events } from '../events';
 import { WKExecutionContext, EVALUATION_SCRIPT_URL } from './wkExecutionContext';
 import { WKNetworkManager } from './wkNetworkManager';
@@ -42,7 +42,7 @@ export class WKPage implements PageDelegate {
   readonly rawKeyboard: RawKeyboardImpl;
   _session: WKSession;
   readonly _page: Page;
-  private readonly _pageProxySession: WKPageProxySession;
+  private readonly _pageProxySession: WKSession;
   private readonly _networkManager: WKNetworkManager;
   private readonly _workers: WKWorkers;
   private readonly _contextIdToContext: Map<number, dom.FrameExecutionContext>;
@@ -50,7 +50,7 @@ export class WKPage implements PageDelegate {
   private _sessionListeners: RegisteredListener[] = [];
   private readonly _bootstrapScripts: string[] = [];
 
-  constructor(browserContext: BrowserContext, pageProxySession: WKPageProxySession) {
+  constructor(browserContext: BrowserContext, pageProxySession: WKSession) {
     this._pageProxySession = pageProxySession;
     this.rawKeyboard = new RawKeyboardImpl(pageProxySession);
     this.rawMouse = new RawMouseImpl(pageProxySession);
@@ -134,6 +134,10 @@ export class WKPage implements PageDelegate {
       this._page._didClose();
   }
 
+  didDisconnect() {
+    this._page._didDisconnect();
+  }
+
   _addSessionListeners() {
     this._sessionListeners = [
       helper.addEventListener(this._session, 'Page.frameNavigated', event => this._onFrameNavigated(event.frame, false)),
@@ -147,7 +151,6 @@ export class WKPage implements PageDelegate {
       helper.addEventListener(this._session, 'Console.messageAdded', event => this._onConsoleMessage(event)),
       helper.addEventListener(this._pageProxySession, 'Dialog.javascriptDialogOpening', event => this._onDialog(event)),
       helper.addEventListener(this._session, 'Page.fileChooserOpened', event => this._onFileChooserOpened(event)),
-      helper.addEventListener(this._session, WKSessionEvents.Disconnected, event => this._page._didDisconnect()),
     ];
   }
 
@@ -393,7 +396,7 @@ export class WKPage implements PageDelegate {
 
   async setBackgroundColor(color?: { r: number; g: number; b: number; a: number; }): Promise<void> {
     // TODO: line below crashes, sort it out.
-    this._session.send('Page.setDefaultBackgroundColorOverride', { color });
+    await this._session.send('Page.setDefaultBackgroundColorOverride', { color });
   }
 
   async takeScreenshot(format: string, options: types.ScreenshotOptions, viewport: types.Viewport): Promise<platform.BufferType> {

src/webkit/wkPageProxy.ts

@@ -5,7 +5,7 @@
 import { BrowserContext } from '../browserContext';
 import { Page } from '../page';
 import { Protocol } from './protocol';
-import { WKPageProxySession, WKSession } from './wkConnection';
+import { WKSession } from './wkConnection';
 import { WKPage } from './wkPage';
 import { RegisteredListener, helper, assert, debugError } from '../helper';
 import { Events } from '../events';
@@ -16,7 +16,7 @@ import { Events } from '../events';
 const provisionalMessagesSymbol = Symbol('provisionalMessages');
 
 export class WKPageProxy {
-  private readonly _pageProxySession: WKPageProxySession;
+  private readonly _pageProxySession: WKSession;
   readonly _browserContext: BrowserContext;
   private _pagePromise: Promise<Page> | null = null;
   private _wkPage: WKPage | null = null;
@@ -25,8 +25,8 @@ export class WKPageProxy {
   private readonly _sessions = new Map<string, WKSession>();
   private readonly _eventListeners: RegisteredListener[];
 
-  constructor(session: WKPageProxySession, browserContext: BrowserContext) {
+  constructor(pageProxySession: WKSession, browserContext: BrowserContext) {
-    this._pageProxySession = session;
+    this._pageProxySession = pageProxySession;
     this._browserContext = browserContext;
     this._firstTargetPromise = new Promise(r => this._firstTargetCallback = r);
     this._eventListeners = [
@@ -38,18 +38,30 @@ export class WKPageProxy {
 
     // Intercept provisional targets during cross-process navigation.
     this._pageProxySession.send('Target.setPauseOnStart', { pauseOnStart: true }).catch(e => {
-      if (this._pageProxySession.isClosed())
+      if (this._pageProxySession.isDisposed())
         return;
       debugError(e);
       throw e;
     });
   }
 
+  didClose() {
+    if (this._wkPage)
+      this._wkPage.didClose(false);
+  }
+
   dispose() {
+    this._pageProxySession.dispose();
     helper.removeEventListeners(this._eventListeners);
     for (const session of this._sessions.values())
       session.dispose();
     this._sessions.clear();
+    if (this._wkPage)
+      this._wkPage.didDisconnect();
+  }
+
+  dispatchMessageToSession(message: any) {
+    this._pageProxySession.dispatchMessage(message);
   }
 
   async page(): Promise<Page> {
@@ -87,7 +99,7 @@ export class WKPageProxy {
 
   private _onTargetCreated(event: Protocol.Target.targetCreatedPayload) {
     const { targetInfo } = event;
-    const session = new WKSession(this._pageProxySession._connection, targetInfo.targetId, `The ${targetInfo.type} has been closed.`, (message: any) => {
+    const session = new WKSession(this._pageProxySession.connection, targetInfo.targetId, `The ${targetInfo.type} has been closed.`, (message: any) => {
       this._pageProxySession.send('Target.sendMessageToTarget', {
         message: JSON.stringify(message), targetId: targetInfo.targetId
       }).catch(e => {
@@ -114,9 +126,7 @@ export class WKPageProxy {
     if (session)
       session.dispose();
     this._sessions.delete(targetId);
-    if (!this._wkPage)
+    if (this._wkPage && this._wkPage._session === session && crashed)
-      return;
-    if (this._wkPage._session === session)
       this._wkPage.didClose(crashed);
   }