TS SDK: delete a task instance's existing XComs when a try starts

ts-sdk/src/coordinator/client.ts

@@ -29,6 +29,7 @@ import type {
   DeleteVariable,
   GetXCom,
   SetXCom,
+  DeleteXCom,
   GetConnection,
   SkipDownstreamTasks,
   TriggerDagRun,
@@ -74,7 +75,7 @@ export interface XComEntry {
 const XCOM_ABSENT: XComEntry = { found: false, value: null };
 
 /**
- * A task's {@link TaskClient} plus the reads only the runtime itself makes.
+ * A task's {@link TaskClient} plus the calls only the runtime itself makes.
  * Handlers are typed against `TaskClient`, so nothing added here is public API.
  */
 export interface CoordinatorClient extends TaskClient {
@@ -87,6 +88,16 @@ export interface CoordinatorClient extends TaskClient {
    */
   getXComEntry(opts: GetXComOpts): Promise<XComEntry>;
 
+  /**
+   * Delete the running task instance's XCom under `key`.
+   *
+   * Sends `map_index` as null for an unmapped task instance, as
+   * {@link TaskClient.setXCom} does. The supervisor then leaves `map_index` out
+   * of the Execution API call, and the API deletes the XCom stored under
+   * map_index -1.
+   */
+  deleteXCom(key: string): Promise<void>;
+
   /** Mark direct downstream tasks of the running task as skipped; none is a no-op. */
   skipDownstreamTasks(taskIds: readonly string[]): Promise<void>;
 
@@ -224,6 +235,20 @@ export function createCoordinatorClient(
       await rpc("SetXCom", null, msg, () => undefined, "throw");
     },
 
+    async deleteXCom(key: string): Promise<void> {
+      const msg: DeleteXCom = {
+        type: "DeleteXCom",
+        key,
+        dag_id: ctx.dagId,
+        task_id: ctx.taskId,
+        run_id: ctx.runId,
+        map_index: resolveWireMapIndex(undefined, ctx.mapIndex),
+      };
+      // The supervisor answers DeleteXCom with an empty frame, not with the
+      // OKResponse it sends for DeleteVariable.
+      await rpc("DeleteXCom", null, msg, () => undefined, "throw");
+    },
+
     // ---- Control flow ----
 
     async skipDownstreamTasks(taskIds: readonly string[]): Promise<void> {

ts-sdk/src/coordinator/protocol.ts

@@ -65,6 +65,7 @@ export type {
   DeleteVariable,
   GetXCom,
   SetXCom,
+  DeleteXCom,
   GetConnection,
   SkipDownstreamTasks,
   TriggerDagRun,

ts-sdk/src/coordinator/runtime.ts

@@ -37,7 +37,7 @@
 //        - StartupDetails      → run task, respond Succeed or Fail, exit
 //
 import { resolveArgs, type BoundArgs } from "./arg-binding.js";
-import { createCoordinatorClient } from "./client.js";
+import { createCoordinatorClient, type CoordinatorClient } from "./client.js";
 import { CommChannel } from "./comm-channel.js";
 import { LogChannel } from "./log-channel.js";
 import {
@@ -348,6 +348,8 @@ async function handleTask(
   if (trigger) {
     const ctx = buildContext(details, signal);
     const client = createCoordinatorClient(comm, ctx, clientLogs);
+    const deleteFailure = await deleteListedXComs(details, client, logs);
+    if (deleteFailure) return deleteFailure;
     const fail = (message: string) => {
       logs.error("Task failed", { task_id: ctx.taskId, error: message });
       return buildFailureResponse(details, message);
@@ -378,6 +380,8 @@ async function handleTask(
 
   const ctx = buildContext(details, signal);
   const client = createCoordinatorClient(comm, ctx, clientLogs);
+  const deleteFailure = await deleteListedXComs(details, client, logs);
+  if (deleteFailure) return deleteFailure;
 
   let bound: BoundArgs;
   try {
@@ -433,6 +437,43 @@ async function handleTask(
   }
 }
 
+/**
+ * Delete the XComs that `ti_context.xcom_keys_to_clear` lists.
+ *
+ * When a try starts, Airflow lists the XComs that the try already has and
+ * leaves deleting them to the runtime. The Python task runner deletes the
+ * listed XComs before it renders templates and fails the try if a delete
+ * fails. This runtime calls this function before a triggerDagRun task runs
+ * and before a task handler's arguments are bound. Those calls match the point
+ * where the Python task runner deletes the listed XComs. The list is empty
+ * when a task resumes from a deferral, so this function does not delete the
+ * XComs that a triggerDagRun task wrote before the task deferred.
+ *
+ * Returns the response that fails the try when a delete fails, and null when
+ * every listed XCom is deleted.
+ */
+async function deleteListedXComs(
+  details: StartupDetails,
+  client: CoordinatorClient,
+  logs: LogChannel,
+): Promise<RuntimeRetryTask | RuntimeTaskState | null> {
+  for (const key of details.ti_context?.xcom_keys_to_clear ?? []) {
+    logs.debug("Clearing XCom with key", { key });
+    try {
+      await client.deleteXCom(key);
+    } catch (err) {
+      const message = (err as Error).message ?? String(err);
+      logs.error("Cannot clear XCom before running the task", {
+        task_id: details.ti.task_id,
+        key,
+        error: message,
+      });
+      return buildFailureResponse(details, message);
+    }
+  }
+  return null;
+}
+
 async function sendSupervisorResponse(
   id: number,
   body: unknown,

ts-sdk/tests/coordinator/client.test.ts

@@ -19,7 +19,7 @@
 
 import { describe, it, expect } from "vitest";
 import { ConnectionNotFoundError } from "../../src/sdk/client.js";
-import { createCoordinatorClient } from "../../src/coordinator/client.js";
+import { createCoordinatorClient, type CoordinatorClient } from "../../src/coordinator/client.js";
 import type { CommChannel } from "../../src/coordinator/comm-channel.js";
 import type { TaskClient } from "../../src/sdk/client.js";
 import type { TaskContext } from "../../src/sdk/task.js";
@@ -157,6 +157,7 @@ describe("writes do not read a supervisor 404 as absence", () => {
     ["setVariable", "PutVariable", (c: TaskClient) => c.setVariable("k", "v")],
     ["deleteVariable", "DeleteVariable", (c: TaskClient) => c.deleteVariable("k")],
     ["setXCom", "SetXCom", (c: TaskClient) => c.setXCom({ key: "k", value: 1 })],
+    ["deleteXCom", "DeleteXCom", (c: CoordinatorClient) => c.deleteXCom("k")],
   ])("%s rejects", async (_name, op, call) => {
     const c = client([
       { body: null, error: { error: "API_SERVER_ERROR", detail: { status_code: 404 } } },

ts-sdk/tests/coordinator/integration.test.ts

@@ -82,8 +82,8 @@ type Responder = (
   body: Record<string, unknown>,
 ) => { body: unknown; error?: unknown } | null;
 
-function frameBytes(id: number, body: unknown, isResponse: boolean): Buffer {
-  const arr = isResponse ? [id, body, null] : [id, body];
+function frameBytes(id: number, body: unknown, isResponse: boolean, error: unknown = null): Buffer {
+  const arr = isResponse ? [id, body, error] : [id, body];
   const payload = Buffer.from(encode(arr));
   const header = Buffer.alloc(4);
   header.writeUInt32BE(payload.length, 0);
@@ -112,6 +112,7 @@ function makeStartupDetails(
   dagId = "test_dag",
   runId = "r1",
   tiContext: Record<string, unknown> = {},
+  mapIndex = -1,
 ): unknown {
   return {
     type: "StartupDetails",
@@ -122,7 +123,7 @@ function makeStartupDetails(
       dag_id: dagId,
       run_id: runId,
       try_number: 1,
-      map_index: -1,
+      map_index: mapIndex,
       hostname: "test-host",
       queue: "default",
     },
@@ -204,7 +205,7 @@ async function driveSupervisor(initialFrame: unknown, responder?: Responder): Pr
           // runtime never hangs on auto-generated RPCs (e.g. the
           // return_value XCom push).
           { body: null };
-        commSock.write(frameBytes(f.id, reply.body, true));
+        commSock.write(frameBytes(f.id, reply.body, true, reply.error ?? null));
         continue;
       }
       // Arity-3: a response from the runtime. The terminal frame
@@ -1286,4 +1287,194 @@ describe("coordinator runtime integration", () => {
       ).toBe(true);
     });
   });
+
+  describe("xcom_keys_to_clear", () => {
+    function deleteXCom(
+      key: string,
+      mapIndex: number | null = null,
+      { dagId = "py_dag", taskId = "extract" } = {},
+    ): Record<string, unknown> {
+      return {
+        type: "DeleteXCom",
+        key,
+        dag_id: dagId,
+        run_id: "r1",
+        task_id: taskId,
+        map_index: mapIndex,
+      };
+    }
+
+    // The supervisor reports a failed Execution API call with an empty body
+    // and an ErrorResponse in the response frame's error field.
+    const failDeletes: Responder = (msgType) =>
+      msgType === "DeleteXCom"
+        ? {
+            body: null,
+            error: {
+              type: "ErrorResponse",
+              error: "API_SERVER_ERROR",
+              detail: { status_code: 500, message: "Server error", detail: null },
+            },
+          }
+        : null;
+
+    it.each([
+      { label: "an unmapped task instance", mapIndex: -1, wireMapIndex: null },
+      { label: "map index 0", mapIndex: 0, wireMapIndex: 0 },
+      { label: "map index 2", mapIndex: 2, wireMapIndex: 2 },
+    ])(
+      "deletes each listed XCom before the task sends anything, for $label",
+      async ({ mapIndex, wireMapIndex }) => {
+        bundle.register(
+          new TaskHandler("py_dag", "extract", async () => {
+            await getClient().setXCom({ key: "progress", value: 1 });
+            return "rows";
+          }),
+        );
+
+        const result = await driveSupervisor(
+          makeStartupDetails(
+            "extract",
+            "py_dag",
+            "r1",
+            { xcom_keys_to_clear: ["return_value", "skipmixin_key"] },
+            mapIndex,
+          ),
+        );
+
+        expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" });
+        expect(result.runtimeRequests.map((r) => r.body)).toEqual([
+          deleteXCom("return_value", wireMapIndex),
+          deleteXCom("skipmixin_key", wireMapIndex),
+          expect.objectContaining({ type: "SetXCom", key: "progress" }),
+          expect.objectContaining({ type: "SetXCom", key: "return_value" }),
+        ]);
+      },
+    );
+
+    it("deletes nothing when ti_context lists no keys", async () => {
+      bundle.register(new TaskHandler("py_dag", "extract", async () => "rows"));
+
+      const result = await driveSupervisor(
+        makeStartupDetails("extract", "py_dag", "r1", { xcom_keys_to_clear: [] }),
+      );
+
+      expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" });
+      expect(result.runtimeRequests.map((r) => r.type)).toEqual(["SetXCom"]);
+    });
+
+    it.each([
+      {
+        shouldRetry: true,
+        terminal: { type: "RetryTask", retry_reason: "DeleteXCom failed: API_SERVER_ERROR" },
+      },
+      { shouldRetry: false, terminal: { type: "TaskState", state: "failed" } },
+    ])(
+      "fails the try without running the task when an XCom cannot be deleted (should_retry=$shouldRetry)",
+      async ({ shouldRetry, terminal }) => {
+        let ran = false;
+        bundle.register(
+          new TaskHandler("py_dag", "extract", async () => {
+            ran = true;
+          }),
+        );
+
+        const result = await driveSupervisor(
+          makeStartupDetails("extract", "py_dag", "r1", {
+            xcom_keys_to_clear: ["return_value", "skipmixin_key"],
+            should_retry: shouldRetry,
+          }),
+          failDeletes,
+        );
+
+        expect(ran).toBe(false);
+        expect(result.firstResponse!.body).toMatchObject(terminal);
+        // The first failed delete stops the rest.
+        expect(result.runtimeRequests.map((r) => r.body)).toEqual([deleteXCom("return_value")]);
+      },
+    );
+
+    it("deletes the listed XComs before binding arguments, so they are gone even when binding fails", async () => {
+      let ran = false;
+      bundle.register(
+        new TaskHandler("py_dag", "extract", async () => {
+          ran = true;
+        }),
+      );
+
+      const result = await driveSupervisor(
+        makeStartupDetails("extract", "py_dag", "r1", {
+          xcom_keys_to_clear: ["skipmixin_key"],
+          arg_bindings: [{ name: "region_code", kind: "bogus" }],
+          should_retry: true,
+        }),
+      );
+
+      expect(ran).toBe(false);
+      expect(result.firstResponse!.body).toMatchObject({
+        type: "RetryTask",
+        retry_reason: expect.stringContaining('has binding kind "bogus"'),
+      });
+      expect(result.runtimeRequests.map((r) => r.body)).toEqual([deleteXCom("skipmixin_key")]);
+    });
+
+    it("deletes nothing for a task the bundle has no handler for", async () => {
+      const result = await driveSupervisor(
+        makeStartupDetails("missing_task", "py_dag", "r1", {
+          xcom_keys_to_clear: ["return_value"],
+        }),
+      );
+
+      expect(result.firstResponse!.body).toMatchObject({ type: "TaskState", state: "removed" });
+      expect(result.runtimeRequests).toEqual([]);
+    });
+
+    it("deletes each listed XCom before a triggerDagRun task sends anything", async () => {
+      testDag.task(triggerDagRun({ dagId: "downstream", failWhenDagIsPaused: true }), {
+        taskId: "trigger",
+      })();
+      const replies: Record<string, unknown> = {
+        GetDag: { type: "DagResult", dag_id: "downstream", is_paused: false },
+        TriggerDagRun: { type: "OKResponse", ok: true },
+      };
+
+      const result = await driveSupervisor(
+        makeStartupDetails("trigger", "test_dag", "r1", {
+          xcom_keys_to_clear: ["trigger_run_id", "_link_TriggerDagRunLink"],
+        }),
+        (msgType) => (msgType in replies ? { body: replies[msgType] } : null),
+      );
+
+      const trigger = { dagId: "test_dag", taskId: "trigger" };
+      expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" });
+      expect(result.runtimeRequests.map((r) => r.body)).toEqual([
+        deleteXCom("trigger_run_id", null, trigger),
+        deleteXCom("_link_TriggerDagRunLink", null, trigger),
+        { type: "GetDag", dag_id: "downstream" },
+        expect.objectContaining({ type: "SetXCom", key: "_link_TriggerDagRunLink" }),
+        expect.objectContaining({ type: "TriggerDagRun", dag_id: "downstream" }),
+        expect.objectContaining({ type: "SetXCom", key: "trigger_run_id" }),
+      ]);
+    });
+
+    it("fails the try without triggering a run when a triggerDagRun task cannot delete an XCom", async () => {
+      testDag.task(triggerDagRun({ dagId: "downstream" }), { taskId: "trigger" })();
+
+      const result = await driveSupervisor(
+        makeStartupDetails("trigger", "test_dag", "r1", {
+          xcom_keys_to_clear: ["trigger_run_id", "_link_TriggerDagRunLink"],
+          should_retry: true,
+        }),
+        failDeletes,
+      );
+
+      expect(result.firstResponse!.body).toMatchObject({
+        type: "RetryTask",
+        retry_reason: "DeleteXCom failed: API_SERVER_ERROR",
+      });
+      expect(result.runtimeRequests.map((r) => r.body)).toEqual([
+        deleteXCom("trigger_run_id", null, { dagId: "test_dag", taskId: "trigger" }),
+      ]);
+    });
+  });
 });