feat: agent token

app/drizzle/20260311132215_fair_rogue/migration.sql

@@ -0,0 +1,29 @@
+CREATE TABLE `agent_tokens` (
+	`id` text PRIMARY KEY,
+	`name` text NOT NULL,
+	`token_hash` text NOT NULL,
+	`token_prefix` text NOT NULL,
+	`agent_id` text NOT NULL,
+	`created_by` text NOT NULL,
+	`last_used_at` integer,
+	`revoked_at` integer,
+	`created_at` integer DEFAULT (unixepoch() * 1000) NOT NULL,
+	CONSTRAINT `fk_agent_tokens_agent_id_agents_id_fk` FOREIGN KEY (`agent_id`) REFERENCES `agents`(`id`) ON DELETE CASCADE,
+	CONSTRAINT `fk_agent_tokens_created_by_users_table_id_fk` FOREIGN KEY (`created_by`) REFERENCES `users_table`(`id`) ON DELETE CASCADE
+);
+--> statement-breakpoint
+CREATE TABLE `agents` (
+	`id` text PRIMARY KEY,
+	`name` text NOT NULL,
+	`organization_id` text NOT NULL,
+	`created_by` text NOT NULL,
+	`last_seen_at` integer,
+	`created_at` integer DEFAULT (unixepoch() * 1000) NOT NULL,
+	CONSTRAINT `fk_agents_organization_id_organization_id_fk` FOREIGN KEY (`organization_id`) REFERENCES `organization`(`id`) ON DELETE CASCADE,
+	CONSTRAINT `fk_agents_created_by_users_table_id_fk` FOREIGN KEY (`created_by`) REFERENCES `users_table`(`id`) ON DELETE CASCADE,
+	CONSTRAINT `agents_name_org_uidx` UNIQUE(`name`,`organization_id`)
+);
+--> statement-breakpoint
+CREATE UNIQUE INDEX `agent_tokens_token_hash_uidx` ON `agent_tokens` (`token_hash`);--> statement-breakpoint
+CREATE INDEX `agent_tokens_agent_id_idx` ON `agent_tokens` (`agent_id`);--> statement-breakpoint
+CREATE INDEX `agents_organization_id_idx` ON `agents` (`organization_id`);

app/server/db/relations.ts

@@ -95,6 +95,29 @@ export const relations = defineRelations(schema, (r) => ({
 			to: r.usersTable.id,
 		}),
 	},
