feat: add platform TaskRunner seam for caller-controlled tool fan-out (#1050)

Add a platform.TaskRunner seam -- the concurrency companion to the
WithTimeProvider / WithUUIDProvider seams -- so a caller can control how ADK
runs its internal tool-call fan-out, and route handleFunctionCalls' parallel
tool execution through it.

- platform.RunTasks(ctx, tasks) runs a batch of independent tasks and blocks
  until all complete; with no runner installed it runs one goroutine per task
  (the unchanged default).
- platform.WithTaskRunner(ctx, runner) lets a caller install its own runner --
  for example to bound concurrency, dispatch tasks onto an existing executor,
  or run them sequentially in a single goroutine.
- Each task is invoked with its own context.Context, so a runner can derive a
  distinct per-task context. The default runner passes the parent ctx to every
  task.
- internal/llminternal tool fan-out (handleFunctionCalls) now executes via
  platform.RunTasks.

Author: DABH

internal/llminternal/base_flow.go

@@ -1057,14 +1057,13 @@ func (f *Flow) handleFunctionCalls(ctx agent.InvocationContext, toolsDict map[st
 	}
 
 	fnResponseEvents := make([]*session.Event, len(fnCalls))
-	var wg sync.WaitGroup
 
+	// Tool calls run via the context's task runner: concurrent goroutines by
+	// default, or a caller-installed runner (platform.WithTaskRunner).
+	tasks := make([]func(context.Context), len(fnCalls))
 	for i, fnCall := range fnCalls {
-		wg.Add(1)
-		go func(i int, fnCall *genai.FunctionCall) {
-			defer wg.Done()
-
-			sctx, span := telemetry.StartExecuteToolSpan(ctx, telemetry.StartExecuteToolSpanParams{
+		tasks[i] = func(taskCtx context.Context) {
+			sctx, span := telemetry.StartExecuteToolSpan(taskCtx, telemetry.StartExecuteToolSpanParams{
 				ToolName: fnCall.Name,
 				Args:     fnCall.Args,
 			})
@@ -1210,9 +1209,9 @@ func (f *Flow) handleFunctionCalls(ctx agent.InvocationContext, toolsDict map[st
 			})
 
 			fnResponseEvents[i] = ev
-		}(i, fnCall)
+		}
 	}
-	wg.Wait()
+	platform.RunTasks(ctx, tasks)
 	mergedEvent, err = mergeParallelFunctionResponseEvents(fnResponseEvents)
 	if err != nil {
 		return mergedEvent, err

platform/exec.go

@@ -0,0 +1,95 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package platform
+
+import (
+	"context"
+	"sync"
+)
+
+// TaskRunner runs a batch of independent tasks, invoking each one exactly
+// once and passing it a context.Context. An implementation must block until
+// every task has completed.
+//
+// Each task receives a context so the runner can give it its own per-task
+// context rather than sharing one. The default runner passes the same ctx it
+// was given to every task; a custom runner may instead derive a distinct
+// context per task.
+//
+// The default runner (used when none is installed on the context) runs the
+// tasks concurrently, one goroutine per task. A caller that wants to control
+// how the fan-out executes — for example to bound concurrency, dispatch onto
+// its own executor, or run the tasks sequentially in a single goroutine — can
+// install its own runner with WithTaskRunner.
+type TaskRunner func(ctx context.Context, tasks []func(context.Context))
+
+// taskRunnerKey is the context key under which a TaskRunner is stored.
+type taskRunnerKey struct{}
+
+// WithTaskRunner returns a copy of ctx that carries runner. RunTasks called
+// with the returned context, or any context derived from it, uses runner
+// instead of the default goroutine-based runner. A nil runner is ignored by
+// RunTasks, which then falls back to the default.
+//
+// WithTaskRunner is the concurrency analog of WithTimeProvider and
+// WithUUIDProvider: it lets a host environment substitute its own execution
+// strategy for ADK's internal fan-out without the ADK runtime taking a
+// dependency on that environment.
+func WithTaskRunner(ctx context.Context, runner TaskRunner) context.Context {
+	return context.WithValue(ctx, taskRunnerKey{}, runner)
+}
+
+// RunTasks runs every task in tasks and blocks until all of them complete. If
+// ctx carries a TaskRunner installed with WithTaskRunner, that runner is used;
+// otherwise RunTasks runs the tasks concurrently on one goroutine each.
+//
+// Each task is invoked with a context.Context. Under the default runner every
+// task receives ctx; an installed runner may instead pass each task its own
+// per-task context. Under the default runner, the tasks run concurrently and
+// must be safe to do so. An empty or nil slice returns without doing anything.
+//
+// RunTasks enforces the completion barrier itself rather than relying on the
+// runner: it wraps each task to signal a sync.WaitGroup on return and waits on
+// that group after the runner returns. Callers may therefore read per-task
+// results as soon as RunTasks returns, even if a custom runner does not block
+// until its tasks finish.
+func RunTasks(ctx context.Context, tasks []func(context.Context)) {
+	if len(tasks) == 0 {
+		return
+	}
+	var wg sync.WaitGroup
+	wg.Add(len(tasks))
+	if ctx != nil {
+		if r, ok := ctx.Value(taskRunnerKey{}).(TaskRunner); ok && r != nil {
+			wrapped := make([]func(context.Context), len(tasks))
+			for i, task := range tasks {
+				wrapped[i] = func(c context.Context) {
+					defer wg.Done()
+					task(c)
+				}
+			}
+			r(ctx, wrapped)
+			wg.Wait()
+			return
+		}
+	}
+	for _, task := range tasks {
+		go func() {
+			defer wg.Done()
+			task(ctx)
+		}()
+	}
+	wg.Wait()
+}

platform/exec_test.go

@@ -0,0 +1,166 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package platform_test
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"google.golang.org/adk/v2/platform"
+)
+
+func TestRunTasksDefaultRunsAllTasks(t *testing.T) {
+	const n = 16
+	var count atomic.Int64
+	seen := make([]bool, n)
+	var mu sync.Mutex
+
+	tasks := make([]func(context.Context), n)
+	for i := range tasks {
+		tasks[i] = func(context.Context) {
+			count.Add(1)
+			mu.Lock()
+			seen[i] = true
+			mu.Unlock()
+		}
+	}
+	platform.RunTasks(context.Background(), tasks)
+
+	if got := count.Load(); got != n {
+		t.Fatalf("ran %d tasks, want %d", got, n)
+	}
+	for i, ok := range seen {
+		if !ok {
+			t.Errorf("task %d did not run", i)
+		}
+	}
+}
+
+func TestRunTasksEmptyIsNoOp(t *testing.T) {
+	// Neither a nil nor an empty slice should block or panic.
+	platform.RunTasks(context.Background(), nil)
+	platform.RunTasks(context.Background(), []func(context.Context){})
+}
+
+func TestWithTaskRunnerIsUsed(t *testing.T) {
+	const n = 5
+	var order []int
+	used := false
+
+	// A sequential runner: runs every task inline, in slice order, on the
+	// caller's goroutine. This is the shape a caller installs to keep ADK's
+	// tool fan-out off goroutines, e.g. to bound concurrency or run on its own
+	// executor.
+	seq := func(ctx context.Context, tasks []func(context.Context)) {
+		used = true
+		for _, task := range tasks {
+			task(ctx)
+		}
+	}
+
+	ctx := platform.WithTaskRunner(context.Background(), seq)
+	tasks := make([]func(context.Context), n)
+	for i := range tasks {
+		tasks[i] = func(context.Context) {
+			order = append(order, i) // safe: no concurrency under the sequential runner
+		}
+	}
+	platform.RunTasks(ctx, tasks)
+
+	if !used {
+		t.Fatal("installed TaskRunner was not used")
+	}
+	for i := 0; i < n; i++ {
+		if order[i] != i {
+			t.Fatalf("sequential runner ran out of order: got %v", order)
+		}
+	}
+}
+
+func TestWithTaskRunnerNilFallsBack(t *testing.T) {
+	var count atomic.Int64
+	ctx := platform.WithTaskRunner(context.Background(), nil)
+	tasks := []func(context.Context){
+		func(context.Context) { count.Add(1) },
+		func(context.Context) { count.Add(1) },
+		func(context.Context) { count.Add(1) },
+	}
+	platform.RunTasks(ctx, tasks)
+	if got := count.Load(); got != 3 {
+		t.Fatalf("ran %d tasks, want 3", got)
+	}
+}
+
+// sentinelKey is a context key used by the per-task context tests.
+type sentinelKey struct{}
+
+func TestRunTasksDefaultPassesContext(t *testing.T) {
+	// With no installed runner, the default runner must invoke every task with
+	// a context derived from (carrying the values of) the parent ctx.
+	const n = 8
+	const want = "sentinel-value"
+	parent := context.WithValue(context.Background(), sentinelKey{}, want)
+
+	var mu sync.Mutex
+	got := make([]any, n)
+	tasks := make([]func(context.Context), n)
+	for i := range tasks {
+		tasks[i] = func(taskCtx context.Context) {
+			v := taskCtx.Value(sentinelKey{})
+			mu.Lock()
+			got[i] = v
+			mu.Unlock()
+		}
+	}
+	platform.RunTasks(parent, tasks)
+
+	for i, v := range got {
+		if v != want {
+			t.Errorf("task %d received sentinel %v, want %q", i, v, want)
+		}
+	}
+}
+
+func TestWithTaskRunnerReceivesPerTaskContext(t *testing.T) {
+	// An installed runner can hand each task its own distinct context, e.g. to
+	// scope per-task cancellation, deadlines, or values to each fanned-out task
+	// instead of sharing one.
+	const n = 6
+
+	// perTask gives task i a context carrying the value i.
+	perTask := func(ctx context.Context, tasks []func(context.Context)) {
+		for i, task := range tasks {
+			task(context.WithValue(ctx, sentinelKey{}, i))
+		}
+	}
+
+	ctx := platform.WithTaskRunner(context.Background(), perTask)
+	got := make([]any, n)
+	tasks := make([]func(context.Context), n)
+	for i := range tasks {
+		tasks[i] = func(taskCtx context.Context) {
+			got[i] = taskCtx.Value(sentinelKey{}) // safe: sequential, no concurrency
+		}
+	}
+	platform.RunTasks(ctx, tasks)
+
+	for i, v := range got {
+		if v != i {
+			t.Errorf("task %d observed per-task value %v, want %d", i, v, i)
+		}
+	}
+}