diff --git a/e2e/agent-token.drizzle.test.ts b/e2e/agent-token.drizzle.test.ts index 07f16f1..1d75a4b 100644 --- a/e2e/agent-token.drizzle.test.ts +++ b/e2e/agent-token.drizzle.test.ts @@ -14,10 +14,12 @@ import { runAgentTokenMigrations } from "../src/migrations"; import { hashAgentToken, mintAgentToken, + mintRunScopedToken, revokeAgentToken, verifyAgentToken, } from "../src/tokens"; import { mountAgentTokens } from "../src/mount"; +import { mountRunScopedTokens } from "../src/mount-run-scoped"; import { createAgentTokenVerifier, requireAgentToken } from "../src/middleware"; import type { WorkflowRunScopeEnv } from "../src/workflow-run-scope"; @@ -128,6 +130,7 @@ describeIfDb("agent tokens", () => { id: minted.id, tenantId, definitionId: "def_artifacts", + runId: null, }); expect(await verifyAgentToken(db, "not-a-token")).toBeUndefined(); @@ -330,6 +333,7 @@ describeIfDb("agent tokens", () => { id: minted.id, tenantId, definitionId: "def_artifacts", + runId: null, }); expect( await identityFor(`Bearer ${minted.token}`, otherTenantId), @@ -340,6 +344,7 @@ describeIfDb("agent tokens", () => { id: minted.id, tenantId, definitionId: "def_artifacts", + runId: null, }); expect(await identityFor(`Basic ${minted.token}`)).toBeNull(); expect(await identityFor("Bearer")).toBeNull(); @@ -366,4 +371,172 @@ describeIfDb("agent tokens", () => { await close(); } }); + + test("run-scoped mint generates the token server-side and replaces on re-register", async () => { + const { db, close } = createDB({ ...target, schema: SCHEMA }); + try { + const tenantId = `tnt_run_${randomUUID().slice(0, 8)}`; + await seedTenant(db, tenantId); + + const first = await mintRunScopedToken(db, { + tenantId, + definitionId: "def_artifacts", + name: "workflow", + runId: "run_1", + }); + expect(first.token.length).toBeGreaterThan(32); + expect(first.runId).toBe("run_1"); + + const verified = await verifyAgentToken(db, first.token); + expect(verified).toEqual({ + id: first.id, + tenantId, + definitionId: "def_artifacts", + runId: "run_1", + }); + + // Re-registering for the same run replaces the old token. + const second = await mintRunScopedToken(db, { + tenantId, + definitionId: "def_artifacts", + name: "workflow", + runId: "run_1", + }); + expect(second.id).not.toBe(first.id); + expect(await verifyAgentToken(db, first.token)).toBeUndefined(); + expect(await verifyAgentToken(db, second.token)).toEqual({ + id: second.id, + tenantId, + definitionId: "def_artifacts", + runId: "run_1", + }); + } finally { + await close(); + } + }); + + test("the run-scoped mount enforces the grant, validates the run, and server-generates the token", async () => { + const { db, close } = createDB({ ...target, schema: SCHEMA }); + try { + const tenantId = `tnt_rsc_${randomUUID().slice(0, 8)}`; + await seedTenant(db, tenantId); + + const runs: Record< + string, + { definitionId: string; runId: string; tenantId: string } + > = { + "runs/abc": { + definitionId: "def_artifacts", + runId: "run_abc", + tenantId, + }, + }; + + function runScopedApp(gate: "allow" | "deny" = "allow") { + const app = new Hono(); + app.use(tenantFrom(db, tenantId)); + mountRunScopedTokens(app, { + db, + requireGrant: async (c, next) => { + if (gate === "deny") return c.json({ error: "forbidden" }, 403); + await next(); + return undefined; + }, + resolveRun: (addr) => runs[addr] ?? null, + }); + return app; + } + + // Without the grant → 403. + const denied = await runScopedApp("deny").request( + "/workflow-run-tokens", + { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ runAddress: "runs/abc" }), + }, + ); + expect(denied.status).toBe(403); + + // Unknown run → 404. + const app = runScopedApp(); + const unknown = await app.request("/workflow-run-tokens", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ runAddress: "runs/missing" }), + }); + expect(unknown.status).toBe(404); + + // Happy path → 201, server-generated token. + const ok = await app.request("/workflow-run-tokens", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ runAddress: "runs/abc" }), + }); + expect(ok.status).toBe(201); + const body = MintResponse.assert(await ok.json()); + expect(body.token.token.length).toBeGreaterThan(32); + + // The token verifies with the correct run. + const identity = await verifyAgentToken(db, body.token.token); + expect(identity).toEqual({ + id: body.token.id, + tenantId, + definitionId: "def_artifacts", + runId: "run_abc", + }); + + // Re-registering replaces the old token. + const second = await app.request("/workflow-run-tokens", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ runAddress: "runs/abc" }), + }); + expect(second.status).toBe(201); + const secondBody = MintResponse.assert(await second.json()); + expect(await verifyAgentToken(db, body.token.token)).toBeUndefined(); + expect(await verifyAgentToken(db, secondBody.token.token)).toEqual({ + id: secondBody.token.id, + tenantId, + definitionId: "def_artifacts", + runId: "run_abc", + }); + } finally { + await close(); + } + }); + + test("the run-scoped mount rejects a run from another tenant", async () => { + const { db, close } = createDB({ ...target, schema: SCHEMA }); + try { + const tenantId = `tnt_cross_${randomUUID().slice(0, 8)}`; + const otherTenantId = `tnt_cross_other_${randomUUID().slice(0, 8)}`; + await seedTenant(db, tenantId); + await seedTenant(db, otherTenantId); + + const app = new Hono(); + app.use(tenantFrom(db, tenantId)); + mountRunScopedTokens(app, { + db, + requireGrant: async (_c, next) => { + await next(); + return undefined; + }, + resolveRun: () => ({ + tenantId: otherTenantId, + definitionId: "def_artifacts", + runId: "run_other", + }), + }); + + const res = await app.request("/workflow-run-tokens", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ runAddress: "runs/cross" }), + }); + expect(res.status).toBe(404); + } finally { + await close(); + } + }); }); diff --git a/e2e/upgrade-from-0.1.0.test.ts b/e2e/upgrade-from-0.1.0.test.ts index 4a7c92c..df499d5 100644 --- a/e2e/upgrade-from-0.1.0.test.ts +++ b/e2e/upgrade-from-0.1.0.test.ts @@ -85,6 +85,7 @@ describeIfDb("upgrading a 0.1.0 database", () => { id: live.id, tenantId: TENANT_ID, definitionId: "def_artifacts", + runId: null, }); expect(await verifyAgentToken(db, revoked.token)).toBeUndefined(); const foreignKeys = await db.execute(sql` diff --git a/migrations/0002_run_scoped_token.sql b/migrations/0002_run_scoped_token.sql new file mode 100644 index 0000000..f8bf2db --- /dev/null +++ b/migrations/0002_run_scoped_token.sql @@ -0,0 +1,5 @@ +ALTER TABLE "agent_token"."token" ADD COLUMN IF NOT EXISTS "run_id" text; +--> statement-breakpoint +CREATE UNIQUE INDEX IF NOT EXISTS "agent_token_token_run_id_active_idx" + ON "agent_token"."token" ("tenant_id", "run_id") + WHERE "revoked_at" IS NULL AND "run_id" IS NOT NULL; diff --git a/src/index.ts b/src/index.ts index 088b37e..5bd0ac8 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,5 +1,6 @@ export { mintAgentToken, + mintRunScopedToken, revokeAgentToken, verifyAgentToken, type AgentTokenDb, @@ -11,6 +12,12 @@ export { type MountAgentTokensOpts, type ResolveDefinition, } from "./mount.js"; +export { + mountRunScopedTokens, + type MountRunScopedTokensOpts, + type ResolveRunForToken, + type RunScopeForToken, +} from "./mount-run-scoped.js"; export { createAgentTokenVerifier, requireAgentToken, diff --git a/src/middleware.ts b/src/middleware.ts index 44866bb..2b037ee 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -14,6 +14,7 @@ export type AgentTokenContext = { tenantId: string; definitionId: string; id: string; + runId: string | null; }; /** The context variable this middleware sets, for a host to widen its Env @@ -45,6 +46,7 @@ async function verifyAuthorization>( id: identity.id, tenantId: identity.tenantId, definitionId: identity.definitionId, + runId: identity.runId, }; } diff --git a/src/mount-run-scoped.ts b/src/mount-run-scoped.ts new file mode 100644 index 0000000..9198d4d --- /dev/null +++ b/src/mount-run-scoped.ts @@ -0,0 +1,75 @@ +// Mint a run-scoped agent token. The server generates the token; the caller +// only names the run. Re-registering for the same run revokes the old token. +// Every route runs the host's `requireGrant`. +import { type } from "arktype"; +import type { TenantEnv } from "@intx/hub-api"; +import type { Context, Hono, MiddlewareHandler } from "hono"; + +import { mintRunScopedToken, type AgentTokenDb } from "./tokens.js"; + +/** What the host's run resolver returns: which tenant and definition a run + * address resolves to. */ +export type RunScopeForToken = { + tenantId: string; + definitionId: string; + runId: string; +}; + +/** `(runAddress) => scope | null`: the host's lookup from the request's run + * address to the run it names. Returns `null` for unknown or inaccessible + * runs. */ +export type ResolveRunForToken = ( + runAddress: string, +) => Promise | RunScopeForToken | null; + +export type MountRunScopedTokensOpts< + E extends TenantEnv, + TSchema extends Record, +> = { + db: AgentTokenDb; + /** + * The host's own authority check, run as middleware on every route. The + * host gates this the same way it gates credential creation. + */ + requireGrant: MiddlewareHandler; + /** Maps a run address to the tenant, definition, and run ID it names. + * Returns `null` for unknown or inaccessible runs. */ + resolveRun: ResolveRunForToken; +}; + +const RunMintBody = type({ + runAddress: "0 < string <= 256", + "name?": "0 < string <= 256", +}); + +/** Mount `POST /workflow-run-tokens` onto the host's app. */ +export function mountRunScopedTokens< + E extends TenantEnv, + TSchema extends Record, +>(app: Hono, opts: MountRunScopedTokensOpts): Hono { + const { db, requireGrant, resolveRun } = opts; + + app.post("/workflow-run-tokens", requireGrant, async (c: Context) => { + const tenantId = c.get("tenant").id; + const parsed = RunMintBody(await c.req.json().catch(() => undefined)); + if (parsed instanceof type.errors) { + return c.json({ error: "invalid_body", detail: parsed.summary }, 400); + } + + const scope = await resolveRun(parsed.runAddress); + if (scope === null || scope.tenantId !== tenantId) { + return c.json({ error: "not_found" }, 404); + } + + const minted = await mintRunScopedToken(db, { + tenantId, + definitionId: scope.definitionId, + name: parsed.name ?? "workflow", + runId: scope.runId, + }); + + return c.json({ token: minted }, 201); + }); + + return app; +} diff --git a/src/schema.ts b/src/schema.ts index 6fd8e26..bc2be01 100644 --- a/src/schema.ts +++ b/src/schema.ts @@ -14,6 +14,7 @@ export const agentTokenTable = agentTokenSchema.table("token", { .references(() => hostTenant.id, { onDelete: "cascade" }), definitionId: text("definition_id").notNull(), name: text("name").notNull(), + runId: text("run_id"), // Only the digest is stored; the plaintext is returned once at mint time. tokenHash: text("token_hash").notNull().unique(), createdAt: timestamp("created_at", { withTimezone: true }) diff --git a/src/tokens.ts b/src/tokens.ts index 889666b..228e0ac 100644 --- a/src/tokens.ts +++ b/src/tokens.ts @@ -19,6 +19,7 @@ export type AgentTokenIdentity = { id: string; tenantId: string; definitionId: string; + runId: string | null; }; export type MintedAgentToken = AgentTokenIdentity & { @@ -67,6 +68,7 @@ export async function mintAgentToken>( id: row.id, tenantId: row.tenantId, definitionId: row.definitionId, + runId: row.runId, name: row.name, createdAt: row.createdAt, token, @@ -91,7 +93,12 @@ export async function verifyAgentToken>( .limit(1); if (row === undefined) return undefined; if (!agentTokenHashEquals(row.tokenHash, presented)) return undefined; - return { id: row.id, tenantId: row.tenantId, definitionId: row.definitionId }; + return { + id: row.id, + tenantId: row.tenantId, + definitionId: row.definitionId, + runId: row.runId, + }; } export async function revokeAgentToken>( @@ -112,6 +119,60 @@ export async function revokeAgentToken>( return row !== undefined; } +/** Revokes any active token for `runId` on the same tenant, then mints a + * replacement. The caller never chooses the token value. The revoke and + * insert run inside a transaction so concurrent calls for the same run + * cannot interleave past the partial unique index. */ +export async function mintRunScopedToken< + TSchema extends Record, +>( + db: AgentTokenDb, + input: { + tenantId: string; + definitionId: string; + name: string; + runId: string; + }, +): Promise { + const token = generateAgentToken(); + const row = await db.transaction(async (tx) => { + const now = new Date(); + await tx + .update(agentTokenTable) + .set({ revokedAt: now }) + .where( + and( + eq(agentTokenTable.tenantId, input.tenantId), + eq(agentTokenTable.runId, input.runId), + isNull(agentTokenTable.revokedAt), + ), + ); + const [inserted] = await tx + .insert(agentTokenTable) + .values({ + id: randomUUID(), + tenantId: input.tenantId, + definitionId: input.definitionId, + name: input.name, + runId: input.runId, + tokenHash: hashAgentToken(token), + }) + .returning(); + if (inserted === undefined) + throw new Error("agent token insert returned no row"); + return inserted; + }); + return { + id: row.id, + tenantId: row.tenantId, + definitionId: row.definitionId, + runId: row.runId, + name: row.name, + createdAt: row.createdAt, + token, + }; +} + /** Reads a `Bearer` credential out of an `Authorization` header value. */ export function bearerFromAuthorization( header: string | undefined,