+	agentsTable: {
+		organization: r.one.organization({
+			from: r.agentsTable.organizationId,
+			to: r.organization.id,
+			optional: false,
+		}),
+		createdByUser: r.one.usersTable({
+			from: r.agentsTable.createdBy,
+			to: r.usersTable.id,
+		}),
+		tokens: r.many.agentTokensTable(),
+	},
+	agentTokensTable: {
+		agent: r.one.agentsTable({
+			from: r.agentTokensTable.agentId,
+			to: r.agentsTable.id,
+			optional: false,
+		}),
+		createdByUser: r.one.usersTable({
+			from: r.agentTokensTable.createdBy,
+			to: r.usersTable.id,
+		}),
+	},
 	organization: {
 		users: r.many.usersTable({
 			alias: "usersTable_id_organization_id_via_member",
@@ -106,6 +129,7 @@ export const relations = defineRelations(schema, (r) => ({
 		members: r.many.member(),
 		invitations: r.many.invitation(),
 		ssoProviders: r.many.ssoProvider(),
+		agents: r.many.agentsTable(),
 	},
 	ssoProvider: {
 		user: r.one.usersTable({

app/server/db/schema.ts

@@ -413,6 +413,61 @@ export const appMetadataTable = sqliteTable("app_metadata", {
 });
 export type AppMetadata = typeof appMetadataTable.$inferSelect;
 
+/**
+ * Agents Table
+ */
+export const agentsTable = sqliteTable(
+	"agents",
+	{
+		id: text("id").primaryKey(),
+		name: text("name").notNull(),
+		organizationId: text("organization_id")
+			.notNull()
+			.references(() => organization.id, { onDelete: "cascade" }),
+		createdBy: text("created_by")
+			.notNull()
+			.references(() => usersTable.id, { onDelete: "cascade" }),
+		lastSeenAt: int("last_seen_at", { mode: "number" }),
+		createdAt: int("created_at", { mode: "number" })
+			.notNull()
+			.default(sql`(unixepoch() * 1000)`),
+	},
+	(table) => [
+		index("agents_organization_id_idx").on(table.organizationId),
+		unique("agents_name_org_uidx").on(table.name, table.organizationId),
+	],
+);
+export type Agent = typeof agentsTable.$inferSelect;
+
+/**
+ * Agent Tokens Table
+ */
+export const agentTokensTable = sqliteTable(
+	"agent_tokens",
+	{
+		id: text("id").primaryKey(),
+		name: text("name").notNull(),
+		tokenHash: text("token_hash").notNull(),
+		tokenPrefix: text("token_prefix").notNull(),
+		agentId: text("agent_id")
+			.notNull()
+			.references(() => agentsTable.id, { onDelete: "cascade" }),
+		createdBy: text("created_by")
+			.notNull()
+			.references(() => usersTable.id, { onDelete: "cascade" }),
+		lastUsedAt: int("last_used_at", { mode: "number" }),
+		revokedAt: int("revoked_at", { mode: "number" }),
+		createdAt: int("created_at", { mode: "number" })
+			.notNull()
+			.default(sql`(unixepoch() * 1000)`),
+	},
+	(table) => [
+		uniqueIndex("agent_tokens_token_hash_uidx").on(table.tokenHash),
+		index("agent_tokens_agent_id_idx").on(table.agentId),
+	],
+);
+export type AgentToken = typeof agentTokensTable.$inferSelect;
+
 export const twoFactor = sqliteTable(
 	"two_factor",
 	{

app/server/modules/agents/agent-tokens.ts

@@ -0,0 +1,87 @@
+import crypto from "node:crypto";
+import { eq } from "drizzle-orm";
+import { db } from "~/server/db/db";
+import { agentTokensTable } from "~/server/db/schema";
+import { cryptoUtils } from "~/server/utils/crypto";
+
+export const generateToken = () => {
+	return `zbk_${crypto.randomBytes(32).toString("hex")}`;
+};
+
+export const hashToken = (token: string) => {
+	return crypto.createHash("sha256").update(token).digest("hex");
+};
+
+export const deriveLocalAgentToken = async () => {
+	const derived = await cryptoUtils.deriveSecret("zerobyte:local-agent-token");
+	return `zbk_${derived}`;
+};
+
+export const createAgentToken = async ({
+	name,
+	agentId,
+	createdBy,
+}: {
+	name: string;
+	agentId: string;
+	createdBy: string;
+}) => {
+	const plaintext = generateToken();
+	const tokenHash = hashToken(plaintext);
+	const tokenPrefix = plaintext.slice(0, 12);
+
+	const id = Bun.randomUUIDv7();
+	await db.insert(agentTokensTable).values({
+		id,
+		name,
+		tokenHash,
+		tokenPrefix,
+		agentId,
+		createdBy,
+	});
+
+	return { id, name, tokenPrefix, plaintext };
+};
+
+export const validateAgentToken = async (token: string) => {
+	const localToken = await deriveLocalAgentToken();
+	if (token === localToken) {
+		return { agentId: "local", organizationId: null, agentName: "local" };
+	}
+
+	const tokenHash = hashToken(token);
+
+	const record = await db.query.agentTokensTable.findFirst({
+		where: { tokenHash, revokedAt: { isNull: true } },
+		with: { agent: true },
+	});
+
+	if (!record) return null;
+
+	await db.update(agentTokensTable).set({ lastUsedAt: Date.now() }).where(eq(agentTokensTable.id, record.id));
+
+	return {
+		agentId: record.agentId,
+		organizationId: record.agent.organizationId,
+		agentName: record.name,
+	};
+};
+
+export const revokeAgentToken = async (tokenId: string, agentId: string) => {
+	const token = await db.query.agentTokensTable.findFirst({
+		where: { id: tokenId, agentId, revokedAt: { isNull: true } },
+	});
+
+	if (!token) return false;
+
+	await db.update(agentTokensTable).set({ revokedAt: Date.now() }).where(eq(agentTokensTable.id, tokenId));
+
+	return true;
+};
+
+export const listAgentTokens = async (agentId: string) => {
+	return db.query.agentTokensTable.findMany({
+		where: { agentId },
+		columns: { id: true, name: true, tokenPrefix: true, lastUsedAt: true, revokedAt: true, createdAt: true },
+	});
+};

app/server/modules/agents/agents-manager.ts

@@ -7,10 +7,13 @@ import {
 	type BackupCommandPayload,
 } from "@zerobyte/contracts/agent-protocol";
 import { logger } from "@zerobyte/core/node";
+import { validateAgentToken, deriveLocalAgentToken } from "./agent-tokens";
 
 type AgentConnectionData = {
 	id: string;
-	agentId?: string;
+	agentId: string;
+	organizationId: string | null;
+	agentName: string;
 };
 
 type AgentSocket = Bun.ServerWebSocket<AgentConnectionData>;
@@ -41,13 +44,15 @@ const setServer = (server: AgentServer | null) => {
 	delete globalState[AGENT_SERVER_KEY];
 };
 
-export const spawnLocalAgent = () => {
+export const spawnLocalAgent = async () => {
 	const agentEntryPoint = path.join(process.cwd(), "apps", "agent", "src", "index.ts");
+	const agentToken = await deriveLocalAgentToken();
 
 	const localAgent = spawn("bun", ["run", agentEntryPoint], {
 		env: {
 			PATH: process.env.PATH,
 			ZEROBYTE_CONTROLLER_URL: "ws://localhost:3001",
+			ZEROBYTE_AGENT_TOKEN: agentToken,
 		},
 		stdio: ["ignore", "pipe", "pipe"],
 	});
@@ -79,57 +84,75 @@ export const agentManager = {
 		logger.info("Starting Agent Manager...");
 		const server = Bun.serve<AgentConnectionData>({
 			port: 3001,
-			fetch(req, srv) {
-				const upgraded = srv.upgrade(req, { data: { id: Bun.randomUUIDv7() } });
+			async fetch(req, srv) {
+				const url = new URL(req.url);
+				const token = url.searchParams.get("token");
+
+				if (!token) {
+					return new Response("Missing token", { status: 401 });
+				}
+
+				const result = await validateAgentToken(token);
+				if (!result) {
+					return new Response("Invalid or revoked token", { status: 401 });
+				}
+
+				const upgraded = srv.upgrade(req, {
+					data: {
+						id: Bun.randomUUIDv7(),
+						agentId: result.agentId,
+						organizationId: result.organizationId,
+						agentName: result.agentName,
+					},
+				});
 				if (upgraded) return undefined;
-				return new Response("Agent WebSocket endpoint", { status: 200 });
+				return new Response("WebSocket upgrade failed", { status: 400 });
 			},
 			websocket: {
-				open: (ws) => logger.info(`WebSocket opened with id: ${ws.data.id}`),
+				open: (ws) => {
+					getAgentSockets().set(ws.data.agentId, ws);
+					logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) connected on ${ws.data.id}`);
+				},
 				message: (ws, data) => {
 					if (typeof data !== "string") {
-						logger.warn(`Ignoring non-text message from agent connection ${ws.data.id}`);
+						logger.warn(`Ignoring non-text message from agent ${ws.data.agentId}`);
 						return;
 					}
 
 					const parsed = parseAgentMessage(data);
 
 					if (parsed === null) {
-						logger.warn(`Invalid JSON from agent connection ${ws.data.id}`);
+						logger.warn(`Invalid JSON from agent ${ws.data.agentId}`);
 						return;
 					}
 
 					if (!parsed.success) {
-						logger.warn(`Invalid agent message on connection ${ws.data.id}: ${parsed.error.message}`);
+						logger.warn(`Invalid agent message from ${ws.data.agentId}: ${parsed.error.message}`);
 						return;
 					}
 
 					switch (parsed.data.type) {
 						case "agent.ready": {
-							ws.data.agentId = parsed.data.payload.agentId;
-							getAgentSockets().set(parsed.data.payload.agentId, ws);
-							logger.info(`Backup agent ${parsed.data.payload.agentId} is ready on connection ${ws.data.id}`);
+							logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) is ready`);
 							break;
 						}
 						case "backup.started": {
-							logger.info(
-								`Backup started on agent ${ws.data.agentId ?? ws.data.id} for schedule ${parsed.data.payload.scheduleId}`,
-							);
+							logger.info(`Backup started on agent ${ws.data.agentId} for schedule ${parsed.data.payload.scheduleId}`);
 							break;
 						}
 					}
 				},
 				close: (ws) => {
-					if (ws.data.agentId && getAgentSockets().get(ws.data.agentId) === ws) {
+					if (getAgentSockets().get(ws.data.agentId) === ws) {
 						getAgentSockets().delete(ws.data.agentId);
 					}
 
-					logger.info(`WebSocket closed for agent ${ws.data.agentId ?? ws.data.id}`);
+					logger.info(`Agent "${ws.data.agentName}" (${ws.data.agentId}) disconnected`);
 				},
 			},
 		});
-		setServer(server);
 
+		setServer(server);
 		logger.info(`Agent Manager listening on port ${server.port}`);
 	},
 	sendBackup: (agentId: string, payload: BackupCommandPayload) => {

app/server/modules/lifecycle/bootstrap.ts

@@ -10,7 +10,7 @@ const runBootstrap = async () => {
 	await runMigrations();
 	await startup();
 	agentManager.start();
-	spawnLocalAgent();
+	await spawnLocalAgent();
 };
 
 export const bootstrapApplication = async () => {

apps/agent/src/index.ts

@@ -2,25 +2,29 @@ import { createAgentMessage, parseControllerMessage, sendAgentMessage } from "@z
 import { logger } from "@zerobyte/core/node";
 
 const controllerUrl = process.env.ZEROBYTE_CONTROLLER_URL;
+const agentToken = process.env.ZEROBYTE_AGENT_TOKEN;
 
 class Agent {
 	private ws: WebSocket | null = null;
 
-	constructor(public id: string) {
-		this.connect();
-	}
-
-	private connect() {
+	connect() {
 		if (!controllerUrl) {
 			throw new Error("Env variable ZEROBYTE_CONTROLLER_URL is not set");
 		}
 
-		this.ws = new WebSocket(controllerUrl);
+		if (!agentToken) {
+			throw new Error("Env variable ZEROBYTE_AGENT_TOKEN is not set");
+		}
+
+		const url = new URL(controllerUrl);
+		url.searchParams.set("token", agentToken);
+
+		this.ws = new WebSocket(url.toString());
 		this.ws.onopen = () => {
-			logger.info(`Agent ${this.id} connected to controller`);
+			logger.info("Agent connected to controller");
 
 			if (this.ws) {
-				sendAgentMessage(this.ws, createAgentMessage("agent.ready", { agentId: this.id }));
+				sendAgentMessage(this.ws, createAgentMessage("agent.ready", { agentId: "" }));
 			}
 		};
 
@@ -28,18 +32,18 @@ class Agent {
 			const parsed = parseControllerMessage(event.data);
 
 			if (parsed === null) {
-				console.error(`Agent ${this.id} received invalid JSON`);
+				console.error("Agent received invalid JSON");
 				return;
 			}
 
 			if (!parsed.success) {
-				console.error(`Agent ${this.id} received an invalid message: ${parsed.error.message}`);
+				console.error(`Agent received an invalid message: ${parsed.error.message}`);
 				return;
 			}
 
 			switch (parsed.data.type) {
 				case "backup":
-					logger.info(`Agent ${this.id} starting backup for schedule ${parsed.data.payload.scheduleId}`);
+					logger.info(`Starting backup for schedule ${parsed.data.payload.scheduleId}`);
 					if (this.ws) {
 						sendAgentMessage(
 							this.ws,
@@ -50,12 +54,13 @@ class Agent {
 			}
 		};
 		this.ws.onclose = () => {
-			logger.info(`Agent ${this.id} disconnected from controller`);
+			logger.info("Agent disconnected from controller");
 		};
 		this.ws.onerror = (error) => {
-			logger.error(`Agent ${this.id} encountered an error:`, error);
+			logger.error("Agent encountered an error:", error);
 		};
 	}
 }
 
-new Agent(Bun.randomUUIDv7());
+const agent = new Agent();
+agent.connect();