Surface hasMappedDependants on ts-sdk TaskContext

task-sdk/src/airflow/sdk/execution_time/schema/schema.json

@@ -5077,6 +5077,11 @@
           "title": "Should Retry",
           "type": "boolean"
         },
+        "has_mapped_dependants": {
+          "default": false,
+          "title": "Has Mapped Dependants",
+          "type": "boolean"
+        },
         "start_date": {
           "anyOf": [
             {

ts-sdk/src/coordinator/runtime.ts

@@ -376,6 +376,7 @@ function buildContext(details: StartupDetails, signal: AbortSignal): TaskContext
     runId: details.ti.run_id,
     tryNumber: details.ti.try_number,
     mapIndex: details.ti.map_index ?? -1,
+    hasMappedDependants: details.ti_context.has_mapped_dependants ?? false,
     signal,
   };
 }

ts-sdk/src/generated/supervisor.ts

@@ -244,6 +244,7 @@ export type NextKwargs1 =
   | null;
 export type XcomKeysToClear = string[];
 export type ShouldRetry = boolean;
+export type HasMappedDependants = boolean;
 export type StartDate2 = string | null;
 export type ArgBindings = TaskArgBinding[] | null;
 /**
@@ -1030,6 +1031,7 @@ export interface TIRunContext {
   next_kwargs?: NextKwargs1;
   xcom_keys_to_clear?: XcomKeysToClear;
   should_retry?: ShouldRetry;
+  has_mapped_dependants?: HasMappedDependants;
   start_date?: StartDate2;
   arg_bindings?: ArgBindings;
 }

ts-sdk/src/sdk/task.ts

@@ -33,6 +33,16 @@ export interface TaskContext {
   readonly tryNumber: number;
   /** -1 for non-mapped tasks, 0..N-1 for mapped instances. */
   readonly mapIndex: number;
+  /**
+   * Whether this task's return value feeds a downstream `.expand()`.
+   *
+   * When `true`, returning an array from this handler causes the
+   * supervisor to record `mapped_length = value.length` on the
+   * return-value XCom automatically, so the scheduler can expand the
+   * mapped dependants. When `false`, no `mapped_length` is recorded
+   * (no behaviour change from today).
+   */
+  readonly hasMappedDependants: boolean;
   /**
    * AbortSignal that fires when Airflow terminates the task subprocess
    * with SIGTERM or SIGINT.

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

@@ -36,6 +36,7 @@ const FAKE_CTX: TaskContext = {
   runId: "r",
   tryNumber: 1,
   mapIndex: -1,
+  hasMappedDependants: false,
   signal: new AbortController().signal,
 };
 

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

@@ -579,4 +579,39 @@ describe("coordinator runtime integration", () => {
     const setXComReqs = result.runtimeRequests.filter((r) => r.type === "SetXCom");
     expect(setXComReqs).toHaveLength(0);
   });
+
+  it("surfaces hasMappedDependants=true on TaskContext when the server sets it", async () => {
+    let observedCtx: unknown = null;
+    registerTask({ dagId: "test_dag", taskId: "mapped_producer" }, async ({ ctx }) => {
+      observedCtx = ctx;
+      return ["a", "b", "c"];
+    });
+
+    const responder: Responder = (msgType) => (msgType === "SetXCom" ? { body: null } : null);
+
+    const result = await driveSupervisor(
+      makeStartupDetails("mapped_producer", "test_dag", "r1", {
+        has_mapped_dependants: true,
+        max_tries: 1,
+      }),
+      responder,
+    );
+
+    expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" });
+    expect(observedCtx).toMatchObject({ hasMappedDependants: true });
+  });
+
+  it("defaults hasMappedDependants to false when the server omits it", async () => {
+    let observedCtx: unknown = null;
+    registerTask({ dagId: "test_dag", taskId: "plain_task" }, async ({ ctx }) => {
+      observedCtx = ctx;
+    });
+
+    const result = await driveSupervisor(
+      makeStartupDetails("plain_task", "test_dag", "r1", { max_tries: 1 }),
+    );
+
+    expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" });
+    expect(observedCtx).toMatchObject({ hasMappedDependants: false });
+  });
 });

ts-sdk/tests/public-api.test.ts

@@ -78,6 +78,7 @@ describe("public API", () => {
       readonly runId: string;
       readonly tryNumber: number;
       readonly mapIndex: number;
+      readonly hasMappedDependants: boolean;
       readonly signal: AbortSignal;
     }>();
     expectTypeOf<GetXComOpts>().toEqualTypeOf<{