Skip to content

Commit b40eeca

Browse files
d-csd-cs
committed
feat(run-engine): source completed-waitpoint envelope fields from both coordinator arms
The resume path only has id, status, type and completedAfter per edge, which is nine fields short of a completion envelope. Add one coordinator method that sources the rest, implemented by both arms so the record build never branches on residency. The store arm reads wp:{id} alone: both halves live under that key, so one pipelined HMGET per id needs no run-scoped key and cannot span two cluster slots. An id with no record, or a record with no completion, is omitted rather than defaulted.
1 parent 755feb2 commit b40eeca

4 files changed

Lines changed: 332 additions & 0 deletions

File tree

‎internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ import type {
1616
CreateWaitpointResult,
1717
RegisterBlocksLocklessParams,
1818
RegisterBlocksParams,
19+
CompletionEnvelopeSource,
20+
ReadCompletionEnvelopesParams,
1921
RunBlockEdge,
2022
WaitpointCoordinator,
2123
} from "./types.js";
@@ -82,6 +84,73 @@ export class LegacyPostgresWaitpointCoordinator implements WaitpointCoordinator
8284
);
8385
}
8486

87+
/**
88+
* Source the envelope fields from the waitpoint rows.
89+
*
90+
* `runId` is unused here and that is correct: a routing store needs it to pick the owning
91+
* database, a single store does not. It stays in the signature so both arms share one
92+
* shape and the caller never branches.
93+
*
94+
* A row whose `outputType` is already a store reference carries `outputRef`, so the
95+
* record build never re-offloads a value that object storage already holds.
96+
*/
97+
async readCompletionEnvelopes({
98+
waitpointIds,
99+
}: ReadCompletionEnvelopesParams): Promise<CompletionEnvelopeSource[]> {
100+
if (waitpointIds.length === 0) {
101+
return [];
102+
}
103+
104+
const rows = await this.runStore.findManyWaitpoints(
105+
{
106+
where: { id: { in: boundedIn(waitpointIds) } },
107+
select: {
108+
id: true,
109+
friendlyId: true,
110+
type: true,
111+
completedAt: true,
112+
output: true,
113+
outputType: true,
114+
outputIsError: true,
115+
completedByTaskRunId: true,
116+
completedByBatchId: true,
117+
completedAfter: true,
118+
idempotencyKey: true,
119+
userProvidedIdempotencyKey: true,
120+
inactiveIdempotencyKey: true,
121+
},
122+
},
123+
this.prisma
124+
);
125+
126+
return rows.map((row) => {
127+
const isRef = row.outputType === "application/store";
128+
129+
return {
130+
id: row.id,
131+
friendlyId: row.friendlyId,
132+
type: row.type,
133+
// A completed waitpoint always has this set. The fallback keeps the shape total
134+
// rather than emitting an invalid Date, and mirrors the same fallback the snapshot
135+
// hydration already applies.
136+
completedAt: row.completedAt ?? new Date(),
137+
outputType: row.outputType,
138+
outputIsError: row.outputIsError,
139+
...(row.output !== null && row.output !== undefined
140+
? isRef
141+
? { outputRef: row.output }
142+
: { output: row.output }
143+
: {}),
144+
...(row.completedByTaskRunId && { completedByTaskRunId: row.completedByTaskRunId }),
145+
...(row.completedByBatchId && { completedByBatchId: row.completedByBatchId }),
146+
...(row.completedAfter && { completedAfter: row.completedAfter }),
147+
...(row.userProvidedIdempotencyKey && !row.inactiveIdempotencyKey && row.idempotencyKey
148+
? { idempotencyKey: row.idempotencyKey }
149+
: {}),
150+
};
151+
});
152+
}
153+
85154
async registerBlocks({
86155
client,
87156
...edge

‎internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1835,3 +1835,140 @@ describe("genuine concurrency", () => {
18351835
}
18361836
);
18371837
});
1838+
1839+
describe("readCompletionEnvelopes", () => {
1840+
redisTest("returns the completion and the immutable half together", async ({ redisOptions }) => {
1841+
const store = coordinator(redisOptions);
1842+
try {
1843+
await store.createIfAbsent({
1844+
record: record("w_env", {
1845+
idempotencyKey: "user-key",
1846+
userProvidedIdempotencyKey: true,
1847+
}),
1848+
status: "PENDING",
1849+
});
1850+
await store.complete({ waitpointId: "w_env", completion: completion() });
1851+
1852+
const envelopes = await store.readCompletionEnvelopes({
1853+
runId: "run_env",
1854+
waitpointIds: ["w_env"],
1855+
});
1856+
1857+
expect(envelopes).toEqual([
1858+
{
1859+
id: "w_env",
1860+
friendlyId: "waitpoint_w_env",
1861+
type: "MANUAL",
1862+
completedAt: new Date(NOW),
1863+
outputType: "application/json",
1864+
outputIsError: false,
1865+
output: '{"ok":true}',
1866+
idempotencyKey: "user-key",
1867+
},
1868+
]);
1869+
} finally {
1870+
await store.quit();
1871+
}
1872+
});
1873+
1874+
redisTest("carries an offloaded value as a ref, not as an inline value", async ({
1875+
redisOptions,
1876+
}) => {
1877+
const store = coordinator(redisOptions);
1878+
try {
1879+
await store.createIfAbsent({ record: record("w_ref"), status: "PENDING" });
1880+
await store.complete({
1881+
waitpointId: "w_ref",
1882+
completion: completion({
1883+
outputType: "application/store",
1884+
output: { ref: "store-key-1" },
1885+
}),
1886+
});
1887+
1888+
const [envelope] = await store.readCompletionEnvelopes({
1889+
runId: "run_env",
1890+
waitpointIds: ["w_ref"],
1891+
});
1892+
1893+
expect(envelope?.outputRef).toBe("store-key-1");
1894+
expect(envelope?.output).toBeUndefined();
1895+
} finally {
1896+
await store.quit();
1897+
}
1898+
});
1899+
1900+
// The omission is the contract. A pending waitpoint has no envelope, and defaulting one
1901+
// here would hand the resolver a record it must not have. The caller's coverage check is
1902+
// what turns the gap into a loud failure.
1903+
redisTest("omits a waitpoint that is not completed", async ({ redisOptions }) => {
1904+
const store = coordinator(redisOptions);
1905+
try {
1906+
await store.createIfAbsent({ record: record("w_pending"), status: "PENDING" });
1907+
1908+
const envelopes = await store.readCompletionEnvelopes({
1909+
runId: "run_env",
1910+
waitpointIds: ["w_pending"],
1911+
});
1912+
1913+
expect(envelopes).toEqual([]);
1914+
} finally {
1915+
await store.quit();
1916+
}
1917+
});
1918+
1919+
redisTest("omits an id that has no record at all", async ({ redisOptions }) => {
1920+
const store = coordinator(redisOptions);
1921+
try {
1922+
const envelopes = await store.readCompletionEnvelopes({
1923+
runId: "run_env",
1924+
waitpointIds: ["w_absent"],
1925+
});
1926+
1927+
expect(envelopes).toEqual([]);
1928+
} finally {
1929+
await store.quit();
1930+
}
1931+
});
1932+
1933+
redisTest("suppresses an idempotency key the user did not provide", async ({ redisOptions }) => {
1934+
const store = coordinator(redisOptions);
1935+
try {
1936+
await store.createIfAbsent({
1937+
record: record("w_internal", {
1938+
idempotencyKey: "internal-key",
1939+
userProvidedIdempotencyKey: false,
1940+
}),
1941+
status: "PENDING",
1942+
});
1943+
await store.complete({ waitpointId: "w_internal", completion: completion() });
1944+
1945+
const [envelope] = await store.readCompletionEnvelopes({
1946+
runId: "run_env",
1947+
waitpointIds: ["w_internal"],
1948+
});
1949+
1950+
expect(envelope?.idempotencyKey).toBeUndefined();
1951+
} finally {
1952+
await store.quit();
1953+
}
1954+
});
1955+
1956+
redisTest("reads many ids in one pass", async ({ redisOptions }) => {
1957+
const store = coordinator(redisOptions);
1958+
try {
1959+
for (const id of ["w_m1", "w_m2", "w_m3"]) {
1960+
await store.createIfAbsent({ record: record(id), status: "PENDING" });
1961+
await store.complete({ waitpointId: id, completion: completion() });
1962+
}
1963+
1964+
const envelopes = await store.readCompletionEnvelopes({
1965+
runId: "run_env",
1966+
waitpointIds: ["w_m1", "w_m2", "w_m3"],
1967+
});
1968+
1969+
expect(envelopes.map((e) => e.id).sort()).toEqual(["w_m1", "w_m2", "w_m3"]);
1970+
} finally {
1971+
await store.quit();
1972+
}
1973+
});
1974+
});

