refactor: gate local agent behind feature flag

This commit is contained in:
Nicolas Meienberger 2026-04-09 20:54:01 +02:00
parent 78758a6c2c
commit a85f00171d
No known key found for this signature in database
12 changed files with 76 additions and 35 deletions

View file

@ -58,7 +58,7 @@ export const createApp = () => {
limit: 1000, limit: 1000,
keyGenerator: (c) => c.req.header("x-forwarded-for") ?? "", keyGenerator: (c) => c.req.header("x-forwarded-for") ?? "",
skip: () => { skip: () => {
return config.disableRateLimiting; return config.flags.disableRateLimiting;
}, },
}), }),
); );

View file

@ -43,6 +43,7 @@ const envSchema = z
APP_SECRET_FILE: z.string().optional(), APP_SECRET_FILE: z.string().optional(),
BASE_URL: z.string(), BASE_URL: z.string(),
ENABLE_DEV_PANEL: z.string().default("false"), ENABLE_DEV_PANEL: z.string().default("false"),
ENABLE_LOCAL_AGENT: z.string().default("false"),
PROVISIONING_PATH: z.string().optional(), PROVISIONING_PATH: z.string().optional(),
}) })
.transform((s, ctx) => { .transform((s, ctx) => {
@ -125,11 +126,14 @@ const envSchema = z
appVersion: s.APP_VERSION, appVersion: s.APP_VERSION,
trustedOrigins: trustedOrigins, trustedOrigins: trustedOrigins,
trustProxy: s.TRUST_PROXY === "true", trustProxy: s.TRUST_PROXY === "true",
disableRateLimiting: s.DISABLE_RATE_LIMITING === "true" || s.NODE_ENV === "test",
appSecret: appSecret ?? "", appSecret: appSecret ?? "",
baseUrl, baseUrl,
isSecure: baseUrl.startsWith("https://"), isSecure: baseUrl.startsWith("https://"),
enableDevPanel: s.ENABLE_DEV_PANEL === "true", flags: {
disableRateLimiting: s.DISABLE_RATE_LIMITING === "true" || s.NODE_ENV === "test",
enableDevPanel: s.ENABLE_DEV_PANEL === "true",
enableLocalAgent: s.ENABLE_LOCAL_AGENT === "true",
},
provisioningPath: s.PROVISIONING_PATH, provisioningPath: s.PROVISIONING_PATH,
allowedHosts, allowedHosts,
}; };

View file

@ -30,13 +30,13 @@ export const auth = betterAuth({
}, },
trustedOrigins: config.trustedOrigins, trustedOrigins: config.trustedOrigins,
rateLimit: { rateLimit: {
enabled: !config.disableRateLimiting, enabled: !config.flags.disableRateLimiting,
}, },
advanced: { advanced: {
cookiePrefix: "zerobyte", cookiePrefix: "zerobyte",
useSecureCookies: config.isSecure, useSecureCookies: config.isSecure,
ipAddress: { ipAddress: {
disableIpTracking: config.disableRateLimiting, disableIpTracking: config.flags.disableRateLimiting,
}, },
}, },
onAPIError: { onAPIError: {

View file

@ -47,18 +47,19 @@ export function createAgentManagerRuntime() {
let backupHandlers: AgentBackupEventHandlers = {}; let backupHandlers: AgentBackupEventHandlers = {};
let runtimeScope: Scope.CloseableScope | null = null; let runtimeScope: Scope.CloseableScope | null = null;
const closeSession = (sessionHandle: ControllerAgentSessionHandle) => { const closeSession = (sessionHandle: ControllerAgentSessionHandle) =>
Effect.runSync(Fiber.interrupt(sessionHandle.runFiber)); Effect.gen(function* () {
Effect.runSync(Scope.close(sessionHandle.scope, Exit.succeed(undefined))); yield* Fiber.interrupt(sessionHandle.runFiber);
}; yield* Scope.close(sessionHandle.scope, Exit.succeed(undefined));
});
const closeAllSessions = () => { const closeAllSessions = Effect.gen(function* () {
const currentSessions = sessions; const currentSessions = sessions;
sessions = new Map(); sessions = new Map();
for (const sessionHandle of currentSessions.values()) { for (const sessionHandle of currentSessions.values()) {
closeSession(sessionHandle); yield* closeSession(sessionHandle);
} }
}; });
const getSessionHandle = (agentId: string) => sessions.get(agentId); const getSessionHandle = (agentId: string) => sessions.get(agentId);
@ -105,7 +106,9 @@ export function createAgentManagerRuntime() {
const setSession = (agentId: string, sessionHandle: ControllerAgentSessionHandle) => { const setSession = (agentId: string, sessionHandle: ControllerAgentSessionHandle) => {
const existingSession = getSessionHandle(agentId); const existingSession = getSessionHandle(agentId);
if (existingSession) { if (existingSession) {
closeSession(existingSession); void Effect.runPromise(closeSession(existingSession)).catch((error) => {
logger.error(`Failed to close existing agent session for ${agentId}: ${toMessage(error)}`);
});
} }
sessions.set(agentId, sessionHandle); sessions.set(agentId, sessionHandle);
@ -118,7 +121,9 @@ export function createAgentManagerRuntime() {
} }
sessions.delete(agentId); sessions.delete(agentId);
closeSession(sessionHandle); void Effect.runPromise(closeSession(sessionHandle)).catch((error) => {
logger.error(`Failed to close agent session for ${agentId}: ${toMessage(error)}`);
});
}; };
const acquireServer = Effect.acquireRelease( const acquireServer = Effect.acquireRelease(
@ -180,7 +185,7 @@ export function createAgentManagerRuntime() {
}), }),
), ),
(server) => (server) =>
Effect.sync(closeAllSessions).pipe( closeAllSessions.pipe(
Effect.andThen( Effect.andThen(
Effect.tryPromise({ Effect.tryPromise({
try: () => server.stop(true), try: () => server.stop(true),

View file

@ -12,6 +12,10 @@ type LocalAgentState = {
export async function spawnLocalAgentProcess(runtime: LocalAgentState) { export async function spawnLocalAgentProcess(runtime: LocalAgentState) {
await stopLocalAgentProcess(runtime); await stopLocalAgentProcess(runtime);
if (!config.flags.enableLocalAgent) {
return;
}
const sourceEntryPoint = path.join(process.cwd(), "apps", "agent", "src", "index.ts"); const sourceEntryPoint = path.join(process.cwd(), "apps", "agent", "src", "index.ts");
const productionEntryPoint = path.join(process.cwd(), ".output", "agent", "index.mjs"); const productionEntryPoint = path.join(process.cwd(), ".output", "agent", "index.mjs");

View file

@ -1,4 +1,4 @@
import { afterEach, describe, expect, mock, spyOn, test } from "bun:test"; import { afterEach, describe, expect, test, vi } from "vitest";
import { Scheduler } from "../../../core/scheduler"; import { Scheduler } from "../../../core/scheduler";
import * as backendModule from "../../backends/backend"; import * as backendModule from "../../backends/backend";
import type { VolumeBackend } from "../../backends/backend"; import type { VolumeBackend } from "../../backends/backend";
@ -12,19 +12,19 @@ const loadShutdownModule = async () => {
}; };
afterEach(() => { afterEach(() => {
mock.restore(); vi.restoreAllMocks();
}); });
describe("shutdown", () => { describe("shutdown", () => {
test("stops the agent runtime before unmounting mounted volumes", async () => { test("stops the agent runtime before unmounting mounted volumes", async () => {
const events: string[] = []; const events: string[] = [];
const stopScheduler = mock(async () => { const stopScheduler = vi.fn(async () => {
events.push("scheduler.stop"); events.push("scheduler.stop");
}); });
const stopApplicationRuntime = mock(async () => { const stopApplicationRuntime = vi.fn(async () => {
events.push("agents.stop"); events.push("agents.stop");
}); });
const unmountVolume = mock(async () => { const unmountVolume = vi.fn(async () => {
events.push("backend.unmount"); events.push("backend.unmount");
return { status: "unmounted" as const }; return { status: "unmounted" as const };
}); });
@ -38,9 +38,9 @@ describe("shutdown", () => {
status: "mounted", status: "mounted",
}); });
spyOn(Scheduler, "stop").mockImplementation(stopScheduler); vi.spyOn(Scheduler, "stop").mockImplementation(stopScheduler);
spyOn(bootstrapModule, "stopApplicationRuntime").mockImplementation(stopApplicationRuntime); vi.spyOn(bootstrapModule, "stopApplicationRuntime").mockImplementation(stopApplicationRuntime);
spyOn(backendModule, "createVolumeBackend").mockImplementation( vi.spyOn(backendModule, "createVolumeBackend").mockImplementation(
() => () =>
({ ({
mount: async () => ({ status: "mounted" as const }), mount: async () => ({ status: "mounted" as const }),

View file

@ -8,7 +8,7 @@ let bootstrapPromise: Promise<void> | undefined;
const runBootstrap = async () => { const runBootstrap = async () => {
await runDbMigrations(); await runDbMigrations();
await runMigrations(); await runMigrations();
agentManager.start(); await agentManager.start();
await spawnLocalAgent(); await spawnLocalAgent();
await startup(); await startup();
}; };

View file

@ -107,7 +107,7 @@ const setRegistrationEnabled = async (enabled: boolean) => {
logger.info(`Registration enabled set to: ${enabled}`); logger.info(`Registration enabled set to: ${enabled}`);
}; };
const isDevPanelEnabled = () => config.enableDevPanel; const isDevPanelEnabled = () => config.flags.enableDevPanel;
export const systemService = { export const systemService = {
getSystemInfo, getSystemInfo,

View file

@ -4,7 +4,8 @@
"type": "module", "type": "module",
"module": "src/index.ts", "module": "src/index.ts",
"scripts": { "scripts": {
"tsc": "tsc --noEmit" "tsc": "tsc --noEmit",
"test": "bunx --bun vitest run --config ./vitest.config.ts"
}, },
"dependencies": { "dependencies": {
"@zerobyte/contracts": "workspace:*", "@zerobyte/contracts": "workspace:*",

View file

@ -1,4 +1,4 @@
import { afterEach, expect, mock, spyOn, test } from "bun:test"; import { afterEach, expect, test, vi } from "vitest";
import { Effect } from "effect"; import { Effect } from "effect";
import waitForExpect from "wait-for-expect"; import waitForExpect from "wait-for-expect";
import { fromPartial } from "@total-typescript/shoehorn"; import { fromPartial } from "@total-typescript/shoehorn";
@ -7,11 +7,11 @@ import * as resticServer from "@zerobyte/core/restic/server";
import { createControllerSession } from "../controller-session"; import { createControllerSession } from "../controller-session";
afterEach(() => { afterEach(() => {
mock.restore(); vi.restoreAllMocks();
}); });
test("emits backup.failed when a backup command hits a restic error", async () => { test("emits backup.failed when a backup command hits a restic error", async () => {
spyOn(resticServer, "createRestic").mockReturnValue( vi.spyOn(resticServer, "createRestic").mockReturnValue(
fromPartial({ fromPartial({
backup: () => Effect.fail("source path missing"), backup: () => Effect.fail("source path missing"),
}), }),

View file

@ -1,4 +1,4 @@
import { afterEach, expect, mock, spyOn, test } from "bun:test"; import { afterEach, expect, test, vi } from "vitest";
import { Effect } from "effect"; import { Effect } from "effect";
import waitForExpect from "wait-for-expect"; import waitForExpect from "wait-for-expect";
import { fromPartial } from "@total-typescript/shoehorn"; import { fromPartial } from "@total-typescript/shoehorn";
@ -18,17 +18,19 @@ const createDeferred = <T>() => {
}; };
afterEach(() => { afterEach(() => {
mock.restore(); vi.restoreAllMocks();
}); });
test("waits for running-job registration before returning to the processor loop", async () => { test("waits for running-job registration before returning to the processor loop", async () => {
const outboundMessages: string[] = []; const outboundMessages: string[] = [];
const runningJobs = new Map<string, RunningJob>(); const runningJobs = new Map<string, RunningJob>();
const setRunningJobGate = createDeferred<void>(); const setRunningJobGate = createDeferred<void>();
const processorLoopGate = createDeferred<void>();
const commandCompleted = createDeferred<void>();
const backupGate = createDeferred<{ exitCode: number; result: null; warningDetails: null }>(); const backupGate = createDeferred<{ exitCode: number; result: null; warningDetails: null }>();
let registeredAbortController: AbortController | undefined; let registeredAbortController: AbortController | undefined;
spyOn(resticServer, "createRestic").mockReturnValue( vi.spyOn(resticServer, "createRestic").mockReturnValue(
fromPartial({ fromPartial({
backup: () => backup: () =>
Effect.async<{ exitCode: number; result: null; warningDetails: null }, never>((resume) => { Effect.async<{ exitCode: number; result: null; warningDetails: null }, never>((resume) => {
@ -82,11 +84,21 @@ test("waits for running-job registration before returning to the processor loop"
scheduleId: "schedule-1", scheduleId: "schedule-1",
}); });
const runPromise = Effect.runPromise(handleBackupRunCommand(context, runPayload)); const processorLoopPromise = Effect.runPromise(
Effect.gen(function* () {
yield* handleBackupRunCommand(context, runPayload);
commandCompleted.resolve(undefined);
yield* Effect.async<void, never>((resume) => {
void processorLoopGate.promise.then(() => {
resume(Effect.void);
});
});
}),
);
try { try {
const returnedBeforeRegistration = await Promise.race([ const returnedBeforeRegistration = await Promise.race([
runPromise.then(() => true), commandCompleted.promise.then(() => true),
new Promise<false>((resolve) => { new Promise<false>((resolve) => {
setTimeout(() => resolve(false), 0); setTimeout(() => resolve(false), 0);
}), }),
@ -95,7 +107,7 @@ test("waits for running-job registration before returning to the processor loop"
expect(returnedBeforeRegistration).toBe(false); expect(returnedBeforeRegistration).toBe(false);
setRunningJobGate.resolve(undefined); setRunningJobGate.resolve(undefined);
await runPromise; await commandCompleted.promise;
await Effect.runPromise(handleBackupCancelCommand(context, cancelPayload)); await Effect.runPromise(handleBackupCancelCommand(context, cancelPayload));
expect(registeredAbortController?.signal.aborted).toBe(true); expect(registeredAbortController?.signal.aborted).toBe(true);
@ -111,7 +123,9 @@ test("waits for running-job registration before returning to the processor loop"
expect(runningJobs.has("job-1")).toBe(false); expect(runningJobs.has("job-1")).toBe(false);
}); });
} finally { } finally {
processorLoopGate.resolve(undefined);
setRunningJobGate.resolve(undefined); setRunningJobGate.resolve(undefined);
backupGate.resolve({ exitCode: 0, result: null, warningDetails: null }); backupGate.resolve({ exitCode: 0, result: null, warningDetails: null });
await processorLoopPromise;
} }
}); });

View file

@ -0,0 +1,13 @@
import { defineConfig } from "vitest/config";
export default defineConfig({
test: {
environment: "node",
server: {
deps: {
inline: ["zod"],
},
},
include: ["src/**/*.test.ts", "src/**/*.spec.ts"],
},
});