From 3fb90ec3f23b9518b0bf968bebb9d9efeb316496 Mon Sep 17 00:00:00 2001 From: Shin <128954611+shin4141@users.noreply.github.com> Date: Sat, 15 Aug 2026 12:41:39 +0900 Subject: [PATCH] fix(world-postgres): make step creation atomic Signed-off-by: Shin <128954611+shin4141@users.noreply.github.com> --- .changeset/atomic-step-created-event.md | 5 + packages/world-postgres/src/storage.ts | 117 +++++++++++++------ packages/world-postgres/test/storage.test.ts | 100 +++++++++++++++- 3 files changed, 188 insertions(+), 34 deletions(-) create mode 100644 .changeset/atomic-step-created-event.md diff --git a/.changeset/atomic-step-created-event.md b/.changeset/atomic-step-created-event.md new file mode 100644 index 0000000000..449ad03e88 --- /dev/null +++ b/.changeset/atomic-step-created-event.md @@ -0,0 +1,5 @@ +--- +'@workflow/world-postgres': patch +--- + +Commit step entities and their `step_created` events atomically. diff --git a/packages/world-postgres/src/storage.ts b/packages/world-postgres/src/storage.ts index 7d6c10f86d..a2a03d651b 100644 --- a/packages/world-postgres/src/storage.ts +++ b/packages/world-postgres/src/storage.ts @@ -1472,31 +1472,6 @@ export function createEventsStorage(drizzle: Drizzle): Storage['events'] { storedEventData = undefined; } - // Handle step_created event: create step entity - if (data.eventType === 'step_created') { - const eventData = (data as any).eventData as { - stepName: string; - input: any; - }; - const [stepValue] = await drizzle - .insert(Schema.steps) - .values({ - runId: effectiveRunId, - stepId: data.correlationId!, - stepName: eventData.stepName, - input: eventData.input as SerializedContent, - status: 'pending', - attempt: 0, - // Propagate specVersion from the event to the step entity - specVersion: effectiveSpecVersion, - }) - .onConflictDoNothing() - .returning(); - if (stepValue) { - step = deserializeStepError(compact(stepValue)); - } - } - let value: { createdAt: Date } | undefined; // Handle step_started event: increment attempt and set the step to @@ -2064,14 +2039,90 @@ export function createEventsStorage(drizzle: Drizzle): Storage['events'] { try { if (!value) { - const inserted = await insertEventRow(drizzle, { - runId: effectiveRunId, - eventId: await getEventId(), - correlationId: data.correlationId, - eventType: data.eventType, - eventData: storedEventData, - specVersion: effectiveSpecVersion, - }); + let inserted: Awaited>; + if (data.eventType === 'step_created') { + const eventData = data.eventData; + const created = await drizzle.transaction(async (tx) => { + let [stepValue] = await tx + .insert(Schema.steps) + .values({ + runId: effectiveRunId, + stepId: data.correlationId, + stepName: eventData.stepName, + input: eventData.input as SerializedContent, + status: 'pending', + attempt: 0, + specVersion: effectiveSpecVersion, + }) + .onConflictDoNothing() + .returning(); + if (!stepValue) { + const [existingEvent] = await tx + .select({ eventId: Schema.events.eventId }) + .from(Schema.events) + .where( + and( + eq(Schema.events.runId, effectiveRunId), + eq(Schema.events.correlationId, data.correlationId), + eq(Schema.events.eventType, 'step_created') + ) + ) + .limit(1); + if (existingEvent) { + throw new EntityConflictError( + `step_created for correlationId "${data.correlationId}" already exists in run "${effectiveRunId}"` + ); + } + + // A row without its matching event was left by the old + // non-transactional path. Keep the row and complete the + // missing event inside this transaction so existing orphans + // remain recoverable while new partial writes cannot escape. + [stepValue] = await tx + .select() + .from(Schema.steps) + .where( + and( + eq(Schema.steps.runId, effectiveRunId), + eq(Schema.steps.stepId, data.correlationId) + ) + ) + .limit(1); + if (!stepValue) { + throw new EntityConflictError( + `step_created for correlationId "${data.correlationId}" already exists in run "${effectiveRunId}"` + ); + } + } + + const eventValue = await insertEventRow(tx, { + runId: effectiveRunId, + eventId: await getEventId(tx), + correlationId: data.correlationId, + eventType: data.eventType, + eventData: storedEventData, + specVersion: effectiveSpecVersion, + }); + if (!eventValue) { + throw new EntityConflictError( + `step_created for run "${effectiveRunId}" could not be created` + ); + } + return { eventValue, stepValue }; + }, SLOT_INSERT_TRANSACTION); + + step = deserializeStepError(compact(created.stepValue)); + inserted = created.eventValue; + } else { + inserted = await insertEventRow(drizzle, { + runId: effectiveRunId, + eventId: await getEventId(), + correlationId: data.correlationId, + eventType: data.eventType, + eventData: storedEventData, + specVersion: effectiveSpecVersion, + }); + } if (inserted) { eventId = inserted.eventId; value = { createdAt: inserted.createdAt }; diff --git a/packages/world-postgres/test/storage.test.ts b/packages/world-postgres/test/storage.test.ts index 4882e17c00..9a6bea0299 100644 --- a/packages/world-postgres/test/storage.test.ts +++ b/packages/world-postgres/test/storage.test.ts @@ -2134,7 +2134,105 @@ describe('Storage (Postgres integration)', () => { stepName: 'test-step', input: new Uint8Array(), }) - ).rejects.toMatchObject({ name: 'EntityConflictError' }); + ).rejects.toMatchObject({ + name: 'EntityConflictError', + message: `step_created for correlationId "step_seq_dup" already exists in run "${testRunId}"`, + }); + }); + + it('recovers an orphaned step row before a plain step_started event', async () => { + const stepId = 'step_orphan'; + await drizzle.insert(DrizzleSchema.steps).values({ + runId: testRunId, + stepId, + stepName: 'test-step', + input: new Uint8Array(), + status: 'pending', + attempt: 0, + specVersion: SPEC_VERSION_CURRENT, + }); + + const recovered = await createStep(events, testRunId, { + stepId, + stepName: 'test-step', + input: new Uint8Array(), + }); + expect(recovered.stepId).toBe(stepId); + + await updateStep(events, testRunId, stepId, 'step_started'); + + const evts = await events.list({ + runId: testRunId, + pagination: {}, + }); + expect( + evts.data + .filter((event) => event.correlationId === stepId) + .map((event) => event.eventType) + ).toEqual(['step_created', 'step_started']); + }); + + it('rolls back the step entity when the matching event insert fails', async () => { + await pool.query(` + CREATE FUNCTION workflow.reject_step_created_event_for_test() + RETURNS trigger + LANGUAGE plpgsql + AS $$ + BEGIN + IF NEW.type = 'step_created' + AND NEW.correlation_id = 'step_partial_write' + THEN + RAISE EXCEPTION 'forced step_created event insert failure'; + END IF; + RETURN NEW; + END; + $$; + + CREATE TRIGGER reject_step_created_event_for_test + BEFORE INSERT ON workflow.workflow_events + FOR EACH ROW + EXECUTE FUNCTION workflow.reject_step_created_event_for_test(); + `); + + try { + await expect( + createStep(events, testRunId, { + stepId: 'step_partial_write', + stepName: 'test-step', + input: new Uint8Array(), + }) + ).rejects.toMatchObject({ + cause: { + message: expect.stringMatching( + /forced step_created event insert failure/ + ), + }, + }); + + const stepRows = await drizzle + .select({ stepId: DrizzleSchema.steps.stepId }) + .from(DrizzleSchema.steps) + .where(eq(DrizzleSchema.steps.stepId, 'step_partial_write')); + expect(stepRows).toEqual([]); + + const evts = await events.list({ + runId: testRunId, + pagination: {}, + }); + expect( + evts.data.filter( + (event) => + event.eventType === 'step_created' && + event.correlationId === 'step_partial_write' + ) + ).toHaveLength(0); + } finally { + await pool.query(` + DROP TRIGGER reject_step_created_event_for_test + ON workflow.workflow_events; + DROP FUNCTION workflow.reject_step_created_event_for_test(); + `); + } }); it('should reject duplicate correlated workflow attr_set events', async () => {