‎internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import {
1010
watcherField,
1111
} from "./keys.js";
1212
import { registerWaitpointCommands } from "./scripts.js";
13+
import type { CompletionEnvelopeSource, ReadCompletionEnvelopesParams } from "./types.js";
1314

1415
/** The values written into a record's `status` field. Uppercase, and never a token. */
1516
export type WaitpointStatus = "PENDING" | "COMPLETED";
@@ -505,6 +506,62 @@ export class WaitpointStoreCoordinator {
505506
return { pendingIds, deliveredIds, edges };
506507
}
507508

509+
/**
510+
* Source the envelope fields for a run's COMPLETED waitpoints.
511+
*
512+
* Reads `wp:{id}` and nothing else. Both halves live under that one key — `r` holds the
513+
* immutable record, `c` holds the completion — so this needs no run-scoped key and no
514+
* script. One HMGET per id, pipelined: each command touches a single key, so nothing can
515+
* span two cluster slots and there is no #call guard to route through.
516+
*
517+
* The run's delivered hash carries the same envelope, but `wp:{id}` is the record of
518+
* origin, and reading it keeps this independent of whether the run's edges were already
519+
* reconciled.
520+
*
521+
* An id with no record, or a record with no completion, is OMITTED rather than defaulted.
522+
* The omission is the contract: the caller's coverage check turns a gap into a loud
523+
* failure, which a defaulted envelope would hide.
524+
*/
525+
async readCompletionEnvelopes({
526+
waitpointIds,
527+
}: ReadCompletionEnvelopesParams): Promise<CompletionEnvelopeSource[]> {
528+
if (waitpointIds.length === 0) {
529+
return [];
530+
}
531+
532+
const pipeline = this.redis.pipeline();
533+
for (const id of waitpointIds) {
534+
pipeline.hmget(waitpointKeys(id).record, "r", "c");
535+
}
536+
const replies = await pipeline.exec();
537+
538+
const out: CompletionEnvelopeSource[] = [];
539+
540+
for (let i = 0; i < waitpointIds.length; i++) {
541+
const id = waitpointIds[i]!;
542+
const reply = replies?.[i];
543+
544+
// A pipelined command reports its own error in slot 0. Surface it rather than reading
545+
// slot 1, because an errored command's value is not a result.
546+
const error = reply?.[0];
547+
if (error) {
548+
throw error;
549+
}
550+
551+
const fields = reply?.[1] as (string | null)[] | undefined;
552+
const record = parseJson<WaitpointRecordInput>(fields?.[0] ?? undefined);
553+
const completion = parseJson<WaitpointCompletion>(fields?.[1] ?? undefined);
554+
555+
if (!record || !completion) {
556+
continue;
557+
}
558+
559+
out.push(toEnvelopeSource(id, record, completion));
560+
}
561+
562+
return out;
563+
}
564+
508565
/**
509566
* Drain one cycle's edges, or clear the run entirely when no edge ids are given.
510567
*
@@ -536,3 +593,37 @@ export class WaitpointStoreCoordinator {
536593
return { outcome: reply[0] as "cleared" | "drained" };
537594
}
538595
}
596+
597+
/**
598+
* Map the store's two halves onto the arm-independent source shape.
599+
*
600+
* The idempotency key is suppressed unless the user provided it, matching the rule the
601+
* snapshot hydration applies today. The store never sets an inactive flag, so
602+
* `userProvidedIdempotencyKey` alone decides it here.
603+
*/
604+
function toEnvelopeSource(
605+
id: string,
606+
record: WaitpointRecordInput,
607+
completion: WaitpointCompletion
608+
): CompletionEnvelopeSource {
609+
const output = completion.output;
610+
const inline = output && "inline" in output ? output.inline : undefined;
611+
const ref = output && "ref" in output ? output.ref : undefined;
612+
613+
return {
614+
id,
615+
friendlyId: record.friendlyId,
616+
type: record.type,
617+
completedAt: new Date(completion.completedAt),
618+
outputType: completion.outputType,
619+
outputIsError: completion.outputIsError,
620+
...(inline !== undefined && { output: inline }),
621+
...(ref !== undefined && { outputRef: ref }),
622+
...(record.completedByTaskRunId && { completedByTaskRunId: record.completedByTaskRunId }),
623+
...(record.completedByBatchId && { completedByBatchId: record.completedByBatchId }),
624+
...(record.completedAfter && { completedAfter: new Date(record.completedAfter) }),
625+
...(record.userProvidedIdempotencyKey && record.idempotencyKey
626+
? { idempotencyKey: record.idempotencyKey }
627+
: {}),
628+
};
629+
}

