diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts index 0efc5f8672a35..d31b5f2fe8dfe 100644 --- a/ts-sdk/src/sdk/dag.ts +++ b/ts-sdk/src/sdk/dag.ts @@ -667,7 +667,7 @@ export class Dag { const taskId = condition.taskId; const branches: ConditionRecord = {}; this.#conditions.set(taskId, branches); - this.#wrapDecider(taskId, async (held: unknown) => { + this.#wrapDecider("dag.if", taskId, async (held: unknown) => { if (typeof held !== "boolean") { throw new Error( `Condition "${taskId}" of Dag "${this.dagId}" returned ${describeValue(held)} ` + @@ -730,12 +730,20 @@ export class Dag { } #wrapDecider( + verb: "dag.if" | "dag.switch", taskId: string, decide: (returned: unknown) => Promise<{ skip: string[]; result: unknown }>, ): void { - const { task, spec, fn } = this.#tasks.get(taskId)!; - // A decider is always declared from a handler, never from an operator. - const inner = fn!; + const record = this.#tasks.get(taskId)!; + // Only an operator task has no handler. + if (record.fn === undefined) { + const { label, operatorName } = record.operator; + throw new Error( + `A ${label ?? operatorName} task of Dag "${this.dagId}" was given to ${verb}, ` + + "which needs a handler to decide the branch with", + ); + } + const { task, spec, fn: inner } = record; const wrapped: TaskFunction = async (args) => { const { skip, result } = await decide(await inner(args as never)); if (skip.length > 0) await skipTasks(getClient() as CoordinatorClient, skip); @@ -775,7 +783,7 @@ export class Dag { const taskId = decider.taskId; const candidates: TaskRef[] = []; this.#branches.set(taskId, candidates); - this.#wrapDecider(taskId, async (chosen: unknown) => { + this.#wrapDecider("dag.switch", taskId, async (chosen: unknown) => { const known = candidates.map((ref) => ref.taskId); const picked = isTaskRef(chosen) && chosen.dagId === this.dagId ? chosen.taskId : undefined; if (picked === undefined || !known.includes(picked)) { diff --git a/ts-sdk/tests/hitl/operator.test.ts b/ts-sdk/tests/hitl/operator.test.ts index f6afb9761d709..7dbb8809aa6c3 100644 --- a/ts-sdk/tests/hitl/operator.test.ts +++ b/ts-sdk/tests/hitl/operator.test.ts @@ -58,6 +58,17 @@ describe("a HITL task in a Dag", () => { expect(dag.taskIds).toEqual(["release.sign_off"]); }); + it.each(["if", "switch"] as const)("is rejected by dag.%s, which needs a handler", (decider) => { + const dag = new Dag("d"); + const declare = dag[decider].bind(dag) as (...args: unknown[]) => unknown; + + expect(() => + declare(hitl({ subject: "s", options: ["a"] }), undefined, { taskId: "choose" }), + ).toThrowError( + `A human-in-the-loop task of Dag "d" was given to dag.${decider}, which needs a handler`, + ); + }); + it("serializes as the operator it mirrors, with its argument bindings", () => { const dag = new Dag("d"); const build = dag.task("build", async () => 1); diff --git a/ts-sdk/tests/sdk/trigger-dag-run.test.ts b/ts-sdk/tests/sdk/trigger-dag-run.test.ts index 7fc3c42d5c9f6..ae3c3712cf165 100644 --- a/ts-sdk/tests/sdk/trigger-dag-run.test.ts +++ b/ts-sdk/tests/sdk/trigger-dag-run.test.ts @@ -237,6 +237,19 @@ describe("triggerDagRun", () => { ); }); + it.each(["if", "switch"] as const)("a trigger task given to dag.%s", (decider) => { + const dag = new Dag("d"); + const declare = dag[decider].bind(dag) as (...args: unknown[]) => unknown; + + expect(() => + declare(triggerDagRun({ dagId: "downstream_etl" }), undefined, { + taskId: "trigger_downstream", + }), + ).toThrowError( + `A triggerDagRun task of Dag "d" was given to dag.${decider}, which needs a handler`, + ); + }); + it("a trigger task left uncalled, when the Dag is read", () => { const dag = new Dag("d"); dag.task(triggerDagRun({ dagId: "downstream_etl" }), { taskId: "trigger_downstream" });