fix(engine): keep prior agent token authenticatable during rotate-token grace

packages/engine/src/__tests__/conformance/registerOrRotateRace.test.ts

@@ -0,0 +1,116 @@
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { createWorkspace, makeNodeStack, registerAgent, type TestStack } from './harness.js';
+
+/**
+ * Regression coverage for relay#1542. Two clients that concurrently reclaim the
+ * same agent name each receive a fresh token from POST /agents/:name/rotate-token.
+ * The stored `token_hash` is a single slot, so the later rotate would invalidate
+ * the earlier caller's token silently — the caller was handed a 200 response
+ * plus a credential that stopped working microseconds later.
+ *
+ * The fix keeps the previous credential live for a short grace window so both
+ * callers can present the token they were handed and be recognised as the agent
+ * they registered under. Once the grace window elapses, only the most recent
+ * token authenticates and everything older stays revoked.
+ */
+describe('registerOrRotate concurrency', () => {
+  let stack: TestStack;
+
+  beforeEach(() => { stack = makeNodeStack(); });
+  afterEach(() => stack.close());
+
+  async function rotate(workspaceKey: string, name: string): Promise<Response> {
+    return stack.app.request(`/v1/agents/${name}/rotate-token`, {
+      method: 'POST',
+      headers: {
+        authorization: `Bearer ${workspaceKey}`,
+        'content-type': 'application/json',
+      },
+      body: '{}',
+    });
+  }
+
+  async function tokenFrom(res: Response): Promise<string> {
+    const body = await res.json() as { data?: { token?: string } };
+    return body.data?.token ?? '';
+  }
+
+  async function authenticate(agentToken: string): Promise<Response> {
+    return stack.app.request('/v1/agent', {
+      headers: { authorization: `Bearer ${agentToken}` },
+    });
+  }
+
+  it('both concurrent rotations yield tokens that authenticate', async () => {
+    const workspace = await createWorkspace(stack.app, 'race-both-authenticate');
+    await registerAgent(stack.app, workspace.workspaceKey, 'chief');
+
+    const [rotateA, rotateB] = await Promise.all([
+      rotate(workspace.workspaceKey, 'chief'),
+      rotate(workspace.workspaceKey, 'chief'),
+    ]);
+
+    expect(rotateA.status).toBe(200);
+    expect(rotateB.status).toBe(200);
+
+    const tokenA = await tokenFrom(rotateA);
+    const tokenB = await tokenFrom(rotateB);
+    expect(tokenA).toMatch(/^at_live_[0-9a-f]{32}$/);
+    expect(tokenB).toMatch(/^at_live_[0-9a-f]{32}$/);
+    expect(tokenA).not.toBe(tokenB);
+
+    // MUST-FIRE: both callers keep working credentials after the race.
+    const [authA, authB] = await Promise.all([
+      authenticate(tokenA),
+      authenticate(tokenB),
+    ]);
+    expect(authA.status).toBe(200);
+    expect(authB.status).toBe(200);
+  });
+
+  it('rejects a revoked (deleted-agent) token even during a grace window', async () => {
+    const workspace = await createWorkspace(stack.app, 'race-revoked-token');
+    const initial = await registerAgent(stack.app, workspace.workspaceKey, 'ephemeral');
+    expect((await authenticate(initial.token)).status).toBe(200);
+
+    // Rotate once first so the grace slot (`previous_token_hash`) is actually
+    // populated. Without this, the delete path is trivially satisfied by a
+    // null slot and the test would pass even if release never cleared grace.
+    const rotateRes = await rotate(workspace.workspaceKey, 'ephemeral');
+    expect(rotateRes.status).toBe(200);
+    const currentToken = await tokenFrom(rotateRes);
+    // Both slots hold a live credential before the delete: the initial token
+    // is now in the grace slot; `currentToken` is in the current slot.
+    expect((await authenticate(initial.token)).status).toBe(200);
+    expect((await authenticate(currentToken)).status).toBe(200);
+
+    const del = await stack.app.request('/v1/agents/ephemeral', {
+      method: 'DELETE',
+      headers: { authorization: `Bearer ${workspace.workspaceKey}` },
+    });
+    expect(del.status).toBe(204);
+
+    // MUST-NOT-FIRE: a revoked identity is not rescued by either slot — the
+    // current token AND the token still inside its grace window are both
+    // rejected. If release stops clearing grace, this second assertion goes
+    // red immediately.
+    expect((await authenticate(currentToken)).status).toBe(401);
+    expect((await authenticate(initial.token)).status).toBe(401);
+  });
+
+  it('only the two most recent tokens stay valid under chained rotations', async () => {
+    const workspace = await createWorkspace(stack.app, 'race-chained-rotations');
+    const initial = await registerAgent(stack.app, workspace.workspaceKey, 'chained');
+
+    const rotateA = await rotate(workspace.workspaceKey, 'chained');
+    const tokenA = await tokenFrom(rotateA);
+    const rotateB = await rotate(workspace.workspaceKey, 'chained');
+    const tokenB = await tokenFrom(rotateB);
+
+    // The token from before the first rotation is fully retired.
+    expect((await authenticate(initial.token)).status).toBe(401);
+    // Both the intermediate and current tokens still authenticate during the grace window.
+    expect((await authenticate(tokenA)).status).toBe(200);
+    expect((await authenticate(tokenB)).status).toBe(200);
+  });
+});

