diff --git a/doc/idle-domain-deadlines.md b/doc/idle-domain-deadlines.md index 73b537ee223..1ec36c14301 100644 --- a/doc/idle-domain-deadlines.md +++ b/doc/idle-domain-deadlines.md @@ -52,9 +52,10 @@ ordinary routine/monitor proof, exact domain rows and plugin worker identities. Any changed state, cancellation, moved deadline, owner replacement or restart invalidates the old proof. Unknown or malformed states fail closed. +The timer owner batches UTC-day run counts, exact decimal cost sums and actionable-work inputs. Credential ownership and selected current/previous version metadata use bounded chunks; they never read credential material or value hashes. Managed GitHub OAuth user and agent grants retain their owner deadlines. Other grant shapes grant no continuity exemption. + The additional domain projection is bounded to 64 data reads, a two-second monotonic budget for -starting further reads, 128 eligible timer candidates and 10,000 rows per -family, while reserving 60 seconds of hold headroom. Overflow refuses sleep; +starting further reads and 10,000 rows per family, while reserving 60 seconds of hold headroom. Overflow refuses sleep; it never truncates the proof or changes a schedule. These limits supplement the transaction's statement timeout. A query already in progress may still take its five-second statement timeout; SQL execution can therefore finish about seven seconds after the projection diff --git a/server/src/__tests__/agent-timer-deadlines.test.ts b/server/src/__tests__/agent-timer-deadlines.test.ts index f925fb8dfd7..d6cc9649ea6 100644 --- a/server/src/__tests__/agent-timer-deadlines.test.ts +++ b/server/src/__tests__/agent-timer-deadlines.test.ts @@ -44,7 +44,7 @@ describe.skipIf(!support.supported)("agent timer owned deadlines", () => { beforeAll(async () => { if (external) await applyPendingMigrations(external); else temporary = await startEmbeddedPostgresTestDatabase("agent-timer-deadlines-"); - db = createDb(external ?? temporary!.connectionString, { idleTimeoutSeconds: 0, connectTimeoutSeconds: 0, maxLifetimeSeconds: 0 }); + db = createDb(external ?? temporary!.connectionString, { maxConnections: 4, idleTimeoutSeconds: 0, connectTimeoutSeconds: 0, maxLifetimeSeconds: 0 }); }, 30_000); beforeEach(async () => { diff --git a/server/src/__tests__/idle-domain-deadlines-postgres.test.ts b/server/src/__tests__/idle-domain-deadlines-postgres.test.ts index 858e6b7072c..31fa4de5c1f 100644 --- a/server/src/__tests__/idle-domain-deadlines-postgres.test.ts +++ b/server/src/__tests__/idle-domain-deadlines-postgres.test.ts @@ -20,8 +20,8 @@ import { startTaskDrain, stopTaskDrain, beginIdleTrackedWork } from "../services import { providerTraceStore } from "../services/provider-trace-store.js"; import { readIdleSleepSafety } from "../services/idle-sleep-safety.js"; import { randomUUID } from "node:crypto"; -import { eq, sql } from "drizzle-orm"; -import { agents, companies, heartbeatRuns, plugins, pluginJobs, pluginJobRuns, pluginState, providerTraceRecords, issues, issueWatchdogs, issueThreadInteractions, toolApplications, toolConnections, connectionGrants, companySecrets, +import { eq, sql, inArray } from "drizzle-orm"; +import { agents, companies, heartbeatRuns, plugins, pluginJobs, pluginJobRuns, pluginState, providerTraceRecords, issues, issueWatchdogs, issueThreadInteractions, toolApplications, toolConnections, connectionGrants, companySecrets, companySecretVersions, applyPendingMigrations, createDb, type Db } from "@paperclipai/db"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; @@ -55,9 +55,12 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu let hold: { ownerId: string; expiresAt: number } = { ownerId: randomUUID(), expiresAt: now + 300_000 }; let pluginReceipt: Awaited>> | undefined; const capability = { deadlineCatalog: "domain_deadlines_v1" } as VerifiedIdleCapabilities; + let projectionReads = 0; const read = (quietPluginIds: string[] = []) => db.transaction(async tx => { + projectionReads = 0; await tx.execute(sql`SET LOCAL statement_timeout = '5s'`); - return readIdleDeadlineProof(tx, hold, { now: Date.now(), verifiedCapability: capability, quietPluginIds, + const counted = { select: (...args: unknown[]) => { projectionReads++; return Reflect.apply(tx.select, tx, args); } } as Pick; + return readIdleDeadlineProof(counted, hold, { now: Date.now(), verifiedCapability: capability, quietPluginIds, quietPluginGeneration: pluginReceipt?.catalogGeneration, quietPluginCatalog: pluginReceipt?.catalog }); }, { isolationLevel: "repeatable read", accessMode: "read only" }); await db.insert(companies).values({ id: companyId, name: "Deadline fixture", issuePrefix: `D${companyId.slice(0, 6)}` }); @@ -104,6 +107,27 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu expect(noAlarm).toMatchObject({ version: 3, backgroundWork: "none", deadlineProof: { catalog: "domain_deadlines_v1" } }); expect(noAlarm.scheduledWake).toBeUndefined(); expect(noAlarm.deadlineProof?.generation).not.toBe(negotiated.deadlineProof?.generation); + const scaleIds = Array.from({ length: 160 }, () => randomUUID()); + const scaleAgent = (id: string, cost = false) => ({ id, companyId, name: "Scale timer", role: "engineer", status: "idle", adapterType: "codex_local", + runtimeConfig: { heartbeat: { enabled: true, intervalSec: 900, maxDailyRuns: 2, ...(cost ? { maxDailyCostCents: 100 } : {}) } }, + createdAt: new Date(now), lastHeartbeatAt: new Date(now) }); + await db.insert(agents).values(scaleIds.slice(0, 80).map(id => scaleAgent(id))); + expect((await read()).scheduledWake?.nextWakeAt).toBe(new Date(now + 900_000).toISOString()); + expect(projectionReads).toBe(5); + await db.insert(agents).values(scaleIds.slice(80).map(id => scaleAgent(id, true))); + const inactiveIds = [randomUUID(), randomUUID(), randomUUID()]; + await db.insert(agents).values(inactiveIds.map((id, i) => ({ ...scaleAgent(id), status: i === 0 ? "paused" : i === 1 ? "terminated" : "idle", + ...(i === 2 ? { runtimeConfig: { heartbeat: { enabled: false } } } : {}) }))); + const scaled = await read(); + expect(scaled.blocked).toBe(false); + expect(scaled.scheduledWake?.nextWakeAt).toBe(new Date(now + 900_000).toISOString()); + expect(projectionReads).toBe(6); + await db.update(agents).set({ lastHeartbeatAt: new Date(now - 900_000) }).where(eq(agents.id, scaleIds[0])); + expect((await read()).blocked).toBe(true); + await db.update(agents).set({ runtimeConfig: { heartbeat: { enabled: false } } }).where(inArray(agents.id, [...scaleIds, ...inactiveIds])); + const cancelledScale = await read(); + expect(cancelledScale.scheduledWake).toBeUndefined(); + expect(cancelledScale.generation).not.toBe(scaled.generation); await db.update(agents).set({ runtimeConfig: { heartbeat: { enabled: true, intervalSec: 0.5 } }, lastHeartbeatAt: new Date(now) }).where(eq(agents.id, agentId)); let checkpointClock = now; const becameDue = await readIdleSleepSafety(db, () => ({ ownerId: hold.ownerId, draining: true, startedAt: new Date(now - 1000), @@ -234,6 +258,27 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu await db.update(connectionGrants).set({ credentialSecretRefs: refs }).where(eq(connectionGrants.id, grant.id)); const vault = secretService(db), credentials = createOAuthCredentialStore(db, vault); await credentials.syncCredentialBindings(connection, refs); + const pinned = await writeConnectionCredential(db, { companyId, connectionName: "GitHub pinned", configPath: "oauth.access_token", + label: "Pinned access", value: "pinned-initial", ownerUserId: owner }); + const pinnedRef = { configPath: "oauth.access_token", secretId: pinned.secret.id, versionSelector: 1 as const, required: true }; + await writeConnectionCredential(db, { companyId, connectionName: "GitHub pinned", configPath: "oauth.access_token", label: "Pinned access", + value: "pinned-rotated", ownerUserId: owner, existingRef: pinnedRef }); + await db.update(connectionGrants).set({ credentialSecretRefs: [pinnedRef, refs[1]] }).where(eq(connectionGrants.id, grant.id)); + await credentials.syncCredentialBindings(connection, [pinnedRef, refs[1]]); + const [pinnedGrant] = await db.select().from(connectionGrants).where(eq(connectionGrants.id, grant.id)); + expect((await credentials.resolveOAuthGrantSecret(connection, pinnedGrant, pinnedRef, + { actorType: "system", actorId: "fixture" }, undefined)).value).toBe("pinned-initial"); + expect((await read()).blocked).toBe(false); + for (const status of ["disabled", "destroyed"]) { + await db.update(companySecretVersions).set({ status }).where(eq(companySecretVersions.secretId, pinned.secret.id)); + expect((await read()).blocked).toBe(true); + } + await db.update(companySecretVersions).set({ status: "previous", revokedAt: new Date() }).where(eq(companySecretVersions.secretId, pinned.secret.id)); + expect((await read()).blocked).toBe(true); + await db.update(companySecretVersions).set({ revokedAt: null }).where(eq(companySecretVersions.secretId, pinned.secret.id)); + expect((await read()).blocked).toBe(false); + await db.update(connectionGrants).set({ credentialSecretRefs: refs }).where(eq(connectionGrants.id, grant.id)); + await credentials.syncCredentialBindings(connection, refs); const baseline = await read(); expect(baseline.blocked).toBe(false); expect(baseline.connectionIds).toContain(connection.id); diff --git a/server/src/__tests__/tool-access-service.test.ts b/server/src/__tests__/tool-access-service.test.ts index aa5afc03b1a..00f40f1c9f1 100644 --- a/server/src/__tests__/tool-access-service.test.ts +++ b/server/src/__tests__/tool-access-service.test.ts @@ -930,14 +930,14 @@ describeEmbeddedPostgres("tool access service", () => { beforeAll(async () => { if (externalDatabaseUrl) { await applyPendingMigrations(externalDatabaseUrl); - db = createDb(externalDatabaseUrl); + db = createDb(externalDatabaseUrl, { maxConnections: 4 }); await instanceSettingsService(db).updateExperimental({ enableMemoryConnectors: true }); return; } tempDb = await startEmbeddedPostgresTestDatabase( "paperclip-tool-access-service-", ); - db = createDb(tempDb.connectionString); + db = createDb(tempDb.connectionString, { maxConnections: 4 }); await instanceSettingsService(db).updateExperimental({ enableMemoryConnectors: true }); }, 20_000); @@ -976,6 +976,7 @@ describeEmbeddedPostgres("tool access service", () => { }); afterAll(async () => { + await db?.$client.end({ timeout: 1 }); await tempDb?.cleanup(); }); diff --git a/server/src/services/connection-credential-state.test.ts b/server/src/services/connection-credential-state.test.ts new file mode 100644 index 00000000000..d0bbcf7dc3a --- /dev/null +++ b/server/src/services/connection-credential-state.test.ts @@ -0,0 +1,67 @@ +import { describe, expect, it } from "vitest"; +import { companySecrets, companySecretVersions, type Db } from "@paperclipai/db"; +import { PgDialect } from "drizzle-orm/pg-core"; +import type { SQL } from "drizzle-orm"; +import { readConnectionGrantCredentialStates, type ConnectionGrantCredentialInput } from "./connection-credentials.js"; + +const input = (index = 0): ConnectionGrantCredentialInput => ({ connection: { id: "connection", companyId: "company" }, + grant: { id: `grant-${index}`, companyId: "company", connectionId: "connection", kind: "user", subjectUserId: "owner", + credentialSecretRefs: [{ configPath: "oauth.access_token", secretId: `secret-${index}`, versionSelector: 1 }] } }); +const metadata = (index = 0) => ({ secret: { id: `secret-${index}`, companyId: "company", scope: "user", ownerUserId: "owner", + userSecretDefinitionId: `definition-${index}`, status: "active", deletedAt: null, latestVersion: 2, updatedAt: new Date(0) }, + definition: { id: `definition-${index}`, companyId: "company", status: "active", deletedAt: null, updatedAt: new Date(0) } }); +const version = (index = 0) => ({ id: `version-${index}`, secretId: `secret-${index}`, version: 1, status: "previous", revokedAt: null as Date | null }); +function metadataDb(rows = [metadata()], versions = [version()], onRead: () => void = () => {}) { + return { select: (columns: Record) => { + onRead(); + for (const selected of [columns, columns.secret ?? {}, columns.definition ?? {}]) { + expect(Object.keys(selected as object)).not.toEqual(expect.arrayContaining(["material"])); + for (const name of ["valueSha256", "fingerprintSha256", "providerMetadata"]) expect(selected).not.toHaveProperty(name); + } + let table: unknown, params: unknown[] = []; + const query = { from: (source: unknown) => { table = source; return query; }, leftJoin: () => query, + where: (where: SQL) => { params = new PgDialect().sqlToQuery(where).params; return query; }, orderBy: () => query, + limit: () => Promise.resolve(table === companySecrets ? rows.filter(row => params.includes(row.secret.id)) + : table === companySecretVersions ? versions.filter(row => params.includes(row.secretId) && params.includes(row.version)) : []) }; + return query; + } } as unknown as Pick; +} +describe("connection grant credential metadata", () => { + it("accepts the exact non-revoked pinned previous version and also selects current latest", async () => { + const request = input(), row = metadata(), saved = version(); + expect((await readConnectionGrantCredentialStates(metadataDb([row], [saved]), [request]))?.size).toBe(1); + request.grant.credentialSecretRefs[0].versionSelector = "latest"; + expect(await readConnectionGrantCredentialStates(metadataDb([row], [saved]), [request])).toBeUndefined(); + saved.version = 2; saved.status = "current"; + expect((await readConnectionGrantCredentialStates(metadataDb([row], [saved]), [request]))?.size).toBe(1); + }); + it("keeps revoked, disabled, destroyed, failed, absent and foreign metadata unproven", async () => { + for (const status of ["disabled", "destroyed", "failed", "unknown"]) { + const saved = { ...version(), status }; + expect(await readConnectionGrantCredentialStates(metadataDb([metadata()], [saved]), [input()])).toBeUndefined(); + } + const revoked = { ...version(), revokedAt: new Date() }; + expect(await readConnectionGrantCredentialStates(metadataDb([metadata()], [revoked]), [input()])).toBeUndefined(); + expect(await readConnectionGrantCredentialStates(metadataDb([], []), [input()])).toBeUndefined(); + expect(await readConnectionGrantCredentialStates(metadataDb([metadata()], []), [input()])).toBeUndefined(); + for (const field of ["companyId", "ownerUserId", "userSecretDefinitionId"] as const) { + const row = metadata(); row.secret[field] = "foreign"; + expect(await readConnectionGrantCredentialStates(metadataDb([row]), [input()])).toBeUndefined(); + } + const row = metadata(); row.definition.status = "disabled"; + expect(await readConnectionGrantCredentialStates(metadataDb([row]), [input()])).toBeUndefined(); + const duplicate = input(); duplicate.grant.credentialSecretRefs.push({ ...duplicate.grant.credentialSecretRefs[0] }); + expect(await readConnectionGrantCredentialStates(metadataDb(), [duplicate])).toBeUndefined(); + }); + it("bounds1001 grants to four metadata reads and checks budget during collection", async () => { + const inputs = Array.from({ length: 1001 }, (_, i) => input(i)); + let reads = 0; + const db = metadataDb(inputs.map((_, i) => metadata(i)), inputs.map((_, i) => version(i)), () => { reads++; }); + expect((await readConnectionGrantCredentialStates(db, inputs))?.size).toBe(1001); + expect(reads).toBe(4); + let checks = 0; + await expect(readConnectionGrantCredentialStates(db, inputs, () => { + if (++checks > 3) throw new Error("budget exhausted"); + })).rejects.toThrow("budget exhausted"); + }); +}); diff --git a/server/src/services/connection-credentials.ts b/server/src/services/connection-credentials.ts index a8e3407705e..4af4f633883 100644 --- a/server/src/services/connection-credentials.ts +++ b/server/src/services/connection-credentials.ts @@ -117,26 +117,36 @@ export async function writeConnectionCredential( return { secret, created: true, definitionId: createdDefinitionId }; } -/** Validate metadata without reading secret values, including for connection pickers. */ -export async function validateConnectionGrantSecretOwnership( - db: Pick, - connection: Pick, - grant: Pick, - ref: ToolCredentialSecretRef, -) { - const [secret] = await db.select().from(companySecrets).where(and( - eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, connection.companyId), - )).limit(1); +const credentialSecretColumns = { + id: companySecrets.id, companyId: companySecrets.companyId, scope: companySecrets.scope, + ownerUserId: companySecrets.ownerUserId, userSecretDefinitionId: companySecrets.userSecretDefinitionId, + status: companySecrets.status, deletedAt: companySecrets.deletedAt, latestVersion: companySecrets.latestVersion, + updatedAt: companySecrets.updatedAt, +}; +type CredentialConnection = Pick; +type CredentialGrantOwner = Pick; +type CredentialSecretOwner = Pick; +function assertConnectionGrantSecretOwnership(connection: CredentialConnection, grant: CredentialGrantOwner, + ref: ToolCredentialSecretRef, secret: CredentialSecretOwner | undefined) { const personal = grant.kind === "user"; if (grant.companyId !== connection.companyId || grant.connectionId !== connection.id || !secret - || (personal + || secret.companyId !== connection.companyId || (personal ? !grant.subjectUserId || secret.scope !== "user" || secret.ownerUserId !== grant.subjectUserId || !secret.userSecretDefinitionId : secret.scope !== "company")) { throw unprocessable("This connection's credential does not belong to the selected identity. Its owner must reconnect it.", { code: "grant_credential_invalid", connectionId: connection.id, grantId: grant.id, credential: ref.configPath, }); } - return secret; +} +export async function validateConnectionGrantSecretOwnership(db: Pick, + connection: CredentialConnection, grant: CredentialGrantOwner, ref: ToolCredentialSecretRef) { + const [secret] = await db.select(credentialSecretColumns).from(companySecrets).where(and( + eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, connection.companyId), + )).limit(1); + assertConnectionGrantSecretOwnership(connection, grant, ref, secret); + return secret!; } /** The same owner/declaration checks apply during setup, testing and execution. */ @@ -166,38 +176,82 @@ export async function resolveConnectionGrantSecret( return { value: resolved.value, latestVersion: secret.latestVersion }; } -export async function readConnectionGrantCredentialState( - db: Pick, - connection: Pick, - grant: Pick, -): Promise { - const refs = grant.credentialSecretRefs.filter(ref => ["oauth.access_token", "oauth.refresh_token"].includes(ref.configPath)); - if (refs.filter(ref => ref.configPath === "oauth.access_token").length !== 1 - || new Set(refs.map(ref => ref.configPath)).size !== refs.length) return; - const state: unknown[] = []; - for (const ref of refs) { - let secret; - try { secret = await validateConnectionGrantSecretOwnership(db, connection, grant, ref); } - catch (error) { - if (error && typeof error === "object" && "status" in error && error.status === 422) return; - throw error; +type CredentialGrant = CredentialGrantOwner & Pick; +export type ConnectionGrantCredentialInput = { connection: CredentialConnection; grant: CredentialGrant }; +export async function readConnectionGrantCredentialStates(db: Pick, + inputs: ConnectionGrantCredentialInput[], checkBudget: () => void = () => {}): Promise | undefined> { + if (inputs.length > 10_000) return; + const states = new Map(); + const requests: { input: ConnectionGrantCredentialInput; ref: ToolCredentialSecretRef }[] = []; + for (const input of inputs) { + checkBudget(); + if (states.has(input.grant.id)) return; + states.set(input.grant.id, []); + const paths = new Set(); + for (const ref of input.grant.credentialSecretRefs) { + checkBudget(); + if (!["oauth.access_token", "oauth.refresh_token"].includes(ref.configPath)) continue; + if (paths.has(ref.configPath)) return; + paths.add(ref.configPath); + requests.push({ input, ref }); } - if (secret.status !== "active" || secret.deletedAt || !Number.isSafeInteger(secret.latestVersion) || secret.latestVersion <= 0) return; - let definitionState: unknown; - if (secret.scope === "user") { - const [definition] = await db.select({ id: userSecretDefinitions.id, companyId: userSecretDefinitions.companyId, - status: userSecretDefinitions.status, deletedAt: userSecretDefinitions.deletedAt, updatedAt: userSecretDefinitions.updatedAt }) - .from(userSecretDefinitions).where(eq(userSecretDefinitions.id, secret.userSecretDefinitionId!)).limit(1); - if (!definition || definition.companyId !== connection.companyId || definition.status !== "active" || definition.deletedAt) return; - definitionState = definition; - } - const version = ref.versionSelector === undefined || ref.versionSelector === "latest" ? secret.latestVersion : ref.versionSelector; - if (typeof version !== "number" || !Number.isSafeInteger(version) || version <= 0) return; - const [saved] = await db.select({ id: companySecretVersions.id, version: companySecretVersions.version, - status: companySecretVersions.status, revokedAt: companySecretVersions.revokedAt }) - .from(companySecretVersions).where(and(eq(companySecretVersions.secretId, secret.id), eq(companySecretVersions.version, version))).limit(1); - if (!saved || saved.status !== "current" || saved.revokedAt) return; - state.push([ref.configPath, secret.id, secret.scope, secret.ownerUserId, secret.latestVersion, secret.updatedAt, definitionState, saved]); + if (!paths.has("oauth.access_token")) return; } - return state; + if (requests.length > 20_000) return; + for (let offset = 0; offset < requests.length; offset += 1000) { + checkBudget(); + const chunk = requests.slice(offset, offset + 1000); + const ids = [...new Set(chunk.map(request => request.ref.secretId))]; + const rows = await db.select({ secret: credentialSecretColumns, definition: { + id: userSecretDefinitions.id, companyId: userSecretDefinitions.companyId, status: userSecretDefinitions.status, + deletedAt: userSecretDefinitions.deletedAt, updatedAt: userSecretDefinitions.updatedAt, + } }).from(companySecrets).leftJoin(userSecretDefinitions, eq(userSecretDefinitions.id, companySecrets.userSecretDefinitionId)) + .where(inArray(companySecrets.id, ids)).orderBy(companySecrets.id).limit(1001); + if (rows.length > ids.length) return; + const secrets = new Map(rows.map(row => [row.secret.id, row])); + if (secrets.size !== rows.length) return; + const selected = []; + const pairs = new Map(); + for (const request of chunk) { + checkBudget(); + const row = secrets.get(request.ref.secretId); + try { assertConnectionGrantSecretOwnership(request.input.connection, request.input.grant, request.ref, row?.secret); } + catch (error) { + if (error && typeof error === "object" && "status" in error && error.status === 422) return; + throw error; + } + if (!row) return; + const { secret, definition } = row; + if (secret.status !== "active" || secret.deletedAt || !Number.isSafeInteger(secret.latestVersion) || secret.latestVersion <= 0) return; + if (secret.scope === "user" && (!definition || definition.id !== secret.userSecretDefinitionId + || definition.companyId !== request.input.connection.companyId || definition.status !== "active" || definition.deletedAt)) return; + const version = request.ref.versionSelector === undefined || request.ref.versionSelector === "latest" + ? secret.latestVersion : request.ref.versionSelector; + if (typeof version !== "number" || !Number.isSafeInteger(version) || version <= 0) return; + const key = JSON.stringify([secret.id, version]); + pairs.set(key, { secretId: secret.id, version }); + selected.push({ request, secret, definition: secret.scope === "user" ? definition : undefined, key }); + } + checkBudget(); + const versions = await db.select({ id: companySecretVersions.id, secretId: companySecretVersions.secretId, + version: companySecretVersions.version, status: companySecretVersions.status, revokedAt: companySecretVersions.revokedAt }) + .from(companySecretVersions).where(or(...[...pairs.values()].map(pair => and( + eq(companySecretVersions.secretId, pair.secretId), eq(companySecretVersions.version, pair.version))))) + .orderBy(companySecretVersions.secretId, companySecretVersions.version).limit(1001); + if (versions.length > pairs.size) return; + const saved = new Map(versions.map(version => [JSON.stringify([version.secretId, version.version]), version])); + if (saved.size !== versions.length || versions.some(version => !pairs.has(JSON.stringify([version.secretId, version.version])))) return; + for (const { request, secret, definition, key } of selected) { + checkBudget(); + const version = saved.get(key); + if (!version || !["current", "previous"].includes(version.status) || version.revokedAt) return; + states.get(request.input.grant.id)!.push([request.ref.configPath, secret.id, secret.scope, secret.ownerUserId, + secret.latestVersion, secret.updatedAt, definition, version]); + } + } + return states; +} +export async function readConnectionGrantCredentialState(db: Pick, + connection: CredentialConnection, grant: CredentialGrant): Promise { + return (await readConnectionGrantCredentialStates(db, [{ connection, grant }]))?.get(grant.id); } diff --git a/server/src/services/heartbeat/idle-timer-schedule.ts b/server/src/services/heartbeat/idle-timer-schedule.ts index ce822586b69..e72f9b3c3c5 100644 --- a/server/src/services/heartbeat/idle-timer-schedule.ts +++ b/server/src/services/heartbeat/idle-timer-schedule.ts @@ -8,20 +8,25 @@ import { idleDeadlineRowLimit, isFutureIdleDeadline, type IdleDeadlineSnapshot } export async function readIdleAgentTimerSchedule(tx: Pick, now: number, env: Record = process.env, checkBudget: () => void = () => {}): Promise { + const flags = [env.HEARTBEAT_SCHEDULER_ENABLED, env.PAPERCLIP_IN_WORKTREE, + env.PAPERCLIP_DATABASE_RESTORE_IN_PROGRESS, env.PAPERCLIP_RESTORE_IN_PROGRESS]; + const state: unknown[] = [flags]; + const blocked = () => ({ blocked: true, state, deadlines: [] }); + if (flags[0] === "false" || flags.slice(1).some(v => ["true", "1", "yes", "on"].includes(v ?? ""))) return blocked(); const rows = await tx.select({ agent: getTableColumns(agents), companyStatus: companies.status, companyUpdatedAt: companies.updatedAt }).from(agents).leftJoin(companies, eq(companies.id, agents.companyId)) .orderBy(agents.id).limit(idleDeadlineRowLimit + 1); - const state: unknown[] = [rows]; - const blocked = () => ({ blocked: true, state, deadlines: [] }); + state.push(rows); if (rows.length > idleDeadlineRowLimit) return blocked(); - const byCompany = new Map(); + const byCompany = new Map>(); for (const { agent } of rows) { - const group = byCompany.get(agent.companyId) ?? []; - group.push(agent); + const group = byCompany.get(agent.companyId) ?? new Map(); + if (group.has(agent.id)) return blocked(); + group.set(agent.id, agent); byCompany.set(agent.companyId, group); } const owner = createHeartbeatTimerPolicy(tx); - let inspectedTimers = 0; + const candidates: typeof agents.$inferSelect[] = []; const deadlines: number[] = []; for (const { agent, companyStatus } of rows) { checkBudget(); @@ -31,17 +36,29 @@ export async function readIdleAgentTimerSchedule(tx: Pick, now: nu const policy = owner.parseHeartbeatPolicy(agent); if (!Number.isFinite(policy.intervalSec)) return blocked(); if (!policy.enabled || policy.intervalSec <= 0 || !supportsUnscopedTimerHeartbeat(agent)) continue; - if (++inspectedTimers > 128) return blocked(); - const invokability = evaluateAgentInvokability(agent, byCompany.get(agent.companyId)!); + const ancestors: typeof agents.$inferSelect[] = []; + const seen = new Set(); + let current: typeof agents.$inferSelect | undefined = agent; + while (current && !seen.has(current.id)) { + checkBudget(); + seen.add(current.id); + ancestors.push(current); + current = current.reportsTo ? byCompany.get(agent.companyId)!.get(current.reportsTo) : undefined; + } + const invokability = evaluateAgentInvokability(agent, ancestors); if (!invokability.invokable) { if (invokability.invalidOrgChain || invokability.reason === "unknown_status") return blocked(); continue; } - const flags = [env.HEARTBEAT_SCHEDULER_ENABLED, env.PAPERCLIP_IN_WORKTREE, - env.PAPERCLIP_DATABASE_RESTORE_IN_PROGRESS, env.PAPERCLIP_RESTORE_IN_PROGRESS]; - state.push(flags); - if (flags[0] === "false" || flags.slice(1).some(v => ["true", "1", "yes", "on"].includes(v ?? ""))) return blocked(); - const snapshot = await owner.readTimerHeartbeatSchedule(agent, new Date(now)); + candidates.push(agent); + } + const batch = await owner.readTimerHeartbeatSchedules(candidates, new Date(now), checkBudget); + if (!batch || batch.schedules.size !== candidates.length) return blocked(); + state.push(batch.state); + for (const agent of candidates) { + checkBudget(); + const snapshot = batch.schedules.get(agent.id); + if (!snapshot) return blocked(); state.push([agent.id, snapshot]); if (snapshot.deadline === null) continue; if (!isFutureIdleDeadline(snapshot.deadline, now)) return blocked(); diff --git a/server/src/services/heartbeat/queue.test.ts b/server/src/services/heartbeat/queue.test.ts index a284c214f09..729504e5663 100644 --- a/server/src/services/heartbeat/queue.test.ts +++ b/server/src/services/heartbeat/queue.test.ts @@ -83,7 +83,7 @@ describe.skipIf(!support.supported)("heartbeat queue database wiring", () => { let db: Db; beforeAll(async () => { if (!external) database = await startEmbeddedPostgresTestDatabase("paperclip-heartbeat-queue-module"); - db = createDb(external ?? database!.connectionString); + db = createDb(external ?? database!.connectionString, { maxConnections: 4 }); }, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS); afterAll(async () => { await db?.$client.end(); diff --git a/server/src/services/heartbeat/timer-policy.ts b/server/src/services/heartbeat/timer-policy.ts index 161a5994836..51b785011ce 100644 --- a/server/src/services/heartbeat/timer-policy.ts +++ b/server/src/services/heartbeat/timer-policy.ts @@ -1,6 +1,6 @@ import { agents, heartbeatRuns, costEvents, issues, type Db } from "@paperclipai/db"; import { AGENT_DEFAULT_MAX_CONCURRENT_RUNS, compareCents } from "@paperclipai/shared"; -import { and, eq, gte, lt, notInArray, sql, isNull, inArray } from "drizzle-orm"; +import { and, eq, gte, lt, notInArray, sql, isNull, inArray, type SQL } from "drizzle-orm"; import { asBoolean, asNumber, parseObject } from "../../adapters/utils.js"; import { isHeartbeatWakeOnDemandEnabled } from "../heartbeat-policy.js"; import { nonIdleSlackIssueCondition } from "../slack-conversation-state.js"; @@ -105,107 +105,148 @@ function currentUtcDayWindow(now = new Date()) { +type TimerAgent = typeof agents.$inferSelect; +type TimerPolicy = ReturnType; +type DailyCap = { reason: string; observed: number; limit: number } | null; +const timerKey = (companyId: string, agentId: string) => JSON.stringify([companyId, agentId]); +function dailyRunConditions(scope: SQL, now: Date) { + const { start, end } = currentUtcDayWindow(now); + return and(scope, gte(heartbeatRuns.startedAt, start), lt(heartbeatRuns.startedAt, end), + notInArray(heartbeatRuns.status, ["queued", "scheduled_retry"])); +} +function dailyCostConditions(scope: SQL, now: Date) { + const { start, end } = currentUtcDayWindow(now); + return and(scope, gte(costEvents.occurredAt, start), lt(costEvents.occurredAt, end)); +} +function actionableConditions(scope: SQL) { + return and(scope, isNull(issues.assigneeUserId), isNull(issues.hiddenAt), + inArray(issues.status, [...TIMER_ACTIONABLE_ISSUE_STATUSES]), isNull(issues.conversationAgentId), nonIdleSlackIssueCondition()); +} +export function heartbeatDailyCapFromTotals(policy: TimerPolicy, totals: { runs?: number; cost?: string }, + checks: { checkRunCap?: boolean; checkCostCap?: boolean } = {}): DailyCap { + if (checks.checkRunCap !== false && policy.maxDailyRuns !== null && (totals.runs ?? 0) >= policy.maxDailyRuns) + return { reason: "heartbeat.daily_run_limit", observed: totals.runs ?? 0, limit: policy.maxDailyRuns }; + if (checks.checkCostCap !== false && policy.maxDailyCostCents !== null + && compareCents(totals.cost ?? "0", policy.maxDailyCostCents) >= 0) + return { reason: "heartbeat.daily_cost_limit", observed: Number(totals.cost ?? "0"), limit: policy.maxDailyCostCents }; + return null; +} +export function projectTimerHeartbeatSchedule(agent: TimerAgent, now: Date, + input: { actionable: boolean | null; cap: DailyCap }) { + const policy = parseHeartbeatPolicy(agent); + if (!policy.enabled || policy.intervalSec <= 0 || policy.maxDailyRuns === 0 || policy.maxDailyCostCents === 0) + return { deadline: null, policy, actionable: null, cap: null }; + if (input.actionable === false) return { deadline: null, policy, actionable: false, cap: null }; + const dueAt = new Date(agent.lastHeartbeatAt ?? agent.createdAt).getTime() + policy.intervalSec * 1000; + return { deadline: input.cap ? Math.max(dueAt, currentUtcDayWindow(now).end.getTime()) : dueAt, + policy, actionable: input.actionable, cap: input.cap }; +} + export function createHeartbeatTimerPolicy(db: Pick) { - async function getHeartbeatDailyCapBlock( - agent: typeof agents.$inferSelect, - policy: ReturnType, - options: { - checkRunCap?: boolean; - checkCostCap?: boolean; - excludeRunId?: string | null; - now?: Date; - } = {}, - client: Pick = db, - ) { - const checkRunCap = options.checkRunCap ?? true; - const checkCostCap = options.checkCostCap ?? true; - const { start, end } = currentUtcDayWindow(options.now); - if (checkRunCap && policy.maxDailyRuns !== null) { - const conditions = [ - eq(heartbeatRuns.companyId, agent.companyId), - eq(heartbeatRuns.agentId, agent.id), - gte(heartbeatRuns.startedAt, start), - lt(heartbeatRuns.startedAt, end), - notInArray(heartbeatRuns.status, ["queued", "scheduled_retry"]), - ]; - if (options.excludeRunId) { - conditions.push(sql`${heartbeatRuns.id} <> ${options.excludeRunId}`); - } - const [row] = await client - .select({ total: sql`count(*)::integer` }) - .from(heartbeatRuns) - .where(and(...conditions)); - const observed = Number(row?.total ?? 0); - if (observed >= policy.maxDailyRuns) { - return { - reason: "heartbeat.daily_run_limit", - observed, - limit: policy.maxDailyRuns, - }; - } + async function getHeartbeatDailyCapBlock(agent: TimerAgent, policy: TimerPolicy, + options: { checkRunCap?: boolean; checkCostCap?: boolean; excludeRunId?: string | null; now?: Date } = {}, + client: Pick = db): Promise { + const now = options.now ?? new Date(); + if (options.checkRunCap !== false && policy.maxDailyRuns !== null) { + const [row] = await client.select({ total: sql`count(*)::integer` }).from(heartbeatRuns) + .where(dailyRunConditions(and(eq(heartbeatRuns.companyId, agent.companyId), eq(heartbeatRuns.agentId, agent.id), + options.excludeRunId ? sql`${heartbeatRuns.id} <> ${options.excludeRunId}` : undefined)!, now)); + const cap = heartbeatDailyCapFromTotals(policy, { runs: Number(row?.total ?? 0) }, { checkCostCap: false }); + if (cap) return cap; } - - if (checkCostCap && policy.maxDailyCostCents !== null) { - const [row] = await client - .select({ - total: sql`coalesce(sum(${costEvents.costCents}), 0)::text`, - }) - .from(costEvents) - .where( - and( - eq(costEvents.companyId, agent.companyId), - eq(costEvents.agentId, agent.id), - gte(costEvents.occurredAt, start), - lt(costEvents.occurredAt, end), - ), - ); - const observed = Number(row?.total ?? 0); - if (compareCents(String(row?.total ?? 0), policy.maxDailyCostCents) >= 0) { - return { - reason: "heartbeat.daily_cost_limit", - observed, - limit: policy.maxDailyCostCents, - }; - } + if (options.checkCostCap !== false && policy.maxDailyCostCents !== null) { + const [row] = await client.select({ total: sql`coalesce(sum(${costEvents.costCents}), 0)::text` }).from(costEvents) + .where(dailyCostConditions(and(eq(costEvents.companyId, agent.companyId), eq(costEvents.agentId, agent.id))!, now)); + return heartbeatDailyCapFromTotals(policy, { cost: String(row?.total ?? 0) }, { checkRunCap: false }); } - return null; } - - async function hasActionableTimerWork(agent: typeof agents.$inferSelect) { - const row = await db - .select({ id: issues.id }) - .from(issues) - .where( - and( - eq(issues.companyId, agent.companyId), - eq(issues.assigneeAgentId, agent.id), - isNull(issues.assigneeUserId), - isNull(issues.hiddenAt), - inArray(issues.status, [...TIMER_ACTIONABLE_ISSUE_STATUSES]), - isNull(issues.conversationAgentId), - nonIdleSlackIssueCondition(), - ), - ) - .limit(1) - .then((rows) => rows[0] ?? null); - return Boolean(row); + async function hasActionableTimerWork(agent: TimerAgent) { + const rows = await db.select({ id: issues.id }).from(issues) + .where(actionableConditions(and(eq(issues.companyId, agent.companyId), eq(issues.assigneeAgentId, agent.id))!)).limit(1); + return rows.length > 0; } - - async function readTimerHeartbeatSchedule(agent: typeof agents.$inferSelect, now = new Date()) { + async function readTimerHeartbeatSchedule(agent: TimerAgent, now = new Date()) { const policy = parseHeartbeatPolicy(agent); if (!policy.enabled || policy.intervalSec <= 0 || policy.maxDailyRuns === 0 || policy.maxDailyCostCents === 0) - return { deadline: null, policy, actionable: null, cap: null }; + return projectTimerHeartbeatSchedule(agent, now, { actionable: null, cap: null }); const actionable = policy.skipTimerWhenNoActionableWork ? await hasActionableTimerWork(agent) : null; - if (actionable === false) return { deadline: null, policy, actionable, cap: null }; - const dueAt = new Date(agent.lastHeartbeatAt ?? agent.createdAt).getTime() + policy.intervalSec * 1000; - const cap = await getHeartbeatDailyCapBlock(agent, policy, { now }); - return { deadline: cap ? Math.max(dueAt, currentUtcDayWindow(now).end.getTime()) : dueAt, - policy, actionable, cap }; + const cap = actionable === false ? null : await getHeartbeatDailyCapBlock(agent, policy, { now }); + return projectTimerHeartbeatSchedule(agent, now, { actionable, cap }); } - - async function nextTimerHeartbeatAt(agent: typeof agents.$inferSelect, now = new Date()): Promise { + async function readTimerHeartbeatSchedules(candidates: TimerAgent[], now: Date, checkBudget: () => void = () => {}) { + if (candidates.length > 10_000) return; + const allowed = new Set(candidates.map(agent => timerKey(agent.companyId, agent.id))); + if (allowed.size !== candidates.length) return; + const eligible = candidates.filter(agent => { + const p = parseHeartbeatPolicy(agent); + return p.enabled && p.intervalSec > 0 && p.maxDailyRuns !== 0 && p.maxDailyCostCents !== 0; + }); + const runIds = eligible.filter(a => parseHeartbeatPolicy(a).maxDailyRuns !== null).map(a => a.id); + const costIds = eligible.filter(a => parseHeartbeatPolicy(a).maxDailyCostCents !== null).map(a => a.id); + const actionableIds = eligible.filter(a => parseHeartbeatPolicy(a).skipTimerWhenNoActionableWork).map(a => a.id); + const runs = new Map(), costs = new Map(), actionable = new Set(); + const state: unknown[] = []; + if (runIds.length) { + checkBudget(); + const rows = await db.select({ companyId: heartbeatRuns.companyId, agentId: heartbeatRuns.agentId, + total: sql`count(*)::integer` }).from(heartbeatRuns) + .where(dailyRunConditions(inArray(heartbeatRuns.agentId, runIds), now)) + .groupBy(heartbeatRuns.companyId, heartbeatRuns.agentId).orderBy(heartbeatRuns.companyId, heartbeatRuns.agentId).limit(10_001); + if (rows.length > 10_000) return; + for (const row of rows) { + checkBudget(); + const key = timerKey(row.companyId, row.agentId); + if (!allowed.has(key) || runs.has(key) || typeof row.total !== "number" || !Number.isSafeInteger(row.total) || row.total < 0) return; + runs.set(key, row.total); + } + state.push(rows); + } + if (costIds.length) { + checkBudget(); + const rows = await db.select({ companyId: costEvents.companyId, agentId: costEvents.agentId, + total: sql`coalesce(sum(${costEvents.costCents}), 0)::text` }).from(costEvents) + .where(dailyCostConditions(inArray(costEvents.agentId, costIds), now)) + .groupBy(costEvents.companyId, costEvents.agentId).orderBy(costEvents.companyId, costEvents.agentId).limit(10_001); + if (rows.length > 10_000) return; + for (const row of rows) { + checkBudget(); + if (!row.agentId) return; + const key = timerKey(row.companyId, row.agentId); + if (!allowed.has(key) || costs.has(key) || typeof row.total !== "string" || !/^-?\d+(?:\.\d+)?$/.test(row.total)) return; + costs.set(key, row.total); + } + state.push(rows); + } + if (actionableIds.length) { + checkBudget(); + const rows = await db.select({ companyId: issues.companyId, agentId: issues.assigneeAgentId }).from(issues) + .where(actionableConditions(inArray(issues.assigneeAgentId, actionableIds))) + .groupBy(issues.companyId, issues.assigneeAgentId).orderBy(issues.companyId, issues.assigneeAgentId).limit(10_001); + if (rows.length > 10_000) return; + for (const row of rows) { + checkBudget(); + if (!row.agentId) return; + const key = timerKey(row.companyId, row.agentId); + if (!allowed.has(key) || actionable.has(key)) return; + actionable.add(key); + } + state.push(rows); + } + const schedules = new Map>(); + for (const agent of candidates) { + checkBudget(); + const policy = parseHeartbeatPolicy(agent), key = timerKey(agent.companyId, agent.id); + schedules.set(agent.id, projectTimerHeartbeatSchedule(agent, now, { + actionable: policy.skipTimerWhenNoActionableWork ? actionable.has(key) : null, + cap: heartbeatDailyCapFromTotals(policy, { runs: runs.get(key), cost: costs.get(key) }), + })); + } + return { schedules, state }; + } + async function nextTimerHeartbeatAt(agent: TimerAgent, now = new Date()): Promise { return (await readTimerHeartbeatSchedule(agent, now)).deadline; } - return { parseHeartbeatPolicy, getHeartbeatDailyCapBlock, hasActionableTimerWork, nextTimerHeartbeatAt, readTimerHeartbeatSchedule }; + return { parseHeartbeatPolicy, getHeartbeatDailyCapBlock, hasActionableTimerWork, nextTimerHeartbeatAt, + readTimerHeartbeatSchedule, readTimerHeartbeatSchedules }; } diff --git a/server/src/services/idle-deadline-proof.test.ts b/server/src/services/idle-deadline-proof.test.ts index deb0142779f..4f3a7bef96c 100644 --- a/server/src/services/idle-deadline-proof.test.ts +++ b/server/src/services/idle-deadline-proof.test.ts @@ -1,5 +1,8 @@ import { pluginIdleManifestGeneration } from "./plugin-idle-catalog.js"; -import type { Db } from "@paperclipai/db"; +import { agents, heartbeatRuns, costEvents, issues, type Db } from "@paperclipai/db"; +import { getTableName } from "drizzle-orm"; +import { readIdleAgentTimerSchedule } from "./heartbeat/idle-timer-schedule.js"; +import { createHeartbeatTimerPolicy } from "./heartbeat/timer-policy.js"; import { describe, expect, it } from "vitest"; import { composeIdleDeadlineProof, createIdleDeadlineInspectionBudget, deadlineWorkReplacement, readIdleDeadlineProof, IdleDeadlineInspectionBudgetExceeded, type IdleDeadlineProofOptions } from "./idle-deadline-proof.js"; @@ -90,7 +93,7 @@ it("trace expiry transfers actual retained data; capture/import/deletion/cross-c expect(projectIdleProviderTraceSchedule([expired], now)).toMatchObject({ blocked: false, deadlines: [] }); }); type ContinuityRow = Parameters[0][number]; -const github = (): ContinuityRow => ({ companyStatus: "active", connection: { id: "connection", companyId: "company", status: "active", enabled: true, config: { sourceTemplateKey: "github", oauth: { connectorProfile: "github.code" } } } as unknown as ContinuityRow["connection"], grant: { id: "grant", companyId: "company", credentialSecretRefs: [{ configPath: "oauth.access_token", secretId: "secret" }], providerTenant: { oauth: { accessTokenExpiresAt: new Date(now + 8 * 3_600_000).toISOString(), refreshedAt: new Date(now).toISOString() }, github: { userId: "user", installationIds: ["101"], lastAccessRefreshAt: new Date(now).toISOString() } } } as ContinuityRow["grant"] }); +const github = (): ContinuityRow => ({ companyStatus: "active", connection: { id: "connection", companyId: "company", status: "active", enabled: true, authKind: "oauth", credentialSource: "paperclip_vault", config: { sourceTemplateKey: "github", oauth: { strategy: "paperclip_cloud_connector", connectorProfile: "github.code" } } } as unknown as ContinuityRow["connection"], grant: { id: "grant", companyId: "company", connectionId: "connection", kind: "user", subjectUserId: "user", credentialSecretRefs: [{ configPath: "oauth.access_token", secretId: "secret" }], providerTenant: { oauth: { accessTokenExpiresAt: new Date(now + 8 * 3_600_000).toISOString(), refreshedAt: new Date(now).toISOString() }, github: { userId: "user", installationIds: ["101"], lastAccessRefreshAt: new Date(now).toISOString() } } } as ContinuityRow["grant"] }); it("GitHub enrollment alone never bypasses its credential continuity deadline or unresolved authority", () => { expect(projectIdleGitHubGrantContinuity([github()], now)).toMatchObject({ blocked: false, deadlines: [now + 7 * 3_600_000], connectionIds: ["connection"] }); const nonexpiring = github(); delete nonexpiring.grant.providerTenant!.oauth; @@ -104,22 +107,112 @@ it("GitHub enrollment alone never bypasses its credential continuity deadline or }); -it("stops owner read orchestration at its query budget and rejects row overflow", async () => { - const agentRows = Array.from({ length: 80 }, (_, i) => ({ agent: { id: `agent-${i}`, companyId: "company", name: "Timer", - status: "idle", reportsTo: null, adapterType: "codex_local", runtimeConfig: { heartbeat: { enabled: true, intervalSec: 3600, maxDailyRuns: 1 } }, - createdAt: new Date(now), lastHeartbeatAt: null }, companyStatus: "active", companyUpdatedAt: new Date(now) })); - let reads = 0; - const tx = { select: (columns: object) => { - reads++; - const rows = "agent" in columns ? agentRows : [{ total: 0 }]; - const query = { from: () => query, leftJoin: () => query, orderBy: () => query, where: () => query, - limit: () => Promise.resolve(rows), then: (resolve: (rows: unknown[]) => unknown) => Promise.resolve(rows).then(resolve) }; +function timerAgent(index: number): typeof agents.$inferSelect { + return { id: `agent-${index}`, companyId: "company", name: "Timer", status: "idle", reportsTo: null, + adapterType: "codex_local", adapterConfig: {}, runtimeConfig: { heartbeat: { enabled: true, intervalSec: 3600, maxDailyRuns: 1 } }, + createdAt: new Date(now), lastHeartbeatAt: new Date(now) } as unknown as typeof agents.$inferSelect; +} +function fixtureDb(tables: Record, onRead: () => void = () => {}) { + return { select: () => { + onRead(); + let table = ""; + const query = { from: (source: Parameters[0]) => { table = getTableName(source); return query; }, + leftJoin: () => query, orderBy: () => query, where: () => query, groupBy: () => query, + limit: () => Promise.resolve(tables[table] ?? []), + then: (resolve: (rows: unknown[]) => unknown) => Promise.resolve(tables[table] ?? []).then(resolve) }; return query; } } as unknown as Pick; - expect(await readIdleDeadlineProof(tx, hold, options)).toEqual({ blocked: true }); - expect(reads).toBe(64); - agentRows.push(...Array.from({ length: 10_001 }, () => agentRows[0])); - reads = 0; +} +it("batches more than128 eligible timers and the old80 run-cap shape within a fixed query budget", async () => { + const rows = Array.from({ length: 160 }, (_, i) => ({ agent: timerAgent(i), companyStatus: "active", companyUpdatedAt: new Date(now) })); + rows.push(...["paused", "terminated", "idle"].map((status, i) => ({ agent: { ...timerAgent(200 + i), status, + runtimeConfig: { heartbeat: { enabled: false } } }, companyStatus: "active", companyUpdatedAt: new Date(now) }))); + let reads = 0; + const tables: Record = { agents: rows }; + const tx = fixtureDb(tables, () => { reads++; }); + const future = await readIdleDeadlineProof(tx, hold, options); + expect(future.blocked).toBe(false); + expect(future.scheduledWake?.nextWakeAt).toBe(new Date(now + 3_600_000).toISOString()); + expect(reads).toBe(5); + tables.agents = rows.slice(0, 80); reads = 0; + expect((await readIdleDeadlineProof(tx, hold, options)).blocked).toBe(false); + expect(reads).toBe(5); + rows[0].agent.lastHeartbeatAt = new Date(now - 3_600_000); tables.agents = rows; reads = 0; + expect((await readIdleDeadlineProof(tx, hold, options)).blocked).toBe(true); + expect(reads).toBe(2); + rows[0].agent.lastHeartbeatAt = new Date(now); + tables.agents = Array.from({ length: 10_001 }, () => rows[0]); reads = 0; expect(await readIdleDeadlineProof(tx, hold, options)).toEqual({ blocked: true }); expect(reads).toBe(1); + tables.agents = rows; + for (const total of [null, "", true, false, "0", NaN, Infinity, -1, 1.5]) { + tables.heartbeat_runs = [{ companyId: rows[0].agent.companyId, agentId: rows[0].agent.id, total }]; + expect((await readIdleDeadlineProof(tx, hold, options)).blocked).toBe(true); + } + tables.heartbeat_runs = [{ companyId: "foreign", agentId: rows[0].agent.id, total: 0 }]; + expect((await readIdleDeadlineProof(tx, hold, options)).blocked).toBe(true); +}); +it("single and batch owner schedules share cap, actionable, UTC and decimal semantics", async () => { + for (const heartbeat of [ + { enabled: false }, { enabled: true, intervalSec: 600, maxDailyRuns: 0 }, + { enabled: true, intervalSec: 600, maxDailyRuns: 2 }, + { enabled: true, intervalSec: 600, maxDailyCostCents: 1 }, + { enabled: true, intervalSec: 600, maxDailyRuns: 2, maxDailyCostCents: 1, skipTimerWhenNoActionableWork: true }, + ]) for (const actionable of [false, true]) for (const runTotal of [0, 2]) for (const cost of ["0.999999", "1.000001"]) { + const agent = { ...timerAgent(0), runtimeConfig: { heartbeat } }; + const owner = createHeartbeatTimerPolicy(fixtureDb({ + [getTableName(heartbeatRuns)]: [{ companyId: agent.companyId, agentId: agent.id, total: runTotal }], + [getTableName(costEvents)]: [{ companyId: agent.companyId, agentId: agent.id, total: cost }], + [getTableName(issues)]: actionable ? [{ companyId: agent.companyId, agentId: agent.id, id: "issue" }] : [], + })); + const single = await owner.readTimerHeartbeatSchedule(agent, new Date(now)); + const batch = await owner.readTimerHeartbeatSchedules([agent], new Date(now)); + expect(batch?.schedules.get(agent.id)).toEqual(single); + } +}); +it("checks global suppression even without eligible timers and retains owner org-chain checks", async () => { + for (const rows of [[], [timerAgent(1)], [{ ...timerAgent(1), status: "paused" }]]) { + if (rows[0]?.status === "idle") rows[0].runtimeConfig = { heartbeat: { enabled: false } }; + for (const env of [{ PAPERCLIP_IN_WORKTREE: "true" }, { PAPERCLIP_DATABASE_RESTORE_IN_PROGRESS: "1" }, + { PAPERCLIP_RESTORE_IN_PROGRESS: "yes" }, { HEARTBEAT_SCHEDULER_ENABLED: "false" }]) { + let reads = 0; + expect((await readIdleAgentTimerSchedule(fixtureDb({ agents: rows }, () => { reads++; }), now, env)).blocked).toBe(true); + expect(reads).toBe(0); + } + } + const parent = timerAgent(1), child = { ...timerAgent(2), reportsTo: parent.id }; + const rows = [parent, child].map(agent => ({ agent, companyStatus: "active", companyUpdatedAt: new Date(now) })); + const tx = fixtureDb({ agents: rows }); + expect((await readIdleAgentTimerSchedule(tx, now, {})).blocked).toBe(false); + parent.status = "paused"; + expect((await readIdleAgentTimerSchedule(tx, now, {})).blocked).toBe(false); + parent.status = "terminated"; + expect((await readIdleAgentTimerSchedule(tx, now, {})).blocked).toBe(true); + parent.status = "idle"; parent.reportsTo = child.id; + expect((await readIdleAgentTimerSchedule(tx, now, {})).blocked).toBe(true); + let checks = 0; + await expect(readIdleAgentTimerSchedule(tx, now, {}, () => { + if (++checks > 1) throw new IdleDeadlineInspectionBudgetExceeded(); + })).rejects.toThrow(IdleDeadlineInspectionBudgetExceeded); +}); +it("does not turn unrelated GitHub grant shapes into continuity authority", () => { + const unrelated = github(); unrelated.connection!.authKind = "api_key"; + expect(projectIdleGitHubGrantContinuity([unrelated], now)).toMatchObject({ blocked: false, deadlines: [], connectionIds: [] }); + unrelated.connection!.authKind = "oauth"; unrelated.connection!.config = { sourceTemplateKey: "other" }; + expect(projectIdleGitHubGrantContinuity([unrelated], now)).toMatchObject({ blocked: false, connectionIds: [] }); + for (const patch of [{ oauth: { connectorProfile: "github.code", strategy: "unknown_strategy" } }, + { oauth: { connectorProfile: "github.code" } }, { oauth: null }, { oauth: { connectorProfile: 123 } }]) { + const malformed = github(); malformed.connection!.config = { sourceTemplateKey: "github", ...patch }; + expect(projectIdleGitHubGrantContinuity([malformed], now).blocked).toBe(true); + } + for (const patch of [{ authKind: "unknown_auth" }, { credentialSource: "unknown_source" }]) { + const malformed = github(); Object.assign(malformed.connection!, patch); + expect(projectIdleGitHubGrantContinuity([malformed], now).blocked).toBe(true); + } + const persisted = github(); persisted.connection!.config.oauth = { connectorProfile: "github.code", strategy: "paperclip_id_connector" }; + expect(projectIdleGitHubGrantContinuity([persisted], now).blocked).toBe(false); + const agent = github(); agent.grant.kind = "agent"; agent.grant.subjectUserId = null; agent.grant.subjectAgentId = "dedicated-agent"; + expect(projectIdleGitHubGrantContinuity([agent], now).connectionIds).toEqual(["connection"]); + const orphan = github(); orphan.connection = null; + expect(projectIdleGitHubGrantContinuity([orphan], now).blocked).toBe(true); }); diff --git a/server/src/services/idle-deadline-proof.ts b/server/src/services/idle-deadline-proof.ts index d9176c05698..089502cca7f 100644 --- a/server/src/services/idle-deadline-proof.ts +++ b/server/src/services/idle-deadline-proof.ts @@ -80,7 +80,7 @@ export async function readIdleDeadlineProof(tx: Pick, hold: { owne const traces = await readIdleProviderTraceSchedule(guarded, options.now); budget.check(); if (traces.blocked) return { blocked: true }; - const github = await readIdleGitHubGrantContinuity(guarded, options.now); + const github = await readIdleGitHubGrantContinuity(guarded, options.now, budget.check); budget.check(); const proof = composeIdleDeadlineProof(hold, options, { timers, jobs, traces, github }); budget.check(); diff --git a/server/src/services/tool-access/github-grant-continuity.ts b/server/src/services/tool-access/github-grant-continuity.ts index 8132f7ed0c5..5786254b5ea 100644 --- a/server/src/services/tool-access/github-grant-continuity.ts +++ b/server/src/services/tool-access/github-grant-continuity.ts @@ -1,15 +1,15 @@ -import { readConnectionGrantCredentialState } from "../connection-credentials.js"; +import { readConnectionGrantCredentialStates } from "../connection-credentials.js"; import { randomUUID } from "node:crypto"; import { notifyGitHubCredentialChange } from "../github-connection-notifications.js"; import type { ActorInfo } from "./types.js"; import { loadGitHubGrantMetadata } from "./github-access-discovery.js"; import { measureRunPhase } from "../execution-preparation-timing.js"; import { measureRunTransaction } from "../execution-transaction-timing.js"; -import { and, eq, sql } from "drizzle-orm"; +import { and, eq, sql, or, isNull, ne } from "drizzle-orm"; import type { Db } from "@paperclipai/db"; import { companies, connectionGrants, toolConnections } from "@paperclipai/db"; import { conflict, HttpError, notFound, unprocessable } from "../../errors.js"; -import { PaperclipCloudConnectorError, type PaperclipCloudConnector } from "../paperclip-cloud-connector.js"; +import { isPaperclipCloudConnectorStrategy, PaperclipCloudConnectorError, type PaperclipCloudConnector } from "../paperclip-cloud-connector.js"; import type { createOAuthGrantRefresh } from "./oauth/grant-refresh.js"; import type { OAuthCredentialStore } from "./oauth/credential-store.js"; @@ -397,36 +397,54 @@ export function githubContinuityRevision(grant: Pick, now: number) { +export async function readIdleGitHubGrantContinuity(tx: Pick, now: number, + checkBudget: () => void = () => {}) { const rows = await tx.select({ grant: connectionGrants, connection: toolConnections, companyStatus: companies.status }) .from(connectionGrants).leftJoin(toolConnections, eq(toolConnections.id, connectionGrants.connectionId)) .leftJoin(companies, eq(companies.id, connectionGrants.companyId)) - .where(and(eq(connectionGrants.status, "active"), sql`${connectionGrants.providerTenant}->'github' IS NOT NULL`)) + .where(and(eq(connectionGrants.status, "active"), sql`${connectionGrants.providerTenant}->'github' IS NOT NULL`, + or(isNull(toolConnections.id), ne(toolConnections.companyId, connectionGrants.companyId), + sql`${toolConnections.config}->>'sourceTemplateKey' = 'github'`))) .orderBy(connectionGrants.id).limit(10_001); - const proof = projectIdleGitHubGrantContinuity(rows, now); + const proof = projectIdleGitHubGrantContinuity(rows, now, checkBudget); if (proof.blocked) return proof; - const credentials: unknown[] = []; const eligible = new Set(proof.connectionIds); - for (const row of rows) { - if (!row.connection || !eligible.has(row.connection.id)) continue; - const state = await readConnectionGrantCredentialState(tx, row.connection, row.grant); - if (!state) return { ...proof, blocked: true, connectionIds: [], deadlines: [] }; - credentials.push([row.grant.id, state]); - } - return { ...proof, state: [rows, credentials] }; + const inputs = rows.flatMap(row => row.connection && eligible.has(row.connection.id) + && managedGitHubContinuityConnection(row.connection) === "managed" ? [{ connection: row.connection, grant: row.grant }] : []); + const credentials = await readConnectionGrantCredentialStates(tx, inputs, checkBudget); + if (!credentials) return { ...proof, blocked: true, connectionIds: [], deadlines: [] }; + return { ...proof, state: [rows, [...credentials]] }; +} +function managedGitHubContinuityConnection(connection: typeof toolConnections.$inferSelect): "managed" | "unrelated" | "invalid" { + const config = asRecord(connection.config), oauth = asRecord(config.oauth); + if (config.sourceTemplateKey !== "github" || connection.authKind !== "oauth" + || connection.credentialSource === "vercel_connect") return "unrelated"; + if (typeof oauth.connectorProfile !== "string" || !oauth.connectorProfile) return "invalid"; + if (oauth.connectorProfile !== "github.code") return "unrelated"; + return isPaperclipCloudConnectorStrategy(oauth.strategy) ? "managed" : "invalid"; } export function projectIdleGitHubGrantContinuity(rows: { grant: typeof connectionGrants.$inferSelect; connection: typeof toolConnections.$inferSelect | null; companyStatus: string | null; -}[], now: number) { +}[], now: number, checkBudget: () => void = () => {}) { const deadlines: number[] = []; const connectionIds = new Set(); const blocked = () => ({ blocked: true, state: rows, deadlines: [], connectionIds: [] as string[] }); if (rows.length > 10_000) return blocked(); for (const { grant, connection, companyStatus } of rows) { - if (!connection || connection.companyId !== grant.companyId || !companyStatus) return blocked(); + checkBudget(); + if (!connection || connection.companyId !== grant.companyId || !companyStatus + || !["active", "paused", "archived"].includes(companyStatus) + || !["draft", "active", "disabled", "archived"].includes(connection.status) + || !["oauth", "api_key", "none"].includes(connection.authKind) + || !["paperclip_vault", "vercel_connect"].includes(connection.credentialSource)) return blocked(); if (companyStatus !== "active" || connection.status !== "active" || !connection.enabled) continue; + const ownership = managedGitHubContinuityConnection(connection); + if (ownership === "invalid") return blocked(); + if (ownership === "unrelated") continue; + if (grant.connectionId !== connection.id || (grant.kind === "user" ? !grant.subjectUserId + : grant.kind === "agent" ? !grant.subjectAgentId : true)) return blocked(); const oauth = asRecord(grant.providerTenant?.oauth); const github = grant.providerTenant?.github; if (!github || github.bindingSync || github.continuityRetry || oauth.refreshLease) return blocked();