shivaam commented on code in PR #74484:
URL: https://github.com/apache/airflow/pull/74484#discussion_r4235811646
##########
ts-sdk/tests/coordinator/integration.test.ts:
##########
@@ -1305,8 +1307,757 @@ describe("coordinator runtime integration", () => {
retry_reason: "Trigger timeout",
});
expect(
- result.logRecords.some((r) => String(r["event"]).includes("Trigger
failed:\nTraceback")),
+ result.logRecords.some((r) =>
+ String(r["event"]).includes("Task could not be resumed:\nTraceback"),
+ ),
+ ).toBe(true);
+ });
+ });
+ describe("human-in-the-loop", () => {
+ const CREATED = { body: { type: "HITLDetailRequestResult", ti_id: "ti-1" }
};
+ // What `handle_event_submit` stores: what the task parked with, plus the
response as `event`,
+ // with `responded_at` in serde's datetime encoding.
+ const RESPONDED_AT = {
+ __classname__: "pendulum.datetime.DateTime",
+ __version__: 2,
+ __data__: {
+ timestamp: 1791462615.123456,
+ tz: {
+ __classname__: "builtins.tuple",
+ __version__: 1,
+ __data__: ["UTC", "pendulum.tz.timezone.Timezone", 1, true],
+ },
+ },
+ };
+
+ function parkedWith(taskId: string, tiContext: Record<string, unknown> =
{}) {
+ return driveSupervisor(
+ makeStartupDetails(taskId, "test_dag", "r1", tiContext),
+ (msgType, body) => {
+ if (msgType === "CreateHITLDetailPayload") return CREATED;
+ if (msgType === "GetXCom") {
+ return { body: { type: "XComResult", key: body["key"], value: {
version: "1.4" } } };
+ }
+ return null;
+ },
+ );
+ }
+
+ /** The task parks with empty `next_kwargs`, so a resume carries only the
response. */
+ function resumedWith(taskId: string, event: unknown, tiContext:
Record<string, unknown> = {}) {
+ return driveSupervisor(
+ makeStartupDetails(taskId, "test_dag", "r1", {
+ next_method: "execute_complete",
+ next_kwargs: { event },
+ ...tiContext,
+ }),
+ );
+ }
+
+ function answer(
+ chosen: string[],
+ user: unknown = { id: "1", name: "Ada" },
+ paramsInput: Record<string, unknown> = {},
+ ) {
+ return {
+ chosen_options: chosen,
+ params_input: paramsInput,
+ responded_at: RESPONDED_AT,
+ responded_by_user: user,
+ timedout: user === null,
+ };
+ }
+
+ function requestTypes(result: MockResult) {
+ return result.runtimeRequests.map((r) => r.type);
+ }
+
+ it("writes the request with rendered text and parks in awaiting_input",
async () => {
+ const build = testDag.task("build", async () => ({ version: "1.4" }));
+ testDag.task(
+ "sign_off",
+ approval({
+ subject: ({ report }: { report: { version: string } }) => `Ship
${report.version}?`,
+ body: async ({ report }: { report: { version: string } }) =>
`**${report.version}**`,
+ assignedUsers: [{ id: "ada", name: "Ada" }],
+ responseTimeout: 3600,
+ }),
+ )({ report: build() });
+
+ const result = await parkedWith("sign_off", {
+ arg_bindings: [{ name: "report", kind: "xcom", task_id: "build" }],
+ });
+
+ const [created] = requestsOf(result, "CreateHITLDetailPayload");
+ expect(created).toEqual({
+ type: "CreateHITLDetailPayload",
+ ti_id: "ti-1",
+ options: ["Approve", "Reject"],
+ subject: "Ship 1.4?",
+ body: "**1.4**",
+ defaults: null,
+ multiple: false,
+ params: {},
+ assigned_users: [{ id: "ada", name: "Ada" }],
+ });
+ expect(result.firstResponse!.body).toEqual({
+ type: "AwaitInputTask",
+ state: "awaiting_input",
+ timeout: "PT3600S",
+ next_method: "execute_complete",
+ next_kwargs: {},
+ });
+ expect(requestsOf(result, "SetXCom")).toEqual([]);
+ });
+
+ it("sends its defaults and multiple, with no body and no timeout as null",
async () => {
+ testDag.task(
+ "choose",
+ hitl({
+ subject: "Pick",
+ options: ["us", "eu"],
+ defaults: ["us", "eu"],
+ multiple: true,
+ }),
+ )();
+
+ const result = await parkedWith("choose");
+
+ expect(requestsOf(result, "CreateHITLDetailPayload")[0]).toMatchObject({
+ body: null,
+ defaults: ["us", "eu"],
+ multiple: true,
+ });
+ expect(result.firstResponse!.body).toMatchObject({ type:
"AwaitInputTask", timeout: null });
+ });
+
+ it("writes its params as Param.serialize does, with the source always
task", async () => {
+ testDag.task(
+ "choose",
+ hitl({
+ subject: "Pick",
+ options: ["go"],
+ params: {
+ region: {
+ value: "us",
+ description: "Where to ship",
+ schema: { type: "string", enum: ["us", "eu"] },
+ },
+ retries: { value: 3 },
+ },
+ }),
+ )();
+
+ const result = await parkedWith("choose");
+
+ expect(requestsOf(result, "CreateHITLDetailPayload")[0]).toMatchObject({
+ params: {
+ region: {
+ value: "us",
+ description: "Where to ship",
+ schema: { type: "string", enum: ["us", "eu"] },
+ source: "task",
+ },
+ retries: { value: 3, description: null, schema: {}, source: "task" },
+ },
+ });
+ });
+
+ it("renders the subject and the body with the renames of both", async ()
=> {
+ testDag.task(
+ "choose",
+ hitl({
+ subject: withArgNames(
+ { version: "release_version" },
+ ({ version }: { version: string }) => `Ship ${version}?`,
+ ),
+ body: withArgNames(
+ { notes: "release_notes" },
+ ({ notes }: { notes: string }) => `Notes: ${notes}`,
+ ),
+ options: ["go"],
+ }),
+ )({ version: "1.4", notes: "Faster" });
+
+ const result = await parkedWith("choose", {
+ arg_bindings: [
+ { name: "release_version", kind: "literal", value: "1.4" },
+ { name: "release_notes", kind: "literal", value: "Faster" },
+ ],
+ });
+
+ expect(requestsOf(result, "CreateHITLDetailPayload")[0]).toMatchObject({
+ subject: "Ship 1.4?",
+ body: "Notes: Faster",
+ });
+ });
+
+ it.each<[string, () => unknown, string]>([
+ [
+ "throws",
+ () => {
+ throw new Error("no report");
+ },
+ "no report",
+ ],
+ [
+ "returns an empty string",
+ () => "",
+ 'The subject of task "sign_off" has to be a non-empty string',
+ ],
+ ["returns a number", () => 7, 'The subject of task "sign_off" has to be
a non-empty string'],
+ ])(
+ "fails without writing the request when its subject function %s",
+ async (_label, subject, reason) => {
+ testDag.task("sign_off", approval({ subject: subject as () => string
}))();
+
+ const result = await parkedWith("sign_off", { should_retry: true });
+
+ expect(requestsOf(result, "CreateHITLDetailPayload")).toEqual([]);
+ expect(result.firstResponse!.body).toMatchObject({
+ type: "RetryTask",
+ retry_reason: reason,
+ });
+ },
+ );
+
+ it("resumes with the response without writing the request again", async ()
=> {
+ let rendered = 0;
+ testDag.task(
+ "choose",
+ hitl({
+ subject: () => `Pick ${++rendered}`,
+ options: ["us", "eu", "apac"],
+ multiple: true,
+ }),
+ )();
+ const parked = await parkedWith("choose");
+ const { next_kwargs } = parked.firstResponse!.body as {
+ next_kwargs: Record<string, unknown>;
+ };
+
+ const result = await driveSupervisor(
+ makeStartupDetails("choose", "test_dag", "r1", {
+ next_method: "execute_complete",
+ next_kwargs: { ...next_kwargs, event: answer(["eu", "apac"]) },
+ }),
+ );
+
+ expect(rendered).toBe(1);
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask"
});
+ expect(requestTypes(result)).toEqual(["SetXCom"]);
+ expect(requestsOf(result, "SetXCom")[0]).toMatchObject({
+ key: "return_value",
+ task_id: "choose",
+ value: {
+ chosen_options: ["eu", "apac"],
+ params_input: {},
+ responded_at: "2026-10-08T12:30:15.123Z",
+ responded_by_user: { id: "1", name: "Ada" },
+ timedout: false,
+ },
+ });
+ });
+
+ it.each([
+ ["an approval answered Approve", approval({ subject: "Go?" }),
"Approve"],
+ [
+ "a generic choice answered Reject",
+ hitl({ subject: "?", options: ["Approve", "Reject"] }),
+ "Reject",
+ ],
+ ])("skips nothing for %s and pushes the response", async (_label,
operator, chosen) => {
+ testDag
+ .task("decide", operator)()
+ .before(testDag.task("publish", async () => {})());
+
+ const result = await resumedWith("decide", answer([chosen]));
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask"
});
+ expect(requestTypes(result)).toEqual(["SetXCom"]);
+ expect(requestsOf(result, "SetXCom")[0]).toMatchObject({ key:
"return_value" });
+ });
+
+ it("on Reject skips the tasks directly downstream and, as SkipMixin does,
pushes no response", async () => {
+ const decide = testDag.task("decide", approval({ subject: "Go?" }))();
+ const publish = testDag.task("publish", async () => {})();
+ decide.before(publish);
+ publish.before(testDag.task("announce", async () => {})());
+
+ const result = await resumedWith("decide", answer(["Reject"]));
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask"
});
+ expect(requestTypes(result)).toEqual(["SetXCom", "SkipDownstreamTasks"]);
+ expect(requestsOf(result, "SetXCom")[0]).toMatchObject({
+ key: "skipmixin_key",
+ value: { skipped: ["publish"] },
+ });
+ expect(requestsOf(result, "SkipDownstreamTasks")[0]).toMatchObject({
tasks: ["publish"] });
+ });
+
+ it("on Reject skips a task group directly downstream by its first tasks",
async () => {
+ const decide = testDag.task("decide", approval({ subject: "Go?" }))();
+ const stage = testDag.taskGroup("stage");
+ const publish = stage.task("publish", async () => {})();
+ publish.before(stage.task("announce", async () => {})());
+ decide.before(stage);
+
+ const result = await resumedWith("decide", answer(["Reject"]));
+
+ expect(requestsOf(result, "SkipDownstreamTasks")[0]).toMatchObject({
+ tasks: ["stage.publish"],
+ });
+ });
+
+ it("on Reject with ignoreDownstreamTriggerRules skips every task
downstream", async () => {
+ const decide = testDag.task(
+ "decide",
+ approval({ subject: "Go?", ignoreDownstreamTriggerRules: true }),
+ )();
+ const publish = testDag.task("publish", async () => {})();
+ decide.before(publish);
+ publish.before(testDag.task("announce", async () => {})());
+
+ const result = await resumedWith("decide", answer(["Reject"]));
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask"
});
+ expect(requestTypes(result)).toEqual(["SetXCom", "SkipDownstreamTasks"]);
+ expect(requestsOf(result, "SkipDownstreamTasks")[0]).toMatchObject({
+ tasks: ["announce", "publish"],
+ });
+ });
+
+ it("on Reject with nothing downstream pushes the response", async () => {
+ testDag.task("decide", approval({ subject: "Go?" }))();
+
+ const result = await resumedWith("decide", answer(["Reject"]));
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask"
});
+ expect(requestsOf(result, "SetXCom")).toEqual([
+ expect.objectContaining({
+ key: "return_value",
+ value: expect.objectContaining({ chosen_options: ["Reject"] }),
+ }),
+ ]);
+ });
+
+ it("on Reject fails with failOnReject, before pushing anything", async ()
=> {
+ testDag.task("decide", approval({ subject: "Go?", failOnReject: true
}))();
+
+ const result = await resumedWith("decide", answer(["Reject"]));
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "TaskState",
state: "failed" });
+ expect(result.runtimeRequests).toEqual([]);
+ });
+
+ it("succeeds with the defaults applied on timeout", async () => {
+ testDag.task(
+ "decide",
+ approval({ subject: "Go?", defaults: "Approve", responseTimeout: 60 }),
+ )();
+
+ const result = await resumedWith("decide", answer(["Approve"], null));
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask"
});
+ expect(requestsOf(result, "SetXCom")[0]).toMatchObject({
+ value: { responded_by_user: null, timedout: true },
+ });
+ });
+
+ it.each([
+ ["retries", true, { type: "RetryTask", retry_reason: "Response timed
out: no response" }],
+ ["fails", false, { type: "TaskState", state: "failed" }],
+ ])("%s when the timeout passes without defaults", async (_label,
shouldRetry, outcome) => {
+ testDag.task("decide", approval({ subject: "Go?", responseTimeout: 60
}))();
+
+ const result = await resumedWith(
+ "decide",
+ { error: "no response", error_type: "timeout" },
+ { should_retry: shouldRetry },
+ );
+
+ expect(result.firstResponse!.body).toMatchObject(outcome);
+ });
+
+ it.each<[string, Record<string, unknown>]>([
+ ["an option it never offered", { event: answer(["Maybe"]) }],
+ ["an unreadable response time", { event: { ...answer(["Approve"]),
responded_at: "soon" } }],
+ ["an unreadable responder", { event: answer(["Approve"], { id: 1 }) }],
+ ["no event", {}],
+ ])("fails on resume with %s", async (_label, nextKwargs) => {
+ testDag.task("decide", approval({ subject: "Go?" }))();
+
+ const result = await driveSupervisor(
+ makeStartupDetails("decide", "test_dag", "r1", {
+ next_method: "execute_complete",
+ next_kwargs: nextKwargs,
+ }),
+ );
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "TaskState",
state: "failed" });
+ expect(result.runtimeRequests).toEqual([]);
+ });
+
+ describe("params input", () => {
+ const withParams = () =>
+ hitl({
+ subject: "Pick",
+ options: ["go"],
+ params: { region: { value: "us" }, retries: { value: 3 } },
+ });
+
+ it("is pushed with the response when it answers every param", async ()
=> {
+ testDag.task("choose", withParams())();
+
+ const result = await resumedWith(
+ "choose",
+ answer(["go"], undefined, { region: "eu", retries: 5 }),
+ );
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask"
});
+ expect(requestsOf(result, "SetXCom")[0]).toMatchObject({
+ key: "return_value",
+ value: { chosen_options: ["go"], params_input: { region: "eu",
retries: 5 } },
+ });
+ });
+
+ it.each<[string, Record<string, unknown>]>([
+ ["only the timeout defaults", { region: "us", retries: 3 }],
+ ["no answers when the form was not filled in", {}],
+ ])("is accepted with %s", async (_label, paramsInput) => {
+ testDag.task("choose", withParams())();
+
+ const result = await resumedWith("choose", answer(["go"], undefined,
paramsInput));
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask"
});
+ });
+
+ it("is not checked for a task that declares no params", async () => {
+ testDag.task("choose", hitl({ subject: "Pick", options: ["go"] }))();
+
+ const result = await resumedWith("choose", answer(["go"], undefined, {
region: "eu" }));
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask"
});
+ expect(requestsOf(result, "SetXCom")[0]).toMatchObject({
+ value: { params_input: { region: "eu" } },
+ });
+ });
+
+ it.each<[string, Record<string, unknown>, string]>([
+ ["a param the task never declared", { region: "eu", other: 1 },
"region,other"],
+ ["only some of the params", { region: "eu" }, "region"],
+ ["other names than the params", { zone: "eu", retries: 5 },
"zone,retries"],
+ ])("fails the task with %s", async (_label, paramsInput, received) => {
+ testDag.task("choose", withParams())();
+
+ const result = await resumedWith("choose", answer(["go"], undefined,
paramsInput), {
+ should_retry: true,
+ });
+
+ expect(result.firstResponse!.body).toMatchObject({
+ type: "RetryTask",
+ retry_reason:
+ `params_input ${JSON.stringify(received.split(","))} does not
match params ` +
+ '["region","retries"]',
+ });
+ expect(result.runtimeRequests).toEqual([]);
+ });
+ });
+
+ it("fails on resume with a next_method it does not know", async () => {
+ testDag.task("decide", approval({ subject: "Go?" }))();
+
+ const result = await driveSupervisor(
+ makeStartupDetails("decide", "test_dag", "r1", {
+ next_method: "something_else",
+ next_kwargs: { event: answer(["Approve"]) },
+ }),
+ );
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "TaskState",
state: "failed" });
+ expect(result.runtimeRequests).toEqual([]);
+ });
+
+ it("fails on resume through __fail__", async () => {
+ testDag.task("decide", approval({ subject: "Go?" }))();
+
+ const result = await driveSupervisor(
+ makeStartupDetails("decide", "test_dag", "r1", {
+ next_method: "__fail__",
+ next_kwargs: { error: "Could not resume the task", traceback:
["Traceback", "Boom"] },
+ }),
+ );
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "TaskState",
state: "failed" });
+ expect(
+ result.logRecords.some((r) =>
+ String(r["event"]).includes("Task could not be
resumed:\nTraceback\nBoom"),
+ ),
).toBe(true);
+ expect(result.logRecords.some((r) =>
String(r["event"]).includes("Trigger failed"))).toBe(
+ false,
+ );
+ });
+
+ // The agent-approval Dag of `example/src/ai-approval.ts`, run one task
per process.
+ describe("agent approval", () => {
+ const DRAFTED_AMOUNT = 120;
+ const EDITED_AMOUNT = 80;
+
+ interface ToolCall {
+ type: "tool-call";
+ toolCallId: string;
+ toolName: string;
+ input: { orderId: string; amount: number };
+ }
+ interface ApprovalRequest {
+ type: "tool-approval-request";
+ approvalId: string;
+ toolCall: ToolCall;
+ }
+ interface ApprovalResponse {
+ type: "tool-approval-response";
+ approvalId: string;
+ approved: boolean;
+ reason?: string;
+ }
+ type ModelMessage =
+ | { role: "user"; content: string }
+ | { role: "assistant"; content: ApprovalRequest[] }
+ | { role: "tool"; content: ApprovalResponse[] };
+ interface AgentStep {
+ messages: ModelMessage[];
+ pending: { id: string; toolName: string; input: ToolCall["input"] }[];
+ }
+
+ let modelCalls: number;
+
+ async function generateText(opts: { messages: ModelMessage[] }) {
+ modelCalls += 1;
+ const response = opts.messages
+ .flatMap((message) => (message.role === "tool" ? message.content :
[]))
+ .at(-1);
+ if (response === undefined) {
+ const request: ApprovalRequest = {
+ type: "tool-approval-request",
+ approvalId: "approval-1",
+ toolCall: {
+ type: "tool-call",
+ toolCallId: "call-1",
+ toolName: "issue_refund",
+ input: { orderId: "ord_1042", amount: DRAFTED_AMOUNT },
+ },
+ };
+ const messages: ModelMessage[] = [{ role: "assistant", content:
[request] }];
+ return { text: "", content: [request], response: { messages } };
+ }
+ const { toolCall } = opts.messages
+ .flatMap((message) => (message.role === "assistant" ?
message.content : []))
+ .find((request) => request.approvalId === response.approvalId)!;
+ const text = `Refunded $${toolCall.input.amount} on order
${toolCall.input.orderId}.`;
+ return { text, content: [], response: { messages: [] } };
+ }
+
+ beforeEach(() => {
+ modelCalls = 0;
+ const agentStep = testDag.task("agent_step", async ():
Promise<AgentStep> => {
+ const messages: ModelMessage[] = [{ role: "user", content: "Refund
order ord_1042." }];
+ const result = await generateText({ messages });
+ return {
+ messages: [...messages, ...result.response.messages],
+ pending: result.content.map((request) => ({
+ id: request.approvalId,
+ toolName: request.toolCall.toolName,
+ input: request.toolCall.input,
+ })),
+ };
+ })();
+ const review = testDag.task(
+ "review",
+ hitl({
+ subject: ({ draft }: { draft: AgentStep }) => `Approve
${draft.pending[0]!.toolName}?`,
+ body: ({ draft }: { draft: AgentStep }) =>
+ [
+ `The agent wants to call \`${draft.pending[0]!.toolName}\`
with:`,
+ "",
+ "```json",
+ JSON.stringify(draft.pending[0]!.input, null, 2),
+ "```",
+ ].join("\n"),
+ options: ["Approve", "Reject"],
+ params: {
+ amount: {
+ value: DRAFTED_AMOUNT,
+ description: "Amount to refund",
+ schema: { type: "number" },
+ },
+ },
+ }),
+ )({ draft: agentStep });
+ testDag.task(
+ "agent_continue",
+ async ({ draft, decision }: { draft: AgentStep; decision: HITLResult
}) => {
+ const edited = decision.params_input["amount"];
+ const amount = typeof edited === "number" ? edited :
draft.pending[0]!.input.amount;
+ const approvalId = draft.pending[0]!.id;
+ const messages: ModelMessage[] = [
+ ...draft.messages.map((message): ModelMessage => {
+ if (message.role !== "assistant") return message;
+ return {
+ role: "assistant",
+ content: message.content.map((request) =>
+ request.approvalId === approvalId
+ ? {
+ ...request,
+ toolCall: {
+ ...request.toolCall,
+ input: { ...request.toolCall.input, amount },
+ },
+ }
+ : request,
+ ),
+ };
+ }),
+ {
+ role: "tool",
+ content: [
+ {
+ type: "tool-approval-response",
+ approvalId,
+ approved: decision.chosen_options[0] === "Approve",
+ reason: `Approved at $${amount}`,
+ },
+ ],
+ },
+ ];
+ const result = await generateText({ messages });
+ return { text: result.text, modelCalls: 1, approvedAmount: amount
};
Review Comment:
Replaced it with a small purpose-built Dag in c3ecdedf59. It checks that a
resume with arg_bindings pulls no inputs and renders nothing again; the e2e
test still runs the real example.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
