Reject invalid Go SDK task handler registrations up front

go-sdk/airflow/bundle.go

@@ -21,13 +21,17 @@ import (
 	"fmt"
 	"slices"
 	"sync"
+	"sync/atomic"
 
 	"github.com/apache/airflow/go-sdk/internal/bundle"
 )
 
 // BundleRef holds the task handlers that this executable runs for Airflow.
 // [Bundle] returns an empty one.
 type BundleRef struct {
+	// closed ends registration for everything the bundle can hold, so a kind added later
+	// is covered without a flag of its own. Serve sets it; Register reads it.
+	closed       atomic.Bool
 	taskHandlers taskHandlerMap
 }
 
@@ -66,8 +70,14 @@ type Registerable interface{ registerable() }
 //
 //	bundle.Register(reports.Handlers()...)
 //
-// Register panics if a task handler with the same dag_id and task_id is already registered.
+// Register panics if a task handler with the same dag_id and task_id is already registered,
+// and if [BundleRef.Serve] has already been called: registration closes when serving starts.
 func (b *BundleRef) Register(items ...Registerable) {
+	if b.closed.Load() {
+		panic(
+			"airflow.BundleRef.Register: Serve has already been called; register everything before Serve",
+		)
+	}
 	for _, item := range items {
 		switch item := item.(type) {
 		case *taskHandler:

go-sdk/airflow/bundle_test.go

@@ -19,6 +19,7 @@ package airflow
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"os/exec"
 	"reflect"
@@ -116,6 +117,38 @@ func TestRegisterRejectsDuplicateTask(t *testing.T) {
 	)
 }
 
+func TestRegisterAfterServePanics(t *testing.T) {
+	b := Bundle()
+	b.Register(TaskHandler("py_etl", "transform", noop))
+	require.NoError(t, b.serve([]string{"--airflow-metadata"}, io.Discard))
+
+	assert.PanicsWithValue(
+		t,
+		"airflow.BundleRef.Register: Serve has already been called; register everything before Serve",
+		func() { b.Register(TaskHandler("py_etl", "load", noop)) },
+	)
+}
+
+// The flag lives on the bundle, not on the task-handler map, so a kind of item added to
+// Register later is covered without a flag of its own.
+func TestRegisterAfterServeRejectsEveryKindOfItem(t *testing.T) {
+	b := Bundle()
+	require.NoError(t, b.serve([]string{"--airflow-metadata"}, io.Discard))
+
+	var nilItem Registerable
+	assert.Panics(t, func() { b.Register(nilItem) },
+		"the closed check runs before Register looks at what the item is")
+}
+
+// Serve closes registration whatever the run does, so a bundle that only printed its usage
+// still refuses a late Register.
+func TestRegisterAfterAFailedServePanics(t *testing.T) {
+	b := Bundle()
+	require.NoError(t, b.serve([]string{"--help"}, io.Discard))
+
+	assert.Panics(t, func() { b.Register(TaskHandler("py_etl", "transform", noop)) })
+}
+
 func TestRegisterIsSafeForConcurrentUse(t *testing.T) {
 	const workers, perWorker = 8, 100
 

go-sdk/airflow/serve.go

@@ -51,6 +51,9 @@ const (
 
 // Serve runs the bundle. Call it as the last statement of main.
 //
+// Serve closes registration: a [BundleRef.Register] call that reaches the bundle afterwards
+// panics rather than changing what the running bundle answers for.
+//
 // The command-line flags of the executable decide what Serve does.
 // With --airflow-metadata it prints the bundle's manifest and returns, which is how
 // airflow-go-pack reads the registered Dag and task ids.
@@ -68,6 +71,10 @@ func (b *BundleRef) Serve() error {
 }
 
 func (b *BundleRef) serve(args []string, stdout io.Writer) error {
+	// Registration closes here, whatever this run turns out to do, so that a Register left
+	// below Serve in main is reported as the mistake it is rather than racing the runtime.
+	b.closed.Store(true)
+
 	// The flags go on their own FlagSet. On pflag.CommandLine, every program that imports this
 	// package would get them, and one that defines its own --format there would panic.
 	flags := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)

go-sdk/airflow/task_handler_test.go

@@ -120,6 +120,11 @@ func TestTaskHandlerPanicsOnBadHandler(t *testing.T) {
 			fn:   func(Context, chan int) error { return nil },
 			want: "type chan int cannot receive a task argument",
 		},
+		{
+			name: "variadic",
+			fn:   func(Context, ...string) error { return nil },
+			want: "is variadic; a task argument cannot fill a ... parameter",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

go-sdk/internal/bundle/task.go

@@ -156,6 +156,15 @@ func (f *taskFunction) validateFn(fnType reflect.Type) error {
 		return fmt.Errorf("expected a func as input but was %s", fnType.Kind())
 	}
 
+	// Execute calls the function with Call, which passes a variadic parameter its slice
+	// rather than spreading it, so a variadic task function would panic at execution.
+	if fnType.IsVariadic() {
+		return fmt.Errorf(
+			"task function %s is variadic; a task argument cannot fill a ... parameter",
+			f.fullName,
+		)
+	}
+
 	// Return values
 	//     `<result>, error`,  or just `error`
 	if fnType.NumOut() < 1 || fnType.NumOut() > 2 {