diff --git a/src/mgr/kafka-outbox-relay.ts b/src/mgr/kafka-outbox-relay.ts index 875d02a..b469692 100644 --- a/src/mgr/kafka-outbox-relay.ts +++ b/src/mgr/kafka-outbox-relay.ts @@ -13,44 +13,60 @@ export interface KafkaOutboxRelayOptions { config: AgentRunKafkaConfig; } +let lastRelayCycle: JsonRecord | null = null; + +export function kafkaOutboxRelayStatus(): JsonRecord { + return { lastCycle: lastRelayCycle, valuesPrinted: false }; +} + export async function relayKafkaEventOutboxOnce(input: { store: AgentRunStore; options: KafkaOutboxRelayOptions; publish?: (config: AgentRunKafkaConfig, message: AgentRunKafkaMessage) => Promise }): Promise { if (!input.options.enabled) return { action: "kafka-event-outbox-relay", enabled: false, claimedCount: 0, deliveredCount: 0, retryCount: 0, staleCount: 0, valuesPrinted: false }; const claimed = await input.store.claimKafkaEventOutbox({ owner: input.options.owner, leaseMs: input.options.leaseMs, limit: input.options.batchSize }); const items = claimed.slice().sort((a, b) => a.outboxSeq - b.outboxSeq); - const blockedKeys = new Set(); + const startedAt = Date.now(); + const groups = groupByPartitionKey(items); let deliveredCount = 0; let retryCount = 0; let staleCount = 0; + let firstError: JsonRecord | null = null; - for (const item of items) { - if (blockedKeys.has(item.partitionKey)) { - if (await retry(input.store, input.options, item, { code: "partition-key-blocked", message: "an earlier outbox record for this partition key failed", valuesPrinted: false })) retryCount++; - else staleCount++; - continue; + await Promise.all(groups.map(async (group) => { + let blocked = false; + for (const item of group) { + if (blocked) { + if (await retry(input.store, input.options, item, { code: "partition-key-blocked", message: "an earlier outbox record for this partition key failed", valuesPrinted: false })) retryCount++; + else staleCount++; + continue; + } + try { + await (input.publish ?? publishAgentRunKafkaMessage)(input.options.config, { + topic: item.topic, + key: item.partitionKey, + value: item.value, + headers: stringHeaders(item.headers), + }); + await input.store.completeKafkaEventOutbox(item); + deliveredCount++; + } catch (error) { + blocked = true; + const typedError = { + code: "kafka-produce-failed", + message: redactText(error instanceof Error ? error.message : String(error)).slice(0, 500), + topic: item.topic, + partitionKey: item.partitionKey, + outboxSeq: item.outboxSeq, + valuesPrinted: false, + }; + firstError ??= typedError; + const retried = await retry(input.store, input.options, item, typedError); + if (retried) retryCount++; + else staleCount++; + } } - try { - await (input.publish ?? publishAgentRunKafkaMessage)(input.options.config, { - topic: item.topic, - key: item.partitionKey, - value: item.value, - headers: stringHeaders(item.headers), - }); - await input.store.completeKafkaEventOutbox(item); - deliveredCount++; - } catch (error) { - blockedKeys.add(item.partitionKey); - const retried = await retry(input.store, input.options, item, { - code: "kafka-produce-failed", - message: redactText(error instanceof Error ? error.message : String(error)).slice(0, 500), - topic: item.topic, - outboxSeq: item.outboxSeq, - valuesPrinted: false, - }); - if (retried) retryCount++; - else staleCount++; - } - } - return { action: "kafka-event-outbox-relay", enabled: true, claimedCount: items.length, deliveredCount, retryCount, staleCount, firstOutboxSeq: items[0]?.outboxSeq ?? null, lastOutboxSeq: items.at(-1)?.outboxSeq ?? null, valuesPrinted: false }; + })); + const durationMs = Math.max(1, Date.now() - startedAt); + lastRelayCycle = { at: new Date().toISOString(), claimedCount: items.length, partitionKeyCount: groups.length, deliveredCount, retryCount, staleCount, durationMs, claimRatePerSecond: Number((items.length * 1_000 / durationMs).toFixed(2)), publishRatePerSecond: Number((deliveredCount * 1_000 / durationMs).toFixed(2)), firstError, firstOutboxSeq: items[0]?.outboxSeq ?? null, lastOutboxSeq: items.at(-1)?.outboxSeq ?? null, valuesPrinted: false }; + return { action: "kafka-event-outbox-relay", enabled: true, ...lastRelayCycle }; } export function startKafkaEventOutboxRelay(input: { store: AgentRunStore; options: KafkaOutboxRelayOptions; onError?: (error: unknown) => void }): () => Promise { @@ -96,3 +112,14 @@ async function retry(store: AgentRunStore, options: KafkaOutboxRelayOptions, ite function stringHeaders(headers: JsonRecord): Record { return Object.fromEntries(Object.entries(headers).map(([key, value]) => [key, typeof value === "string" ? value : JSON.stringify(value)])); } + +function groupByPartitionKey(items: KafkaEventOutboxRecord[]): KafkaEventOutboxRecord[][] { + const groups = new Map(); + for (const item of items) { + const key = `${item.topic}\u0000${item.partitionKey}`; + const group = groups.get(key) ?? []; + group.push(item); + groups.set(key, group); + } + return Array.from(groups.values()); +} diff --git a/src/mgr/postgres-durable-queues.ts b/src/mgr/postgres-durable-queues.ts index 0ba8d30..69e3e98 100644 --- a/src/mgr/postgres-durable-queues.ts +++ b/src/mgr/postgres-durable-queues.ts @@ -51,6 +51,12 @@ CREATE INDEX IF NOT EXISTS agentrun_kafka_event_outbox_due_idx ON agentrun_kafka_event_outbox (delivered_at, next_attempt_at, lease_expires_at); `; +export const kafkaOutboxPendingKeyOrderMigrationSql = ` +CREATE INDEX IF NOT EXISTS agentrun_kafka_event_outbox_pending_key_order_idx + ON agentrun_kafka_event_outbox (topic, partition_key, outbox_seq) + WHERE delivered_at IS NULL; +`; + export const runnerDispatchOutcomeMigrationSql = ` ALTER TABLE agentrun_runner_dispatch_intents ADD COLUMN IF NOT EXISTS dispatch_outcome text; ALTER TABLE agentrun_runner_dispatch_intents ADD COLUMN IF NOT EXISTS actual_runner_job_id text; diff --git a/src/mgr/postgres-store.ts b/src/mgr/postgres-store.ts index 4ddaff6..a6cf049 100644 --- a/src/mgr/postgres-store.ts +++ b/src/mgr/postgres-store.ts @@ -11,7 +11,7 @@ import { backendCapabilitiesSqlValues, mergeBackendCapability } from "../common/ import { normalizeRunEventPayload, requireEventType, requireExternallyAppendableEventType, userMessagePayloadForCommand } from "../common/events.js"; import { buildKafkaEventOutboxRecord, canonicalAgentRunEventPartitionKey } from "./event-outbox.js"; import { assertRunnerDispatchReplay, attachRunnerDispatchIntent, newRunnerDispatchIntent } from "./runner-dispatch-intent.js"; -import { durableDispatchAndKafkaOutboxMigrationSql, durableQueueStatusFromRow, kafkaEventOutboxFromRow, runnerDispatchIntentFromRow, runnerDispatchOutcomeMigrationSql } from "./postgres-durable-queues.js"; +import { durableDispatchAndKafkaOutboxMigrationSql, durableQueueStatusFromRow, kafkaEventOutboxFromRow, kafkaOutboxPendingKeyOrderMigrationSql, runnerDispatchIntentFromRow, runnerDispatchOutcomeMigrationSql } from "./postgres-durable-queues.js"; import { fenceActiveRunnerDispatchIntents, fenceTerminalRunnerDispatchIntents, lockRunnerDispatchClaim, staleClaimError } from "./postgres-runner-dispatch.js"; interface PostgresStoreOptions { @@ -497,6 +497,11 @@ const postgresMigrations: MigrationDefinition[] = [ checksum: checksumSql(queueAttemptsMigrationSql), sql: queueAttemptsMigrationSql, }, + { + id: "015_v02_kafka_outbox_pending_key_order", + checksum: checksumSql(kafkaOutboxPendingKeyOrderMigrationSql), + sql: kafkaOutboxPendingKeyOrderMigrationSql, + }, ]; export function postgresMigrationContract(): JsonRecord { @@ -514,6 +519,12 @@ export function postgresMigrationContract(): JsonRecord { failureVisibility: ["failure_kind", "failure_message"], valuesPrinted: false, }, + kafkaOutbox: { + pendingKeyOrderIndexMigrationId: "015_v02_kafka_outbox_pending_key_order", + pendingKeyOrderIndex: "agentrun_kafka_event_outbox_pending_key_order_idx", + relayParallelismBound: "AGENTRUN_KAFKA_OUTBOX_BATCH_SIZE", + valuesPrinted: false, + }, checksums: Object.fromEntries(postgresMigrations.map((migration) => [migration.id, migration.checksum])), }; } @@ -816,8 +827,9 @@ CREATE TABLE IF NOT EXISTS agentrun_schema_migrations ( async claimKafkaEventOutbox(input: { owner: string; leaseMs: number; limit: number }): Promise { const leaseExpiresAt = new Date(Date.now() + input.leaseMs).toISOString(); const result = await this.pool.query( - `WITH due AS ( - SELECT candidate.id FROM agentrun_kafka_event_outbox candidate + `WITH heads AS MATERIALIZED ( + SELECT candidate.id, candidate.topic, candidate.partition_key + FROM agentrun_kafka_event_outbox candidate WHERE candidate.delivered_at IS NULL AND candidate.next_attempt_at <= now() AND (candidate.lease_expires_at IS NULL OR candidate.lease_expires_at <= now()) @@ -831,6 +843,29 @@ CREATE TABLE IF NOT EXISTS agentrun_schema_migrations ( ORDER BY candidate.outbox_seq ASC FOR UPDATE OF candidate SKIP LOCKED LIMIT $1 + ), ranked AS MATERIALIZED ( + SELECT candidate.id, + row_number() OVER (PARTITION BY candidate.topic, candidate.partition_key ORDER BY candidate.outbox_seq ASC) AS key_offset + FROM agentrun_kafka_event_outbox candidate + JOIN heads ON heads.topic = candidate.topic AND heads.partition_key = candidate.partition_key + WHERE candidate.delivered_at IS NULL + AND candidate.next_attempt_at <= now() + AND (candidate.lease_expires_at IS NULL OR candidate.lease_expires_at <= now()) + AND NOT EXISTS ( + SELECT 1 FROM agentrun_kafka_event_outbox blocker + WHERE blocker.topic = candidate.topic + AND blocker.partition_key = candidate.partition_key + AND blocker.delivered_at IS NULL + AND blocker.outbox_seq < candidate.outbox_seq + AND (blocker.next_attempt_at > now() OR blocker.lease_expires_at > now()) + ) + ), due AS MATERIALIZED ( + SELECT candidate.id + FROM agentrun_kafka_event_outbox candidate + JOIN ranked ON ranked.id = candidate.id + ORDER BY ranked.key_offset ASC, candidate.outbox_seq ASC + FOR UPDATE OF candidate SKIP LOCKED + LIMIT $1 ) UPDATE agentrun_kafka_event_outbox o SET attempt_count = o.attempt_count + 1, lease_owner = $2, lease_expires_at = $3, updated_at = now() @@ -871,10 +906,11 @@ CREATE TABLE IF NOT EXISTS agentrun_schema_migrations ( count(*) FILTER (WHERE delivered_at IS NULL)::int AS backlog_count, count(*) FILTER (WHERE delivered_at IS NULL AND attempt_count > 0)::int AS retrying_count, min(created_at) FILTER (WHERE delivered_at IS NULL) AS oldest_pending_at, - COALESCE(max(attempt_count), 0)::int AS max_attempt_count + COALESCE(max(attempt_count), 0)::int AS max_attempt_count, + (array_agg(last_error ORDER BY outbox_seq ASC) FILTER (WHERE delivered_at IS NULL AND last_error IS NOT NULL))[1] AS first_error FROM agentrun_kafka_event_outbox`, ); - return { ...durableQueueStatusFromRow(result.rows[0], "kafka-event-outbox"), enabled: this.eventOutboxConfig.enabled }; + return { ...durableQueueStatusFromRow(result.rows[0], "kafka-event-outbox"), firstError: result.rows[0]?.first_error ?? null, enabled: this.eventOutboxConfig.enabled }; } async registerRunner(input: Partial): Promise { diff --git a/src/mgr/server.ts b/src/mgr/server.ts index 6f546c6..89c1802 100644 --- a/src/mgr/server.ts +++ b/src/mgr/server.ts @@ -16,7 +16,7 @@ import { buildRunResult } from "./result.js"; import { runnerJobStatusSummary } from "./runner-job-status.js"; import { reconcileRunnerJobsOnce, startRunnerJobReconciler } from "./runner-reconciler.js"; import { assertRunnerDispatcherAttemptTiming, dispatchRunnerIntentsOnce, startRunnerDispatcher, type RunnerDispatcherOptions } from "./runner-dispatcher.js"; -import { relayKafkaEventOutboxOnce, startKafkaEventOutboxRelay, type KafkaOutboxRelayOptions } from "./kafka-outbox-relay.js"; +import { kafkaOutboxRelayStatus, relayKafkaEventOutboxOnce, startKafkaEventOutboxRelay, type KafkaOutboxRelayOptions } from "./kafka-outbox-relay.js"; import { createSessionPvc, deleteSessionPvc, getSessionPvcSummary, refreshSessionPvcSummary, runSessionStorageGc } from "./session-pvc.js"; import type { SessionPvcSummary } from "./session-pvc.js"; import type { SessionPvcOptions } from "./session-pvc.js"; @@ -646,11 +646,11 @@ async function route({ method, url, body, store, sourceCommit, authSummary, diag if (method === "GET" && (path === "/health" || path === "/health/live" || path === "/health/readiness")) { const [database, runnerDispatch, kafkaEventOutbox] = await Promise.all([store.health(), store.runnerDispatchStatus(), store.kafkaEventOutboxStatus()]); const ready = path === "/health/live" ? true : database.ready; - return { serviceId: "agentrun-mgr", live: true, ready, database, sourceCommit, auth: authSummary ?? null, runnerWorkReady: staticWorkReadyCapabilitySummary(), durableQueues: { runnerDispatch: { ...runnerDispatch, enabled: runnerDispatcherOptions.enabled }, kafkaEventOutbox }, secretRefs: { databaseUrl: database.adapter === "postgres" ? "redacted" : "not-used", valuesPrinted: false } }; + return { serviceId: "agentrun-mgr", live: true, ready, database, sourceCommit, auth: authSummary ?? null, runnerWorkReady: staticWorkReadyCapabilitySummary(), durableQueues: { runnerDispatch: { ...runnerDispatch, enabled: runnerDispatcherOptions.enabled }, kafkaEventOutbox: { ...kafkaEventOutbox, relay: kafkaOutboxRelayStatus() } }, secretRefs: { databaseUrl: database.adapter === "postgres" ? "redacted" : "not-used", valuesPrinted: false } }; } if (method === "GET" && path === "/api/v1/operations/durable-queues") { const [runnerDispatch, kafkaEventOutbox] = await Promise.all([store.runnerDispatchStatus(), store.kafkaEventOutboxStatus()]); - return { runnerDispatch: { ...runnerDispatch, enabled: runnerDispatcherOptions.enabled }, kafkaEventOutbox, relayEnabled: kafkaOutboxRelayOptions.enabled, valuesPrinted: false }; + return { runnerDispatch: { ...runnerDispatch, enabled: runnerDispatcherOptions.enabled }, kafkaEventOutbox: { ...kafkaEventOutbox, relay: kafkaOutboxRelayStatus() }, relayEnabled: kafkaOutboxRelayOptions.enabled, valuesPrinted: false }; } if (method === "GET" && path === "/api/v1/backends") return await listBackendCapabilities(providerProfileDefaults) as JsonValue; if (method === "GET" && path === "/api/v1/tool-credentials") return await listToolCredentials(toolCredentialDefaults) as JsonValue; diff --git a/src/mgr/store.ts b/src/mgr/store.ts index 7db8634..af4c827 100644 --- a/src/mgr/store.ts +++ b/src/mgr/store.ts @@ -516,11 +516,27 @@ export class MemoryAgentRunStore implements AgentRunStore { claimKafkaEventOutbox(input: { owner: string; leaseMs: number; limit: number }): KafkaEventOutboxRecord[] { const now = Date.now(); const all = Array.from(this.kafkaEventOutbox.values()); - const claimed = all - .filter((item) => !item.deliveredAt && Date.parse(item.nextAttemptAt) <= now && (!item.leaseExpiresAt || Date.parse(item.leaseExpiresAt) <= now)) - .filter((item) => !all.some((earlier) => !earlier.deliveredAt && earlier.topic === item.topic && earlier.partitionKey === item.partitionKey && earlier.outboxSeq < item.outboxSeq)) - .sort((a, b) => a.outboxSeq - b.outboxSeq) - .slice(0, Math.max(1, Math.min(input.limit, 500))); + const limit = Math.max(1, Math.min(input.limit, 500)); + const claimed: KafkaEventOutboxRecord[] = []; + const pendingByKey = new Map(); + for (const item of all.filter((candidate) => !candidate.deliveredAt).sort((a, b) => a.outboxSeq - b.outboxSeq)) { + const key = `${item.topic}\u0000${item.partitionKey}`; + const pending = pendingByKey.get(key) ?? []; + pending.push(item); + pendingByKey.set(key, pending); + } + const eligible = Array.from(pendingByKey.values()) + .filter((items) => isDueKafkaOutboxItem(items[0], now)) + .flatMap((items) => { + const prefix: KafkaEventOutboxRecord[] = []; + for (const item of items) { + if (!isDueKafkaOutboxItem(item, now)) break; + prefix.push(item); + } + return prefix.map((item, keyOffset) => ({ item, keyOffset })); + }) + .sort((a, b) => a.keyOffset - b.keyOffset || a.item.outboxSeq - b.item.outboxSeq); + claimed.push(...eligible.slice(0, limit).map(({ item }) => item)); return claimed.map((item) => { const updated: KafkaEventOutboxRecord = { ...item, attemptCount: item.attemptCount + 1, leaseOwner: input.owner, leaseExpiresAt: new Date(now + input.leaseMs).toISOString(), updatedAt: nowIso() }; this.kafkaEventOutbox.set(item.id, updated); @@ -546,7 +562,9 @@ export class MemoryAgentRunStore implements AgentRunStore { } kafkaEventOutboxStatus(): JsonRecord { - return { ...workQueueStatus(Array.from(this.kafkaEventOutbox.values()), "kafka-event-outbox", (item) => Boolean((item as KafkaEventOutboxRecord).deliveredAt)), enabled: this.eventOutboxConfig.enabled }; + const items = Array.from(this.kafkaEventOutbox.values()); + const firstError = items.filter((item) => !item.deliveredAt && item.lastError).sort((a, b) => a.outboxSeq - b.outboxSeq)[0]?.lastError ?? null; + return { ...workQueueStatus(items, "kafka-event-outbox", (item) => Boolean((item as KafkaEventOutboxRecord).deliveredAt)), firstError, enabled: this.eventOutboxConfig.enabled }; } getCommand(commandId: string): CommandRecord { @@ -1905,6 +1923,10 @@ function workQueueStatus(items: T[], component: string, co }; } +function isDueKafkaOutboxItem(item: KafkaEventOutboxRecord | undefined, now: number): boolean { + return Boolean(item && Date.parse(item.nextAttemptAt) <= now && (!item.leaseExpiresAt || Date.parse(item.leaseExpiresAt) <= now)); +} + export function isSessionOutputEvent(event: RunEvent): boolean { return event.type === "assistant_progress" || event.type === "assistant_message" || event.type === "command_output" || event.type === "diff" || event.type === "error" || event.type === "terminal_status"; } diff --git a/src/selftest/cases/00-redaction-postgres.ts b/src/selftest/cases/00-redaction-postgres.ts index 9d1b0b2..898e109 100644 --- a/src/selftest/cases/00-redaction-postgres.ts +++ b/src/selftest/cases/00-redaction-postgres.ts @@ -3,6 +3,7 @@ import { openAgentRunStoreFromEnv } from "../../mgr/store.js"; import { postgresMigrationContract } from "../../mgr/postgres-store.js"; import { redactText } from "../../common/redaction.js"; import { AgentRunError } from "../../common/errors.js"; +import type { JsonRecord } from "../../common/types.js"; import type { SelfTestCase } from "../harness.js"; const selfTest: SelfTestCase = async () => { @@ -13,7 +14,7 @@ const selfTest: SelfTestCase = async () => { (error) => error instanceof AgentRunError && error.failureKind === "infra-failed" && error.message.includes("DATABASE_URL is required"), ); const postgresContract = postgresMigrationContract(); - assert.equal(postgresContract.latestMigrationId, "014_v02_queue_attempt_history_retry"); + assert.equal(postgresContract.latestMigrationId, "015_v02_kafka_outbox_pending_key_order"); assert.equal((postgresContract.migrationIds as string[]).includes("008_v01_dsflash_go_backend_profile"), true); assert.equal((postgresContract.migrationIds as string[]).includes("009_v01_dsflash_go_model_catalog"), true); assert.equal((postgresContract.migrationIds as string[]).includes("010_v01_queue_session_ref"), true); @@ -21,6 +22,7 @@ const selfTest: SelfTestCase = async () => { assert.equal((postgresContract.migrationIds as string[]).includes("012_v02_durable_dispatch_kafka_outbox"), true); assert.equal((postgresContract.migrationIds as string[]).includes("013_v02_runner_dispatch_outcome"), true); assert.equal((postgresContract.migrationIds as string[]).includes("014_v02_queue_attempt_history_retry"), true); + assert.equal((postgresContract.migrationIds as string[]).includes("015_v02_kafka_outbox_pending_key_order"), true); assert.ok(typeof (postgresContract.checksums as Record)["008_v01_dsflash_go_backend_profile"] === "string" && (postgresContract.checksums as Record)["008_v01_dsflash_go_backend_profile"].length > 0); assert.ok(typeof (postgresContract.checksums as Record)["009_v01_dsflash_go_model_catalog"] === "string" && (postgresContract.checksums as Record)["009_v01_dsflash_go_model_catalog"].length > 0); assert.ok(typeof (postgresContract.checksums as Record)["010_v01_queue_session_ref"] === "string" && (postgresContract.checksums as Record)["010_v01_queue_session_ref"].length > 0); @@ -28,7 +30,11 @@ const selfTest: SelfTestCase = async () => { assert.ok(typeof (postgresContract.checksums as Record)["012_v02_durable_dispatch_kafka_outbox"] === "string" && (postgresContract.checksums as Record)["012_v02_durable_dispatch_kafka_outbox"].length > 0); assert.ok(typeof (postgresContract.checksums as Record)["013_v02_runner_dispatch_outcome"] === "string" && (postgresContract.checksums as Record)["013_v02_runner_dispatch_outcome"].length > 0); assert.ok(typeof (postgresContract.checksums as Record)["014_v02_queue_attempt_history_retry"] === "string" && (postgresContract.checksums as Record)["014_v02_queue_attempt_history_retry"].length > 0); + assert.ok(typeof (postgresContract.checksums as Record)["015_v02_kafka_outbox_pending_key_order"] === "string" && (postgresContract.checksums as Record)["015_v02_kafka_outbox_pending_key_order"].length > 0); assert.equal((postgresContract.checksums as Record)["002_v01_backend_profiles"], "928b5c490cc4539cb64ecef34784557601b2724fa2870570f16a53576804e49c"); + assert.equal((postgresContract.checksums as Record)["012_v02_durable_dispatch_kafka_outbox"], "48c552aae6d9a70eb8f1c0638aeda6ac24dcb6efa4d9a9c56e471dbc17367dc5"); + assert.equal((postgresContract.kafkaOutbox as JsonRecord).pendingKeyOrderIndexMigrationId, "015_v02_kafka_outbox_pending_key_order"); + assert.equal((postgresContract.kafkaOutbox as JsonRecord).relayParallelismBound, "AGENTRUN_KAFKA_OUTBOX_BATCH_SIZE"); assert.ok(Array.isArray(postgresContract.requiredTables)); assert.ok(postgresContract.requiredTables.includes("agentrun_schema_migrations")); assert.ok(postgresContract.requiredTables.includes("agentrun_runs")); @@ -41,7 +47,7 @@ const selfTest: SelfTestCase = async () => { assert.ok(postgresContract.requiredTables.includes("agentrun_cancel_requests")); assert.ok(postgresContract.requiredTables.includes("agentrun_runner_dispatch_intents")); assert.ok(postgresContract.requiredTables.includes("agentrun_kafka_event_outbox")); - return { name: "redaction-postgres", tests: ["redaction", "postgres-store-contract"] }; + return { name: "redaction-postgres", tests: ["redaction", "postgres-store-contract", "historical-migration-checksum"] }; }; export default selfTest; diff --git a/src/selftest/cases/35-kafka-durable-outbox.ts b/src/selftest/cases/35-kafka-durable-outbox.ts index b69b9fe..a117e0e 100644 --- a/src/selftest/cases/35-kafka-durable-outbox.ts +++ b/src/selftest/cases/35-kafka-durable-outbox.ts @@ -91,7 +91,7 @@ const selfTest: SelfTestCase = async () => { const normalizedDispatch = validateCreateCommand({ type: "turn", payload: {}, dispatch: { kind: "kubernetes-job", input: { transientEnv: [{ name: "HWLAB_RUNTIME_LANE", value: "v02", sensitive: false }] } } }).dispatch?.input; assert.deepEqual(normalizedDispatch, { transientEnv: [{ name: "HWLAB_RUNTIME_LANE", value: "v02", sensitive: false }] }); - const firstClaim = store.claimKafkaEventOutbox({ owner: "first", leaseMs: 60_000, limit: 100 }); + const firstClaim = store.claimKafkaEventOutbox({ owner: "first", leaseMs: 60_000, limit: 1 }); assert.equal(firstClaim.length, 1); assert.equal(firstClaim[0]?.outboxSeq, 1); store.retryKafkaEventOutbox(firstClaim[0]!, { code: "broker-unavailable" }, new Date(Date.now() + 20).toISOString()); @@ -103,9 +103,10 @@ const selfTest: SelfTestCase = async () => { const firstCycle = await relayKafkaEventOutboxOnce({ store, options: relayOptions, publish }); const secondCycle = await relayKafkaEventOutboxOnce({ store, options: relayOptions, publish }); const thirdCycle = await relayKafkaEventOutboxOnce({ store, options: relayOptions, publish }); - assert.equal(firstCycle.deliveredCount, 1); - assert.equal(secondCycle.deliveredCount, 1); - assert.equal(thirdCycle.deliveredCount, 1); + assert.equal(firstCycle.deliveredCount, 3); + assert.equal(firstCycle.partitionKeyCount, 1); + assert.equal(secondCycle.deliveredCount, 0); + assert.equal(thirdCycle.deliveredCount, 0); assert.deepEqual(published.map((item) => item.outboxSeq), [1, 2, 3]); assert.deepEqual(published.map((item) => item.sessionId), ["ses_kafka_durable", "ses_kafka_durable", "ses_kafka_durable"]); assert.equal(published[2]?.schema, "agentrun.event.v1"); @@ -132,6 +133,8 @@ const selfTest: SelfTestCase = async () => { assert.equal(stdioMetadata.frameSeq, 7); assert.equal(stdioMetadata.runId, run.id); assert.equal(stdioMetadata.commandId, command.id); + await assertRelayParallelismAndFailureIsolation(); + await assertRelayParallelismBoundedByBatchSize(); assertWarmRunCommandTraceAuthority(); const claimedIntent = store.claimRunnerDispatchIntents({ owner: "terminalizer", leaseMs: 60_000, limit: 1 })[0]; @@ -171,7 +174,7 @@ const selfTest: SelfTestCase = async () => { await assertStaleClaimsCannotOverwrite(); await assertProducerRecreatedAfterDisconnected(); - return { name: "kafka-durable-outbox", tests: ["explicit-enable", "canonical-event", "warm-run-command-trace-authority", "run-level-trace-fallback", "stdio-command-trace-context", "kafka-tail-correlation-summary", "durable-head-of-line", "atomic-dispatch-terminal", "terminal-command-idempotent-replay", "terminal-command-replay-no-new-events-or-outbox", "terminal-command-same-key-payload-conflict", "terminal-command-different-key-fails-closed", "stale-claim-fencing", "producer-reconnect"] }; + return { name: "kafka-durable-outbox", tests: ["explicit-enable", "canonical-event", "warm-run-command-trace-authority", "run-level-trace-fallback", "stdio-command-trace-context", "kafka-tail-correlation-summary", "durable-key-batch", "cross-key-parallelism", "batch-bounded-parallelism", "failed-key-isolation", "atomic-dispatch-terminal", "terminal-command-idempotent-replay", "terminal-command-replay-no-new-events-or-outbox", "terminal-command-same-key-payload-conflict", "terminal-command-different-key-fails-closed", "stale-claim-fencing", "producer-reconnect"] }; }; function assertWarmRunCommandTraceAuthority(): void { @@ -321,6 +324,70 @@ function staleClaim(error: unknown): boolean { return error instanceof AgentRunError && error.failureKind === "runner-lease-conflict"; } +async function assertRelayParallelismAndFailureIsolation(): Promise { + const store = new MemoryAgentRunStore({ eventOutbox: { enabled: true, topic: config.agentrunEventTopic, source: config.clientId } }); + const blockedRun = store.createRun(runInput("ses_kafka_blocked")); + store.appendEvent(blockedRun.id, "backend_status", { phase: "blocked-1" }); + store.appendEvent(blockedRun.id, "backend_status", { phase: "blocked-2" }); + store.appendEvent(blockedRun.id, "backend_status", { phase: "blocked-3" }); + const healthyRun = store.createRun(runInput("ses_kafka_healthy")); + store.appendEvent(healthyRun.id, "backend_status", { phase: "healthy-1" }); + store.appendEvent(healthyRun.id, "backend_status", { phase: "healthy-2" }); + store.appendEvent(healthyRun.id, "backend_status", { phase: "healthy-3" }); + const publishedByKey = new Map(); + let active = 0; + let maxActive = 0; + const result = await relayKafkaEventOutboxOnce({ + store, + options: { ...relayOptions, batchSize: 6 }, + publish: async (_config, message) => { + active++; + maxActive = Math.max(maxActive, active); + try { + await delay(10); + const key = String(message.value.sessionId); + publishedByKey.set(key, [...(publishedByKey.get(key) ?? []), Number(message.value.outboxSeq)]); + if (key === "ses_kafka_blocked") throw new Error("selftest broker rejection"); + } finally { + active--; + } + }, + }); + assert.equal(result.claimedCount, 6); + assert.equal(result.partitionKeyCount, 2); + assert.equal(result.deliveredCount, 3); + assert.equal(result.retryCount, 3); + assert.equal((result.firstError as JsonRecord).code, "kafka-produce-failed"); + assert.equal(maxActive, 2); + assert.deepEqual(publishedByKey.get("ses_kafka_blocked"), [1]); + assert.deepEqual(publishedByKey.get("ses_kafka_healthy"), [4, 5, 6]); + assert.equal(store.kafkaEventOutboxStatus().backlogCount, 3); +} + +async function assertRelayParallelismBoundedByBatchSize(): Promise { + const store = new MemoryAgentRunStore({ eventOutbox: { enabled: true, topic: config.agentrunEventTopic, source: config.clientId } }); + for (let index = 0; index < 20; index++) { + const run = store.createRun(runInput(`ses_kafka_parallel_${index}`)); + store.appendEvent(run.id, "backend_status", { phase: `parallel-${index}` }); + } + let active = 0; + let maxActive = 0; + const result = await relayKafkaEventOutboxOnce({ + store, + options: { ...relayOptions, batchSize: 8 }, + publish: async () => { + active++; + maxActive = Math.max(maxActive, active); + await delay(10); + active--; + }, + }); + assert.equal(result.claimedCount, 8); + assert.equal(result.partitionKeyCount, 8); + assert.equal(maxActive, 8); + assert.equal(store.kafkaEventOutboxStatus().backlogCount, 12); +} + function delay(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } diff --git a/src/selftest/integration/postgres-kafka-order.ts b/src/selftest/integration/postgres-kafka-order.ts index a0db713..634bda3 100644 --- a/src/selftest/integration/postgres-kafka-order.ts +++ b/src/selftest/integration/postgres-kafka-order.ts @@ -2,7 +2,8 @@ import assert from "node:assert/strict"; import { Pool, type PoolClient } from "pg"; import type { CreateRunInput, JsonRecord, KafkaEventOutboxRecord } from "../../common/types.js"; import { validateCreateCommand } from "../../common/validation.js"; -import { createPostgresAgentRunStore } from "../../mgr/postgres-store.js"; +import { relayKafkaEventOutboxOnce, type KafkaOutboxRelayOptions } from "../../mgr/kafka-outbox-relay.js"; +import { createPostgresAgentRunStore, postgresMigrationContract } from "../../mgr/postgres-store.js"; import { runCreatedEventPayload } from "../../mgr/store.js"; const connectionString = process.env.AGENTRUN_SELFTEST_DATABASE_URL?.trim(); @@ -12,7 +13,7 @@ const sessionId = "ses_pg_kafka_commit_order"; const topic = "agentrun.event.v1"; const gate = "746173865219047"; const controlPool = new Pool({ connectionString, application_name: "agentrun-pg-order-control", max: 2 }); -const store = await createPostgresAgentRunStore({ +let store = await createPostgresAgentRunStore({ connectionString, poolMax: 6, eventOutbox: { enabled: true, topic, source: "agentrun-pg-order-test" }, @@ -58,6 +59,9 @@ try { assert.equal(marker(second), "B"); await store.completeKafkaEventOutbox(second); + await assertPostgresBatchRelay(store, controlPool); + store = await assertForwardMigrationFromExistingV02Ledger(store, controlPool); + const userRun = await store.createRun(runInput("ses_pg_user_message_order")); const userCommandInput = validateCreateCommand({ type: "turn", @@ -182,6 +186,79 @@ function runInput(targetSessionId: string): CreateRunInput { }; } +async function assertPostgresBatchRelay(store: Awaited>, pool: Pool): Promise { + await pool.query("DELETE FROM agentrun_kafka_event_outbox"); + await pool.query("ALTER SEQUENCE agentrun_kafka_event_outbox_seq RESTART WITH 1"); + const blockedRun = await store.createRun(runInput("ses_pg_kafka_blocked")); + const healthyRun = await store.createRun(runInput("ses_pg_kafka_healthy")); + for (let index = 1; index <= 100; index++) { + await store.appendEvent(blockedRun.id, "backend_status", { phase: `blocked-${index}` }); + await store.appendEvent(healthyRun.id, "backend_status", { phase: `healthy-${index}` }); + } + let active = 0; + let maxActive = 0; + const publishedByKey = new Map(); + const options: KafkaOutboxRelayOptions = { + enabled: true, + intervalMs: 1_000, + batchSize: 200, + leaseMs: 60_000, + retryBackoffMs: 1, + owner: "postgres-batch-relay", + config: { brokers: ["127.0.0.1:9092"], clientId: "postgres-batch-relay", agentrunEventTopic: topic, codexStdioTopic: "codex-stdio.raw.v1", enabled: true, valuesPrinted: false }, + }; + const result = await relayKafkaEventOutboxOnce({ + store, + options, + publish: async (_config, message) => { + active++; + maxActive = Math.max(maxActive, active); + try { + await delay(1); + const key = String(message.value.sessionId); + publishedByKey.set(key, [...(publishedByKey.get(key) ?? []), Number(message.value.outboxSeq)]); + if (key === "ses_pg_kafka_blocked") throw new Error("postgres selftest broker rejection"); + } finally { + active--; + } + }, + }); + assert.equal(result.claimedCount, 200); + assert.equal(result.partitionKeyCount, 2); + assert.equal(result.deliveredCount, 100); + assert.equal(result.retryCount, 100); + assert.equal(maxActive, 2); + assert.deepEqual(publishedByKey.get("ses_pg_kafka_blocked"), [1]); + assert.deepEqual(publishedByKey.get("ses_pg_kafka_healthy"), Array.from({ length: 100 }, (_, index) => (index + 1) * 2)); + const status = await store.kafkaEventOutboxStatus(); + assert.equal(status.backlogCount, 100); + assert.equal((status.firstError as JsonRecord).code, "kafka-produce-failed"); + await delay(5); + const retryResult = await relayKafkaEventOutboxOnce({ store, options, publish: async () => undefined }); + assert.equal(retryResult.claimedCount, 100); + assert.equal(retryResult.deliveredCount, 100); + assert.equal((await store.kafkaEventOutboxStatus()).backlogCount, 0); +} + +async function assertForwardMigrationFromExistingV02Ledger(store: Awaited>, pool: Pool): Promise>> { + const contract = postgresMigrationContract(); + const checksums = contract.checksums as Record; + const legacyIds = ["012_v02_durable_dispatch_kafka_outbox", "013_v02_runner_dispatch_outcome", "014_v02_queue_attempt_history_retry"]; + const legacyLedger = await pool.query<{ id: string; checksum: string }>("SELECT id, checksum FROM agentrun_schema_migrations WHERE id = ANY($1::text[]) ORDER BY id", [legacyIds]); + assert.deepEqual(legacyLedger.rows.map((row) => row.id), legacyIds); + assert.deepEqual(legacyLedger.rows.map((row) => row.checksum), legacyIds.map((id) => checksums[id])); + await store.close(); + await pool.query("DROP INDEX IF EXISTS agentrun_kafka_event_outbox_pending_key_order_idx"); + await pool.query("DELETE FROM agentrun_schema_migrations WHERE id = $1", ["015_v02_kafka_outbox_pending_key_order"]); + const reopened = await createPostgresAgentRunStore({ connectionString, poolMax: 6, eventOutbox: { enabled: true, topic, source: "agentrun-pg-order-test" } }); + const applied = await pool.query<{ checksum: string }>("SELECT checksum FROM agentrun_schema_migrations WHERE id = $1", ["015_v02_kafka_outbox_pending_key_order"]); + assert.equal(applied.rows[0]?.checksum, checksums["015_v02_kafka_outbox_pending_key_order"]); + const index = await pool.query<{ present: boolean }>("SELECT to_regclass('agentrun_kafka_event_outbox_pending_key_order_idx') IS NOT NULL AS present"); + assert.equal(index.rows[0]?.present, true); + assert.equal((await reopened.health()).migrationId, "015_v02_kafka_outbox_pending_key_order"); + return reopened; +} + async function installGateTrigger(pool: Pool, advisoryGate: string): Promise { await removeGateTrigger(pool); await pool.query(`