‎internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@ import type { ShardKey } from "@trigger.dev/core/v3/isomorphic";
1717
export type WaitpointCoordinator = {
1818
clearRunBlockState(params: ClearRunBlockStateParams): Promise<{ count: number }>;
1919
readRunBlockState(runId: string): Promise<RunBlockEdge[]>;
20+
readCompletionEnvelopes(
21+
params: ReadCompletionEnvelopesParams
22+
): Promise<CompletionEnvelopeSource[]>;
2023
registerBlocks(params: RegisterBlocksParams): Promise<{ pendingCount: number }>;
2124
registerBlocksLockless(params: RegisterBlocksLocklessParams): Promise<void>;
2225
complete(params: CompleteParams): Promise<CompleteResult>;
@@ -34,6 +37,38 @@ export type WaitpointCoordinator = {
3437
}): Promise<Waitpoint>;
3538
};
3639

40+
export type ReadCompletionEnvelopesParams = {
41+
runId: string;
42+
/** The DISTINCT completed waitpoint ids to source. Result order is not meaningful. */
43+
waitpointIds: string[];
44+
};
45+
46+
/**
47+
* One completed waitpoint's fields, sourced from whichever arm owns it.
48+
*
49+
* Deliberately NOT the frozen record type. This is the raw material; the record build
50+
* decides which output variant a record carries. Both arms return this same shape, so the
51+
* record build never branches on residency, which is what makes a mixed wait work.
52+
*
53+
* `output` is the literal stored value. `outputRef` is set instead when the value was
54+
* already offloaded to object storage. At most one of the two is set.
55+
*/
56+
export type CompletionEnvelopeSource = {
57+
id: string;
58+
friendlyId: string;
59+
type: "RUN" | "BATCH" | "DATETIME" | "MANUAL";
60+
completedAt: Date;
61+
outputType: string;
62+
outputIsError: boolean;
63+
output?: string;
64+
outputRef?: string;
65+
completedByTaskRunId?: string;
66+
completedByBatchId?: string;
67+
completedAfter?: Date;
68+
/** Already resolved by the arm: userProvidedIdempotencyKey && !inactiveIdempotencyKey. */
69+
idempotencyKey?: string;
70+
};
71+
3772
export type ClearRunBlockStateParams = {
3873
runId: string;
3974
/** Edge ids to delete. Omit to clear every edge for the run. */

0 commit comments

Comments
 (0)