diff --git a/testing/marionette/driver.js b/testing/marionette/driver.js
index 658ff4feb5e8..2a1237c533c6 100644
--- a/testing/marionette/driver.js
+++ b/testing/marionette/driver.js
@@ -2535,8 +2535,6 @@ GeckoDriver.prototype.clearImportedScripts = function(cmd, resp) {
  *     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();
@@ -2579,7 +2577,10 @@ GeckoDriver.prototype.takeScreenshot = function(cmd, resp) {
       break;
 
     case Context.CONTENT:
-      return this.listener.takeScreenshot(id, highlights, full);
+      resp.body.value = yield this.listener.takeScreenshot({
+        id: cmd.parameters.id,
+        highlights: cmd.parameters.highlights,
+        full: cmd.parameters.full});
       break;
   }
 };
diff --git a/testing/marionette/listener.js b/testing/marionette/listener.js
index 65086e7aa01c..1d0fb244aa51 100644
--- a/testing/marionette/listener.js
+++ b/testing/marionette/listener.js
@@ -216,7 +216,6 @@ var getElementValueOfCssPropertyFn = dispatch(getElementValueOfCssProperty);
 var switchToShadowRootFn = dispatch(switchToShadowRoot);
 var getCookiesFn = dispatch(getCookies);
 var singleTapFn = dispatch(singleTap);
-var takeScreenshotFn = dispatch(takeScreenshot);
 
 /**
  * Start all message listeners
@@ -261,7 +260,7 @@ function startListeners() {
   addMessageListenerId("Marionette:importScript", importScript);
   addMessageListenerId("Marionette:getAppCacheStatus", getAppCacheStatus);
   addMessageListenerId("Marionette:setTestName", setTestName);
-  addMessageListenerId("Marionette:takeScreenshot", takeScreenshotFn);
+  addMessageListenerId("Marionette:takeScreenshot", takeScreenshot);
   addMessageListenerId("Marionette:addCookie", addCookie);
   addMessageListenerId("Marionette:getCookies", getCookiesFn);
   addMessageListenerId("Marionette:deleteAllCookies", deleteAllCookies);
@@ -365,7 +364,7 @@ function deleteSession(msg) {
   removeMessageListenerId("Marionette:importScript", importScript);
   removeMessageListenerId("Marionette:getAppCacheStatus", getAppCacheStatus);
   removeMessageListenerId("Marionette:setTestName", setTestName);
-  removeMessageListenerId("Marionette:takeScreenshot", takeScreenshotFn);
+  removeMessageListenerId("Marionette:takeScreenshot", takeScreenshot);
   removeMessageListenerId("Marionette:addCookie", addCookie);
   removeMessageListenerId("Marionette:getCookies", getCookiesFn);
   removeMessageListenerId("Marionette:deleteAllCookies", deleteAllCookies);
@@ -1997,35 +1996,44 @@ function importScript(msg) {
  * <code>msg.json.highlights</code>, a red box will be painted around
  * them to highlight their position.
  */
-function takeScreenshot(id, highlights, full) {
-  let node;
-  if (id) {
-    node = elementManager.getKnownElement(id, curContainer);
-  } else {
+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 {
     node = curContainer.frame;
   }
+  let highlights = msg.json.highlights;
 
-  let document = curContainer.frame.document;
-  let rect, win, width, height, left, top;
-
+  var document = curContainer.frame.document;
+  var 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 (full) {
+    if (msg.json.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();
@@ -2035,16 +2043,15 @@ function takeScreenshot(id, highlights, full) {
     left = rect.left;
   }
 
-  let canvas = document.createElementNS(
-      "http://www.w3.org/1999/xhtml", "canvas");
+  var canvas = document.createElementNS("http://www.w3.org/1999/xhtml",
+                                        "canvas");
   canvas.width = width;
   canvas.height = height;
-  let ctx = canvas.getContext("2d");
-
-  // draws the DOM contents of the window to the canvas
+  var 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";
@@ -2052,26 +2059,25 @@ function takeScreenshot(id, highlights, full) {
     ctx.save();
 
     for (var i = 0; i < highlights.length; ++i) {
-      let elem = elementManager.getKnownElement(highlights[i], curContainer);
+      var elem = elementManager.getKnownElement(highlights[i], curContainer);
       rect = elem.getBoundingClientRect();
 
-      let offsetY = -top;
-      let offsetX = -left;
+      var offsetY = -top;
+      var 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
-  let dataUrl = canvas.toDataURL("image/png", "");
-  let encoded = dataUrl.substring(dataUrl.indexOf(",") + 1);
-  return encoded;
+  // 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);
 }
 
 // Call register self when we get loaded