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..e371f7f5 --- /dev/null +++ b/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 { + 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); + + // 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); + }); +}); 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, }; }