diff --git a/lib/RuntimeBabel.js b/lib/RuntimeBabel.js
index 1962a6c..f0f0d62 100644
--- a/lib/RuntimeBabel.js
+++ b/lib/RuntimeBabel.js
@@ -234,9 +234,9 @@ module.exports = function(S) {
               process.env[key] = envVars[key];
             }
 
-            exports.handler = (event, context) => {
+            exports.handler = (event, context, cb) => {
               try {
-                const result = require('./${handlerFile}')['${handlerMethod}'](event, context);
+                const result = require('./${handlerFile}')['${handlerMethod}'](event, context, cb);
 
                 if (result && typeof result.then == 'function') {
                   result.then(context.succeed).catch(context.fail);