Make in-page callback survive navigations

This patch makes in-page callbacks survive navigations via
running in-page harness code on page load.

lib/Page.js

@@ -50,12 +50,13 @@ class Page extends EventEmitter {
         this._client = client;
         this._screenDPI = screenDPI;
         this._extraHeaders = {};
-        /** @type {!Map<string, function()>} */
+        /** @type {!Map<string, !InPageCallback>} */
+        this._sourceURLToPageCallback = new Map();
+        /** @type {!Map<string, !InPageCallback>} */
         this._scriptIdToPageCallback = new Map();
-        /** @type {!Map<string, string>} */
-        this._scriptIdToCallbackName = new Map();
 
         client.on('Debugger.paused', event => this._onDebuggerPaused(event));
+        client.on('Debugger.scriptParsed', event => this._onScriptParsed(event));
         client.on('Network.responseReceived', event => this.emit(Page.Events.ResponseReceived, event.response));
         client.on('Network.loadingFailed', event => this.emit(Page.Events.ResourceLoadingFailed, event));
         client.on('Runtime.consoleAPICalled', event => this._onConsoleAPI(event));
@@ -103,15 +104,13 @@ class Page extends EventEmitter {
             throw new Error(`Failed to set in-page callback with name ${name}: window['${name}'] already exists!`);
 
         var sourceURL = '__in_page_callback__' + name;
-        // Ensure debugger is enabled.
-        await this._client.send('Debugger.enable', {});
-        var scriptPromise = helpers.waitForScriptWithURL(this._client, sourceURL);
-        helpers.evaluate(this._client, inPageCallback, [name], false /* awaitPromise */, sourceURL);
-        var script = await scriptPromise;
-        if (!script)
-            throw new Error(`Failed to set in-page callback with name "${name}"`);
-        this._scriptIdToPageCallback.set(script.scriptId, callback);
-        this._scriptIdToCallbackName.set(script.scriptId, name);
+        this._sourceURLToPageCallback.set(sourceURL, new InPageCallback(name, callback));
+        var text = helpers.evaluationString(inPageCallback, [name], false /* awaitPromise */, sourceURL);
+        await Promise.all([
+            this._client.send('Debugger.enable', {}),
+            this._client.send('Page.addScriptToEvaluateOnLoad', { scriptSource: text }),
+            helpers.evaluateText(this._client, text, false /* awaitPromise */)
+        ]);
 
         function inPageCallback(callbackName) {
             window[callbackName] = (...args) => {
@@ -124,11 +123,11 @@ class Page extends EventEmitter {
     }
 
     /**
-     * @param {string} scriptId
+     * @param {!InPageCallback} inPageCallback
      */
-    async _handleInPageCallback(scriptId) {
-        var name = /** @type {string} */ (this._scriptIdToCallbackName.get(scriptId));
-        var callback = /** @type {function()} */ (this._scriptIdToPageCallback.get(scriptId));
+    async _handleInPageCallback(inPageCallback) {
+        var name = inPageCallback.name;
+        var callback = inPageCallback.callback;
         var args = await this.evaluate(callbackName => window[callbackName].__args, name);
         var result = callback.apply(null, args);
         await this.evaluate(assignResult, name, result);
@@ -145,13 +144,20 @@ class Page extends EventEmitter {
 
     _onDebuggerPaused(event) {
         var location = event.callFrames[0] ? event.callFrames[0].location : null;
-        if (location && this._scriptIdToPageCallback.has(location.scriptId)) {
-            this._handleInPageCallback(location.scriptId);
+        var inPageCallback = location ? this._scriptIdToPageCallback.get(location.scriptId) : null;
+        if (inPageCallback) {
+            this._handleInPageCallback(inPageCallback);
             return;
         }
         this._client.send('Debugger.resume');
     }
 
+    _onScriptParsed(event) {
+        var inPageCallback = this._sourceURLToPageCallback.get(event.url);
+        if (inPageCallback)
+            this._scriptIdToPageCallback.set(event.scriptId, inPageCallback);
+    }
+
     /**
      * @param {!Object} headers
      * @return {!Promise}
@@ -392,6 +398,17 @@ class Page extends EventEmitter {
     }
 }
 
+class InPageCallback {
+    /**
+     * @param {string} name
+     * @param {function(?):?} callback
+     */
+    constructor(name, callback) {
+        this.name = name;
+        this.callback = callback;
+    }
+}
+
 /** @enum {string} */
 Page.ScreenshotTypes = {
     PNG: "png",

test/test.js

@@ -1,6 +1,9 @@
+var path = require('path');
 var Browser = require('../lib/Browser');
 
-describe('Page', function() {
+var EMPTY_PAGE = 'file://' + path.join(__dirname, 'assets', 'empty.html');
+
+describe('Puppeteer', function() {
     var browser;
     var page;
 
@@ -29,6 +32,30 @@ describe('Page', function() {
         var result = await page.evaluateAsync(() => Promise.resolve(8 * 7));
         expect(result).toBe(56);
     }));
+
+    describe('Page.setInPageCallback', function() {
+        it('should work', SX(async function() {
+            await page.setInPageCallback('callController', function(a, b) {
+                return a * b;
+            });
+
+            var result = await page.evaluate(function() {
+                return callController(9, 4);
+            });
+            expect(result).toBe(36);
+        }));
+        it('should survive navigation', SX(async function() {
+            await page.setInPageCallback('callController', function(a, b) {
+                return a * b;
+            });
+
+            await page.navigate(EMPTY_PAGE);
+            var result = await page.evaluate(function() {
+                return callController(9, 4);
+            });
+            expect(result).toBe(36);
+        }));
+    });
 });
 
 // Since Jasmine doesn't like async functions, they should be wrapped