diff --git a/src/firmware/sagas.test.ts b/src/firmware/sagas.test.ts
index 98ab7dd0..0abfbc67 100644
--- a/src/firmware/sagas.test.ts
+++ b/src/firmware/sagas.test.ts
@@ -6,7 +6,6 @@ import {
     FirmwareReaderError,
     FirmwareReaderErrorCode,
 } from '@pybricks/firmware';
-import { mock } from 'jest-mock-extended';
 import JSZip from 'jszip';
 import { AsyncSaga } from '../../test';
 import {
diff --git a/src/firmware/sagas.ts b/src/firmware/sagas.ts
index 12b49676..8f0ffcbc 100644
--- a/src/firmware/sagas.ts
+++ b/src/firmware/sagas.ts
@@ -65,16 +65,6 @@ import {
     flashFirmware,
 } from './actions';
 
-/**
- * Function that returns the next unused message ID.
- */
-type NextMessageIdFunc = () => number;
-
-/**
- * Partial saga context type for context used in the firmware sagas.
- */
-export type FirmwareSagaContext = { nextMessageId: NextMessageIdFunc };
-
 const firmwareZipMap = new Map<HubType, string>([
     [HubType.CityHub, cityHubZip],
     [HubType.TechnicHub, technicHubZip],
@@ -305,7 +295,7 @@ function* handleFlashFirmware(action: ReturnType<typeof flashFirmware>): Generat
             return;
         }
 
-        const nextMessageId = yield* getContext<NextMessageIdFunc>('nextMessageId');
+        const nextMessageId = yield* getContext<() => number>('nextMessageId');
 
         const infoAction = yield* put(infoRequest(nextMessageId()));
         const { info } = yield* all({
diff --git a/src/sagas.ts b/src/sagas.ts
index 32ab3dc9..7eebb3bb 100644
--- a/src/sagas.ts
+++ b/src/sagas.ts
@@ -10,7 +10,7 @@ import editor from './editor/sagas';
 import errorLog from './error-log/sagas';
 import explorer from './explorer/sagas';
 import fileStorage from './fileStorage/sagas';
-import flashFirmware, { FirmwareSagaContext } from './firmware/sagas';
+import flashFirmware from './firmware/sagas';
 import hub from './hub/sagas';
 import lwp3BootloaderProtocol from './lwp3-bootloader/sagas';
 import lwp3BootloaderBle from './lwp3-bootloader/sagas-ble';
@@ -42,6 +42,7 @@ export default function* (): Generator {
 /**
  * Combined type for all saga contexts.
  */
-export type RootSagaContext = FirmwareSagaContext &
-    NotificationSagaContext &
+export type RootSagaContext = {
+    nextMessageId: () => number;
+} & NotificationSagaContext &
     TerminalSagaContext;