diff --git a/js/background.js b/js/background.js
index c48c3bcc3..d4cbb8a2b 100644
--- a/js/background.js
+++ b/js/background.js
@@ -14,7 +14,7 @@
 ;(function() {
     'use strict';
 
-    const { IdleDetector } = Signal;
+    const { IdleDetector, MessageMigration } = Signal;
     const { Errors, Message } = window.Signal.Types;
     const { upgradeMessageSchema } = window.Signal.Migrations;
 
diff --git a/js/modules/types/message.js b/js/modules/types/message.js
index b4680818f..a43e9c664 100644
--- a/js/modules/types/message.js
+++ b/js/modules/types/message.js
@@ -166,8 +166,13 @@ const toVersion3 = exports._withSchemaVersion(
 );
 
 // UpgradeStep
-exports.upgradeSchema = async (message, { writeAttachmentData } = {}) =>
-  toVersion3(
+exports.upgradeSchema = async (message, { writeAttachmentData } = {}) => {
+  if (!isFunction(writeAttachmentData)) {
+    throw new TypeError('`context.writeAttachmentData` is required');
+  }
+
+  return toVersion3(
     await toVersion2(await toVersion1(await toVersion0(message))),
     { writeAttachmentData }
   );
+};