Retry zrok token setup after stale-environment cleanup

packages/server/package.json

@@ -11,6 +11,7 @@
     "scripts": {
         "build": "export NODE_ENV=production && webpack --config ./scripts/webpack.main.prod.config.js",
         "start": "export NODE_ENV=development && webpack --config ./scripts/webpack.main.config.js && electron ./dist/main.js",
+        "test:zrok": "node --test test/zrokManager.test.cjs",
         "lint": "eslint --ext=jsx,js,tsx,ts src",
         "dist": "npm run build && electron-builder build --mac --publish never --config ./scripts/electron-builder-config.js",
         "release": "npm run build && electron-builder build --mac --publish always --config ./scripts/electron-builder-config.js",

packages/server/src/server/managers/zrokManager/index.ts

@@ -8,6 +8,8 @@ import { spawn, ChildProcess } from "child_process";
 import path from "path";
 import { ProcessSpawner, ProcessSpawnerError } from "@server/lib/ProcessSpawner";
 
+const ALREADY_ENABLED_ENVIRONMENT_ERROR = "you already have an enabled environment";
+
 export class ZrokManager extends Loggable {
     tag = "ZrokManager";
 
@@ -158,19 +160,33 @@ export class ZrokManager extends Loggable {
 
     static async setToken(token: string): Promise<string> {
         const logger = getLogger("ZrokManager");
+        const enable = () => ProcessSpawner.executeCommand(this.daemonPath, ["enable", token], {}, "ZrokManager");
+        const throwEnableError = (ex: any | ProcessSpawnerError): never => {
+            const output = ex?.output ?? ex?.message ?? String(ex);
+            if (output.includes("enableUnauthorized")) {
+                throw new Error("Invalid Zrok token!");
+            }
 
+            logger.error(`Failed to set Zrok token! Error: ${output}`);
+            throw new Error("Failed to set Zrok token! Please check your server logs for more information.");
+        };
+
+        logger.info(`Enabling Zrok...`);
         try {
-            logger.info(`Enabling Zrok...`);
-            return await ProcessSpawner.executeCommand(this.daemonPath, ["enable", token], {}, "ZrokManager");
+            return await enable();
         } catch (ex: any | ProcessSpawnerError) {
             const output = ex?.output ?? ex?.message ?? String(ex);
-            if (output.includes("you already have an enabled environment")) {
-                return output;
-            } else if (output.includes("enableUnauthorized")) {
-                throw new Error("Invalid Zrok token!");
-            } else {
-                logger.error(`Failed to set Zrok token! Error: ${output}`);
-                throw new Error("Failed to set Zrok token! Please check your server logs for more information.");
+            if (!output.includes(ALREADY_ENABLED_ENVIRONMENT_ERROR)) {
+                return throwEnableError(ex);
+            }
+
+            logger.info("Zrok already has an enabled environment. Disabling it before retrying...");
+            await this.disable();
+
+            try {
+                return await enable();
+            } catch (retryEx: any | ProcessSpawnerError) {
+                return throwEnableError(retryEx);
             }
         }
     }

packages/server/test/zrokManager.test.cjs

@@ -0,0 +1,141 @@
+const assert = require("node:assert/strict");
+const path = require("node:path");
+const test = require("node:test");
+const babel = require("@babel/core");
+
+const loadTypeScriptModule = (modulePath, overrides = {}) => {
+    const transformed = babel.transformFileSync(modulePath, {
+        presets: [
+            [require.resolve("@babel/preset-env"), { targets: { node: "20" }, modules: "commonjs" }],
+            require.resolve("@babel/preset-typescript")
+        ]
+    });
+    const loadedModule = { exports: {} };
+    const loadModule = new Function("module", "exports", "require", transformed.code);
+    const customRequire = request =>
+        Object.prototype.hasOwnProperty.call(overrides, request) ? overrides[request] : require(request);
+    loadModule(loadedModule, loadedModule.exports, customRequire);
+    return loadedModule.exports;
+};
+
+const zrokManagerPath = path.join(__dirname, "../src/server/managers/zrokManager/index.ts");
+const token = "test-token";
+const alreadyEnabledError = {
+    output: "[ERROR]: unable to enable environment (you already have an enabled environment)"
+};
+
+const loadZrokManager = responses => {
+    const calls = [];
+    const logs = {
+        debug: [],
+        error: [],
+        info: []
+    };
+    const executeCommand = async (...args) => {
+        calls.push(args);
+        const response = responses.shift();
+        if (!response) throw new Error("Unexpected ProcessSpawner call");
+        if ("error" in response) throw response.error;
+        return response.output;
+    };
+    const logger = {
+        debug: (...args) => logs.debug.push(args),
+        error: (...args) => logs.error.push(args),
+        info: (...args) => logs.info.push(args)
+    };
+
+    const { ZrokManager } = loadTypeScriptModule(zrokManagerPath, {
+        axios: { post: async () => undefined },
+        child_process: { spawn: () => undefined },
+        electron: { app: { getVersion: () => "test" } },
+        "@server": { Server: () => undefined },
+        "@server/fileSystem": { FileSystem: { resources: "/resources" } },
+        "@server/helpers/utils": {
+            isEmpty: value => value == null || value.length === 0,
+            isNotEmpty: value => value != null && value.length > 0
+        },
+        "@server/lib/logging/Loggable": {
+            Loggable: class {},
+            getLogger: () => logger
+        },
+        "@server/lib/ProcessSpawner": {
+            ProcessSpawner: { executeCommand }
+        }
+    });
+
+    return { calls, logs, ZrokManager };
+};
+
+const success = output => ({ output });
+const failure = output => ({ error: { output } });
+const expectedCall = (ZrokManager, args) => [ZrokManager.daemonPath, args, {}, "ZrokManager"];
+
+test("returns the output when zrok enables normally", async () => {
+    const { calls, ZrokManager } = loadZrokManager([success("environment enabled")]);
+
+    assert.equal(await ZrokManager.setToken(token), "environment enabled");
+    assert.deepEqual(calls, [expectedCall(ZrokManager, ["enable", token])]);
+});
+
+test("preserves invalid-token errors without disabling or retrying", async () => {
+    const { calls, ZrokManager } = loadZrokManager([failure("[ERROR]: enableUnauthorized")]);
+
+    await assert.rejects(() => ZrokManager.setToken(token), { message: "Invalid Zrok token!" });
+    assert.deepEqual(calls, [expectedCall(ZrokManager, ["enable", token])]);
+});
+
+test("preserves generic enable errors without disabling or retrying", async () => {
+    const { calls, ZrokManager } = loadZrokManager([failure("[ERROR]: network unavailable")]);
+
+    await assert.rejects(() => ZrokManager.setToken(token), {
+        message: "Failed to set Zrok token! Please check your server logs for more information."
+    });
+    assert.deepEqual(calls, [expectedCall(ZrokManager, ["enable", token])]);
+});
+
+test("disables a stale environment and retries enable exactly once", async () => {
+    const { calls, ZrokManager } = loadZrokManager([
+        { error: alreadyEnabledError },
+        success("environment disabled"),
+        success("environment enabled")
+    ]);
+
+    assert.equal(await ZrokManager.setToken(token), "environment enabled");
+    assert.deepEqual(calls, [
+        expectedCall(ZrokManager, ["enable", token]),
+        expectedCall(ZrokManager, ["disable"]),
+        expectedCall(ZrokManager, ["enable", token])
+    ]);
+});
+
+test("does not retry again when the second enable fails", async () => {
+    const { calls, ZrokManager } = loadZrokManager([
+        { error: alreadyEnabledError },
+        success("environment disabled"),
+        { error: alreadyEnabledError }
+    ]);
+
+    await assert.rejects(() => ZrokManager.setToken(token), {
+        message: "Failed to set Zrok token! Please check your server logs for more information."
+    });
+    assert.deepEqual(calls, [
+        expectedCall(ZrokManager, ["enable", token]),
+        expectedCall(ZrokManager, ["disable"]),
+        expectedCall(ZrokManager, ["enable", token])
+    ]);
+});
+
+test("preserves invalid-token errors from the retry", async () => {
+    const { calls, ZrokManager } = loadZrokManager([
+        { error: alreadyEnabledError },
+        success("environment disabled"),
+        failure("[ERROR]: enableUnauthorized")
+    ]);
+
+    await assert.rejects(() => ZrokManager.setToken(token), { message: "Invalid Zrok token!" });
+    assert.deepEqual(calls, [
+        expectedCall(ZrokManager, ["enable", token]),
+        expectedCall(ZrokManager, ["disable"]),
+        expectedCall(ZrokManager, ["enable", token])
+    ]);
+});