@desplega.ai/agent-swarm 2.4.3 → 2.4.5

package/src/be/users.ts

@@ -22,9 +22,10 @@
 import { createHash, randomBytes, randomUUID } from "node:crypto";
 import { scheduleOrgDomainRecompute } from "../telemetry-identity";
 import type { User } from "../types";
-import { unregisterVolatileSecret } from "../utils/secret-scrubber";
+import { registerVolatileSecret, unregisterVolatileSecret } from "../utils/secret-scrubber";
 import { closeSessionTokenConnections } from "../utils/session-token-connections";
 import { getDbClient } from "./db";
+import type { DbExecutor } from "./db-client";
 import { validateRuntimeTokenTtlHours } from "./swarm-config-guard";
 
 // ---------------------------------------------------------------------------
@@ -547,6 +548,12 @@ export async function resolveUserByToken(plaintext: string): Promise<User | null
 
 const SESSION_TOKEN_PREFIX = "aseph_"; // agent-swarm-ephemeral
 
+/**
+ * What an agent-principal session token is bound to: one task (ACP), one
+ * runtime instance (runtime pools), or only its agent (a worker token).
+ */
+export type SessionTokenKind = "task" | "runtime" | "agent";
+
 export class InvalidRuntimeTokenError extends Error {
   constructor() {
     super("Runtime token is no longer valid");
@@ -563,14 +570,21 @@ export async function mintSessionToken(
   agentId: string,
   taskId: string,
   ttlMs: number,
-  parent?: { tokenId: string; runtimeId: string },
+  parent?: { tokenId: string; runtimeId: string | null },
 ): Promise<{ tokenId: string; plaintext: string }> {
   if (parent) {
     return await getDbClient().transaction(async (client) => {
+      // A runtime parent is matched by its runtime; an agent parent has no
+      // binding but its agent, so it is matched by the empty binding columns.
       const row = await client.get<{ expiresAt: string }>(
-        `SELECT expiresAt FROM session_tokens WHERE id = ? AND agentId = ? AND runtimeId = ?
-         AND revokedAt IS NULL AND expiresAt > ?`,
-        [parent.tokenId, agentId, parent.runtimeId, new Date().toISOString()],
+        parent.runtimeId
+          ? `SELECT expiresAt FROM session_tokens WHERE id = ? AND agentId = ? AND runtimeId = ?
+             AND revokedAt IS NULL AND expiresAt > ?`
+          : `SELECT expiresAt FROM session_tokens WHERE id = ? AND agentId = ? AND ${AGENT_TOKEN_BINDING}
+             AND revokedAt IS NULL AND expiresAt > ?`,
+        parent.runtimeId
+          ? [parent.tokenId, agentId, parent.runtimeId, new Date().toISOString()]
+          : [parent.tokenId, agentId, new Date().toISOString()],
       );
       if (!row) throw new InvalidRuntimeTokenError();
       return await insertSessionToken(
@@ -580,6 +594,8 @@ export async function mintSessionToken(
         ttlMs,
         parent.runtimeId,
         row.expiresAt,
+        null,
+        parent.runtimeId ? null : parent.tokenId,
       );
     });
   }
@@ -607,6 +623,7 @@ export async function insertSessionToken(
   parentRuntimeId: string | null = null,
   expiresAtCeiling?: string,
   controllerId: string | null = null,
+  parentTokenId: string | null = null,
 ): Promise<{ tokenId: string; plaintext: string }> {
   const plaintext = `${SESSION_TOKEN_PREFIX}${base62(randomBytes(24))}`;
   const tokenId = randomUUID().replace(/-/g, "");
@@ -619,8 +636,8 @@ export async function insertSessionToken(
 
   await getDbClient().run(
     `INSERT INTO session_tokens
-     (id, tokenHash, tokenPreview, agentId, taskId, runtimeId, parentRuntimeId, expiresAt, createdAt, controllerId)
-     VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+     (id, tokenHash, tokenPreview, agentId, taskId, runtimeId, parentRuntimeId, expiresAt, createdAt, controllerId, parentTokenId)
+     VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
     [
       tokenId,
       hash,
@@ -632,6 +649,7 @@ export async function insertSessionToken(
       expiresAt,
       now.toISOString(),
       controllerId,
+      parentTokenId,
     ],
   );
 
@@ -648,15 +666,31 @@ export async function revokeSessionToken(tokenId: string): Promise<void> {
     if (row?.runtimeId) {
       await revokeRuntimeTokensForRuntime(row.runtimeId);
     } else {
-      await client.run(
-        "UPDATE session_tokens SET revokedAt = ? WHERE id = ? AND revokedAt IS NULL",
-        [new Date().toISOString(), tokenId],
-      );
-      getDbClient().afterCommit(() => closeSessionTokenConnections([tokenId]));
+      await revokeTokensWithChildren(client, [tokenId]);
     }
   });
 }
 
+/** Revoke tokens and the child task tokens an agent token minted. */
+async function revokeTokensWithChildren(client: DbExecutor, tokenIds: string[]) {
+  if (tokenIds.length === 0) return;
+  const placeholders = tokenIds.map(() => "?").join(", ");
+  const children = await client.query<{ id: string }>(
+    `SELECT id FROM session_tokens WHERE parentTokenId IN (${placeholders})`,
+    tokenIds,
+  );
+  const ids = [...tokenIds, ...children.map((row) => row.id)];
+  await client.run(
+    `UPDATE session_tokens SET revokedAt = ?
+     WHERE id IN (${ids.map(() => "?").join(", ")}) AND revokedAt IS NULL`,
+    [new Date().toISOString(), ...ids],
+  );
+  getDbClient().afterCommit(() => {
+    for (const id of tokenIds) unregisterVolatileSecret(`agent-token:${id}`);
+    closeSessionTokenConnections(ids);
+  });
+}
+
 export async function revokeRuntimeTokensForRuntime(runtimeId: string): Promise<void> {
   await getDbClient().transaction(async (client) => {
     const rows = await client.query<{ id: string }>(
@@ -675,6 +709,126 @@ export async function revokeRuntimeTokensForRuntime(runtimeId: string): Promise<
   });
 }
 
+// ---------------------------------------------------------------------------
+// Agent tokens (long-lived, one agent, no runtime)
+// ---------------------------------------------------------------------------
+
+/** SQL predicate for an agent token: bound to its agent and nothing else. */
+const AGENT_TOKEN_BINDING =
+  "taskId IS NULL AND runtimeId IS NULL AND controllerId IS NULL AND parentRuntimeId IS NULL";
+
+/** Stored expiry of a token that does not expire. `expiresAt` is NOT NULL. */
+export const AGENT_TOKEN_NO_EXPIRY = "9999-12-31T23:59:59.999Z";
+
+/**
+ * Default agent-token lifetime in ms. `null` means the token does not expire
+ * and stays valid until it is rotated or revoked. A fixed lifetime (for
+ * example 90 days) is a one-line change here.
+ */
+export const AGENT_TOKEN_DEFAULT_TTL_MS: number | null = null;
+
+export type AgentTokenCredential = {
+  tokenId: string;
+  plaintext: string;
+  /** null when the token does not expire. */
+  expiresAt: string | null;
+};
+
+export class AgentTokenExistsError extends Error {
+  constructor() {
+    super("Agent already has an active token; rotate or revoke it");
+  }
+}
+
+function storedExpiry(ttlMs: number | null): string {
+  return ttlMs === null ? AGENT_TOKEN_NO_EXPIRY : new Date(Date.now() + ttlMs).toISOString();
+}
+
+function publicExpiry(expiresAt: string): string | null {
+  return expiresAt === AGENT_TOKEN_NO_EXPIRY ? null : expiresAt;
+}
+
+async function activeAgentTokenIds(client: DbExecutor, agentId: string) {
+  const rows = await client.query<{ id: string }>(
+    `SELECT id FROM session_tokens WHERE agentId = ? AND ${AGENT_TOKEN_BINDING}
+     AND revokedAt IS NULL AND expiresAt > ?`,
+    [agentId, new Date().toISOString()],
+  );
+  return rows.map((row) => row.id);
+}
+
+async function insertAgentToken(
+  agentId: string,
+  ttlMs: number | null,
+): Promise<AgentTokenCredential> {
+  const expiresAt = storedExpiry(ttlMs);
+  // Pass the expiry as the ceiling so it is stored exactly; a TTL recomputed
+  // from a later clock read could step past AGENT_TOKEN_NO_EXPIRY into year 10000.
+  const token = await insertSessionToken(
+    agentId,
+    null,
+    null,
+    Number.MAX_SAFE_INTEGER,
+    null,
+    expiresAt,
+  );
+  getDbClient().afterCommit(() =>
+    registerVolatileSecret(token.plaintext, `agent-token:${token.tokenId}`),
+  );
+  return { ...token, expiresAt: publicExpiry(expiresAt) };
+}
+
+/**
+ * Mint the agent-bound worker token for `agentId`. A worker runs on it in
+ * place of the swarm API key. Throws `AgentTokenExistsError` when the agent
+ * already holds an active one, so tokens do not accumulate silently.
+ */
+export async function mintAgentToken(
+  agentId: string,
+  ttlMs: number | null = AGENT_TOKEN_DEFAULT_TTL_MS,
+): Promise<AgentTokenCredential> {
+  return await getDbClient().transaction(async (client) => {
+    if ((await activeAgentTokenIds(client, agentId)).length > 0) {
+      throw new AgentTokenExistsError();
+    }
+    return await insertAgentToken(agentId, ttlMs);
+  });
+}
+
+/**
+ * Replace the agent's active token. The old token keeps working for
+ * `graceMs` so a running worker can pick up the new one, then expires. Its
+ * child task tokens expire with it, so no credential it delegated outlives
+ * it. Returns null when the agent holds no active token.
+ */
+export async function rotateAgentToken(
+  agentId: string,
+  graceMs: number,
+  ttlMs: number | null = AGENT_TOKEN_DEFAULT_TTL_MS,
+): Promise<(AgentTokenCredential & { previousExpiresAt: string }) | null> {
+  return await getDbClient().transaction(async (client) => {
+    const previous = await activeAgentTokenIds(client, agentId);
+    if (previous.length === 0) return null;
+    const previousExpiresAt = new Date(Date.now() + graceMs).toISOString();
+    const placeholders = previous.map(() => "?").join(", ");
+    await client.run(
+      `UPDATE session_tokens SET expiresAt = ?
+       WHERE (id IN (${placeholders}) OR parentTokenId IN (${placeholders})) AND expiresAt > ?`,
+      [previousExpiresAt, ...previous, ...previous, previousExpiresAt],
+    );
+    return { ...(await insertAgentToken(agentId, ttlMs)), previousExpiresAt };
+  });
+}
+
+/** Revoke every active agent token of `agentId` and their child task tokens. */
+export async function revokeAgentTokens(agentId: string): Promise<number> {
+  return await getDbClient().transaction(async (client) => {
+    const ids = await activeAgentTokenIds(client, agentId);
+    await revokeTokensWithChildren(client, ids);
+    return ids.length;
+  });
+}
+
 /**
  * Resolve a plaintext `aseph_` token to its principal. Returns null when the
  * token is unknown, revoked, or past its expiry. On a hit, touches `lastUsedAt`
@@ -685,6 +839,7 @@ export async function resolveBySessionToken(
   options: { touch?: boolean } = {},
 ): Promise<{
   tokenId: string;
+  tokenKind: SessionTokenKind;
   agentId: string;
   taskId: string | null;
   runtimeId: string | null;
@@ -704,6 +859,7 @@ export async function resolveSessionToken(
   | {
       kind: "agent";
       tokenId: string;
+      tokenKind: SessionTokenKind;
       agentId: string;
       taskId: string | null;
       runtimeId: string | null;
@@ -758,6 +914,7 @@ export async function resolveSessionToken(
   return {
     kind: "agent",
     tokenId: row.id,
+    tokenKind: row.taskId !== null ? "task" : row.runtimeId !== null ? "runtime" : "agent",
     agentId: row.agentId,
     taskId: row.taskId,
     runtimeId: row.runtimeId,

package/src/cli.tsx

@@ -360,6 +360,25 @@ const COMMAND_HELP: Record<
       `  ${binName} claude-managed-setup --api-url https://swarm.example.com`,
     ].join("\n"),
   },
+  "worker-token": {
+    usage: `${binName} worker-token <create|rotate|revoke> (--agent-id <id> | --name <name>) [options]`,
+    description:
+      "Manage the agent-bound token a worker runs on in place of the swarm API key. The token acts only as its agent. create and rotate print the new token to stdout once; messages go to stderr. Uses your operator swarm API key.",
+    options: [
+      "  --agent-id <id>    Agent the token is bound to",
+      "  --name <name>      Agent name; create makes the agent when no agent has this name",
+      "  --ttl-days <n>     Token lifetime (create, rotate). Default: no expiry",
+      "  --grace-sec <n>    How long the old token keeps working after rotate (default 300)",
+      "  --api-url <url>    Swarm API URL (default: MCP_BASE_URL or http://localhost:3013)",
+      "  --api-key <key>    Operator swarm API key (default: AGENT_SWARM_API_KEY or API_KEY, falling back to 123123)",
+      "  -h, --help         Show this help",
+    ].join("\n"),
+    examples: [
+      `  ${binName} worker-token create --name worker-1`,
+      `  ${binName} worker-token rotate --agent-id <agent-id> --grace-sec 600`,
+      `  ${binName} worker-token revoke --agent-id <agent-id>`,
+    ].join("\n"),
+  },
   e2b: {
     usage: `${binName} e2b <subcommand> [options]`,
     description:
@@ -445,6 +464,7 @@ function printHelp(command?: string) {
     ["docs", "Open documentation (--open to launch in browser)"],
     ["codex-login", "Authenticate Codex via ChatGPT OAuth"],
     ["claude-managed-setup", "Bootstrap Anthropic Managed Agents (agent + env + skills)"],
+    ["worker-token", "Mint, rotate, or revoke an agent-bound worker token"],
     ["e2b", "Build templates and start E2B API/worker sandboxes"],
     ["version", "Show version number"],
     ["help", "Show this help message"],
@@ -803,6 +823,10 @@ if (args.showHelp || args.command === "help" || args.command === undefined) {
   const { runClaudeManagedSetup } = await import("./commands/claude-managed-setup");
   const setupArgs = process.argv.slice(process.argv.indexOf("claude-managed-setup") + 1);
   await runClaudeManagedSetup(setupArgs);
+} else if (args.command === "worker-token") {
+  const { runWorkerToken } = await import("./commands/worker-token");
+  const tokenArgs = process.argv.slice(process.argv.indexOf("worker-token") + 1);
+  process.exit(await runWorkerToken(tokenArgs));
 } else if (args.command === "e2b") {
   const { runE2BCommand } = await import("./commands/e2b");
   const e2bArgs = process.argv.slice(process.argv.indexOf("e2b") + 1);

package/src/commands/worker-token.ts

@@ -0,0 +1,164 @@
+/**
+ * `agent-swarm worker-token create|rotate|revoke`: manage the agent-bound
+ * token a worker runs on in place of the swarm API key.
+ *
+ * Operator-run. It authenticates with the operator's own swarm API key and
+ * talks to the API over HTTP only. The new token goes to stdout, once, so it
+ * can be captured (`API_KEY=$(agent-swarm worker-token create --name w1)`);
+ * everything else goes to stderr.
+ */
+import { getApiKey } from "../utils/api-key";
+
+type Action = "create" | "rotate" | "revoke";
+
+export type WorkerTokenArgs = {
+  action?: Action;
+  agentId?: string;
+  name?: string;
+  apiUrl?: string;
+  apiKey?: string;
+  ttlDays?: number;
+  graceSec?: number;
+  showHelp: boolean;
+  error?: string;
+};
+
+export const WORKER_TOKEN_USAGE = `Usage:
+  agent-swarm worker-token <create|rotate|revoke> (--agent-id <id> | --name <name>) [options]
+
+Options:
+  --agent-id <id>     Agent the token is bound to
+  --name <name>       Agent name. create makes the agent when no agent has this name
+  --ttl-days <n>      Token lifetime (create, rotate). Default: no expiry
+  --grace-sec <n>     How long the old token keeps working after rotate (default 300)
+  --api-url <url>     Swarm API URL (default: MCP_BASE_URL or http://localhost:3013)
+  --api-key <key>     Operator swarm API key (default: AGENT_SWARM_API_KEY or API_KEY, falling back to 123123)
+  -h, --help          Show this help`;
+
+const ACTIONS = new Set<string>(["create", "rotate", "revoke"]);
+
+export function parseWorkerTokenArgs(args: string[]): WorkerTokenArgs {
+  const parsed: WorkerTokenArgs = { showHelp: false };
+  const number = (flag: string, value: string | undefined, min: number) => {
+    const n = Number(value);
+    if (!Number.isInteger(n) || n < min) parsed.error = `${flag} must be an integer >= ${min}`;
+    return n;
+  };
+  for (let i = 0; i < args.length; i++) {
+    const arg = args[i]!;
+    if (arg === "--help" || arg === "-h") parsed.showHelp = true;
+    else if (arg === "--agent-id") parsed.agentId = args[++i];
+    else if (arg === "--name") parsed.name = args[++i];
+    else if (arg === "--api-url") parsed.apiUrl = args[++i];
+    else if (arg === "--api-key") parsed.apiKey = args[++i];
+    else if (arg === "--ttl-days") parsed.ttlDays = number(arg, args[++i], 1);
+    else if (arg === "--grace-sec") parsed.graceSec = number(arg, args[++i], 0);
+    else if (!parsed.action && ACTIONS.has(arg)) parsed.action = arg as Action;
+    else parsed.error = `Unknown argument: ${arg}`;
+  }
+  if (parsed.error || parsed.showHelp) return parsed;
+  if (!parsed.action) parsed.error = "Missing action: create, rotate, or revoke";
+  else if (!parsed.agentId === !parsed.name)
+    parsed.error = "Pass exactly one of --agent-id or --name";
+  return parsed;
+}
+
+type Deps = {
+  env?: Record<string, string | undefined>;
+  fetch?: typeof fetch;
+  stdout?: (line: string) => void;
+  stderr?: (line: string) => void;
+};
+
+async function readError(res: Response): Promise<string> {
+  const text = await res.text().catch(() => "");
+  try {
+    return (JSON.parse(text) as { error?: string }).error ?? text;
+  } catch {
+    return text;
+  }
+}
+
+/** Runs the command. Returns the process exit code. */
+export async function runWorkerToken(args: string[], deps: Deps = {}): Promise<number> {
+  const env = deps.env ?? process.env;
+  const doFetch = deps.fetch ?? fetch;
+  const stdout = deps.stdout ?? ((line: string) => console.log(line));
+  const stderr = deps.stderr ?? ((line: string) => console.error(line));
+  const parsed = parseWorkerTokenArgs(args);
+  if (parsed.showHelp) {
+    stdout(WORKER_TOKEN_USAGE);
+    return 0;
+  }
+  if (parsed.error) {
+    stderr(`${parsed.error}\n\n${WORKER_TOKEN_USAGE}`);
+    return 1;
+  }
+
+  const apiUrl = (parsed.apiUrl ?? env.MCP_BASE_URL ?? "http://localhost:3013").replace(/\/+$/, "");
+  const apiKey = parsed.apiKey ?? (getApiKey(env) || "123123");
+  const headers = { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" };
+
+  let agentId = parsed.agentId;
+  let name: string | undefined;
+  if (!agentId) {
+    const res = await doFetch(`${apiUrl}/api/agents`, { headers });
+    if (!res.ok) {
+      stderr(`Listing agents failed (${res.status}): ${await readError(res)}`);
+      return 1;
+    }
+    const { agents } = (await res.json()) as { agents: { id: string; name: string }[] };
+    const matches = agents.filter((agent) => agent.name === parsed.name);
+    if (matches.length > 1) {
+      stderr(`Several agents are named "${parsed.name}"; pass --agent-id instead`);
+      return 1;
+    }
+    agentId = matches[0]?.id;
+    if (!agentId) {
+      if (parsed.action !== "create") {
+        stderr(`No agent is named "${parsed.name}"`);
+        return 1;
+      }
+      agentId = crypto.randomUUID();
+      name = parsed.name;
+    }
+  }
+
+  const base = `${apiUrl}/api/agents/${encodeURIComponent(agentId)}/token`;
+  const ttlMs = parsed.ttlDays === undefined ? undefined : parsed.ttlDays * 86_400_000;
+  const res =
+    parsed.action === "revoke"
+      ? await doFetch(base, { method: "DELETE", headers })
+      : await doFetch(parsed.action === "rotate" ? `${base}/rotate` : base, {
+          method: "POST",
+          headers,
+          body: JSON.stringify(
+            parsed.action === "rotate" ? { graceSec: parsed.graceSec, ttlMs } : { name, ttlMs },
+          ),
+        });
+  if (!res.ok) {
+    stderr(`worker-token ${parsed.action} failed (${res.status}): ${await readError(res)}`);
+    return 1;
+  }
+
+  if (parsed.action === "revoke") {
+    const { revoked } = (await res.json()) as { revoked: number };
+    stderr(`Revoked ${revoked} worker token(s) of agent ${agentId}.`);
+    return 0;
+  }
+  const token = (await res.json()) as {
+    plaintext: string;
+    expiresAt: string | null;
+    agentCreated?: boolean;
+    previousExpiresAt?: string;
+  };
+  if (token.agentCreated) stderr(`Created agent ${agentId} (${name}).`);
+  stderr(
+    `Worker token for agent ${agentId}, ${token.expiresAt ? `expires ${token.expiresAt}` : "no expiry"}.`,
+  );
+  if (token.previousExpiresAt)
+    stderr(`The previous token stops working at ${token.previousExpiresAt}.`);
+  stderr(`Set it as the worker's API_KEY with AGENT_ID=${agentId}. It is shown only once.`);
+  stdout(token.plaintext);
+  return 0;
+}