diff --git a/src/Netscript/TypeAssertion.ts b/src/Netscript/TypeAssertion.ts
index 27d7c9ac0..fe6709e75 100644
--- a/src/Netscript/TypeAssertion.ts
+++ b/src/Netscript/TypeAssertion.ts
@@ -26,3 +26,7 @@ export const debugType = (v: unknown): string => {
 export function assertString(ctx: NetscriptContext, argName: string, v: unknown): asserts v is string {
   if (typeof v !== "string") throw errorMessage(ctx, `${argName} expected to be a string. ${debugType(v)}`, "TYPE");
 }
+
+export function assertFunction(ctx: NetscriptContext, argName: string, v: unknown): asserts v is () => void {
+  if (typeof v !== "function") throw errorMessage(ctx, `${argName} expected to be a function ${debugType(v)}`, "TYPE");
+}
diff --git a/src/Netscript/killWorkerScript.ts b/src/Netscript/killWorkerScript.ts
index 57e9bb84c..5090ee77b 100644
--- a/src/Netscript/killWorkerScript.ts
+++ b/src/Netscript/killWorkerScript.ts
@@ -47,12 +47,11 @@ function stopAndCleanUpWorkerScript(ws: WorkerScript): void {
   //so the map must be cleared before looping
   ws.atExit = new Map();
 
-  for (const key of atExit.keys()) {
+  for (const [id, callback] of atExit) {
     try {
-      const callback = atExit.get(key);
-      if (typeof callback == "function") callback();
+      callback();
     } catch (e: unknown) {
-      handleUnknownError(e, ws, "Error running atExit function.\n\n");
+      handleUnknownError(e, ws, `Error running atExit function with id ${id}.\n\n`);
     }
   }
 
diff --git a/src/NetscriptFunctions.ts b/src/NetscriptFunctions.ts
index 316c51764..93a247d0f 100644
--- a/src/NetscriptFunctions.ts
+++ b/src/NetscriptFunctions.ts
@@ -108,6 +108,7 @@ import { getRamCost } from "./Netscript/RamCostGenerator";
 import { getEnumHelper } from "./utils/EnumHelper";
 import { setDeprecatedProperties, deprecationWarning } from "./utils/DeprecationHelper";
 import { ServerConstants } from "./Server/data/Constants";
+import { assertFunction } from "./Netscript/TypeAssertion";
 
 export const enums: NSEnums = {
   CityName,
@@ -1708,21 +1709,11 @@ export const ns: InternalAPI<NSFull> = {
     sinceInstall: Object.assign({}, Player.moneySourceA),
     sinceStart: Object.assign({}, Player.moneySourceB),
   }),
-  atExit:
-    (ctx) =>
-    (f, id = "default") => {
-      if (typeof f !== "function") {
-        throw helpers.errorMessage(ctx, "argument should be function");
-      }
-
-      if (typeof id !== "string") {
-        throw helpers.errorMessage(ctx, "id should be a string");
-      }
-
-      ctx.workerScript.atExit.set(id, () => {
-        f();
-      }); // Wrap the user function to prevent WorkerScript leaking as 'this'
-    },
+  atExit: (ctx) => (callback, _id) => {
+    const id = _id ? helpers.string(ctx, "id", _id) : "default";
+    assertFunction(ctx, "callback", callback);
+    ctx.workerScript.atExit.set(id, callback);
+  },
   mv: (ctx) => (_host, _source, _destination) => {
     const hostname = helpers.string(ctx, "host", _host);
     const server = helpers.getServer(ctx, hostname);