Keep taskDef out of Context's constructor and reuse resolveWiredAll

Author: jason810496

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt

@@ -123,9 +123,10 @@ data class TaskInstance(
 data class Context(
   @JvmField val dagRun: DagRun,
   @JvmField val ti: TaskInstance,
-  /** Registration of the executing task; resolves wired data inputs. */
-  internal val taskDef: TaskDef? = null,
 ) {
+  /** Registration of the executing task; resolves wired data inputs. */
+  internal var taskDef: TaskDef? = null
+
   internal companion object {
     fun from(
       request: StartupDetails,
@@ -146,7 +147,6 @@ data class Context(
             )
           },
         ti = with(request.ti) { TaskInstance(dagId, runId, taskId, mapIndex, tryNumber) },
-        taskDef = taskDef,
-      )
+      ).also { it.taskDef = taskDef }
   }
 }

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt

@@ -94,7 +94,7 @@ object ArgValues {
     // is well defined because a TaskInput is a task's only data parameter.
     wiredInputs(context, client)?.let { wired ->
       warnWiredArity(client, type, wired.size)
-      return type.cast(decode(resolveWired(wired[0], client), type))
+      return type.cast(decode(resolveWiredAll(wired.take(1), client).single(), type))
         ?: throw missingInput(wired[0], type.simpleName)
     }
     val input = newInput(type)
@@ -271,15 +271,6 @@ object ArgValues {
       is LiteralArg<*> -> MissingXComException("Input '$target' is wired to a null literal, so there is nothing to bind.")
     }
 
-  private fun resolveWired(
-    input: Arg<*>,
-    client: Client,
-  ): Any? =
-    when (input) {
-      is TaskRef<*> -> client.getXCom(taskId = input.def.id)
-      is LiteralArg<*> -> input.value
-    }
-
   /**
    * Builds the failure for a binding that resolved to nothing where a value is
    * required, naming [target] — the stub argument, or the [TaskInput] field

java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt

@@ -107,5 +107,5 @@ internal fun contextWiredWith(inputs: List<Arg<*>>): Context {
   val def = TaskDef("t", NoopTask::class.java)
   DagDef("d").addTask(def)
   def.inputs += inputs
-  return taskContext().copy(taskDef = def)
+  return taskContext().also { it.taskDef = def }
 }

java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt

@@ -121,7 +121,7 @@ internal class ArgValuesTest {
     val def = TaskDef("consumer", NoopArgTask::class.java)
     dag.addTask(def)
     def.inputs += inputs
-    return contextWithoutTaskDef().copy(taskDef = def)
+    return contextWithoutTaskDef().also { it.taskDef = def }
   }
 
   private fun contextWithoutTaskDef(): Context =