diff --git a/server/src/db/compliance-db.ts b/server/src/db/compliance-db.ts index 42a707bf041..fdd0538c1a9 100644 --- a/server/src/db/compliance-db.ts +++ b/server/src/db/compliance-db.ts @@ -959,15 +959,21 @@ export class ComplianceDatabase { tracksSummary[t.track] = t.status; } - // 4. Upsert the materialized status and capture transition + // 4. Upsert the materialized status and capture transition. + // Owner-initiated runs publish the verdict but leave the independent + // heartbeat schedule and pending requeues untouched (#7680). const statusResult = await client.query( `WITH schedule_next AS ( INSERT INTO agent_registry_metadata (agent_url, next_compliance_check_at) - VALUES ($1, NOW() + INTERVAL '12 hours') + VALUES ($1, CASE WHEN $7 = 'owner_test' THEN NULL ELSE NOW() + INTERVAL '12 hours' END) ON CONFLICT (agent_url) DO UPDATE SET next_compliance_check_at = NOW() + make_interval(hours => agent_registry_metadata.check_interval_hours), compliance_inconclusive_streak = 0, requeued_at = NULL + -- TriggeredBy constrains $7 at compile time. Adding a trigger opts it + -- into cadence advancement unless this predicate and the INSERT + -- CASE above are updated together. + WHERE $7 != 'owner_test' ) INSERT INTO agent_compliance_status ( agent_url, status, last_checked_at, @@ -1024,6 +1030,7 @@ export class ComplianceDatabase { input.headline ?? null, input.requested_compliance_target ?? null, input.adcp_version ?? null, + input.triggered_by ?? 'heartbeat', ], ); diff --git a/server/src/training-agent/idempotency.ts b/server/src/training-agent/idempotency.ts index af625d28395..62f2d68d659 100644 --- a/server/src/training-agent/idempotency.ts +++ b/server/src/training-agent/idempotency.ts @@ -19,7 +19,8 @@ * public sandbox token is in use (otherwise every caller on that token * sees the same oracle). * - `getIdempotencyStore` — returns a process-wide store backed by - * Postgres when a DB pool is available, in-memory otherwise. + * Postgres in production (with lazy pool lookup) or when a DB pool is + * available, in-memory otherwise. */ import { @@ -472,7 +473,12 @@ function validateTtl(seconds: number): number { } function fencedPgBackend(): FencedIdempotencyBackend { - const db = getPool(); + // Tenant registry prewarming starts before initializeDatabase(). Resolve + // the pool only when querying so construction never downgrades production + // to process-local memory; operations still fail closed until DB startup. + const db = { + query: (text: string, values?: unknown[]) => getPool().query(text, values), + }; const base = pgBackend(db); return { ...base, @@ -504,7 +510,7 @@ function fencedPgBackend(): FencedIdempotencyBackend { export function getIdempotencyStore(): OwnedIdempotencyStore { if (storeInstance) return storeInstance; - const backend = isDatabaseInitialized() + const backend = process.env.NODE_ENV === 'production' || isDatabaseInitialized() ? fencedPgBackend() : memoryBackend(); const base = createHashAwareIdempotencyStore({ backend, ttlSeconds: REPLAY_TTL_SECONDS }); diff --git a/server/tests/integration/compliance-run-publication.test.ts b/server/tests/integration/compliance-run-publication.test.ts index ca2e1dc6d4d..a2da6da9a30 100644 --- a/server/tests/integration/compliance-run-publication.test.ts +++ b/server/tests/integration/compliance-run-publication.test.ts @@ -351,6 +351,143 @@ describe.skipIf(!process.env.DATABASE_URL)('compliance publication transaction', } }); + describe('owner-initiated runs leave the heartbeat schedule alone (#7680)', () => { + const urls: string[] = []; + const freshAgent = async (options: { withMetadata?: boolean } = {}) => { + const url = `https://${randomUUID()}.example.test/mcp`; + urls.push(url); + // Metadata alone is not a current-agent source (#7844). Register each + // fixture with the owning profile so due-selection exercises scheduling. + await pool.query( + `UPDATE member_profiles SET agents = agents || $2::jsonb + WHERE workos_organization_id = $1`, + [ownerOrgId, JSON.stringify([{ url }])], + ); + if (options.withMetadata ?? true) await db.upsertRegistryMetadata(url, { lifecycle_stage: 'production' }); + return url; + }; + const schedule = async (url: string) => (await pool.query<{ + next_compliance_check_at: Date | null; requeued_at: Date | null; hours_until_next: number | null; + }>( + `SELECT next_compliance_check_at, requeued_at, + EXTRACT(EPOCH FROM next_compliance_check_at - NOW()) / 3600 AS hours_until_next + FROM agent_registry_metadata WHERE agent_url = $1`, + [url], + )).rows[0]; + const isDue = async (url: string) => + (await db.getAgentsDueForCheck(100000)).some(agent => agent.agent_url === url); + + afterAll(async () => { + for (const table of ['agent_compliance_step_diagnostics', 'agent_storyboard_status', 'agent_verification_badges', + 'agent_compliance_status', 'agent_compliance_runs', 'agent_registry_metadata']) { + await pool.query(`DELETE FROM ${table} WHERE agent_url = ANY($1::text[])`, [urls]); + } + }); + + it('keeps a pending requeue when a full-suite owner test publishes', async () => { + const url = await freshAgent(); + await db.requeueForHeartbeat(url); + const before = await schedule(url); + await db.recordComplianceRun({ ...input, agent_url: url, triggered_by: 'owner_test' }); + + const after = await schedule(url); + expect(after.next_compliance_check_at).toBeNull(); + expect(after.requeued_at).toEqual(before.requeued_at); + expect(after.requeued_at).not.toBeNull(); + expect(await isDue(url)).toBe(true); + // The owner verdict still publishes to the public card. + expect((await db.getComplianceStatus(url))?.status).toBe('passing'); + }); + + it('consumes the requeue when the heartbeat serving it publishes', async () => { + const url = await freshAgent(); + await db.requeueForHeartbeat(url); + await db.recordComplianceRun({ ...input, agent_url: url, triggered_by: 'heartbeat' }); + + const after = await schedule(url); + expect(after.requeued_at).toBeNull(); + expect(Number(after.hours_until_next)).toBeCloseTo(12, 1); + expect(await isDue(url)).toBe(false); + }); + + it.each(['manual', 'webhook'] as const)('advances the cadence and clears the requeue for %s runs', async (triggeredBy) => { + const url = await freshAgent(); + await db.requeueForHeartbeat(url); + await db.recordComplianceRun({ ...input, agent_url: url, triggered_by: triggeredBy }); + + const after = await schedule(url); + expect(after.requeued_at).toBeNull(); + expect(Number(after.hours_until_next)).toBeCloseTo(12, 1); + }); + + it('defaults an unlabelled run to heartbeat scheduling', async () => { + const url = await freshAgent(); + await db.requeueForHeartbeat(url); + expect(input.triggered_by).toBeUndefined(); + await db.recordComplianceRun({ ...input, agent_url: url }); + + const after = await schedule(url); + expect(after.requeued_at).toBeNull(); + expect(Number(after.hours_until_next)).toBeCloseTo(12, 1); + }); + + it('does not push a scheduled heartbeat further out when an owner test publishes', async () => { + const url = await freshAgent(); + await pool.query( + `UPDATE agent_registry_metadata SET next_compliance_check_at = NOW() + INTERVAL '2 hours' WHERE agent_url = $1`, + [url], + ); + const before = await schedule(url); + await db.recordComplianceRun({ ...input, agent_url: url, triggered_by: 'owner_test' }); + + expect((await schedule(url)).next_compliance_check_at).toEqual(before.next_compliance_check_at); + }); + + it('leaves an overdue heartbeat due after an owner test publishes', async () => { + const url = await freshAgent(); + await pool.query( + `UPDATE agent_registry_metadata SET next_compliance_check_at = NOW() - INTERVAL '1 hour' WHERE agent_url = $1`, + [url], + ); + await db.recordComplianceRun({ ...input, agent_url: url, triggered_by: 'owner_test' }); + + expect(await isDue(url)).toBe(true); + }); + + it('leaves a first-ever owner test immediately eligible for an independent heartbeat', async () => { + const url = await freshAgent({ withMetadata: false }); + expect(await schedule(url)).toBeUndefined(); + expect(await isDue(url)).toBe(true); + await db.recordComplianceRun({ ...input, agent_url: url, triggered_by: 'owner_test' }); + + const after = await schedule(url); + expect(after).toBeDefined(); + expect(after.next_compliance_check_at).toBeNull(); + expect(await isDue(url)).toBe(true); + }); + + it('still schedules a first-ever heartbeat on the default cadence', async () => { + const url = await freshAgent({ withMetadata: false }); + await db.recordComplianceRun({ ...input, agent_url: url, triggered_by: 'heartbeat' }); + + expect(Number((await schedule(url)).hours_until_next)).toBeCloseTo(12, 1); + }); + + it('leaves scheduling untouched for a scoped (non-authoritative) owner test', async () => { + const url = await freshAgent(); + await db.requeueForHeartbeat(url); + await db.recordComplianceRun({ + ...input, agent_url: url, triggered_by: 'owner_test', + is_authoritative: false, replace_storyboard_statuses: false, + storyboard_statuses: [{ storyboard_id: 'first', status: 'passing', steps_passed: 1, steps_total: 1 }], + }); + + const after = await schedule(url); + expect(after.next_compliance_check_at).toBeNull(); + expect(after.requeued_at).not.toBeNull(); + }); + }); + it('uses only complete public profiles for version and specialism fallback', async () => { const complete = await db.recordComplianceRun(input); const epoch = Date.parse('2030-01-01T00:00:00Z'); diff --git a/server/tests/integration/training-agent-idempotency-cold-boot.test.ts b/server/tests/integration/training-agent-idempotency-cold-boot.test.ts new file mode 100644 index 00000000000..0fdde4c553e --- /dev/null +++ b/server/tests/integration/training-agent-idempotency-cold-boot.test.ts @@ -0,0 +1,141 @@ +import { randomUUID } from 'node:crypto'; +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'; +import type { Pool } from 'pg'; +import { closeDatabase, initializeDatabase, isDatabaseInitialized } from '../../src/db/client.js'; +import { runMigrations } from '../../src/db/migrate.js'; +import { + clearIdempotencyCache, + getIdempotencyStore, + getSdkIdempotencyStore, + type OwnedIdempotencyStore, +} from '../../src/training-agent/idempotency.js'; + +describe.skipIf(!process.env.DATABASE_URL)('training-agent idempotency across cold boot and stores', () => { + const principal = `cold-boot-${randomUUID()}`; + const keys: string[] = []; + let pool: Pool | undefined; + let earlyStore: OwnedIdempotencyStore; + let secondStore: OwnedIdempotencyStore; + let earlyReadiness: PromiseSettledResult; + let earlyMutation: PromiseSettledResult; + + const request = (suffix: string, budget = 5_000) => { + const key = `cold-boot-${suffix}-${randomUUID()}`; + keys.push(`${principal}\u001F${key}`); + return { principal, key, payload: { budget } }; + }; + + beforeAll(async () => { + vi.stubEnv('NODE_ENV', 'production'); + await clearIdempotencyCache(); + expect(isDatabaseInitialized()).toBe(false); + + // Tenant prewarming resolves the SDK adapter before database startup. + // Construction must succeed, but readiness and mutations must fail closed. + const sdkStore = getSdkIdempotencyStore(); + earlyStore = getIdempotencyStore(); + // Record both outcomes without aborting setup: the original memory backend + // must still reach the real-PG cross-store regressions below. + [earlyReadiness, earlyMutation] = await Promise.allSettled([ + Promise.resolve().then(() => sdkStore.probe!()), + Promise.resolve().then(() => sdkStore.check(request('before-pool'))), + ]); + + pool = initializeDatabase({ connectionString: process.env.DATABASE_URL }); + await runMigrations(); + await earlyStore.probe!(); + + // Retain the earlier adapter while constructing another independent store. + // PostgreSQL backends do not own or close the caller's database pool. + await clearIdempotencyCache(); + secondStore = getIdempotencyStore(); + await secondStore.probe!(); + }, 180_000); + + afterAll(async () => { + try { + if (pool) { + await pool.query('DELETE FROM adcp_idempotency WHERE scoped_key = ANY($1::text[])', [keys]); + } + } finally { + await clearIdempotencyCache(); + await closeDatabase(); + vi.unstubAllEnvs(); + } + }); + + it('rejects readiness before PostgreSQL is initialized', () => { + expect(earlyReadiness.status).toBe('rejected'); + if (earlyReadiness.status === 'rejected') { + expect(earlyReadiness.reason).toBeInstanceOf(Error); + expect(earlyReadiness.reason.message).toMatch(/idempotency backend probe failed/); + } + }); + + it('rejects mutations before PostgreSQL is initialized', () => { + expect(earlyMutation.status).toBe('rejected'); + if (earlyMutation.status === 'rejected') { + expect(earlyMutation.reason).toBeInstanceOf(Error); + expect(earlyMutation.reason.message).toMatch(/database operation failed/); + } + }); + + it('persists an adapter created before initialization and replays from another store', async () => { + const input = request('replay'); + const claim = await earlyStore.check(input); + if (claim.kind !== 'miss') throw new Error('Expected the first request to claim the key'); + expect(await secondStore.check(input)).toMatchObject({ kind: 'in-flight' }); + expect(await secondStore.check({ ...input, payload: { budget: 25_000 } })).toEqual({ kind: 'conflict' }); + + await earlyStore.renew({ ...input, claimToken: claim.claimToken }); + const response = { media_buy_id: 'cold-boot-buy' }; + await earlyStore.save({ ...input, payloadHash: claim.payloadHash, claimToken: claim.claimToken, response }); + expect(await secondStore.check(input)).toEqual({ kind: 'replay', response }); + const rows = await pool!.query('SELECT response FROM adcp_idempotency WHERE scoped_key = $1', [ + `${input.principal}\u001F${input.key}`, + ]); + expect(rows.rows).toEqual([{ response }]); + }); + + it('grants exactly one concurrent claim across independent PostgreSQL stores', async () => { + const input = request('concurrent'); + const stores = [earlyStore, secondStore]; + const results = await Promise.all(stores.map(store => store.check(input))); + expect(results.map(result => result.kind).sort()).toEqual(['in-flight', 'miss']); + const winnerIndex = results.findIndex(result => result.kind === 'miss'); + const claim = results[winnerIndex]!; + if (claim.kind !== 'miss') throw new Error('Expected exactly one claim owner'); + const response = { media_buy_id: 'concurrent-buy' }; + await stores[winnerIndex]!.save({ ...input, payloadHash: claim.payloadHash, claimToken: claim.claimToken, response }); + expect(await stores[1 - winnerIndex]!.check(input)).toEqual({ kind: 'replay', response }); + }); + + it('keeps owner fencing when another store takes over a released claim', async () => { + const input = request('released'); + const original = await earlyStore.check(input); + if (original.kind !== 'miss') throw new Error('Expected the initial claim'); + await earlyStore.release({ ...input, claimToken: original.claimToken }); + const replacement = await secondStore.check(input); + if (replacement.kind !== 'miss') throw new Error('Expected a fresh owner after release'); + + await expect(earlyStore.renew({ ...input, claimToken: original.claimToken })).rejects.toThrow(/ownership was lost/); + await expect(earlyStore.save({ + ...input, payloadHash: original.payloadHash, claimToken: original.claimToken, + response: { media_buy_id: 'stale-owner' }, + })).rejects.toThrow(/ownership was lost/); + const response = { media_buy_id: 'replacement-owner' }; + await secondStore.save({ ...input, payloadHash: replacement.payloadHash, claimToken: replacement.claimToken, response }); + expect(await earlyStore.check(input)).toEqual({ kind: 'replay', response }); + }); + + it('retains PostgreSQL replay when development has an initialized pool', async () => { + const input = request('development'); + const claim = await earlyStore.check(input); + if (claim.kind !== 'miss') throw new Error('Expected the initial claim'); + const response = { media_buy_id: 'development-replay' }; + await earlyStore.save({ ...input, payloadHash: claim.payloadHash, claimToken: claim.claimToken, response }); + vi.stubEnv('NODE_ENV', 'development'); + await clearIdempotencyCache(); + expect(await getIdempotencyStore().check(input)).toEqual({ kind: 'replay', response }); + }); +});