diff --git a/doc/DATABASE.md b/doc/DATABASE.md index b231152f6d4..eca8bea9183 100644 --- a/doc/DATABASE.md +++ b/doc/DATABASE.md @@ -2285,6 +2285,19 @@ named `review_wait_benchmark...` can be supplied through `PAPERCLIP_TEST_DATABASE_URL` to compare another PostgreSQL major version. DataRow bytes exclude response metadata and transport overhead. +An issue list or inbox request reads the waits of all its tasks in one pass. +That pass checks the requester provenance of every waiting run with +`chatRunSourceMessageIds`. The runs that need the same kind of provenance read +share one statement, so the statement count follows the longest requester +lineage, not the number of waits. Each run gets the decision of a separate +`chatRunSourceMessageId` call: both use the same walk and the same reads. The +benchmark counts the statements of one request and measures its wall and CPU +time against a pre-fix commit: + +```sh +pnpm --filter @paperclipai/server exec tsx scripts/benchmark-review-wait-provenance.ts --baseline-ref= +``` + ### Sandbox creation ownership diff --git a/scripts/benchmark-github-review-wait.ts b/scripts/benchmark-github-review-wait.ts index 9485b202c6b..cdc1f437bc5 100644 --- a/scripts/benchmark-github-review-wait.ts +++ b/scripts/benchmark-github-review-wait.ts @@ -14,10 +14,9 @@ assert.ok(baselineRef, "Pass an explicit pre-fix --baseline-ref="); const baselineRevision = execFileSync("git", ["rev-parse", "--verify", `${baselineRef}^{commit}`], { encoding: "utf8" }).trim(); const treatmentRevision = execFileSync("git", ["rev-parse", "HEAD"], { encoding: "utf8" }).trim(); const modulePath = "server/src/services/chat-channels/github/review-wait.ts"; -const helperPath = "server/src/services/chat-channels/run-origin.ts"; +// The baseline module and the treatment module both import the current +// provenance helper (run-origin.ts), so its reads are the same on both sides. const root = fileURLToPath(new URL("../", import.meta.url)); -assert.equal(execFileSync("git", ["show", `${baselineRevision}:${helperPath}`], { encoding: "utf8" }), - await fs.readFile(path.join(root, helperPath), "utf8"), "Baseline provenance helper must be unchanged"); const keepAlive = setInterval(() => {}, 1_000); let temporary: Awaited> | undefined; let db: ReturnType | undefined; diff --git a/server/scripts/benchmark-review-wait-provenance.ts b/server/scripts/benchmark-review-wait-provenance.ts new file mode 100644 index 00000000000..0939a314232 --- /dev/null +++ b/server/scripts/benchmark-review-wait-provenance.ts @@ -0,0 +1,267 @@ +/** + * Local proof for the batched requester-provenance check of inherited GitHub + * review waits. + * + * `listExternalReviewWaits` serves the review waits of an issue list or an + * inbox request. The pre-fix modules checked the requester provenance of every + * waiting run with its own `chatRunSourceMessageId` call inside `Promise.all`, + * so a request with N waits built and sent at least 2N statements at once + * (one identity-context read and one source-message link read for each run, + * and two more for each continuation in a run's lineage). The fix checks the + * runs of one request with a number of statements that depends on the longest + * lineage, not on the number of runs. + * + * The pre-fix modules are loaded from --baseline-ref. The script fails unless + * both versions return the same waits and the same inherited wait for each + * run. + * + * pnpm --filter @paperclipai/server exec tsx scripts/benchmark-review-wait-provenance.ts --baseline-ref= + * + * Optional sizes: WAITS (default 250), CONTINUED_EVERY (10: every tenth wait + * continues an earlier run), CONTINUATIONS (3: the length of that lineage), + * POOL (50), PAIRS (15). + */ +import assert from "node:assert/strict"; +import { execFileSync } from "node:child_process"; +import { createHash, randomUUID } from "node:crypto"; +import { promises as fs } from "node:fs"; +import path from "node:path"; +import { fileURLToPath, pathToFileURL } from "node:url"; +import { chatActions, chatConversations, chatDeliveries, chatEndpointResources, chatEndpoints, + chatExternalPrincipals, chatGitHubReviews, chatIdentityLinks, chatMessageLinks, companies, companyMemberships, + createDb, heartbeatRuns, issueComments, runIdentityContexts, toolApplications, toolConnections, + type Db } from "@paperclipai/db"; +import { defaultGitHubReviewPolicy } from "@paperclipai/shared"; +import { startEmbeddedPostgresTestDatabase } from "../src/__tests__/helpers/embedded-postgres.js"; +import { readInheritedGitHubReviewWaits } from "../src/services/chat-channels/github/review-wait.js"; +import { listExternalReviewWaits } from "../src/services/native-runtime/external-review-wait.js"; + +const size = (name: string, fallback: number) => Number(process.env[name] ?? fallback); +const WAITS = size("WAITS", 250), CONTINUED_EVERY = size("CONTINUED_EVERY", 10), CONTINUATIONS = size("CONTINUATIONS", 3); +const POOL = size("POOL", 50), PAIRS = size("PAIRS", 15); +for (const [name, value] of Object.entries({ WAITS, CONTINUED_EVERY, CONTINUATIONS, POOL, PAIRS })) { + 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 root = fileURLToPath(new URL("../../", import.meta.url)); +const git = (...args: string[]) => execFileSync("git", args, { cwd: root, encoding: "utf8" }); +const baselineRevision = git("rev-parse", "--verify", `${baselineRef}^{commit}`).trim(); +// The modules of the read path. Each pre-fix copy imports the pre-fix copies +// of the others; every other import resolves to the current tree. +const modules = ["server/src/services/native-runtime/external-review-wait.ts", + "server/src/services/chat-channels/github/review-wait.ts", "server/src/services/chat-channels/run-origin.ts"]; + +type Versions = { list: typeof listExternalReviewWaits; inherited: typeof readInheritedGitHubReviewWaits }; +const keepAlive = setInterval(() => {}, 1_000); +let temporary: Awaited> | undefined; +let db: Db | undefined; +let baselineDir: string | undefined; +try { + baselineDir = await fs.mkdtemp(path.join(root, "server/scripts/.review-wait-provenance-baseline-")); + const copyOf = (file: string) => path.join(baselineDir!, path.basename(file)); + for (const file of modules) { + const source = git("show", `${baselineRevision}:${file}`).replace(/from "(\.{1,2}\/[^"]+)"/g, (_match, specifier: string) => { + const target = fileURLToPath(new URL(specifier, pathToFileURL(path.join(root, file)))); + const module = modules.find(candidate => path.join(root, candidate) === target.replace(/\.js$/, ".ts")); + return `from ${JSON.stringify(module ? copyOf(module) : target)}`; + }); + await fs.writeFile(copyOf(file), source); + } + const load = async (file: string) => await import(pathToFileURL(copyOf(file)).href); + const versions: Record = { + baseline: { list: (await load(modules[0]!)).listExternalReviewWaits, + inherited: (await load(modules[1]!)).readInheritedGitHubReviewWaits }, + treatment: { list: listExternalReviewWaits, inherited: readInheritedGitHubReviewWaits }, + }; + + temporary = await startEmbeddedPostgresTestDatabase("review-wait-provenance-benchmark-"); + db = createDb(temporary.connectionString, { maxConnections: POOL }); + const sql = db.$client; + const loadStart = performance.now(); + + // One company, one agent and one GitHub repository with WAITS reviewed pull + // requests. Each pull request has a task in review that waits for the next + // GitHub reply. + const companyId = randomUUID(), agentId = randomUUID(), applicationId = randomUUID(), connectionId = randomUUID(); + const endpointId = randomUUID(), resourceId = randomUUID(), principalId = randomUUID(); + const repository = "synthetic/repository", repositoryId = "42", headSha = "1".repeat(40), user = "operator"; + await db.insert(companies).values({ id: companyId, name: "Synthetic review", issuePrefix: "SYN" }); + // A fixture agent that never runs. It is not created through the agent lifecycle module. + await sql.unsafe("insert into agents (id, company_id, name, adapter_type) values ($1::uuid, $2::uuid, 'Reviewer', 'process')", + [agentId, companyId]); + await db.insert(toolApplications).values({ id: applicationId, companyId, name: "GitHub", type: "chat" }); + await db.insert(toolConnections).values({ id: connectionId, companyId, applicationId, name: "GitHub", + uid: connectionId, transport: "chat_sdk", connectionPurpose: "channel", enabled: true, status: "active" }); + await db.insert(chatEndpoints).values({ id: endpointId, companyId, connectionId, provider: "github", + publicId: endpointId, assignedAgentId: agentId, status: "active", providerAccountId: "synthetic" }); + await db.insert(chatEndpointResources).values({ id: resourceId, companyId, endpointId, type: "repository", + providerResourceId: repository, label: repository, metadata: { providerRepositoryId: repositoryId }, + availability: "available", enabled: true }); + await db.insert(chatExternalPrincipals).values({ id: principalId, companyId, provider: "github", + providerAccountId: "synthetic", externalId: "1", kind: "user" }); + await db.insert(chatIdentityLinks).values({ companyId, endpointId, principalId, paperclipUserId: user, status: "linked" }); + await db.insert(companyMemberships).values({ companyId, principalType: "user", principalId: user, + status: "active", membershipRole: "member" }); + + const waits = Array.from({ length: WAITS }, (_unused, index) => ({ + pullNumber: index + 1, issueId: randomUUID(), runId: randomUUID(), reviewRunId: randomUUID(), reviewId: randomUUID(), + decisionId: randomUUID(), conversationId: randomUUID(), commentId: randomUUID(), deliveryId: randomUUID(), + // The run of every CONTINUED_EVERY-th wait continues earlier runs of its + // task. The first run of that lineage accepted the GitHub comment. + lineage: Array.from({ length: index % CONTINUED_EVERY === CONTINUED_EVERY - 1 ? CONTINUATIONS : 0 }, + () => ({ runId: randomUUID(), contextId: randomUUID() })), + contextId: randomUUID(), + })); + const insert = async (rows: Row[], write: (chunk: Row[]) => Promise) => { + for (let start = 0; start < rows.length; start += 500) await write(rows.slice(start, start + 500)); + }; + await sql.begin(async tx => { + // Skip the work-assessment owner chain of the status decisions. + await tx.unsafe("set local session_replication_role = replica"); + await tx.unsafe(`insert into issues (id, company_id, title, status, assignee_agent_id, issue_number, identifier, last_status_decision_id) + select w.issue_id, $1::uuid, 'wait ' || w.n, 'in_review', $2::uuid, w.n, 'SYN-' || w.n, w.decision_id + from unnest($3::uuid[], $4::uuid[]) with ordinality as w(issue_id, decision_id, n)`, + [companyId, agentId, waits.map(wait => wait.issueId), waits.map(wait => wait.decisionId)]); + await tx.unsafe(`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 w.decision_id, $1::uuid, w.issue_id, w.run_id, gen_random_uuid(), 1, 'v1', 'in_progress', 'in_review', + 'external_chat_response_waiting', '{}'::jsonb, w.decision_id::text, 'applied', now() + from unnest($2::uuid[], $3::uuid[], $4::uuid[]) as w(decision_id, issue_id, run_id)`, + [companyId, waits.map(wait => wait.decisionId), waits.map(wait => wait.issueId), waits.map(wait => wait.runId)]); + }); + await insert(waits, chunk => db!.insert(heartbeatRuns).values(chunk.flatMap(wait => [ + { id: wait.runId, companyId, agentId, status: "succeeded", runtimeMode: "native", nativeIssueId: wait.issueId, + responsibleUserId: user, activeIdentityContextId: wait.contextId, contextSnapshot: { prompt: "x".repeat(4096) } }, + { id: wait.reviewRunId, companyId, agentId, status: "succeeded", nativeIssueId: wait.issueId }, + ...wait.lineage.map(earlier => ({ id: earlier.runId, companyId, agentId, status: "succeeded", + nativeIssueId: wait.issueId, responsibleUserId: user })), + ]))); + await insert(waits, chunk => db!.insert(issueComments).values(chunk.map(wait => ({ id: wait.commentId, companyId, + issueId: wait.issueId, authorType: "user" as const, authorUserId: user, body: "Review" })))); + await insert(waits, chunk => db!.insert(runIdentityContexts).values(chunk.flatMap(wait => { + // The lineage from the first run (which accepted the comment) to the waiting run. + const contexts = [...wait.lineage, { runId: wait.runId, contextId: wait.contextId }]; + return contexts.map((context, index) => ({ id: context.contextId, companyId, runId: context.runId, revision: 1, + responsibleUserId: user, messageId: index === 0 ? wait.commentId : null, + parentContextId: index === 0 ? null : contexts[index - 1]!.contextId, + cause: index === 0 ? "instruction" : "monitor", correlationId: "benchmark", status: "accepted" })); + }))); + await insert(waits, chunk => db!.insert(chatConversations).values(chunk.map(wait => ({ id: wait.conversationId, companyId, + endpointId, resourceId, issueId: wait.issueId, externalConversationId: repository, + externalThreadId: `github:${repository}:${wait.pullNumber}`, externalLabel: "Review", state: "waiting" })))); + await insert(waits, chunk => db!.insert(chatDeliveries).values(chunk.map(wait => ({ id: wait.deliveryId, companyId, + endpointId, conversationId: wait.conversationId, principalId, providerEventId: wait.deliveryId, + deduplicationKey: wait.deliveryId, eventKind: "message" as const, state: "processed" as const, normalizedEvent: {}, processedAt: new Date() })))); + await insert(waits, chunk => db!.insert(chatMessageLinks).values(chunk.map(wait => ({ companyId, endpointId, + conversationId: wait.conversationId, deliveryId: wait.deliveryId, commentId: wait.commentId, + providerMessageId: wait.commentId, direction: "inbound" })))); + await insert(waits, chunk => db!.insert(chatGitHubReviews).values(chunk.map(wait => ({ id: wait.reviewId, companyId, + endpointId, issueId: wait.issueId, runId: wait.reviewRunId, repositoryId, repository, pullNumber: wait.pullNumber, + headSha, deliveryId: wait.reviewId, configurationRevision: 1, state: "completed" as const, + policySnapshot: defaultGitHubReviewPolicy(), + event: { event: "mention" as const, deliveryId: wait.reviewId, repositoryId, repository, pullNumber: wait.pullNumber, + title: "Review", body: "", baseSha: "0".repeat(40), headSha, baseBranch: "main", + author: { id: "1", login: "author", isBot: false }, sender: { id: "1", login: "author" }, draft: false, labels: [] }, + assessment: { reviewedCommit: headSha, score: 5 as const, complete: true, summary: "Reviewed", rationale: "Reviewed", + coverage: { reviewedPaths: [], omittedPaths: [], limitations: [] }, findings: [] } })))); + await insert(waits, chunk => db!.insert(chatActions).values(chunk.map(wait => ({ companyId, endpointId, + kind: "github_review_publication", status: "processed", + providerActionId: `github_publication:${wait.issueId}:assessment:${createHash("sha256").update(JSON.stringify(wait.reviewId)).digest("hex")}` })))); + await sql.unsafe("vacuum analyze"); + const loadMs = performance.now() - loadStart; + const issueIds = waits.map(wait => wait.issueId); + const runs = await db.select().from(heartbeatRuns); + const waitRuns = waits.map(wait => runs.find(run => run.id === wait.runId)!); + + // Both versions return the same waits, and the same inherited wait for each run. + const listed: Record = {}, inherited: Record = {}; + for (const [name, version] of Object.entries(versions)) { + listed[name] = (await version.list(db, companyId, issueIds)).map(wait => wait.issueId).sort(); + inherited[name] = await version.inherited(db, waitRuns); + } + assert.deepEqual(listed.treatment, listed.baseline, "Both versions must return the same waits"); + assert.deepEqual(inherited.treatment, inherited.baseline, "Both versions must return the same inherited waits"); + const published = (inherited.baseline as Array<{ kind?: string } | null>).filter(wait => wait?.kind === "published").length; + assert.equal(published, WAITS, "Every seeded wait must inherit its published review"); + + const round = (value: number) => Number(value.toFixed(2)); + // Statements that one request sends, counted at the driver (prepared + // templates are counted at each execution). + const statementCounts = async (list: Versions["list"]) => { + const client = sql as unknown as { unsafe: (query: string, ...rest: unknown[]) => unknown }; + const original = client.unsafe; + const counts: Record = {}; + let total = 0, sqlBytes = 0; + const batches: Array<{ query: string; params: unknown }> = []; + client.unsafe = function (this: unknown, query: string, ...rest: unknown[]) { + // A batch provenance read selects from the request rows and reads its + // table in a lateral subquery. + const batch = /^select .*? from unnest\(.*? cross join lateral \(select .*? from "([a-z_]+)"/s.exec(query)?.[1]; + if (batch) batches.push({ query, params: rest[0] }); + const key = batch ? `provenance batch: ${batch}` : /^select .*? from "([a-z_]+)"/s.exec(query)?.[1] ?? "other"; + counts[key] = (counts[key] ?? 0) + 1; + total += 1; + sqlBytes += Buffer.byteLength(query); + return original.call(this, query, ...rest); + }; + try { await list(db!, companyId, issueIds); } finally { client.unsafe = original; } + // Server execution and scans of the batch provenance reads of one request. + let executionMs = 0; + const scans = new Set(); + for (const statement of batches) { + 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 }>; + executionMs += plan!["Execution Time"]; + const walk = (node: Record) => { + if (/Scan/.test(String(node["Node Type"])) && node["Relation Name"]) { + scans.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, sqlTextBytes: sqlBytes, byFirstTable: counts, + provenanceBatches: { count: batches.length, summedExecutionMs: round(executionMs), scans: [...scans].sort() } }; + }; + const stats = (samples: number[]) => { + const sorted = [...samples].sort((a, b) => a - b); + return { median: round(sorted[Math.floor(sorted.length / 2)]!), min: round(sorted[0]!), max: round(sorted.at(-1)!) }; + }; + // Wall time and process CPU time (user and system) of one request. + const measure = async (list: Versions["list"]) => { + const cpu = process.cpuUsage(), start = performance.now(); + await list(db!, companyId, issueIds); + const used = process.cpuUsage(cpu); + return { wallMs: performance.now() - start, cpuMs: (used.user + used.system) / 1_000 }; + }; + for (const version of Object.values(versions)) for (let warm = 0; warm < 3; warm++) await measure(version.list); + const samples: Record> = { baseline: [], treatment: [] }; + for (let pair = 0; pair < PAIRS; pair++) { + const order = pair % 2 === 0 ? ["baseline", "treatment"] : ["treatment", "baseline"]; + for (const name of order) samples[name]!.push(await measure(versions[name]!.list)); + } + const report: Record = {}; + for (const [name, version] of Object.entries(versions)) { + report[name] = { + statements: await statementCounts(version.list), + wallMs: stats(samples[name]!.map(sample => sample.wallMs)), + cpuMs: stats(samples[name]!.map(sample => sample.cpuMs)), + }; + } + console.log(JSON.stringify({ + baselineRevision, treatmentRevision: git("rev-parse", "HEAD").trim(), + treatmentDirty: git("status", "--porcelain", "--", ...modules).trim().length > 0, + postgresVersion: (await sql.unsafe("select version()"))[0]!.version, node: process.version, + fixture: { synthetic: true, waits: WAITS, continuedEvery: CONTINUED_EVERY, continuations: CONTINUATIONS, + pool: POOL, pairs: PAIRS, warmups: 3, order: "alternating", loadMs: Math.round(loadMs) }, + correctness: { sameWaits: true, sameInheritedWaits: true, waits: listed.baseline!.length, publishedInheritedWaits: published }, + 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__/github-review-wait.test.ts b/server/src/__tests__/github-review-wait.test.ts index f4874dff9a1..5db20770e8b 100644 --- a/server/src/__tests__/github-review-wait.test.ts +++ b/server/src/__tests__/github-review-wait.test.ts @@ -136,6 +136,21 @@ describe("inherited GitHub review wait projection", () => { expect(await readInheritedGitHubReviewWaits(db, [])).toEqual([]); }); + it("reads the waits of forty runs with the statements of four runs", async () => { + const first = await seedGitHubReviewWait(db); + const runs = [first.run]; + while (runs.length < 40) runs.push((await seedGitHubReviewWait(db, { companyId: first.companyId, agentId: first.agentId })).run); + const few = await withStatements(db, () => readInheritedGitHubReviewWaits(db, runs.slice(0, 4))); + const many = await withStatements(db, () => readInheritedGitHubReviewWaits(db, runs)); + expect(many.result.map(wait => wait?.kind)).toEqual(runs.map(() => "published")); + expect(many.result.slice(0, 4)).toEqual(few.result); + // The requester lineage, the source message link, the reviews and the + // publications: one statement each, for any number of runs. + expect({ few: few.sql.length, many: many.sql.length }).toEqual({ few: 4, many: 4 }); + expect(many.sql.filter(text => /from "run_identity_contexts"/i.test(text))).toHaveLength(1); + expect(many.sql.filter(text => /from "chat_message_links"/i.test(text))).toHaveLength(2); + }); + it("binds each call afresh while reusing the compiled query", async () => { const first = await seedGitHubReviewWait(db), second = await seedGitHubReviewWait(db); const queries: CapturedReviewWaitQuery[] = []; diff --git a/server/src/services/chat-channels/github/review-wait.ts b/server/src/services/chat-channels/github/review-wait.ts index 7f490c327a4..5758d35554c 100644 --- a/server/src/services/chat-channels/github/review-wait.ts +++ b/server/src/services/chat-channels/github/review-wait.ts @@ -1,6 +1,6 @@ import { createHash } from "node:crypto"; import { HttpError } from "../../../errors.js"; -import { chatRunSourceMessageId } from "../run-origin.js"; +import { chatRunSourceMessageId, chatRunSourceMessageIds } from "../run-origin.js"; import { and, desc, eq, or, sql } from "drizzle-orm"; import { chatActions, chatConversations, chatEndpointResources, chatEndpoints, chatGitHubReviews, chatMessageLinks, heartbeatRuns, toolConnections, type Db } from "@paperclipai/db"; @@ -141,6 +141,21 @@ async function inheritedSourceComment(db: Db, run: WaitRun): Promise<{ kind: "re return commentId ? { commentId } : null; } +/** `inheritedSourceComment` for the runs of one pass. The lineage of every + * run is checked with the statements of one batch, not with separate reads + * for each run. */ +async function inheritedSourceComments(db: Db, runs: WaitRun[]) { + const candidates = runs.filter(run => run.activeIdentityContextId && run.nativeIssueId); + const checked = await chatRunSourceMessageIds(db, candidates, { provider: "github" }); + const sources = new Map(candidates.map((run, index) => [run, checked[index]!])); + return runs.map((run): Awaited> => { + const source = sources.get(run); + if (!source) return null; + if ("forbidden" in source) return { kind: "revoked" }; + return source.messageId ? { commentId: source.messageId } : null; + }); +} + function inheritedReviewRevoked(run: WaitRun, review: InheritedReviewRow) { return !review.connectionEnabled || review.connectionStatus !== "active" || review.state !== "completed" || review.reviewCompanyId !== run.companyId || review.reviewAgentId !== run.agentId || @@ -187,11 +202,13 @@ export async function readInheritedGitHubReviewWait(db: Db, run: WaitRun): Promi /** `readInheritedGitHubReviewWait` for the runs of one read-only pass, such * as the waits of an issue page. Each run gets the same result as a separate - * call. The source lineage is still checked for each run. The reviews and the - * publications of all the runs of one company are read with one statement - * each, not with two statements for each run. */ + * call. The source lineage is still checked for each run, with the same + * decision. The lineage reads of all the runs share their statements, and the + * reviews and the publications of all the runs of one company are read with + * one statement each, so the number of statements does not follow the number + * of runs. */ export async function readInheritedGitHubReviewWaits(db: Db, runs: WaitRun[]): Promise { - const sources = await Promise.all(runs.map(run => inheritedSourceComment(db, run))); + const sources = await inheritedSourceComments(db, runs); const results: InheritedGitHubReviewWait[] = sources.map(source => source && !("commentId" in source) ? source : null); const byCompany = new Map>(); runs.forEach((run, index) => { diff --git a/server/src/services/chat-channels/run-origin.test.ts b/server/src/services/chat-channels/run-origin.test.ts index 6289333d6dc..2facf7c2248 100644 --- a/server/src/services/chat-channels/run-origin.test.ts +++ b/server/src/services/chat-channels/run-origin.test.ts @@ -1,10 +1,14 @@ import { randomUUID } from "node:crypto"; import { eq } from "drizzle-orm"; import { afterAll, beforeAll, describe, expect, it } from "vitest"; -import { agents, companies, createDb, heartbeatRuns, issueComments, issues, runIdentityContexts } from "@paperclipai/db"; +import { agents, chatConversations, chatEndpoints, chatMessageLinks, companies, createDb, heartbeatRuns, issueComments, + issues, runIdentityContexts } from "@paperclipai/db"; +import type { ChatProvider } from "@paperclipai/shared"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../__tests__/helpers/embedded-postgres.js"; +import { seedGitHubReviewWait } from "../../__tests__/helpers/github-review-wait.js"; +import { HttpError } from "../../errors.js"; import { boundRunIssueId, boundRunIssueIdSql } from "../run-task-binding.js"; -import { chatRunSourceMessageId } from "./run-origin.js"; +import { chatRunSourceMessageId, chatRunSourceMessageIds } from "./run-origin.js"; const support = await getEmbeddedPostgresTestSupport(); (support.supported ? describe : describe.skip)("channel task provenance", () => { @@ -89,31 +93,243 @@ const support = await getEmbeddedPostgresTestSupport(); expect(await chatRunSourceMessageId(db, run, { provider: "github", allowLegacyWakeComment: true })).toBeNull(); }); + async function breakInheritedAuthority(f: Awaited>, fault: string) { + if (fault === "task") { + const [other] = await db.insert(issues).values({ companyId: f.company.id, title: "Other task" }).returning(); + const [otherRun] = await db.insert(heartbeatRuns).values({ companyId: f.company.id, agentId: f.agent.id, status: "succeeded", nativeIssueId: other.id }).returning(); + await db.update(runIdentityContexts).set({ runId: otherRun.id }).where(eq(runIdentityContexts.id, f.origin.id)); + } + if (fault === "agent") { + const [other] = await db.insert(agents).values({ companyId: f.company.id, name: "Other agent", role: "general", adapterType: "codex_local" }).returning(); + await db.update(heartbeatRuns).set({ agentId: other.id }).where(eq(heartbeatRuns.id, f.source.id)); + } + if (fault === "company") { + const [other] = await db.insert(companies).values({ name: "Other company", issuePrefix: randomUUID().slice(0, 8) }).returning(); + await db.update(runIdentityContexts).set({ companyId: other.id }).where(eq(runIdentityContexts.id, f.origin.id)); + } + if (fault === "requester") await db.update(runIdentityContexts).set({ responsibleUserId: "another-person" }).where(eq(runIdentityContexts.id, f.origin.id)); + if (fault === "missing-parent") await db.update(runIdentityContexts).set({ parentContextId: randomUUID() }).where(eq(runIdentityContexts.id, f.context.id)); + if (fault === "pending-parent") await db.update(runIdentityContexts).set({ status: "pending" }).where(eq(runIdentityContexts.id, f.origin.id)); + if (fault === "pending-current") await db.update(runIdentityContexts).set({ status: "pending" }).where(eq(runIdentityContexts.id, f.context.id)); + if (fault === "cycle") await db.update(runIdentityContexts).set({ messageId: null, parentContextId: f.context.id }).where(eq(runIdentityContexts.id, f.origin.id)); + } + it.each(["task", "agent", "company", "requester", "missing-parent", "pending-parent", "pending-current", "cycle"])( "rejects a %s break in inherited authority", async fault => { const f = await fixture(); - if (fault === "task") { - const [other] = await db.insert(issues).values({ companyId: f.company.id, title: "Other task" }).returning(); - const [otherRun] = await db.insert(heartbeatRuns).values({ companyId: f.company.id, agentId: f.agent.id, status: "succeeded", nativeIssueId: other.id }).returning(); - await db.update(runIdentityContexts).set({ runId: otherRun.id }).where(eq(runIdentityContexts.id, f.origin.id)); - } - if (fault === "agent") { - const [other] = await db.insert(agents).values({ companyId: f.company.id, name: "Other agent", role: "general", adapterType: "codex_local" }).returning(); - await db.update(heartbeatRuns).set({ agentId: other.id }).where(eq(heartbeatRuns.id, f.source.id)); - } - if (fault === "company") { - const [other] = await db.insert(companies).values({ name: "Other company", issuePrefix: randomUUID().slice(0, 8) }).returning(); - await db.update(runIdentityContexts).set({ companyId: other.id }).where(eq(runIdentityContexts.id, f.origin.id)); - } - if (fault === "requester") await db.update(runIdentityContexts).set({ responsibleUserId: "another-person" }).where(eq(runIdentityContexts.id, f.origin.id)); - if (fault === "missing-parent") await db.update(runIdentityContexts).set({ parentContextId: randomUUID() }).where(eq(runIdentityContexts.id, f.context.id)); - if (fault === "pending-parent") await db.update(runIdentityContexts).set({ status: "pending" }).where(eq(runIdentityContexts.id, f.origin.id)); - if (fault === "pending-current") await db.update(runIdentityContexts).set({ status: "pending" }).where(eq(runIdentityContexts.id, f.context.id)); - if (fault === "cycle") await db.update(runIdentityContexts).set({ messageId: null, parentContextId: f.context.id }).where(eq(runIdentityContexts.id, f.origin.id)); + await breakInheritedAuthority(f, fault); await expect(chatRunSourceMessageId(db, f.run)).rejects.toMatchObject({ status: 403 }); }, ); + type Run = Parameters[1]; + type Options = { allowLegacyWakeComment?: boolean; provider?: ChatProvider }; + type Seeded = Awaited>; + const optionSets: Options[] = [{}, { allowLegacyWakeComment: true }, { provider: "github" }, + { provider: "github", allowLegacyWakeComment: true }, { provider: "slack" }]; + const FORBIDDEN = { changed: "Chat requester provenance changed", otherTask: "Chat request belongs to another task", + inherited: "Chat tools cannot be inherited by another task", unverified: "Chat source provenance could not be verified" }; + + /** The decision of the single check: its message, or the 403 that it throws. */ + async function singleDecision(run: Run, options: Options) { + try { return { messageId: await chatRunSourceMessageId(db, run, options) }; } catch (error) { + if (!(error instanceof HttpError)) throw error; + return { status: error.status, message: error.message }; + } + } + async function batchDecisions(runs: Run[], options: Options) { + return (await chatRunSourceMessageIds(db, runs, options)).map(result => "forbidden" in result + ? { status: result.forbidden.status, message: result.forbidden.message } : result); + } + + /** Count the statements that the client sends. */ + async function countStatements(read: () => Promise) { + const client = db.$client as unknown as { unsafe: (query: string, ...rest: unknown[]) => unknown }; + const unsafe = client.unsafe; + let statements = 0; + client.unsafe = (query: string, ...rest: unknown[]) => { statements += 1; return unsafe.call(client, query, ...rest); }; + try { return { result: await read(), statements }; } finally { client.unsafe = unsafe; } + } + + /** A later run of the task that continues the identity context of an earlier run. */ + async function continuation(companyId: string, agentId: string, issueId: string, parentContextId: string, responsibleUserId = "operator") { + const [next] = await db.insert(heartbeatRuns).values({ companyId, agentId, status: "running", nativeIssueId: issueId, responsibleUserId }).returning(); + const [context] = await db.insert(runIdentityContexts).values({ companyId, runId: next.id, revision: 1, responsibleUserId, parentContextId, cause: "monitor", correlationId: "dispatch" }).returning(); + const [run] = await db.update(heartbeatRuns).set({ activeIdentityContextId: context.id }).where(eq(heartbeatRuns.id, next.id)).returning(); + return run; + } + + /** Runs in every provenance state, each with the decision that the GitHub + * check of a separate call must make. The states are a GitHub request with + * an inbound message link, its breaks, and the states of the tests above. */ + async function provenanceStates() { + const states: Array<{ state: string; run: Run; github: string | null | { forbidden: string } }> = []; + + // A GitHub request: the accepted message has an inbound link to the + // conversation of the task, on a GitHub endpoint. + const linked = await seedGitHubReviewWait(db); + states.push({ state: "linked request", run: linked.run, github: linked.commentId }); + states.push({ state: "the same run again", run: linked.run, github: linked.commentId }); + states.push({ state: "continuation of a linked request", github: linked.commentId, + run: await continuation(linked.companyId, linked.agentId, linked.issueId, linked.activeIdentityContextId) }); + const [otherTask] = await db.insert(issues).values({ companyId: linked.companyId, title: "Other task", assigneeAgentId: linked.agentId }).returning(); + states.push({ state: "continuation on another task", github: { forbidden: FORBIDDEN.inherited }, + run: await continuation(linked.companyId, linked.agentId, otherTask.id, linked.activeIdentityContextId) }); + states.push({ state: "continuation by another requester", github: { forbidden: FORBIDDEN.changed }, + run: await continuation(linked.companyId, linked.agentId, linked.issueId, linked.activeIdentityContextId, "another-person") }); + states.push({ state: "run of another requester", github: { forbidden: FORBIDDEN.changed }, + run: { ...linked.run, responsibleUserId: "another-person" } }); + states.push({ state: "run of another company", github: { forbidden: FORBIDDEN.changed }, + run: { ...linked.run, companyId: (await seedGitHubReviewWait(db)).companyId } }); + states.push({ state: "no identity context", github: null, + run: { ...linked.run, activeIdentityContextId: null, contextSnapshot: { wakeCommentId: linked.commentId } } }); + + // A run of another company whose accepted message is the linked message + // of the first company. The link, the conversation and the endpoint of + // another company are never read for this run. + const borrower = await seedGitHubReviewWait(db); + await db.update(runIdentityContexts).set({ messageId: linked.commentId }).where(eq(runIdentityContexts.id, borrower.activeIdentityContextId)); + states.push({ state: "message linked in another company", run: borrower.run, github: { forbidden: FORBIDDEN.changed } }); + + // A run of another task of the same company whose accepted message is + // the linked message of the first task. + const sibling = await seedGitHubReviewWait(db, { companyId: linked.companyId, agentId: linked.agentId }); + await db.update(runIdentityContexts).set({ messageId: linked.commentId }).where(eq(runIdentityContexts.id, sibling.activeIdentityContextId)); + states.push({ state: "message linked to another task", run: sibling.run, github: { forbidden: FORBIDDEN.otherTask } }); + + const moved = await seedGitHubReviewWait(db); + const [elsewhere] = await db.insert(issues).values({ companyId: moved.companyId, title: "Other task" }).returning(); + await db.update(chatConversations).set({ issueId: elsewhere.id }).where(eq(chatConversations.id, moved.conversationId)); + states.push({ state: "conversation moved to another task", run: moved.run, github: { forbidden: FORBIDDEN.otherTask } }); + + const otherProvider = await seedGitHubReviewWait(db); + await db.update(chatEndpoints).set({ provider: "slack" }).where(eq(chatEndpoints.id, otherProvider.endpointId)); + states.push({ state: "link of another provider", run: otherProvider.run, github: null }); + + const rejected = await seedGitHubReviewWait(db); + await db.update(runIdentityContexts).set({ status: "rejected" }).where(eq(runIdentityContexts.id, rejected.activeIdentityContextId)); + states.push({ state: "rejected identity context", run: rejected.run, github: { forbidden: FORBIDDEN.changed } }); + + const reassigned = await seedGitHubReviewWait(db); + await db.update(runIdentityContexts).set({ responsibleUserId: "another-person" }).where(eq(runIdentityContexts.id, reassigned.activeIdentityContextId)); + states.push({ state: "identity context of another requester", run: reassigned.run, github: { forbidden: FORBIDDEN.changed } }); + + // Without an inbound link, the accepted message must be a board message + // of the requester on the task of the run. + const unlinked = async (change: (f: Seeded) => Promise) => { + const f = await seedGitHubReviewWait(db); + await db.delete(chatMessageLinks).where(eq(chatMessageLinks.commentId, f.commentId)); + await change(f); + return f.run; + }; + const message = (f: Seeded, values: Partial) => + db.update(issueComments).set(values).where(eq(issueComments.id, f.commentId)); + states.push({ state: "missing link, board message of the requester", run: await unlinked(async () => {}), github: null }); + states.push({ state: "missing link, message of another person", github: { forbidden: FORBIDDEN.changed }, + run: await unlinked(f => message(f, { authorUserId: "another-person" })) }); + states.push({ state: "missing link, agent message", github: { forbidden: FORBIDDEN.changed }, + run: await unlinked(f => message(f, { authorType: "agent" })) }); + states.push({ state: "missing link, deleted message", github: { forbidden: FORBIDDEN.changed }, + run: await unlinked(f => message(f, { deletedAt: new Date() })) }); + states.push({ state: "missing link, message of another task", github: { forbidden: FORBIDDEN.changed }, + run: await unlinked(async f => { + const [other] = await db.insert(issues).values({ companyId: f.companyId, title: "Other task" }).returning(); + await message(f, { issueId: other.id }); + }) }); + const outbound = await seedGitHubReviewWait(db); + await db.update(chatMessageLinks).set({ direction: "outbound" }).where(eq(chatMessageLinks.commentId, outbound.commentId)); + states.push({ state: "outbound link only", run: outbound.run, github: null }); + + // The states of the single-check tests above: a continuation whose + // accepted message has no link and is not a board message. + states.push({ state: "continuation of an unlinked message", run: (await fixture()).run, github: { forbidden: FORBIDDEN.changed } }); + for (const fault of ["task", "agent", "company", "requester", "missing-parent", "pending-parent", "pending-current", "cycle"]) { + const f = await fixture(); + await breakInheritedAuthority(f, fault); + states.push({ state: `${fault} break in inherited authority`, run: f.run, github: { forbidden: + fault === "task" || fault === "agent" ? FORBIDDEN.inherited : fault === "cycle" ? FORBIDDEN.unverified : FORBIDDEN.changed } }); + } + + // A legacy wake comment, through a verified parent run of the same + // requester, and after a board message of the same requester. + const legacy = await fixture(); + await db.update(runIdentityContexts).set({ messageId: null }).where(eq(runIdentityContexts.id, legacy.origin.id)); + await db.update(heartbeatRuns).set({ contextSnapshot: { issueId: legacy.issue.id, wakeCommentId: legacy.messageId } }).where(eq(heartbeatRuns.id, legacy.source.id)); + states.push({ state: "legacy wake of a verified parent run", run: legacy.run, github: null }); + const board = await fixture(); + const [boardMessage] = await db.insert(issueComments).values({ companyId: board.company.id, issueId: board.issue.id, authorType: "user", authorUserId: "requester", body: "Continue" }).returning(); + await db.update(runIdentityContexts).set({ messageId: null }).where(eq(runIdentityContexts.id, board.origin.id)); + await db.update(heartbeatRuns).set({ contextSnapshot: { issueId: board.issue.id, wakeCommentId: board.messageId } }).where(eq(heartbeatRuns.id, board.source.id)); + await db.update(runIdentityContexts).set({ messageId: boardMessage.id }).where(eq(runIdentityContexts.id, board.context.id)); + const [boardRun] = await db.update(heartbeatRuns).set({ contextSnapshot: { wakeCommentId: boardMessage.id } }).where(eq(heartbeatRuns.id, board.run.id)).returning(); + states.push({ state: "legacy wake after a board message of the requester", run: boardRun, github: null }); + const later = await fixture(); + const [laterMessage] = await db.insert(issueComments).values({ companyId: later.company.id, issueId: later.issue.id, authorType: "user", authorUserId: "new-requester", body: "Continue" }).returning(); + await db.update(runIdentityContexts).set({ messageId: laterMessage.id, responsibleUserId: "new-requester" }).where(eq(runIdentityContexts.id, later.context.id)); + const [laterRun] = await db.update(heartbeatRuns).set({ responsibleUserId: "new-requester", contextSnapshot: { wakeCommentId: later.messageId } }).where(eq(heartbeatRuns.id, later.run.id)).returning(); + states.push({ state: "board message of a new requester", run: laterRun, github: null }); + + // A long lineage: the request is twelve continuations back. + let deep = linked.run; + for (let depth = 0; depth < 12; depth++) deep = await continuation(linked.companyId, linked.agentId, linked.issueId, deep.activeIdentityContextId!); + states.push({ state: "long lineage of a linked request", run: deep, github: linked.commentId }); + return states; + } + + it("gives every run of a batch the decision of its own check, in every provenance state", async () => { + const states = await provenanceStates(); + const runs = states.map(state => state.run); + // The GitHub decision of each state is the documented one. + const github: Array>> = []; + for (const run of runs) github.push(await singleDecision(run, { provider: "github" })); + expect(Object.fromEntries(states.map((state, index) => [state.state, github[index]]))).toEqual( + Object.fromEntries(states.map(state => [state.state, state.github !== null && typeof state.github === "object" + ? { status: 403, message: state.github.forbidden } : { messageId: state.github }]))); + for (const options of optionSets) { + const single: Array>> = []; + for (const run of runs) single.push(await singleDecision(run, options)); + // Every state is decided in one batch, beside every other state. + expect(await batchDecisions(runs, options)).toEqual(single); + // The decision of a run does not depend on its position or on its neighbors. + expect((await batchDecisions([...runs].reverse(), options)).reverse()).toEqual(single); + for (const [index, run] of runs.entries()) expect(await batchDecisions([run], options)).toEqual([single[index]]); + } + // The legacy wake is returned only when the caller allows it. + const legacyWake = states.findIndex(state => state.state === "legacy wake of a verified parent run"); + expect((await batchDecisions(runs, { provider: "github", allowLegacyWakeComment: true }))[legacyWake]) + .toEqual({ messageId: expect.any(String) }); + expect(await chatRunSourceMessageIds(db, [])).toEqual([]); + }, 120_000); + + it("checks the runs of a batch with statements that follow the lineage, not the number of runs", async () => { + const first = await seedGitHubReviewWait(db); + const few = [first.run], many = [first.run]; + for (let wait = 1; wait < 40; wait++) { + const { run } = await seedGitHubReviewWait(db, { companyId: first.companyId, agentId: first.agentId }); + many.push(run); + if (wait < 4) few.push(run); + } + const options = { provider: "github" as const }; + const batchOfFew = await countStatements(() => chatRunSourceMessageIds(db, few, options)); + const batchOfMany = await countStatements(() => chatRunSourceMessageIds(db, many, options)); + // One identity-context read and one message-link read for the whole batch. + expect(batchOfFew.statements).toBe(2); + expect(batchOfMany.statements).toBe(2); + expect(new Set(batchOfMany.result.map(result => "messageId" in result ? result.messageId : null)).size).toBe(many.length); + // Separate checks send both reads for every run. + const separate = await countStatements(() => Promise.all(many.map(run => chatRunSourceMessageId(db, run, options)))); + expect(separate.statements).toBe(2 * many.length); + expect(batchOfMany.result).toEqual(separate.result.map(messageId => ({ messageId }))); + + // One continuation in every lineage adds the two reads of that step to + // the batch, not two reads for every run. + const continued: Run[] = []; + for (const run of many) continued.push(await continuation(first.companyId, first.agentId, run.nativeIssueId!, run.activeIdentityContextId!)); + const batchOfContinued = await countStatements(() => chatRunSourceMessageIds(db, continued, options)); + expect(batchOfContinued.statements).toBe(4); + expect(batchOfContinued.result).toEqual(batchOfMany.result); + }, 120_000); + it.each(["native", "issue-column", "legacy-issue", "legacy-task"])("uses the same %s task binding in SQL and runtime authority", async source => { const f = await fixture(); const values = { diff --git a/server/src/services/chat-channels/run-origin.ts b/server/src/services/chat-channels/run-origin.ts index f73f9dca22c..241055d2b5e 100644 --- a/server/src/services/chat-channels/run-origin.ts +++ b/server/src/services/chat-channels/run-origin.ts @@ -1,51 +1,124 @@ -import { and, eq, isNull, or, sql } from "drizzle-orm"; +import { and, eq, isNull, or, sql, type SQL } from "drizzle-orm"; import { chatConversations, chatEndpoints, chatMessageLinks, heartbeatRuns, issueComments, runIdentityContexts, type Db } from "@paperclipai/db"; import type { ChatProvider } from "@paperclipai/shared"; -import { forbidden } from "../../errors.js"; +import { forbidden, HttpError } from "../../errors.js"; import { boundRunIssueId, boundRunIssueIdSql } from "../run-task-binding.js"; type Run = Pick; type Context = typeof runIdentityContexts.$inferSelect; +type Options = { allowLegacyWakeComment?: boolean; provider?: ChatProvider }; -export async function chatRunSourceMessageId(db: Db, run: Run, options: { context?: Context | null; allowLegacyWakeComment?: boolean; provider?: ChatProvider } = {}): Promise { +/** A value of one provenance read: a literal of the run in a single read, or a + * column of the request row in a batch read. */ +type Value = string | SQL; + +const contextColumns = { + id: runIdentityContexts.id, runId: runIdentityContexts.runId, companyId: runIdentityContexts.companyId, + status: runIdentityContexts.status, responsibleUserId: runIdentityContexts.responsibleUserId, + messageId: runIdentityContexts.messageId, parentContextId: runIdentityContexts.parentContextId, +}; +const readColumns = { + context: contextColumns, + link: { issueId: chatConversations.issueId, provider: chatEndpoints.provider }, + boardMessage: { authorUserId: issueComments.authorUserId }, + parent: contextColumns, + sourceRun: { + issueId: boundRunIssueIdSql().as("bound_issue_id"), + responsibleUserId: heartbeatRuns.responsibleUserId, + wakeCommentId: sql`${heartbeatRuns.contextSnapshot}->'wakeCommentId'`.as("wake_comment_id"), + }, +}; + +/** The reads of the provenance check. Each read returns its first row, if + * any. The single check and the batch check build every read here, so both + * apply the same joins, company scope and filters. */ +const provenanceReads = { + /** The active identity context of the run. */ + context: (db: Db, of: Record<"contextId" | "runId" | "companyId", Value>) => db.select(readColumns.context) + .from(runIdentityContexts).where(and( + eq(runIdentityContexts.id, of.contextId), + eq(runIdentityContexts.runId, of.runId), + eq(runIdentityContexts.companyId, of.companyId), + eq(runIdentityContexts.status, "accepted"), + )).limit(1), + /** The channel conversation that received the accepted message. */ + link: (db: Db, of: Record<"messageId" | "companyId", Value>) => db.select(readColumns.link).from(chatMessageLinks) + .innerJoin(chatConversations, and(eq(chatConversations.id, chatMessageLinks.conversationId), eq(chatConversations.companyId, of.companyId))) + .innerJoin(chatEndpoints, and(eq(chatEndpoints.id, chatConversations.endpointId), eq(chatEndpoints.companyId, of.companyId))) + .where(and(eq(chatMessageLinks.companyId, of.companyId), eq(chatMessageLinks.commentId, of.messageId), + eq(chatMessageLinks.direction, "inbound"))).limit(1), + /** The accepted message as a board message of the task of the run. */ + boardMessage: (db: Db, of: Record<"messageId" | "companyId" | "issueId", Value>) => db.select(readColumns.boardMessage) + .from(issueComments).where(and( + eq(issueComments.id, of.messageId), eq(issueComments.companyId, of.companyId), + eq(issueComments.issueId, of.issueId), + or(eq(issueComments.authorType, "user"), and(isNull(issueComments.authorType), isNull(issueComments.authorAgentId))), + isNull(issueComments.deletedAt), + )).limit(1), + /** The identity context that the current context continues. */ + parent: (db: Db, of: Record<"contextId" | "companyId", Value>) => db.select(readColumns.parent) + .from(runIdentityContexts).where(and( + eq(runIdentityContexts.id, of.contextId), + eq(runIdentityContexts.companyId, of.companyId), + eq(runIdentityContexts.status, "accepted"), + )).limit(1), + /** The run of the parent context, for the same agent. */ + sourceRun: (db: Db, of: Record<"runId" | "companyId" | "agentId", Value>) => db.select(readColumns.sourceRun) + .from(heartbeatRuns).where(and( + eq(heartbeatRuns.id, of.runId), + eq(heartbeatRuns.companyId, of.companyId), + eq(heartbeatRuns.agentId, of.agentId), + )).limit(1), +}; + +type Kind = keyof typeof provenanceReads; +type Row = Awaited>[number]; +type Values = Record[1], string>; +type Read = { [K in Kind]: { kind: K; values: Values } }[Kind]; +type ReadOf = (db: Db, of: Record) => ReturnType; +type Walk = Generator; + +/** Every value of a read is a UUID. A run without a task binding matches no board message. */ +const NO_TASK = "00000000-0000-0000-0000-000000000000"; + +function* read(kind: K, values: Values): Generator | undefined, unknown> { + return (yield { kind, values } as Read) as Row | undefined; +} + +/** The provenance decision for one run. The walk asks for one read at a time + * and decides on the returned rows only, so the single check and the batch + * check make the same decision from the same rows. */ +function* sourceMessageWalk(run: Run, options: Options & { context?: Row<"context"> | null }): Walk { let legacyMessageId = options.allowLegacyWakeComment && typeof run.contextSnapshot?.wakeCommentId === "string" ? run.contextSnapshot.wakeCommentId : null; if (!run.activeIdentityContextId) { return legacyMessageId; } - let origin = options.context ?? (await db.select().from(runIdentityContexts).where(and( - eq(runIdentityContexts.id, run.activeIdentityContextId), - eq(runIdentityContexts.runId, run.id), - eq(runIdentityContexts.companyId, run.companyId), - eq(runIdentityContexts.status, "accepted"), - )).limit(1))[0]; - if (!origin || origin.id !== run.activeIdentityContextId || origin.runId !== run.id || - origin.companyId !== run.companyId || origin.status !== "accepted" || - origin.responsibleUserId !== run.responsibleUserId) { + const active = options.context ?? (yield* read("context", { + contextId: run.activeIdentityContextId, runId: run.id, companyId: run.companyId, + })); + if (!active || active.id !== run.activeIdentityContextId || active.runId !== run.id || + active.companyId !== run.companyId || active.status !== "accepted" || + active.responsibleUserId !== run.responsibleUserId) { throw forbidden("Chat requester provenance changed"); } + let origin: Row<"context"> = active; const seen = new Set(); let acceptedMessageSeen = false; while (true) { if (origin.messageId) { acceptedMessageSeen = true; if (!options.provider) return origin.messageId; - const [link] = await db.select({ issueId: chatConversations.issueId, provider: chatEndpoints.provider }).from(chatMessageLinks) - .innerJoin(chatConversations, and(eq(chatConversations.id, chatMessageLinks.conversationId), eq(chatConversations.companyId, run.companyId))) - .innerJoin(chatEndpoints, and(eq(chatEndpoints.id, chatConversations.endpointId), eq(chatEndpoints.companyId, run.companyId))) - .where(and(eq(chatMessageLinks.companyId, run.companyId), eq(chatMessageLinks.commentId, origin.messageId), eq(chatMessageLinks.direction, "inbound"))).limit(1); + const link = yield* read("link", { messageId: origin.messageId, companyId: run.companyId }); if (link) { if (link.issueId !== boundRunIssueId(run)) throw forbidden("Chat request belongs to another task"); if (link.provider !== options.provider) return null; return origin.messageId; } - const [message] = await db.select({ authorUserId: issueComments.authorUserId }).from(issueComments).where(and( - eq(issueComments.id, origin.messageId), eq(issueComments.companyId, run.companyId), - eq(issueComments.issueId, boundRunIssueId(run) ?? "00000000-0000-0000-0000-000000000000"), - or(eq(issueComments.authorType, "user"), and(isNull(issueComments.authorType), isNull(issueComments.authorAgentId))), - isNull(issueComments.deletedAt), - )).limit(1); + const message = yield* read("boardMessage", { + messageId: origin.messageId, companyId: run.companyId, issueId: boundRunIssueId(run) ?? NO_TASK, + }); if (!run.responsibleUserId || message?.authorUserId !== run.responsibleUserId) throw forbidden("Chat requester provenance changed"); // This verified board message may continue its same-user channel history, // but its own wake metadata is not the earlier channel request. @@ -54,25 +127,13 @@ export async function chatRunSourceMessageId(db: Db, run: Run, options: { contex if (!origin.parentContextId) return legacyMessageId; if (seen.has(origin.id)) throw forbidden("Chat source provenance could not be verified"); seen.add(origin.id); - const [parent] = await db.select().from(runIdentityContexts).where(and( - eq(runIdentityContexts.id, origin.parentContextId), - eq(runIdentityContexts.companyId, run.companyId), - eq(runIdentityContexts.status, "accepted"), - )).limit(1); + const parent: Row<"parent"> | undefined = yield* read("parent", { contextId: origin.parentContextId, companyId: run.companyId }); if (!parent) throw forbidden("Chat requester provenance changed"); if (parent.responsibleUserId !== run.responsibleUserId) { if (acceptedMessageSeen) return null; throw forbidden("Chat requester provenance changed"); } - const [sourceRun] = await db.select({ - issueId: boundRunIssueIdSql(), - responsibleUserId: heartbeatRuns.responsibleUserId, - wakeCommentId: sql`${heartbeatRuns.contextSnapshot}->'wakeCommentId'`, - }).from(heartbeatRuns).where(and( - eq(heartbeatRuns.id, parent.runId), - eq(heartbeatRuns.companyId, run.companyId), - eq(heartbeatRuns.agentId, run.agentId), - )).limit(1); + const sourceRun = yield* read("sourceRun", { runId: parent.runId, companyId: run.companyId, agentId: run.agentId }); if (!sourceRun || !boundRunIssueId(run) || sourceRun.issueId !== boundRunIssueId(run)) { throw forbidden("Chat tools cannot be inherited by another task"); } @@ -81,3 +142,71 @@ export async function chatRunSourceMessageId(db: Db, run: Run, options: { contex origin = parent; } } + +export async function chatRunSourceMessageId(db: Db, run: Run, options: Options & { context?: Context | null } = {}): Promise { + const walk = sourceMessageWalk(run, options); + let step = walk.next(); + while (!step.done) { + const [row] = await (provenanceReads[step.value.kind] as ReadOf)(db, step.value.values); + step = walk.next(row); + } + return step.value; +} + +const request = sql.identifier("chat_run_source_request"); +const ORDINAL = "chatRunSourceRequestOrdinal"; + +/** The reads of one kind for many runs, in one statement. For each request + * row, the lateral subquery is the single read with the request columns in + * place of the literals, so each request gets the row of its single read. */ +async function readRows(db: Db, kind: Kind, requests: Array>): Promise { + const names = Object.keys(requests[0]!); + const column = (name: string) => sql`${request}.${sql.identifier(name)}`; + const found = (provenanceReads[kind] as ReadOf)(db, Object.fromEntries(names.map(name => [name, column(name)]))) + .as("chat_run_source_row"); + const columns = found as unknown as Record; + const rows = await db.select({ + [ORDINAL]: sql`${column("ordinal")}`.mapWith(Number), + ...Object.fromEntries(Object.keys(readColumns[kind]).map(name => [name, columns[name]!])), + }).from(sql`unnest(${sql.join(names.map(name => sql`${sql.param(requests.map(values => values[name]))}::uuid[]`), sql`, `)}) + with ordinality as ${request}(${sql.join([...names, "ordinal"].map(name => sql.identifier(name)), sql`, `)})`) + .crossJoinLateral(found); + const result: unknown[] = new Array(requests.length).fill(undefined); + for (const { [ORDINAL]: ordinal, ...row } of rows) result[(ordinal as number) - 1] = row; + return result; +} + +/** The result of `chatRunSourceMessageId` for one run of a batch: the message + * that the single check returns, or the 403 that it throws. */ +export type ChatRunSourceMessage = { messageId: string | null } | { forbidden: HttpError }; + +/** `chatRunSourceMessageId` for many runs. Each run gets the decision of a + * separate call, from the same walk and the same reads. The runs that wait for + * the same kind of read share one statement, so the number of statements + * follows the longest requester lineage, not the number of runs. A database + * error rejects the whole batch. */ +export async function chatRunSourceMessageIds(db: Db, runs: Run[], options: Options = {}): Promise { + const results = new Array(runs.length); + let waiting: Array<{ index: number; walk: Walk; read: Read }> = []; + const advance = (index: number, walk: Walk, row?: unknown) => { + try { + const step = walk.next(row); + if (step.done) results[index] = { messageId: step.value }; + else waiting.push({ index, walk, read: step.value }); + } catch (error) { + if (!(error instanceof HttpError) || error.status !== 403) throw error; + results[index] = { forbidden: error }; + } + }; + runs.forEach((run, index) => advance(index, sourceMessageWalk(run, options))); + while (waiting.length) { + const round = waiting; + waiting = []; + for (const kind of new Set(round.map(entry => entry.read.kind))) { + const entries = round.filter(entry => entry.read.kind === kind); + const rows = await readRows(db, kind, entries.map(entry => entry.read.values)); + entries.forEach((entry, position) => advance(entry.index, entry.walk, rows[position])); + } + } + return results; +}