packages/engine/src/auth/index.ts

@@ -1,4 +1,4 @@
-import { eq } from 'drizzle-orm';
+import { and, eq, gt } from 'drizzle-orm';
 import { workspaces, agents, nodes } from '../db/schema.js';
 import { sha256Hex } from '../lib/crypto.js';
 import { getActiveObserverTokenByHash } from '../engine/observerToken.js';
@@ -49,7 +49,19 @@ export class SqliteApiKeyAuthProvider implements AuthProvider {
     }
 
     if (parsedToken.kind === 'agent') {
-      const [agent] = await db.select().from(agents).where(eq(agents.tokenHash, hash));
+      // Current slot first — the common case is a token that has not been
+      // rotated out from under this caller.
+      let [agent] = await db.select().from(agents).where(eq(agents.tokenHash, hash));
+      if (!agent) {
+        // Fall back to the previous slot inside its grace window. This is the
+        // credential a caller that lost a `registerOrRotate` race was handed
+        // (relay#1542); it must remain live long enough for that caller to
+        // upgrade to a persistent session.
+        [agent] = await db
+          .select()
+          .from(agents)
+          .where(and(eq(agents.previousTokenHash, hash), gt(agents.previousTokenExpiresAt, new Date())));
+      }
       if (!agent) return unauthorized('Invalid agent token', 'agent_token_invalid');
       const [workspace] = await db.select().from(workspaces).where(eq(workspaces.id, agent.workspaceId));
       if (!workspace) return unauthorized('Workspace not found');

packages/engine/src/db/migrations/0035_agent_token_grace.sql

@@ -0,0 +1,18 @@
+-- relay#1542. Give `POST /agents/:name/rotate-token` a two-slot outcome so
+-- concurrent rotations do not silently strand the earlier caller with a token
+-- that stopped authenticating between the response body and the next request.
+-- The prior credential is retained in `previous_token_hash` until
+-- `previous_token_expires_at`, then the auth path stops honouring it.
+--
+-- Nullable and unbounded so a first-ever rotate on a legacy row is trivial
+-- (both columns stay NULL until the second write moves the current hash into
+-- the previous slot). No UNIQUE constraint on `previous_token_hash`: a random
+-- 256-bit token collision is a non-event, and enforcing global uniqueness
+-- across current+previous would abort otherwise-correct rotations.
+ALTER TABLE agents ADD COLUMN previous_token_hash TEXT;
+ALTER TABLE agents ADD COLUMN previous_token_expires_at INTEGER;
+
+-- The auth path fans out to a second lookup by `previous_token_hash` on a
+-- miss against `token_hash`; without this index every rejected token pays a
+-- full-table scan.
+CREATE INDEX IF NOT EXISTS idx_agents_previous_token ON agents(previous_token_hash);

packages/engine/src/db/schema.ts

@@ -62,6 +62,10 @@ export const agents = sqliteTable(
     name: text('name').notNull(),
     type: text('type').notNull().default('agent'),
     tokenHash: text('token_hash').notNull().unique(),
+    // Superseded credential retained during the rotation grace window. See
+    // migration 0035 for the concurrency defect this closes (relay#1542).
+    previousTokenHash: text('previous_token_hash'),
+    previousTokenExpiresAt: integer('previous_token_expires_at', { mode: 'timestamp' }),
     status: text('status').notNull().default('active'),
     handle: text('handle'),
     persona: text('persona'),
@@ -88,6 +92,7 @@ export const agents = sqliteTable(
     uniqueIndex('agents_workspace_id_unique').on(table.workspaceId, table.id),
     index('idx_agents_workspace').on(table.workspaceId),
     index('idx_agents_token').on(table.tokenHash),
+    index('idx_agents_previous_token').on(table.previousTokenHash),
   ],
 );
 

packages/engine/src/engine/action.ts

@@ -764,6 +764,11 @@ async function dispatchRelease(args: {
           // NOT NULL UNIQUE and cannot be cleared, so rotate it to a value
           // nobody holds; the released agent's old token stops authenticating.
           tokenHash: releasedTokenHash,
+          // Clear the rotation grace slot too, otherwise any token issued by
+          // the last live rotation would keep authenticating for its grace
+          // window on an agent that is supposed to be gone. See 0035_agent_token_grace.
+          previousTokenHash: null,
+          previousTokenExpiresAt: null,
           // Same `release` shape the dispatched path writes, so an audit does
           // not have to know which path released the agent.
           metadata: sql`json_patch(COALESCE(${agents.metadata}, '{}'), ${JSON.stringify({
@@ -1387,6 +1392,10 @@ async function applyReleaseCompletionEffect(
         handle: `@${releasedName}`,
         status: RELEASED_AGENT_STATUS,
         tokenHash: releasedTokenHash,
+        // Same reason as the other release paths — the grace slot survives
+        // `token_hash` rewrites unless we clear it. See 0035_agent_token_grace.
+        previousTokenHash: null,
+        previousTokenExpiresAt: null,
         locationType: 'self_connected',
         locationNodeId: null,
         lastSeen: new Date(),

packages/engine/src/engine/agent.ts

@@ -560,6 +560,12 @@ export async function deleteAgent(db: Db, workspaceId: string, name: string) {
         handle: `@${releasedName}`,
         status: RELEASED_AGENT_STATUS,
         tokenHash: releasedTokenHash,
+        // Clear the rotation grace slot too. `token_hash` alone is not the
+        // whole credential surface since 0035_agent_token_grace; a bare rotate
+        // of the current slot would leave a released agent still reachable via
+        // whatever token was in the previous slot until its grace expired.
+        previousTokenHash: null,
+        previousTokenExpiresAt: null,
         locationType: 'self_connected',
         locationNodeId: null,
         lastSeen: releasedAt,

packages/engine/src/engine/tokenRotate.ts

@@ -1,31 +1,55 @@
-import { eq, and } from 'drizzle-orm';
+import { eq, and, sql } from 'drizzle-orm';
 import type { getDb } from '../db/index.js';
 import { agents } from '../db/schema.js';
 import { randomHex, sha256Hex } from '../lib/crypto.js';
 import { codedError } from '../lib/httpError.js';
 
 type Db = ReturnType<typeof getDb>;
 
+/**
+ * How long a superseded agent token stays authenticatable after a rotation.
+ *
+ * Sized to cover the concurrency envelope reported in relay#1542: the broker
+ * and the MCP layer both fire `registerOrRotate` at agent spawn, and the loser
+ * of that race must have long enough to make the follow-up request that gets
+ * it a persistent WebSocket session (which then carries its own auth state).
+ * Sixty seconds is well past the observed request latencies for that path and
+ * well short of a duration that would functionally weaken a rotate-to-revoke.
+ */
+export const AGENT_TOKEN_GRACE_MS = 60_000;
+
 export async function rotateAgentToken(db: Db, workspaceId: string, agentName: string) {
-  const [agent] = await db
-    .select()
+  const [existing] = await db
+    .select({ id: agents.id })
     .from(agents)
     .where(and(eq(agents.workspaceId, workspaceId), eq(agents.name, agentName)));
 
-  if (!agent) {
+  if (!existing) {
     throw codedError(`Agent "${agentName}" not found`, 'agent_not_found', 404);
   }
 
   const newToken = `at_live_${randomHex(16)}`;
   const newTokenHash = await sha256Hex(newToken);
+  const graceExpiresAtSeconds = Math.floor((Date.now() + AGENT_TOKEN_GRACE_MS) / 1000);
 
+  // SQLite evaluates every SET expression against the row's pre-update values
+  // before writing any of them. That is what makes the current→previous
+  // handoff atomic: two concurrent rotations serialize on the row's write lock,
+  // each captures its predecessor into `previous_token_hash`, and the loser of
+  // the race authenticates against the previous slot instead of being handed a
+  // silently-dead credential. Chained rotations retire the older previous slot
+  // — see the "chained rotations" case in registerOrRotateRace.test.ts.
   await db
     .update(agents)
-    .set({ tokenHash: newTokenHash })
-    .where(eq(agents.id, agent.id));
+    .set({
+      previousTokenHash: sql`${agents.tokenHash}`,
+      previousTokenExpiresAt: sql`${graceExpiresAtSeconds}`,
+      tokenHash: newTokenHash,
+    })
+    .where(eq(agents.id, existing.id));
 
   return {
-    name: agent.name,
+    name: agentName,
     token: newToken,
   };
 }