--- a/testing/marionette/driver.js
+++ b/testing/marionette/driver.js
@@ -2530,16 +2530,18 @@ GeckoDriver.prototype.clearImportedScrip
  *     Reference to a web element.
  * @param {string} highlights
  *     List of web elements to highlight.
  *
  * @return {string}
  *     PNG image encoded as base64 encoded string.
  */
 GeckoDriver.prototype.takeScreenshot = function(cmd, resp) {
+  let {id, highlights, full} = cmd.parameters;
+
   switch (this.context) {
     case Context.CHROME:
       let win = this.getCurrentWindow();
       let canvas = win.document.createElementNS("http://www.w3.org/1999/xhtml", "canvas");
       let doc;
       if (this.appName == "B2G") {
         doc = win.document.body;
       } else {
@@ -2572,20 +2574,17 @@ GeckoDriver.prototype.takeScreenshot = f
       context.scale(scale, scale);
       context.drawWindow(win, 0, 0, width, height, "rgb(255,255,255)", flags);
       let dataUrl = canvas.toDataURL("image/png", "");
       let data = dataUrl.substring(dataUrl.indexOf(",") + 1);
       resp.body.value = data;
       break;
 
     case Context.CONTENT:
-      resp.body.value = yield this.listener.takeScreenshot({
-        id: cmd.parameters.id,
-        highlights: cmd.parameters.highlights,
-        full: cmd.parameters.full});
+      return this.listener.takeScreenshot(id, highlights, full);
       break;
   }
 };
 
 /**
  * Get the current browser orientation.
  *
  * Will return one of the valid primary orientation values

--- a/testing/marionette/listener.js
+++ b/testing/marionette/listener.js
@@ -211,16 +211,17 @@ var findElementContentFn = dispatch(find
 var findElementsContentFn = dispatch(findElementsContent);
 var isElementSelectedFn = dispatch(isElementSelected);
 var clearElementFn = dispatch(clearElement);
 var isElementDisplayedFn = dispatch(isElementDisplayed);
 var getElementValueOfCssPropertyFn = dispatch(getElementValueOfCssProperty);
 var switchToShadowRootFn = dispatch(switchToShadowRoot);
 var getCookiesFn = dispatch(getCookies);
 var singleTapFn = dispatch(singleTap);
+var takeScreenshotFn = dispatch(takeScreenshot);
 
 /**
  * Start all message listeners
  */
 function startListeners() {
   addMessageListenerId("Marionette:receiveFiles", receiveFiles);
   addMessageListenerId("Marionette:newSession", newSession);
   addMessageListenerId("Marionette:executeScript", executeScript);
@@ -255,17 +256,17 @@ function startListeners() {
   addMessageListenerId("Marionette:switchToFrame", switchToFrame);
   addMessageListenerId("Marionette:switchToShadowRoot", switchToShadowRootFn);
   addMessageListenerId("Marionette:deleteSession", deleteSession);
   addMessageListenerId("Marionette:sleepSession", sleepSession);
   addMessageListenerId("Marionette:emulatorCmdResult", emulatorCmdResult);
   addMessageListenerId("Marionette:importScript", importScript);
   addMessageListenerId("Marionette:getAppCacheStatus", getAppCacheStatus);
   addMessageListenerId("Marionette:setTestName", setTestName);
-  addMessageListenerId("Marionette:takeScreenshot", takeScreenshot);
+  addMessageListenerId("Marionette:takeScreenshot", takeScreenshotFn);
   addMessageListenerId("Marionette:addCookie", addCookie);
   addMessageListenerId("Marionette:getCookies", getCookiesFn);
   addMessageListenerId("Marionette:deleteAllCookies", deleteAllCookies);
   addMessageListenerId("Marionette:deleteCookie", deleteCookie);
 }
 
 /**
  * Used during newSession and restart, called to set up the modal dialog listener in b2g
@@ -359,17 +360,17 @@ function deleteSession(msg) {
   removeMessageListenerId("Marionette:switchToFrame", switchToFrame);
   removeMessageListenerId("Marionette:switchToShadowRoot", switchToShadowRootFn);
   removeMessageListenerId("Marionette:deleteSession", deleteSession);
   removeMessageListenerId("Marionette:sleepSession", sleepSession);
   removeMessageListenerId("Marionette:emulatorCmdResult", emulatorCmdResult);
   removeMessageListenerId("Marionette:importScript", importScript);
   removeMessageListenerId("Marionette:getAppCacheStatus", getAppCacheStatus);
   removeMessageListenerId("Marionette:setTestName", setTestName);
-  removeMessageListenerId("Marionette:takeScreenshot", takeScreenshot);
+  removeMessageListenerId("Marionette:takeScreenshot", takeScreenshotFn);
   removeMessageListenerId("Marionette:addCookie", addCookie);
   removeMessageListenerId("Marionette:getCookies", getCookiesFn);
   removeMessageListenerId("Marionette:deleteAllCookies", deleteAllCookies);
   removeMessageListenerId("Marionette:deleteCookie", deleteCookie);
   if (isB2G) {
     content.removeEventListener("mozbrowsershowmodalprompt", modalHandler, false);
   }
   elementManager.reset();
@@ -1991,94 +1992,87 @@ function importScript(msg) {
  * Takes a screen capture of the given web element if <code>id</code>
  * property exists in the message's JSON object, or if null captures
  * the bounding box of the current frame.
  *
  * If given an array of web element references in
  * <code>msg.json.highlights</code>, a red box will be painted around
  * them to highlight their position.
  */
-function takeScreenshot(msg) {
-  let node = null;
-  if (msg.json.id) {
-    try {
-      node = elementManager.getKnownElement(msg.json.id, curContainer)
-    }
-    catch (e) {
-      sendResponse(e.message, e.code, e.stack, msg.json.command_id);
-      return;
-    }
-  }
-  else {
+function takeScreenshot(id, highlights, full) {
+  let node;
+  if (id) {
+    node = elementManager.getKnownElement(id, curContainer);
+  } else {
     node = curContainer.frame;
   }
-  let highlights = msg.json.highlights;
 
-  var document = curContainer.frame.document;
-  var rect, win, width, height, left, top;
+  let document = curContainer.frame.document;
+  let rect, win, width, height, left, top;
+
   // node can be either a window or an arbitrary DOM node
   if (node == curContainer.frame) {
     // node is a window
     win = node;
-    if (msg.json.full) {
+    if (full) {
       // the full window
       width = document.body.scrollWidth;
       height = document.body.scrollHeight;
       top = 0;
       left = 0;
-    }
-    else {
+    } else {
       // only the viewport
       width = document.documentElement.clientWidth;
       height = document.documentElement.clientHeight;
       left = curContainer.frame.pageXOffset;
       top = curContainer.frame.pageYOffset;
     }
-  }
-  else {
+  } else {
     // node is an arbitrary DOM node
     win = node.ownerDocument.defaultView;
     rect = node.getBoundingClientRect();
     width = rect.width;
     height = rect.height;
     top = rect.top;
     left = rect.left;
   }
 
-  var canvas = document.createElementNS("http://www.w3.org/1999/xhtml",
-                                        "canvas");
+  let canvas = document.createElementNS(
+      "http://www.w3.org/1999/xhtml", "canvas");
   canvas.width = width;
   canvas.height = height;
-  var ctx = canvas.getContext("2d");
-  // Draws the DOM contents of the window to the canvas
+  let ctx = canvas.getContext("2d");
+
+  // draws the DOM contents of the window to the canvas
   ctx.drawWindow(win, left, top, width, height, "rgb(255,255,255)");
 
-  // This section is for drawing a red rectangle around each element
+  // this section is for drawing a red rectangle around each element
   // passed in via the highlights array
   if (highlights) {
     ctx.lineWidth = "2";
     ctx.strokeStyle = "red";
     ctx.save();
 
     for (var i = 0; i < highlights.length; ++i) {
-      var elem = elementManager.getKnownElement(highlights[i], curContainer);
+      let elem = elementManager.getKnownElement(highlights[i], curContainer);
       rect = elem.getBoundingClientRect();
 
-      var offsetY = -top;
-      var offsetX = -left;
+      let offsetY = -top;
+      let offsetX = -left;
 
-      // Draw the rectangle
-      ctx.strokeRect(rect.left + offsetX,
-                     rect.top + offsetY,
-                     rect.width,
-                     rect.height);
+      // draw the rectangle
+      ctx.strokeRect(
+          rect.left + offsetX,
+          rect.top + offsetY,
+          rect.width,
+          rect.height);
     }
   }
 
-  // Return the Base64 encoded string back to the client so that it
-  // can save the file to disk if it is required
-  var dataUrl = canvas.toDataURL("image/png", "");
-  var data = dataUrl.substring(dataUrl.indexOf(",") + 1);
-  sendResponse({value: data}, msg.json.command_id);
+  // return the Base64 encoded string back to the client
+  // so that it can save the file to disk if it is required
+  let dataUrl = canvas.toDataURL("image/png", "");
+  let encoded = dataUrl.substring(dataUrl.indexOf(",") + 1);
+  return encoded;
 }
 
 // Call register self when we get loaded
 registerSelf();