diff --git a/scripts/benchmark-chat-bound-authorization.ts b/scripts/benchmark-chat-bound-authorization.ts new file mode 100644 index 00000000000..80ad0be035f --- /dev/null +++ b/scripts/benchmark-chat-bound-authorization.ts @@ -0,0 +1,228 @@ +/** + * Local proof for the batched chat authorization of external review waits. + * + * `listExternalReviewWaits` authorizes the chat conversation of every waiting + * task on an issue page. The pre-fix module ran one bound-run authorization + * per wait at the same time, so a page with N waits sent N copies of the + * message-link statement (and N resource, principal, identity-link and + * membership reads) on N pool connections at once. The fix reads the + * messages of all waits with one statement and shares the policy reads. + * + * The pre-fix modules are loaded from --baseline-ref. The script fails unless + * both versions return the same waits. + * + * pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-chat-bound-authorization.ts --baseline-ref= + * + * Optional sizes: WAITS (default 20), PEOPLE (5), NOISE_DELIVERIES (50000), + * EVENT_KIB (6), POOL (10), PAIRS (15), CONCURRENT_PAGES (8). + */ +import assert from "node:assert/strict"; +import { execFileSync } from "node:child_process"; +import { promises as fs } from "node:fs"; +import path from "node:path"; +import { fileURLToPath, pathToFileURL } from "node:url"; +import { createDb, startEmbeddedPostgresTestDatabase, type Db } from "../packages/db/src/index.js"; +import { listExternalReviewWaits } from "../server/src/services/native-runtime/external-review-wait.js"; +import { seedBoundChatAuthorization } from "../server/src/__tests__/helpers/bound-chat-authorization.js"; + +const size = (name: string, fallback: number) => Number(process.env[name] ?? fallback); +const WAITS = size("WAITS", 20), PEOPLE = size("PEOPLE", 5), NOISE_DELIVERIES = size("NOISE_DELIVERIES", 50_000); +const EVENT_KIB = size("EVENT_KIB", 6), POOL = size("POOL", 10), PAIRS = size("PAIRS", 15); +const CONCURRENT_PAGES = size("CONCURRENT_PAGES", 8); +for (const [name, value] of Object.entries({ WAITS, PEOPLE, NOISE_DELIVERIES, EVENT_KIB, POOL, PAIRS, CONCURRENT_PAGES })) { + assert.ok(Number.isSafeInteger(value) && value > 0, `${name} must be a positive integer`); +} +const baselineRef = process.argv.find(arg => arg.startsWith("--baseline-ref="))?.slice(15); +assert.ok(baselineRef, "Pass an explicit pre-fix --baseline-ref="); +const git = (...args: string[]) => execFileSync("git", args, { encoding: "utf8" }); +const baselineRevision = git("rev-parse", "--verify", `${baselineRef}^{commit}`).trim(); +const root = fileURLToPath(new URL("../", import.meta.url)); +const modules = ["external-review-wait.ts", "chat-attachment-reuse.ts"]; +const moduleDir = path.join(root, "server/src/services/native-runtime"); + +type List = typeof listExternalReviewWaits; +const keepAlive = setInterval(() => {}, 1_000); +let temporary: Awaited> | undefined; +let db: Db | undefined; +let baselineDir: string | undefined; +try { + // Pre-fix copies of both modules. Their other relative imports resolve to + // the current tree, which this change does not touch. + baselineDir = await fs.mkdtemp(path.join(root, "server/scripts/.chat-authorization-baseline-")); + for (const file of modules) { + const source = git("show", `${baselineRevision}:server/src/services/native-runtime/${file}`) + .replace(/from "(\.{1,2}\/[^"]+)"/g, (_match, specifier: string) => + specifier === "./chat-attachment-reuse.js" ? `from "./chat-attachment-reuse.ts"` + : `from ${JSON.stringify(fileURLToPath(new URL(specifier, pathToFileURL(path.join(moduleDir, file)))))}`); + await fs.writeFile(path.join(baselineDir, file), source); + } + const baseline = (await import(pathToFileURL(path.join(baselineDir, modules[0]!)).href) as + { listExternalReviewWaits: List }).listExternalReviewWaits; + + temporary = await startEmbeddedPostgresTestDatabase("chat-authorization-benchmark-"); + db = createDb(temporary.connectionString, { maxConnections: POOL }); + const sql = db.$client; + const f = await seedBoundChatAuthorization(db); + const loadStart = performance.now(); + // People who write in the waiting conversations, each linked and a member. + await sql.unsafe(`insert into chat_external_principals (id, company_id, provider, provider_account_id, external_id, kind) + select gen_random_uuid(), $1::uuid, 'slack', 'guild-1', 'person-' || g, 'user' from generate_series(1, $2::int) g`, + [f.companyId, PEOPLE]); + await sql.unsafe(`insert into chat_identity_links (company_id, endpoint_id, principal_id, paperclip_user_id, status) + select $1::uuid, $2::uuid, id, 'user-' || external_id, 'linked' from chat_external_principals + where company_id = $1::uuid and external_id like 'person-%'`, [f.companyId, f.endpointId]); + await sql.unsafe(`insert into company_memberships (company_id, principal_type, principal_id, status, membership_role) + select $1::uuid, 'user', 'user-' || external_id, 'active', 'member' from chat_external_principals + where company_id = $1::uuid and external_id like 'person-%'`, [f.companyId]); + // WAITS tasks in review, each waiting for an external chat response. + // Skip the work-assessment owner chain of the status decisions. + await sql.begin(tx => tx.unsafe("set local session_replication_role = replica").then(() => tx.unsafe(`with w as ( + select g, gen_random_uuid() issue_id, gen_random_uuid() run_id, gen_random_uuid() decision_id, + gen_random_uuid() conversation_id, gen_random_uuid() comment_id, gen_random_uuid() delivery_id, + (select id from chat_external_principals where company_id = $1::uuid and external_id = 'person-' || (1 + g % $5::int)) principal_id + from generate_series(1, $4::int) g), + i as (insert into issues (id, company_id, title, status, assignee_agent_id, issue_number, identifier, last_status_decision_id) + select issue_id, $1::uuid, 'wait ' || g, 'in_review', $2::uuid, 100 + g, 'RCB-' || (100 + g), decision_id from w), + r as (insert into heartbeat_runs (id, company_id, agent_id, status, native_issue_id, context_snapshot) + select run_id, $1::uuid, $2::uuid, 'succeeded', issue_id, jsonb_build_object('source', 'chat:slack', + 'paperclipHarnessCheckedOut', true, 'wakeCommentIds', jsonb_build_array(comment_id), + 'prompt', repeat('x', 4096)) from w), + d as (insert into status_decisions (id, company_id, issue_id, run_id, assessment_id, decision_version, policy_version, + from_status, to_status, reason_code, decision_json, decision_digest, application_state, applied_at) + select decision_id, $1::uuid, issue_id, run_id, gen_random_uuid(), 1, 'v1', 'in_progress', 'in_review', + 'external_chat_response_waiting', '{}'::jsonb, md5(g::text), 'applied', now() from w), + c as (insert into chat_conversations (id, company_id, endpoint_id, resource_id, issue_id, external_conversation_id, + external_thread_id, external_label, state) + select conversation_id, $1::uuid, $3::uuid, $6::uuid, issue_id, 'channel-1', 'wait-' || g, 'wait', 'waiting' from w), + m as (insert into issue_comments (id, company_id, issue_id, author_type, author_user_id, body) + select comment_id, $1::uuid, issue_id, 'user', 'u', 'Please check' from w), + l as (insert into chat_deliveries (id, company_id, endpoint_id, conversation_id, principal_id, provider_event_id, + deduplication_key, event_kind, state, normalized_event, processed_at) + select delivery_id, $1::uuid, $3::uuid, conversation_id, principal_id, 'wait-' || g, 'wait-' || g, 'message', + 'processed', jsonb_build_object('message', jsonb_build_object('text', repeat('y', $7::int * 1024))), now() from w) + insert into chat_message_links (company_id, endpoint_id, conversation_id, delivery_id, comment_id, provider_message_id, direction) + select $1::uuid, $3::uuid, conversation_id, delivery_id, comment_id, 'wait-' || g, 'inbound' from w`, + [f.companyId, f.agentId, f.endpointId, WAITS, PEOPLE, f.resourceId, EVENT_KIB]))); + // Chat history of the same endpoint on other tasks. + await sql.unsafe(`insert into issues (company_id, title, status, assignee_agent_id, issue_number, identifier) + select $1::uuid, 'history ' || g, 'done', $2::uuid, 10000 + g, 'RCB-' || (10000 + g) from generate_series(1, 2000) g`, + [f.companyId, f.agentId]); + await sql.unsafe(`insert into chat_conversations (company_id, endpoint_id, resource_id, issue_id, external_conversation_id, + external_thread_id, external_label, state) + select $1::uuid, $2::uuid, $3::uuid, i.id, 'channel-1', 'history-' || i.id, 'history', 'completed' + from issues i where i.company_id = $1::uuid and i.title like 'history %'`, [f.companyId, f.endpointId, f.resourceId]); + await sql.unsafe(`with c as (select id, issue_id, row_number() over (order by id) rn, count(*) over () n + from chat_conversations where company_id = $1::uuid and external_label = 'history') + insert into chat_deliveries (company_id, endpoint_id, conversation_id, principal_id, provider_event_id, + deduplication_key, event_kind, state, normalized_event, processed_at) + select $1::uuid, $2::uuid, c.id, $3::uuid, 'history-' || g, 'history-' || g, 'message', 'processed', + jsonb_build_object('message', jsonb_build_object('text', (select string_agg(md5(random()::text || s.i || g), ' ') + from generate_series(1, $4::int) s(i)))), now() + from generate_series(1, $5::int) g join c on c.rn = 1 + (g % c.n)`, + [f.companyId, f.endpointId, f.principalId, EVENT_KIB * 31, NOISE_DELIVERIES]); + await sql.unsafe(`insert into issue_comments (id, company_id, issue_id, author_type, author_user_id, body) + select d.id, $1::uuid, c.issue_id, 'user', 'u', 'history' from chat_deliveries d + join chat_conversations c on c.id = d.conversation_id where d.provider_event_id like 'history-%'`, [f.companyId]); + await sql.unsafe(`insert into chat_message_links (company_id, endpoint_id, conversation_id, delivery_id, comment_id, + provider_message_id, direction) + select $1::uuid, endpoint_id, conversation_id, id, id, provider_event_id, 'inbound' + from chat_deliveries where provider_event_id like 'history-%'`, [f.companyId]); + await sql.unsafe("vacuum analyze"); + const loadMs = performance.now() - loadStart; + const issueIds = (await sql.unsafe(`select id from issues where company_id = $1::uuid and title like 'wait %' order by title`, + [f.companyId])).map(row => row.id as string); + // One waiting task that is no longer authorized: its person lost membership. + await sql.unsafe(`update company_memberships set status = 'suspended' where company_id = $1::uuid and principal_id = 'user-person-1'`, + [f.companyId]); + + const round = (value: number) => Number(value.toFixed(2)); + // Statements that one page sends, counted at the driver (prepared + // templates are counted at each execution). + const statementCounts = async (list: List) => { + const client = sql as unknown as { unsafe: (query: string, ...rest: unknown[]) => unknown }; + const original = client.unsafe; + const counts: Record = {}; + let total = 0; + const links: Array<{ query: string; params: unknown }> = []; + client.unsafe = function (this: unknown, query: string, ...rest: unknown[]) { + const table = /^select .*? from "([a-z_]+)"/s.exec(query)?.[1] ?? "other"; + if (table === "chat_message_links") links.push({ query, params: rest[0] }); + counts[table] = (counts[table] ?? 0) + 1; + total += 1; + return original.call(this, query, ...rest); + }; + try { await list(db!, f.companyId, issueIds); } finally { client.unsafe = original; } + // Server execution of the message-link statements of one page. + let linkExecutionMs = 0; + const nodes = new Set(); + for (const statement of links) { + const [row] = await sql.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${statement.query}`, statement.params as never[]); + const [plan] = row!["QUERY PLAN"] as Array<{ Plan: Record; "Execution Time": number }>; + linkExecutionMs += plan!["Execution Time"]; + const walk = (node: Record) => { + if (/Scan/.test(String(node["Node Type"]))) nodes.add(`${node["Node Type"]} ${node["Index Name"] ?? ""} on ${node["Relation Name"]}`.replace(" ", " ")); + for (const child of (node.Plans as Array> | undefined) ?? []) walk(child); + }; + walk(plan!.Plan); + } + return { total, byTable: counts, linkStatements: { count: links.length, + summedExecutionMs: round(linkExecutionMs), scans: [...nodes].sort() } }; + }; + + const stats = (samples: number[]) => { + const sorted = [...samples].sort((a, b) => a - b); + return { medianMs: round(sorted[Math.floor(sorted.length / 2)]!), minMs: round(sorted[0]!), + maxMs: round(sorted.at(-1)!), samplesMs: samples.map(round) }; + }; + const time = async (run: () => Promise) => { + const start = performance.now(); + await run(); + return performance.now() - start; + }; + const versions: Record = { baseline, treatment: listExternalReviewWaits }; + const page = (list: List) => () => list(db!, f.companyId, issueIds); + const pages = (list: List) => () => Promise.all(Array.from({ length: CONCURRENT_PAGES }, page(list))); + + const results = Object.fromEntries(await Promise.all(Object.entries(versions).map(async ([name, list]) => + [name, (await list(db!, f.companyId, issueIds)).map(wait => wait.issueId).sort()]))); + assert.deepEqual(results.treatment, results.baseline, "Both versions must return the same waits"); + // Every size works: count the waits of the suspended person from the data. + const [{ denied }] = await sql.unsafe(`select count(*)::int denied from chat_deliveries d + join chat_external_principals p on p.id = d.principal_id + where d.company_id = $1::uuid and d.provider_event_id like 'wait-%' and p.external_id = 'person-1'`, [f.companyId]); + assert.equal((results.baseline as string[]).length, WAITS - (denied as number), + "Waits of the suspended person must be excluded"); + + for (const list of Object.values(versions)) { await page(list)(); await pages(list)(); } + const samples: Record = { + baseline: { page: [], concurrentPages: [] }, treatment: { page: [], concurrentPages: [] } }; + for (let pair = 0; pair < PAIRS; pair++) { + const order = pair % 2 === 0 ? ["baseline", "treatment"] : ["treatment", "baseline"]; + for (const name of order) samples[name]!.page.push(await time(page(versions[name]!))); + for (const name of order) samples[name]!.concurrentPages.push(await time(pages(versions[name]!))); + } + const report: Record = {}; + for (const [name, list] of Object.entries(versions)) { + report[name] = { + statements: await statementCounts(list), + page: stats(samples[name]!.page), + concurrentPages: stats(samples[name]!.concurrentPages), + }; + } + console.log(JSON.stringify({ + baselineRevision, treatmentRevision: git("rev-parse", "HEAD").trim(), + postgresVersion: (await sql.unsafe("select version()"))[0]!.version, + fixture: { synthetic: true, waits: WAITS, people: PEOPLE, noiseDeliveries: NOISE_DELIVERIES, eventKiB: EVENT_KIB, + pool: POOL, concurrentPages: CONCURRENT_PAGES, pairs: PAIRS, loadMs: Math.round(loadMs), + tableSizes: await sql.unsafe(`select relname as table, n_live_tup::int as rows, pg_size_pretty(pg_total_relation_size(relid)) as size + from pg_stat_user_tables where relname in ('chat_message_links', 'chat_deliveries', 'chat_conversations') order by relname`) }, + correctness: { sameWaits: true, waits: (results.baseline as string[]).length, deniedWaits: WAITS - (results.baseline as string[]).length }, + results: report, + }, null, 2)); +} finally { + try { + await db?.$client.end({ timeout: 0 }); + await temporary?.cleanup(); + if (baselineDir) await fs.rm(baselineDir, { recursive: true, force: true }); + } finally { clearInterval(keepAlive); } +} diff --git a/server/src/__tests__/batched-chat-authorization.integration.test.ts b/server/src/__tests__/batched-chat-authorization.integration.test.ts new file mode 100644 index 00000000000..834f6ce4315 --- /dev/null +++ b/server/src/__tests__/batched-chat-authorization.integration.test.ts @@ -0,0 +1,194 @@ +import { randomUUID } from "node:crypto"; +import { and, eq, sql } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { chatConversations, chatDeliveries, chatEndpointResources, chatEndpoints, chatExternalPrincipals, + chatIdentityLinks, chatMessageLinks, companyMemberships, createDb, issueComments, type Db } from "@paperclipai/db"; +import { seedBoundChatAuthorization } from "./helpers/bound-chat-authorization.js"; +import { startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; +import { authorizeChatConversationForBoundRun, authorizeChatConversationsForBoundRuns, + isBoundChatAuthorizationDenial, type ChatReuseBinding } from "../services/native-runtime/chat-attachment-reuse.js"; + +type Run = { binding: ChatReuseBinding; contextSnapshot: unknown }; + +describe("batched bound-chat authorization", () => { + let temporary: Awaited>; + let db: Db; + let f: Awaited>; + // A second conversation of the same company and endpoint, on another issue, + // from another person who has no identity link. + let other: { conversationId: string; commentId: string; principalId: string }; + + // The outcome of one separate read-mode call: the conversation, null for a + // policy denial, or the thrown error. + const single = (run: Run) => authorizeChatConversationForBoundRun(db, run.binding, run.contextSnapshot, "read") + .catch((error: unknown) => { + if (isBoundChatAuthorizationDenial(error)) return null; + throw error; + }); + const expectSameOutcomes = async (runs: Run[]) => { + const separate = await Promise.all(runs.map(single)); + const batched = await authorizeChatConversationsForBoundRuns(db, f.companyId, runs); + expect(batched).toEqual(separate); + return batched; + }; + const run = (overrides: { binding?: Partial; context?: Record } = {}): Run => ({ + binding: { ...f.binding, ...overrides.binding }, + contextSnapshot: { ...f.context, ...overrides.context }, + }); + const otherRun = (context: Record = {}) => + run({ binding: { issueId: f.otherIssueId, runId: randomUUID() }, context: { wakeCommentIds: [other.commentId], ...context } }); + + beforeAll(async () => { + temporary = await startEmbeddedPostgresTestDatabase("batched-chat-"); + db = createDb(temporary.connectionString); + f = await seedBoundChatAuthorization(db); + other = { conversationId: randomUUID(), commentId: randomUUID(), principalId: randomUUID() }; + const deliveryId = randomUUID(); + await db.insert(chatConversations).values({ id: other.conversationId, companyId: f.companyId, + endpointId: f.endpointId, resourceId: f.resourceId, issueId: f.otherIssueId, + externalConversationId: "channel-1", externalThreadId: "thread-2", externalLabel: "Other thread", state: "waiting" }); + await db.insert(chatExternalPrincipals).values({ id: other.principalId, companyId: f.companyId, + provider: "slack", providerAccountId: "guild-1", externalId: "user-2", kind: "user" }); + await db.insert(issueComments).values({ id: other.commentId, companyId: f.companyId, issueId: f.otherIssueId, + authorType: "user", authorUserId: "other-user", body: "Another request" }); + await db.insert(chatDeliveries).values({ id: deliveryId, companyId: f.companyId, endpointId: f.endpointId, + conversationId: other.conversationId, principalId: other.principalId, providerEventId: deliveryId, + deduplicationKey: deliveryId, eventKind: "message", normalizedEvent: {}, state: "processed", processedAt: new Date() }); + await db.insert(chatMessageLinks).values({ companyId: f.companyId, endpointId: f.endpointId, + conversationId: other.conversationId, deliveryId, commentId: other.commentId, + providerMessageId: other.commentId, direction: "inbound" }); + }); + afterAll(async () => { + try { await db?.$client.end({ timeout: 1 }); } + finally { await temporary?.cleanup(); } + }); + + it("gives every run the outcome of a separate read, including each denial", async () => { + const authorized = { conversationId: f.conversationId, endpointId: f.endpointId }; + const runs = [ + run(), + run({ context: { wakeCommentIds: [f.commentId, f.commentId] } }), + run({ context: { wakeCommentIds: [f.commentId, randomUUID()] } }), + run({ context: { wakeCommentIds: [] } }), + run({ context: { paperclipHarnessCheckedOut: false } }), + run({ context: { paperclipHarnessCheckedOut: false, paperclipExternalChatExecutionBound: true } }), + run({ context: { source: "chat:discord" } }), + run({ context: { source: "not-a-chat" } }), + run({ binding: { agentId: randomUUID() } }), + run({ binding: { issueId: f.otherIssueId } }), + // The comment of one run never authorizes another run's issue. + run({ context: { wakeCommentIds: [other.commentId] } }), + run({ context: { wakeCommentIds: [f.commentId, other.commentId] } }), + // The other person has no identity link and the endpoint does not + // allow unlinked people. + otherRun(), + ]; + const outcomes = await expectSameOutcomes(runs); + expect(outcomes).toEqual([authorized, authorized, null, null, null, authorized, + null, null, null, null, null, null, null]); + expect(await expectSameOutcomes([])).toEqual([]); + }); + + it("does not share a principal decision between the people of different runs", async () => { + await db.insert(chatIdentityLinks).values({ companyId: f.companyId, endpointId: f.endpointId, + principalId: other.principalId, paperclipUserId: "other-user", status: "linked" }); + const otherAuthorized = { conversationId: other.conversationId, endpointId: f.endpointId }; + try { + // No membership yet: the first person is authorized and the second is not. + expect(await expectSameOutcomes([run(), otherRun(), run()])) + .toEqual([{ conversationId: f.conversationId, endpointId: f.endpointId }, null, + { conversationId: f.conversationId, endpointId: f.endpointId }]); + await db.insert(companyMemberships).values({ companyId: f.companyId, principalType: "user", + principalId: "other-user", status: "active", membershipRole: "viewer" }); + expect((await expectSameOutcomes([otherRun(), run()]))[0]).toBeNull(); + await db.update(companyMemberships).set({ membershipRole: "member" }) + .where(and(eq(companyMemberships.companyId, f.companyId), eq(companyMemberships.principalId, "other-user"))); + expect(await expectSameOutcomes([otherRun(), run()])) + .toEqual([otherAuthorized, { conversationId: f.conversationId, endpointId: f.endpointId }]); + await db.update(chatIdentityLinks).set({ status: "revoked" }) + .where(eq(chatIdentityLinks.principalId, f.principalId)); + expect(await expectSameOutcomes([run(), otherRun()])).toEqual([null, otherAuthorized]); + } finally { + await db.update(chatIdentityLinks).set({ status: "linked" }).where(eq(chatIdentityLinks.principalId, f.principalId)); + await db.delete(chatIdentityLinks).where(eq(chatIdentityLinks.principalId, other.principalId)); + await db.delete(companyMemberships).where(eq(companyMemberships.principalId, "other-user")); + } + }); + + it("applies current destination, endpoint and delivery policy to every run", async () => { + const both = () => [run(), otherRun()]; + const restore = async () => { + await db.update(chatEndpointResources).set({ enabled: true, availability: "available" }) + .where(eq(chatEndpointResources.id, f.resourceId)); + await db.update(chatEndpoints).set({ status: "active" }).where(eq(chatEndpoints.id, f.endpointId)); + await db.update(chatConversations).set({ state: "active" }).where(eq(chatConversations.id, f.conversationId)); + await db.update(chatDeliveries).set({ state: "processed" }).where(eq(chatDeliveries.id, f.deliveryId)); + }; + try { + await db.update(chatEndpointResources).set({ enabled: false }).where(eq(chatEndpointResources.id, f.resourceId)); + expect(await expectSameOutcomes(both())).toEqual([null, null]); + await restore(); + await db.update(chatEndpointResources).set({ availability: "unavailable" }).where(eq(chatEndpointResources.id, f.resourceId)); + expect(await expectSameOutcomes(both())).toEqual([null, null]); + await restore(); + await db.update(chatEndpoints).set({ status: "paused" }).where(eq(chatEndpoints.id, f.endpointId)); + expect(await expectSameOutcomes(both())).toEqual([null, null]); + await restore(); + await db.update(chatConversations).set({ state: "completed" }).where(eq(chatConversations.id, f.conversationId)); + expect((await expectSameOutcomes(both()))[0]).toBeNull(); + await restore(); + await db.update(chatDeliveries).set({ state: "failed" }).where(eq(chatDeliveries.id, f.deliveryId)); + expect((await expectSameOutcomes(both()))[0]).toBeNull(); + await restore(); + // A delivery without a person is never authorized. + await db.update(chatDeliveries).set({ principalId: null }).where(eq(chatDeliveries.id, f.deliveryId)); + expect((await expectSameOutcomes(both()))[0]).toBeNull(); + await db.update(chatDeliveries).set({ principalId: f.principalId }).where(eq(chatDeliveries.id, f.deliveryId)); + // A verifying endpoint authorizes only deliveries from its current test. + await db.update(chatEndpoints).set({ status: "verifying", setup: sql`${JSON.stringify({ step: "test", + testStartedAt: new Date(Date.now() - 60_000).toISOString(), runtimeGeneration: 3 })}::jsonb` }) + .where(eq(chatEndpoints.id, f.endpointId)); + expect(await expectSameOutcomes(both())).toEqual([null, null]); + } finally { + await db.update(chatEndpoints).set({ setup: sql`'{}'::jsonb` }).where(eq(chatEndpoints.id, f.endpointId)); + await restore(); + } + }); + + it("throws the errors that a separate read throws", async () => { + const invalid = run({ context: { wakeCommentIds: ["not-a-uuid"] } }); + await expect(single(invalid)).rejects.toThrow(); + await expect(authorizeChatConversationsForBoundRuns(db, f.companyId, [run(), invalid])).rejects.toThrow(); + await expect(authorizeChatConversationsForBoundRuns(db, randomUUID(), [run()])) + .rejects.toThrow("paperclip_runner_chat_authorization_batch_company_mismatch"); + }); + + it("reads the messages of all runs with one statement", async () => { + const statements: string[] = []; + const counting = new Proxy(db, { get(target, key) { + if (key !== "select") return Reflect.get(target, key); + return (...fields: Parameters) => { + const builder = target.select(...fields), from = builder.from.bind(builder); + builder.from = ((table: Parameters[0]) => { + statements.push((table as unknown as Record)[Symbol.for("drizzle:Name")]); + return from(table); + }) as typeof builder.from; + return builder; + }; + } }) as Db; + await db.insert(chatIdentityLinks).values({ companyId: f.companyId, endpointId: f.endpointId, + principalId: other.principalId, paperclipUserId: f.userId, status: "linked" }); + try { + const runs = [run(), otherRun(), run(), otherRun(), run()]; + const outcomes = await authorizeChatConversationsForBoundRuns(counting, f.companyId, runs); + expect(outcomes.every(outcome => outcome !== null)).toBe(true); + // One message read; one resource read; per person a principal, link + // and membership read. + expect(statements.filter(name => name === "chat_message_links")).toHaveLength(1); + expect(statements.filter(name => name === "chat_endpoint_resources")).toHaveLength(1); + expect(statements.filter(name => name === "chat_external_principals")).toHaveLength(2); + } finally { + await db.delete(chatIdentityLinks).where(eq(chatIdentityLinks.principalId, other.principalId)); + } + }); +}); diff --git a/server/src/services/native-runtime/chat-attachment-reuse.ts b/server/src/services/native-runtime/chat-attachment-reuse.ts index b037d8b1371..7878583ca01 100644 --- a/server/src/services/native-runtime/chat-attachment-reuse.ts +++ b/server/src/services/native-runtime/chat-attachment-reuse.ts @@ -1,7 +1,8 @@ import { notifyIssueRecovery } from "../issue-recovery-notifications.js"; import { createHash } from "node:crypto"; -import { and, desc, eq, inArray, isNull, sql } from "drizzle-orm"; +import { and, desc, eq, inArray, isNull, sql, type SQL } from "drizzle-orm"; +import type { SelectedFields } from "drizzle-orm/pg-core"; import type { Db } from "@paperclipai/db"; import { assets, @@ -457,39 +458,47 @@ async function principalAuthorized( return endpoint.sponsorUserId ? activeMember(endpoint.sponsorUserId) : true; } -function prepareBoundRunLinks(tx: Db, lockMode: AuthorizationLockMode) { - const linksQuery = tx - .select({ - commentId: chatMessageLinks.commentId, - conversation: { - id: chatConversations.id, - resourceId: chatConversations.resourceId, - isDirectMessage: chatConversations.isDirectMessage, - }, - endpoint: { - id: chatEndpoints.id, - companyId: chatEndpoints.companyId, - provider: chatEndpoints.provider, - providerAccountId: chatEndpoints.providerAccountId, - status: chatEndpoints.status, - allowDirectMessages: chatEndpoints.allowDirectMessages, - allowGroupChats: chatEndpoints.allowGroupChats, - allowUnlinkedPeople: chatEndpoints.allowUnlinkedPeople, - sponsorUserId: chatEndpoints.sponsorUserId, - setup: sql`jsonb_build_object( - 'step', ${chatEndpoints.setup}->'step', - 'testStartedAt', ${chatEndpoints.setup}->'testStartedAt', - 'runtimeGeneration', ${chatEndpoints.setup}->'runtimeGeneration' - )`, - }, - principalId: chatDeliveries.principalId, - normalizedEvent: sql`jsonb_build_object('runtimeContext', jsonb_build_object( - 'generation', ${chatDeliveries.normalizedEvent}->'runtimeContext'->'generation', - 'credentialFingerprint', ${chatDeliveries.normalizedEvent}->'runtimeContext'->'credentialFingerprint' - ))`, - processedAt: chatDeliveries.processedAt, - receivedAt: chatDeliveries.receivedAt, - }) +const boundRunLinkFields = { + commentId: chatMessageLinks.commentId, + conversation: { + id: chatConversations.id, + resourceId: chatConversations.resourceId, + isDirectMessage: chatConversations.isDirectMessage, + }, + endpoint: { + id: chatEndpoints.id, + companyId: chatEndpoints.companyId, + provider: chatEndpoints.provider, + providerAccountId: chatEndpoints.providerAccountId, + status: chatEndpoints.status, + allowDirectMessages: chatEndpoints.allowDirectMessages, + allowGroupChats: chatEndpoints.allowGroupChats, + allowUnlinkedPeople: chatEndpoints.allowUnlinkedPeople, + sponsorUserId: chatEndpoints.sponsorUserId, + setup: sql`jsonb_build_object( + 'step', ${chatEndpoints.setup}->'step', + 'testStartedAt', ${chatEndpoints.setup}->'testStartedAt', + 'runtimeGeneration', ${chatEndpoints.setup}->'runtimeGeneration' + )`, + }, + principalId: chatDeliveries.principalId, + normalizedEvent: sql`jsonb_build_object('runtimeContext', jsonb_build_object( + 'generation', ${chatDeliveries.normalizedEvent}->'runtimeContext'->'generation', + 'credentialFingerprint', ${chatDeliveries.normalizedEvent}->'runtimeContext'->'credentialFingerprint' + ))`, + processedAt: chatDeliveries.processedAt, + receivedAt: chatDeliveries.receivedAt, +}; + +/** The one join and policy filter for the inbound messages of a bound run. + * The caller adds the run's company, comments, issue, provider and agent. */ +function boundRunLinksQuery( + tx: Db, + extraFields: T, + bindingFilter: SQL, +) { + return tx + .select({ ...boundRunLinkFields, ...extraFields }) .from(chatMessageLinks) .innerJoin( chatDeliveries, @@ -517,17 +526,23 @@ function prepareBoundRunLinks(tx: Db, lockMode: AuthorizationLockMode) { ) .where( and( - eq(chatMessageLinks.companyId, sql.placeholder("companyId")), + bindingFilter, eq(chatMessageLinks.direction, "inbound"), - sql`${chatMessageLinks.commentId} = any(${sql.placeholder("commentIds")}::uuid[])`, eq(chatDeliveries.state, "processed"), - eq(chatConversations.issueId, sql.placeholder("issueId")), inArray(chatConversations.state, ["active", "waiting"]), - eq(chatEndpoints.provider, sql.placeholder("provider")), - eq(chatEndpoints.assignedAgentId, sql.placeholder("agentId")), inArray(chatEndpoints.status, ["active", "verifying"]), ), ); +} + +function prepareBoundRunLinks(tx: Db, lockMode: AuthorizationLockMode) { + const linksQuery = boundRunLinksQuery(tx, {}, and( + eq(chatMessageLinks.companyId, sql.placeholder("companyId")), + sql`${chatMessageLinks.commentId} = any(${sql.placeholder("commentIds")}::uuid[])`, + eq(chatConversations.issueId, sql.placeholder("issueId")), + eq(chatEndpoints.provider, sql.placeholder("provider")), + eq(chatEndpoints.assignedAgentId, sql.placeholder("agentId")), + )!); return (lockMode === "read" ? linksQuery : lockMode === "nonblocking" @@ -535,6 +550,8 @@ function prepareBoundRunLinks(tx: Db, lockMode: AuthorizationLockMode) { : linksQuery.for("update")).prepare("chat_bound_run_links"); } +type BoundRunLink = Awaited["execute"]>>[number]; + const boundRunLinksTemplates = new WeakMap< Db, Partial>> @@ -549,13 +566,25 @@ function boundRunLinksTemplate(tx: Db, lockMode: AuthorizationLockMode) { return templates[lockMode] ??= prepareBoundRunLinks(tx, lockMode); } -/** Caller must independently verify the current run/issue execution owner. */ -export async function authorizeChatConversationForBoundRun( +const boundChatAuthorizationDenials = new Set([ + "paperclip_runner_chat_attachment_binding_denied", + "paperclip_runner_chat_attachment_destination_denied", + "paperclip_runner_chat_attachment_principal_denied", +]); + +/** True for the three policy denials of bound-run chat authorization. Database + * errors, lock contention and malformed input are not denials. */ +export function isBoundChatAuthorizationDenial(error: unknown): boolean { + return error instanceof Error && boundChatAuthorizationDenials.has(error.message); +} + +/** The provider and wake comments that the saved context of a run binds. */ +async function boundRunAuthorizationRequest( tx: Db, binding: ChatReuseBinding, contextSnapshot: unknown, - lockMode: AuthorizationLockMode = "blocking", -): Promise { + lockMode: AuthorizationLockMode, +): Promise<{ provider: string; commentIds: string[] }> { let context = record(contextSnapshot); if (context.source === "issue.interaction.respond") { const answer = await resolveExternalChatQuestionResponse(tx, binding, context, lockMode); @@ -573,13 +602,75 @@ export async function authorizeChatConversationForBoundRun( ) { throw new Error("paperclip_runner_chat_attachment_binding_denied"); } - const links = await boundRunLinksTemplate(tx, lockMode).execute({ - companyId: binding.companyId, - issueId: binding.issueId, - agentId: binding.agentId, - provider, - commentIds, - }); + return { provider, commentIds }; +} + +type BoundRunPolicyReads = { + resource: (endpoint: AuthorizationEndpoint, resourceId: string) => Promise; + principal: (endpoint: AuthorizationEndpoint, principalId: string) => Promise; +}; + +function boundRunPolicyReads( + tx: Db, + companyId: string, + lockMode: AuthorizationLockMode, +): BoundRunPolicyReads { + return { + resource: (endpoint, resourceId) => { + const query = tx + .select({ + availability: chatEndpointResources.availability, + type: chatEndpointResources.type, + enabled: chatEndpointResources.enabled, + }) + .from(chatEndpointResources) + .where( + and( + eq(chatEndpointResources.id, resourceId), + eq(chatEndpointResources.companyId, companyId), + eq(chatEndpointResources.endpointId, endpoint.id), + ), + ); + return lockMode === "read" + ? query.then((rows) => rows[0] ?? null) + : lockMode === "nonblocking" + ? query + .for("update", { noWait: true }) + .then((rows) => rows[0] ?? null) + : query.for("update").then((rows) => rows[0] ?? null); + }, + principal: (endpoint, principalId) => + principalAuthorized(tx, endpoint, principalId, lockMode), + }; +} + +/** Read-only policy reads that the runs of one batch share. The batch reads + * each distinct resource and principal one time. */ +function sharedBoundRunPolicyReads(reads: BoundRunPolicyReads): BoundRunPolicyReads { + const resources = new Map>(); + const principals = new Map>(); + const once = (cache: Map>, key: string, read: () => Promise) => { + let value = cache.get(key); + if (!value) { + value = read(); + cache.set(key, value); + } + return value; + }; + return { + resource: (endpoint, resourceId) => + once(resources, `${endpoint.id}:${resourceId}`, () => reads.resource(endpoint, resourceId)), + principal: (endpoint, principalId) => + once(principals, `${endpoint.id}:${principalId}`, () => reads.principal(endpoint, principalId)), + }; +} + +/** Apply the bound-run policy to the inbound message rows of one run. */ +async function authorizeBoundRunLinks( + links: BoundRunLink[], + commentIds: string[], + reads: BoundRunPolicyReads, +): Promise { const linkedCommentIds = new Set(links.map((row) => row.commentId)); const conversationIds = new Set(links.map((row) => row.conversation.id)); const endpointIds = new Set(links.map((row) => row.endpoint.id)); @@ -598,41 +689,93 @@ export async function authorizeChatConversationForBoundRun( throw new Error("paperclip_runner_chat_attachment_binding_denied"); } const resource = conversation.resourceId - ? await (() => { - const query = tx - .select({ - availability: chatEndpointResources.availability, - type: chatEndpointResources.type, - enabled: chatEndpointResources.enabled, - }) - .from(chatEndpointResources) - .where( - and( - eq(chatEndpointResources.id, conversation.resourceId!), - eq(chatEndpointResources.companyId, binding.companyId), - eq(chatEndpointResources.endpointId, endpoint.id), - ), - ); - return lockMode === "read" - ? query.then((rows) => rows[0] ?? null) - : lockMode === "nonblocking" - ? query - .for("update", { noWait: true }) - .then((rows) => rows[0] ?? null) - : query.for("update").then((rows) => rows[0] ?? null); - })() + ? await reads.resource(endpoint, conversation.resourceId) : null; if (!destinationAllowed(endpoint, conversation, resource)) { throw new Error("paperclip_runner_chat_attachment_destination_denied"); } for (const principalId of new Set(links.map((row) => row.principalId!))) { - if (!(await principalAuthorized(tx, endpoint, principalId, lockMode))) { + if (!(await reads.principal(endpoint, principalId))) { throw new Error("paperclip_runner_chat_attachment_principal_denied"); } } return { conversationId: conversation.id, endpointId: endpoint.id }; } +/** Caller must independently verify the current run/issue execution owner. */ +export async function authorizeChatConversationForBoundRun( + tx: Db, + binding: ChatReuseBinding, + contextSnapshot: unknown, + lockMode: AuthorizationLockMode = "blocking", +): Promise { + const { provider, commentIds } = await boundRunAuthorizationRequest( + tx, binding, contextSnapshot, lockMode, + ); + const links = await boundRunLinksTemplate(tx, lockMode).execute({ + companyId: binding.companyId, + issueId: binding.issueId, + agentId: binding.agentId, + provider, + commentIds, + }); + return authorizeBoundRunLinks( + links, commentIds, boundRunPolicyReads(tx, binding.companyId, lockMode), + ); +} + +/** + * The read mode of `authorizeChatConversationForBoundRun` for several runs of + * one company. One statement reads the inbound messages of all the runs, and + * the runs share their resource and principal reads. Each run gets the same + * outcome as a separate read-mode call: its conversation, or null for a + * policy denial. Other errors are thrown. This takes no locks. + */ +export async function authorizeChatConversationsForBoundRuns( + db: Db, + companyId: string, + runs: Array<{ binding: ChatReuseBinding; contextSnapshot: unknown }>, +): Promise> { + if (runs.some((run) => run.binding.companyId !== companyId)) { + throw new Error("paperclip_runner_chat_authorization_batch_company_mismatch"); + } + const deniedAsNull = (error: unknown) => { + if (isBoundChatAuthorizationDenial(error)) return null; + throw error; + }; + const requests = await Promise.all(runs.map((run) => + boundRunAuthorizationRequest(db, run.binding, run.contextSnapshot, "read").catch(deniedAsNull))); + const pending = runs.flatMap((run, index) => { + const request = requests[index]; + return request ? [{ binding: run.binding, ...request }] : []; + }); + if (pending.length === 0) return runs.map(() => null); + const unique = (values: string[]) => [...new Set(values)]; + const rows = await boundRunLinksQuery(db, { + issueId: chatConversations.issueId, + assignedAgentId: chatEndpoints.assignedAgentId, + }, and( + eq(chatMessageLinks.companyId, companyId), + sql`${chatMessageLinks.commentId} = any(${sql.param(unique(pending.flatMap((run) => run.commentIds)))}::uuid[])`, + sql`${chatConversations.issueId} = any(${sql.param(unique(pending.map((run) => run.binding.issueId)))}::uuid[])`, + sql`${chatEndpoints.provider} = any(${sql.param(unique(pending.map((run) => run.provider)))}::text[])`, + sql`${chatEndpoints.assignedAgentId} = any(${sql.param(unique(pending.map((run) => run.binding.agentId)))}::uuid[])`, + )!); + const reads = sharedBoundRunPolicyReads(boundRunPolicyReads(db, companyId, "read")); + return Promise.all(requests.map(async (request, index) => { + if (!request) return null; + const { binding } = runs[index]!; + const comments = new Set(request.commentIds); + // The equality filters of the single-run statement, for this run's rows. + const links = rows.filter((row) => + row.issueId === binding.issueId && + row.assignedAgentId === binding.agentId && + row.endpoint.provider === request.provider && + row.commentId !== null && comments.has(row.commentId)); + return authorizeBoundRunLinks(links, request.commentIds, reads).catch(deniedAsNull); + })); +} + function externalChatWaitCandidate( contextSnapshot: unknown, binding: ChatReuseBinding, diff --git a/server/src/services/native-runtime/external-review-wait.ts b/server/src/services/native-runtime/external-review-wait.ts index 054113ce9f1..5d05c779cb1 100644 --- a/server/src/services/native-runtime/external-review-wait.ts +++ b/server/src/services/native-runtime/external-review-wait.ts @@ -1,7 +1,7 @@ import { and, eq, inArray, sql } from "drizzle-orm"; import { chatConversations, chatEndpoints, heartbeatRuns, issues, statusDecisions, type Db } from "@paperclipai/db"; import { readInheritedGitHubReviewWait } from "../chat-channels/github/review-wait.js"; -import { authorizeChatConversationForBoundRun } from "./chat-attachment-reuse.js"; +import { authorizeChatConversationsForBoundRuns } from "./chat-attachment-reuse.js"; import { EXTERNAL_CHAT_QUESTION_RESPONSE_KEY } from "./external-chat-question-response.js"; const authorizationContextKeys = [ @@ -61,22 +61,19 @@ export async function listExternalReviewWaits(db: Db, companyId: string, issueId eq(chatEndpoints.status, "active"), )) .where(and(eq(issues.companyId, companyId), inArray(issues.id, issueIds), eq(issues.status, "in_review"))); - const waits = await Promise.all(candidates.map(async ({ runId, agentId, contextSnapshot, responsibleUserId, activeIdentityContextId, ...wait }) => { - try { - const inheritedWait = await readInheritedGitHubReviewWait(db, { - id: runId, companyId, agentId, responsibleUserId, activeIdentityContextId, - nativeIssueId: wait.issueId, issueId: wait.issueId, contextSnapshot, - }); - if (inheritedWait?.kind === "revoked") return null; - await authorizeChatConversationForBoundRun(db, { companyId, issueId: wait.issueId, runId, agentId }, inheritedWait?.context ?? contextSnapshot, "read"); - return wait; - } catch (error) { - if (error instanceof Error && ["paperclip_runner_chat_attachment_binding_denied", - "paperclip_runner_chat_attachment_destination_denied", "paperclip_runner_chat_attachment_principal_denied"].includes(error.message)) return null; - throw error; - } - })); - return waits.filter((wait): wait is NonNullable => wait !== null); + const authorizable = (await Promise.all(candidates.map(async ({ runId, agentId, contextSnapshot, responsibleUserId, activeIdentityContextId, ...wait }) => { + const inheritedWait = await readInheritedGitHubReviewWait(db, { + id: runId, companyId, agentId, responsibleUserId, activeIdentityContextId, + nativeIssueId: wait.issueId, issueId: wait.issueId, contextSnapshot, + }); + if (inheritedWait?.kind === "revoked") return null; + return { wait, binding: { companyId, issueId: wait.issueId, runId, agentId }, + contextSnapshot: inheritedWait?.context ?? contextSnapshot }; + }))).filter((candidate): candidate is NonNullable => candidate !== null); + // One authorization read for every wait of the page, not one per wait: the + // per-wait reads each held a pool connection at the same moment. + const authorized = await authorizeChatConversationsForBoundRuns(db, companyId, authorizable); + return authorizable.filter((_candidate, index) => authorized[index] !== null).map(({ wait }) => wait); } export type ExternalReviewWait = Awaited>[number];