diff --git a/doc/DATABASE.md b/doc/DATABASE.md index b6ab2ec8d59..207e6a20a5b 100644 --- a/doc/DATABASE.md +++ b/doc/DATABASE.md @@ -43,6 +43,23 @@ launch context does not travel with each policy or parent-report check. JSON types, task-key precedence, and company and agent identity checks are preserved. Each decision still reads current data; this is not an authorization cache. +The initial agent JWT run lookup reads only the cancellation state and +conversation flag from saved run JSON. It does not transfer the full context or +result. Run status, responsible-user identity, company and agent boundaries, +and the cancellation response remain unchanged. Later identity refreshes keep +their existing transaction and data reads. + +To measure authentication query responses on disposable synthetic data: + +```sh +pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-agent-authentication.ts --baseline-ref=06689c1b3410000bcd65ff40cb469e600209c4db +``` + +This compares the current query with the original query in that revision. It +checks the complete authentication decision and reports returned JSON bytes, +local query timings, and execution plans. It does not measure request latency +or fleet-wide savings. + To compare the real summary and full history queries with an earlier revision: ```sh diff --git a/scripts/benchmark-agent-authentication.ts b/scripts/benchmark-agent-authentication.ts new file mode 100644 index 00000000000..b50bae740e7 --- /dev/null +++ b/scripts/benchmark-agent-authentication.ts @@ -0,0 +1,85 @@ +import assert from "node:assert/strict"; +import { createHash, randomBytes, randomUUID } from "node:crypto"; +import { execFileSync } from "node:child_process"; +import { readFileSync } from "node:fs"; +import { createRequire } from "node:module"; +const { and, eq } = createRequire(new URL("../server/package.json", import.meta.url))("drizzle-orm"); +import { agents, companies, createDb, heartbeatRuns, startEmbeddedPostgresTestDatabase, type Db } from "../packages/db/src/index.js"; +import { agentRunWritesRevoked } from "../server/src/agent-run-cancellation.js"; +import { loadAgentRunAuthentication } from "../server/src/middleware/agent-run-authentication.js"; + +type Query = { sql: string; params: unknown[] }; +const baselineRef = process.argv.find(arg => arg.startsWith("--baseline-ref="))?.slice(15); +assert.ok(baselineRef, "Supply an explicit pre-fix --baseline-ref"); +const treatmentRevision = execFileSync("git", ["rev-parse", "HEAD"], { encoding: "utf8" }).trim(); +assert.equal(execFileSync("git", ["status", "--porcelain"], { encoding: "utf8" }).trim(), "", + "Commit benchmark and treatment changes before measuring"); +const baselineRevision = execFileSync("git", ["rev-parse", "--verify", `${baselineRef}^{commit}`], { encoding: "utf8" }).trim(); +const source = execFileSync("git", ["show", `${baselineRevision}:server/src/middleware/auth.ts`], { encoding: "utf8" }); +const expression = source.match(/const \[identityRun\] = await (db\.select\([\s\S]*?);\n if/)?.[1]; +assert.ok(expression, "The baseline must contain the original JWT identity query"); +const keepAlive = setInterval(() => {}, 1000); +let database: Awaited> | undefined; +let db: Db | undefined; +try { + database = await startEmbeddedPostgresTestDatabase("agent-authentication-benchmark-"); + db = createDb(database.connectionString); + const companyId = randomUUID(), agentId = randomUUID(), runId = randomUUID(); + await db.insert(companies).values({ id: companyId, name: "Synthetic authentication", issuePrefix: "AUTH" }); + await db.insert(agents).values({ id: agentId, companyId, name: "Fixture", adapterType: "process" }); + await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, status: "running", + activeIdentityContextId: randomUUID(), responsibleUserId: "fixture-user" }); + const baseline = new Function("db", "heartbeatRuns", "and", "eq", "claims", `return ${expression}`)( + db, heartbeatRuns, and, eq, { sub: agentId, company_id: companyId, run_id: runId }, + ).toSQL() as Query; + const treatment = loadAgentRunAuthentication(db, { companyId, agentId, runId }).toSQL(); + + function contract(rows: unknown[][]) { + return rows.map(([activeIdentityContextId, responsibleUserId, status, resultJson, contextSnapshot]) => ({ + activeIdentityContextId, responsibleUserId, status, + writesRevoked: agentRunWritesRevoked({ status: status as string, resultJson: resultJson as Record | null }), + conversation: (contextSnapshot as Record | null)?.conversationMode === true, + })); + } + + const results = []; + for (const contextBytes of [1024, 96 * 1024]) { + const context = { prompt: randomBytes(contextBytes * 3 / 4).toString("base64"), conversationMode: true }; + const result = { summary: randomBytes(3 * 1024).toString("base64") }; + await db.update(heartbeatRuns).set({ contextSnapshot: context, resultJson: result }).where(eq(heartbeatRuns.id, runId)); + await db.$client.unsafe("ANALYZE heartbeat_runs"); + const samples: Record = { baseline: [], treatment: [] }; + const rows: Record = {}; + for (let iteration = 0; iteration < 12; iteration++) { + const order = iteration % 2 ? ["treatment", "baseline"] : ["baseline", "treatment"]; + for (const name of order) { + const query = name === "baseline" ? baseline : treatment; + const start = performance.now(); + rows[name] = await db.$client.unsafe(query.sql, query.params as never[]).values(); + if (iteration >= 2) samples[name].push(performance.now() - start); + } + } + assert.deepEqual(contract(rows.treatment), contract(rows.baseline)); + const measurements = []; + for (const [name, query] of Object.entries({ baseline, treatment })) { + const sorted = [...samples[name]].sort((a, b) => a - b); + const [plan] = await db.$client.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${query.sql}`, query.params as never[]); + measurements.push({ name, samplesMs: samples[name], medianMs: sorted[5], + returnedJsonBytes: Buffer.byteLength(JSON.stringify(rows[name])), plan: plan["QUERY PLAN"] }); + } + results.push({ contextBytes: Buffer.byteLength(JSON.stringify(context)), + resultBytes: Buffer.byteLength(JSON.stringify(result)), contractEqual: true, measurements }); + } + const [version] = await db.$client.unsafe("SELECT version() AS version"); + const settings = await db.$client.unsafe("SELECT name, setting, unit FROM pg_settings WHERE name IN ('shared_buffers','work_mem','effective_cache_size','jit') ORDER BY name"); + console.log(JSON.stringify({ baselineRevision, treatmentRevision, node: process.version, + dependencyLockfileSha256: createHash("sha256").update(readFileSync(new URL("../pnpm-lock.yaml", import.meta.url))).digest("hex"), + postgres: version.version, settings, fixtureRuns: 1, queries: { baseline, treatment }, results }, null, 2)); +} finally { + try { + await db?.$client.end({ timeout: 0 }); + await database?.cleanup(); + } finally { + clearInterval(keepAlive); + } +} diff --git a/server/src/__tests__/agent-run-authentication.test.ts b/server/src/__tests__/agent-run-authentication.test.ts new file mode 100644 index 00000000000..fe2da4eb363 --- /dev/null +++ b/server/src/__tests__/agent-run-authentication.test.ts @@ -0,0 +1,94 @@ +import { randomBytes, randomUUID } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { and, eq } from "drizzle-orm"; +import { agents, companies, createDb, heartbeatRuns, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS } from "@paperclipai/db"; +import { agentRunWritesRevoked } from "../agent-run-cancellation.js"; +import { loadAgentRunAuthentication } from "../middleware/agent-run-authentication.js"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; + +const support = await getEmbeddedPostgresTestSupport(); +(support.supported ? describe : describe.skip)("agent run authentication projection", () => { + let database: Awaited>; + let db: ReturnType; + const companyId = randomUUID(), agentId = randomUUID(), runId = randomUUID(); + const identityContextId = randomUUID(); + const input = { companyId, agentId, runId }; + + beforeAll(async () => { + database = await startEmbeddedPostgresTestDatabase("agent-run-authentication-"); + db = createDb(database.connectionString); + await db.insert(companies).values({ id: companyId, name: "Authentication fixture", issuePrefix: "AUTH" }); + await db.insert(agents).values({ id: agentId, companyId, name: "Worker", adapterType: "process" }); + await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, status: "running", + activeIdentityContextId: identityContextId, responsibleUserId: "fixture-user" }); + }, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS); + + afterAll(async () => { + await db?.$client.end({ timeout: 0 }); + await database?.cleanup(); + }); + + async function baseline() { + const [row] = await db.select({ + activeIdentityContextId: heartbeatRuns.activeIdentityContextId, + responsibleUserId: heartbeatRuns.responsibleUserId, + status: heartbeatRuns.status, + resultJson: heartbeatRuns.resultJson, + contextSnapshot: heartbeatRuns.contextSnapshot, + }).from(heartbeatRuns).where(and(eq(heartbeatRuns.id, runId), + eq(heartbeatRuns.companyId, companyId), eq(heartbeatRuns.agentId, agentId))); + return row; + } + + function contract(row: Awaited>) { + return { + activeIdentityContextId: row.activeIdentityContextId, + responsibleUserId: row.responsibleUserId, + status: row.status, + writesRevoked: agentRunWritesRevoked(row), + conversation: row.contextSnapshot?.conversationMode === true, + }; + } + + it("preserves cancellation and conversation decisions for legacy JSON shapes", async () => { + const results = [null, "null", "false", "42", '"requested"', "[]", "{}", + '{"executionCancellation":null}', '{"executionCancellation":true}', + '{"executionCancellation":"requested"}', '{"executionCancellation":[]}', + '{"executionCancellation":{"state":"requested"}}', + '{"executionCancellation":{"state":"completed"}}', + '{"executionCancellation":{"state":true}}', + '{"executionCancellation":{"state":["requested"]}}']; + const contexts = [null, "null", '"true"', "[]", "{}", '{"conversationMode":true}', + '{"conversationMode":false}', '{"conversationMode":"true"}', '{"conversationMode":1}']; + for (const status of ["running", "cancelled"]) { + for (const result of results) { + for (const context of contexts) { + await db.$client.unsafe("UPDATE heartbeat_runs SET status=$1, result_json=$2::jsonb, context_snapshot=$3::jsonb WHERE id=$4", + [status, result, context, runId]); + const [row] = await loadAgentRunAuthentication(db, input); + expect(contract(row), JSON.stringify({ status, result, context })).toEqual(contract(await baseline())); + } + } + } + }); + + it("omits large unused payloads and reads new cancellation state on every request", async () => { + const payload = randomBytes(72 * 1024).toString("base64"); + await db.update(heartbeatRuns).set({ status: "running", resultJson: { summary: payload }, + contextSnapshot: { prompt: payload, conversationMode: true } }).where(eq(heartbeatRuns.id, runId)); + const [row] = await loadAgentRunAuthentication(db, input); + expect(contract(row)).toEqual(contract(await baseline())); + expect(row.resultJson).toEqual({ executionCancellation: { state: null } }); + expect(row.contextSnapshot).toEqual({ conversationMode: true }); + expect(Buffer.byteLength(JSON.stringify(row))).toBeLessThan(512); + await db.update(heartbeatRuns).set({ resultJson: { executionCancellation: { state: "requested" } } }) + .where(eq(heartbeatRuns.id, runId)); + expect(agentRunWritesRevoked((await loadAgentRunAuthentication(db, input))[0])).toBe(true); + }); + + it("does not return another company, agent, or run", async () => { + for (const key of ["companyId", "agentId", "runId"] as const) { + expect(await loadAgentRunAuthentication(db, { ...input, [key]: randomUUID() })).toEqual([]); + } + }); +}); diff --git a/server/src/middleware/agent-run-authentication.ts b/server/src/middleware/agent-run-authentication.ts new file mode 100644 index 00000000000..d130e9e706b --- /dev/null +++ b/server/src/middleware/agent-run-authentication.ts @@ -0,0 +1,26 @@ +import { and, eq, sql } from "drizzle-orm"; +import { heartbeatRuns, type Db } from "@paperclipai/db"; + +export function loadAgentRunAuthentication(db: Db, input: { + runId: string; + companyId: string; + agentId: string; +}) { + return db.select({ + activeIdentityContextId: heartbeatRuns.activeIdentityContextId, + responsibleUserId: heartbeatRuns.responsibleUserId, + status: heartbeatRuns.status, + resultJson: sql>`jsonb_build_object( + 'executionCancellation', jsonb_build_object( + 'state', ${heartbeatRuns.resultJson}->'executionCancellation'->'state' + ) + )`, + contextSnapshot: sql>`jsonb_build_object( + 'conversationMode', ${heartbeatRuns.contextSnapshot}->'conversationMode' + )`, + }).from(heartbeatRuns).where(and( + eq(heartbeatRuns.id, input.runId), + eq(heartbeatRuns.companyId, input.companyId), + eq(heartbeatRuns.agentId, input.agentId), + )); +} diff --git a/server/src/middleware/auth.ts b/server/src/middleware/auth.ts index 6736f68db1a..ea23e1f3946 100644 --- a/server/src/middleware/auth.ts +++ b/server/src/middleware/auth.ts @@ -22,6 +22,7 @@ import { } from "../services/issue-prefix.js"; import { verifyLocalAgentJwt } from "../agent-auth-jwt.js"; import { agentRunWritesRevoked } from "../agent-run-cancellation.js"; +import { loadAgentRunAuthentication } from "./agent-run-authentication.js"; import { isUuidLike, normalizeAgentApiKeyScope, type DeploymentMode } from "@paperclipai/shared"; import type { BetterAuthSessionResult } from "../auth/better-auth.js"; import { logger } from "./logger.js"; @@ -411,11 +412,11 @@ export function actorMiddleware(db: Db, opts: ActorMiddlewareOptions): RequestHa return; } - const [identityRun] = await db.select({ activeIdentityContextId: heartbeatRuns.activeIdentityContextId, - responsibleUserId: heartbeatRuns.responsibleUserId, status: heartbeatRuns.status, resultJson: heartbeatRuns.resultJson, - contextSnapshot: heartbeatRuns.contextSnapshot }).from(heartbeatRuns).where(and( - eq(heartbeatRuns.id, claims.run_id), eq(heartbeatRuns.companyId, claims.company_id), eq(heartbeatRuns.agentId, claims.sub), - )); + const [identityRun] = await loadAgentRunAuthentication(db, { + runId: claims.run_id, + companyId: claims.company_id, + agentId: claims.sub, + }); if (agentRunWritesRevoked(identityRun) && !["GET", "HEAD", "OPTIONS"].includes(req.method)) { const conversation = identityRun?.contextSnapshot?.conversationMode === true;