From c349a1615f7b4d25aea593b47fc0db63b68cea6e Mon Sep 17 00:00:00 2001 From: kjgbot Date: Mon, 17 Aug 2026 00:39:23 +0200 Subject: [PATCH 1/2] fix(engine): keep prior agent token authenticatable during rotate-token grace (#1542) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `POST /v1/agents/:name/rotate-token` overwrote the single `token_hash` slot, so two concurrent callers each got a 200 with their own new token, but the later rotate invalidated the earlier caller's credential mid-flight. The "loser" of the race was handed a token that had already stopped authenticating. That silent failure was reachable from every code path that runs `registerOrRotate` (SDK, MCP, and node reconnect) and looked indistinguishable from an agent going quiet. Give the agents row a two-slot outcome: rotate moves the current hash into `previous_token_hash` with a bounded grace window (60s) as one atomic UPDATE. SQLite evaluates every SET expression against the pre-update row, so two serialized rotations both preserve the credential they superseded and both callers stay authenticatable long enough to establish a persistent session. Auth accepts either the current or (previous ∧ not-yet-expired) slot, in that order, so a genuinely revoked token still 401s the moment the grace expires or the agent is released. Release paths (`deleteAgent`, dispatched release, node-completed release) clear the previous slot alongside the current `token_hash` rewrite so a released agent's grace token stops working immediately. MUST-FIRE: two concurrent rotations return distinct tokens that both authenticate against `GET /v1/agent`. MUST-NOT-FIRE: a deleted agent's last-issued token authenticates 401. Deferred, filed separately: - Broker WS re-register storm at crates/broker/src/relaycast/ws.rs:129 is the relay-repo aggravator that amplified this defect. - The SDK `registerOrRotate` shape (get + rotateToken) becomes correct with this server change; no SDK patch ships in this PR. - `registerAgentViaNode`'s ON CONFLICT DO UPDATE clobbers `token_hash` under the same shape; it needs the same dual-slot treatment as a follow-up. Session-Id: 6560d879-2098-406f-82cf-4bc2365ca27d --- .../conformance/registerOrRotateRace.test.ts | 102 ++++++++++++++++++ packages/engine/src/auth/index.ts | 16 ++- .../db/migrations/0035_agent_token_grace.sql | 18 ++++ packages/engine/src/db/schema.ts | 5 + packages/engine/src/engine/action.ts | 9 ++ packages/engine/src/engine/agent.ts | 6 ++ packages/engine/src/engine/tokenRotate.ts | 38 +++++-- 7 files changed, 185 insertions(+), 9 deletions(-) create mode 100644 packages/engine/src/__tests__/conformance/registerOrRotateRace.test.ts create mode 100644 packages/engine/src/db/migrations/0035_agent_token_grace.sql diff --git a/packages/engine/src/__tests__/conformance/registerOrRotateRace.test.ts b/packages/engine/src/__tests__/conformance/registerOrRotateRace.test.ts new file mode 100644 index 00000000..5a4000b6 --- /dev/null +++ b/packages/engine/src/__tests__/conformance/registerOrRotateRace.test.ts @@ -0,0 +1,102 @@ +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 { + 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 { + const body = await res.json() as { data?: { token?: string } }; + return body.data?.token ?? ''; + } + + async function authenticate(agentToken: string): Promise { + 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); + + 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 genuinely revoked identity is not rescued by any grace slot. + const auth = await authenticate(initial.token); + expect(auth.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); + }); +}); diff --git a/packages/engine/src/auth/index.ts b/packages/engine/src/auth/index.ts index 26a0ee96..225c70e8 100644 --- a/packages/engine/src/auth/index.ts +++ b/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'); diff --git a/packages/engine/src/db/migrations/0035_agent_token_grace.sql b/packages/engine/src/db/migrations/0035_agent_token_grace.sql new file mode 100644 index 00000000..208f3725 --- /dev/null +++ b/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); diff --git a/packages/engine/src/db/schema.ts b/packages/engine/src/db/schema.ts index 6aaea9e1..718d301c 100644 --- a/packages/engine/src/db/schema.ts +++ b/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), ], ); diff --git a/packages/engine/src/engine/action.ts b/packages/engine/src/engine/action.ts index 59ebb506..12036da6 100644 --- a/packages/engine/src/engine/action.ts +++ b/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(), diff --git a/packages/engine/src/engine/agent.ts b/packages/engine/src/engine/agent.ts index b9fe652e..531d9ddc 100644 --- a/packages/engine/src/engine/agent.ts +++ b/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, diff --git a/packages/engine/src/engine/tokenRotate.ts b/packages/engine/src/engine/tokenRotate.ts index ef55c270..3e42da13 100644 --- a/packages/engine/src/engine/tokenRotate.ts +++ b/packages/engine/src/engine/tokenRotate.ts @@ -1,4 +1,4 @@ -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'; @@ -6,26 +6,50 @@ import { codedError } from '../lib/httpError.js'; type Db = ReturnType; +/** + * 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, }; } From 448047b981a15e91741d05ecc478654e24fd4b6e Mon Sep 17 00:00:00 2001 From: kjgbot Date: Mon, 17 Aug 2026 05:25:39 +0200 Subject: [PATCH 2/2] test(engine): make revoked-token must-not-fire actually populate the grace slot MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The must-not-fire in registerOrRotateRace previously deleted an agent whose `previous_token_hash` had never been populated, so the 401 passed trivially even for an implementation that never cleared the grace slot on delete. Rotate once before delete, then assert BOTH the current and the now-previous (grace-window) tokens return 401. Verified locally by removing the `previousTokenHash: null` / `previousTokenExpiresAt: null` writes in deleteAgent — the new assertion goes red — then restoring and reconfirming green. Session-Id: 66dc7bf3-e321-4b55-b65a-9276081ad12c --- .../conformance/registerOrRotateRace.test.ts | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/packages/engine/src/__tests__/conformance/registerOrRotateRace.test.ts b/packages/engine/src/__tests__/conformance/registerOrRotateRace.test.ts index 5a4000b6..e371f7f5 100644 --- a/packages/engine/src/__tests__/conformance/registerOrRotateRace.test.ts +++ b/packages/engine/src/__tests__/conformance/registerOrRotateRace.test.ts @@ -73,15 +73,29 @@ describe('registerOrRotate concurrency', () => { 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 genuinely revoked identity is not rescued by any grace slot. - const auth = await authenticate(initial.token); - expect(auth.status).toBe(401); + // 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 () => {