fix: make transfer_task concurrency-safe under parallel tool calls

pkg/runtime/agent_delegation.go

@@ -192,19 +192,24 @@ type SubSessionConfig struct {
 type delegationRequest struct {
 	SubSessionConfig
 
+	// CallerAgent is the agent that issued the delegating tool call,
+	// snapshotted by the caller before the dispatcher's parallel fan-out.
+	// runForwarding falls back to session resolution when it is nil.
+	CallerAgent *agent.Agent
+
 	// SwitchCurrentAgent, when true, swaps r.currentAgent to AgentName
 	// for the lifetime of the call and emits AgentSwitching/AgentInfo
 	// events on entry and exit. Used by transfer_task. Mutually
 	// exclusive in spirit with PinAgent: pinning is for concurrent
 	// sub-sessions that must NOT share the runtime's mutable
-	// currentAgent, while switching is for sequential delegations where
-	// the parent loop is blocked anyway.
+	// currentAgent, while switching is for the one delegation that owns it.
 	//
-	// When the parent session is itself pinned (a background agent's
-	// session), runForwarding downgrades the switch to pinning the child
-	// to AgentName instead: the shared current agent belongs to the
-	// concurrent foreground loop and must not be mutated from a
-	// background task (#3886).
+	// runForwarding downgrades the switch to pinning the child to AgentName
+	// in two cases. When the parent session is itself pinned (a background
+	// agent's session), the shared current agent belongs to the concurrent
+	// foreground loop and must not be mutated from a background task (#3886).
+	// When a sibling delegation from the same parallel tool batch already
+	// holds the switch, mutating it would misroute this child's turns (#4156).
 	SwitchCurrentAgent bool
 }
 
@@ -340,11 +345,14 @@ func (r *LocalRuntime) swapCurrentAgent(ctx context.Context, sessionID string, f
 func (r *LocalRuntime) runForwarding(ctx context.Context, parent *session.Session, evts EventSink, req delegationRequest) (*tools.ToolCallResult, error) {
 	span := trace.SpanFromContext(ctx)
 
-	// The caller resolves from the parent session, not the shared current
-	// agent: a nested transfer from a pinned background session must
-	// attribute events, hooks, and completion to the pinned agent, no
-	// matter where the concurrent foreground loop points (#3886).
-	callerAgent := r.resolveSessionAgent(parent)
+	// The caller never resolves from the shared current agent here: a nested
+	// transfer from a pinned background session must attribute events, hooks,
+	// and completion to the pinned agent (#3886), and a sibling call in the
+	// same parallel batch may already have swapped it (#4156).
+	callerAgent := req.CallerAgent
+	if callerAgent == nil {
+		callerAgent = r.resolveSessionAgent(parent)
+	}
 	if callerAgent == nil {
 		return nil, errors.New("no agent resolved for the parent session")
 	}
@@ -354,15 +362,22 @@ func (r *LocalRuntime) runForwarding(ctx context.Context, parent *session.Sessio
 	}
 
 	if req.SwitchCurrentAgent {
-		if parent.AgentName == "" {
-			defer r.swapCurrentAgent(ctx, parent.ID, callerAgent, child, evts)()
-		} else {
+		switch {
+		case parent.AgentName != "":
 			// Pinned parent (background delegation): the shared current
 			// agent belongs to the concurrent foreground loop and must not
 			// be mutated. Pin the child to the target instead — RunStream
 			// resolves pinned sessions directly, so the child still
 			// executes as the target agent, without switch events/hooks.
 			req.PinAgent = true
+		case r.agentSwitchInFlight.CompareAndSwap(false, true):
+			defer r.agentSwitchInFlight.Store(false)
+			defer r.swapCurrentAgent(ctx, parent.ID, callerAgent, child, evts)()
+		default:
+			// A sibling delegation from the same parallel tool batch owns the
+			// shared current agent. Stomping it would make this child's turns
+			// resolve to the wrong agent, so pin the child instead (#4156).
+			req.PinAgent = true
 		}
 	}
 
@@ -420,14 +435,15 @@ func (r *LocalRuntime) runForwarding(ctx context.Context, parent *session.Sessio
 // Unlike runForwarding it does not emit AgentSwitching/AgentInfo events:
 // callers like background agents PinAgent the child session so the
 // runtime never mutates the shared currentAgent state.
-func (r *LocalRuntime) runCollecting(ctx context.Context, parent *session.Session, cfg SubSessionConfig, onContent func(string)) *agenttool.RunResult {
-	// The caller resolves from the parent session, not the shared current
-	// agent: a nested background dispatch from a pinned session must
-	// attribute the child's completion to the pinned agent, no matter
-	// where the concurrent foreground loop points (#3886). Resolved once
-	// up front so the subagent_stop defer below can't drift to a
-	// different agent if the shared current changes mid-run.
-	callerAgent := r.resolveSessionAgent(parent)
+func (r *LocalRuntime) runCollecting(ctx context.Context, parent *session.Session, caller *agent.Agent, cfg SubSessionConfig, onContent func(string)) *agenttool.RunResult {
+	// Resolved by the caller before dispatch, never from the shared current
+	// agent: a nested background dispatch from a pinned session must attribute
+	// the child's completion to the pinned agent (#3886), and the shared field
+	// may be swapped mid-run by a concurrent delegation (#4156).
+	callerAgent := caller
+	if callerAgent == nil {
+		callerAgent = r.resolveSessionAgent(parent)
+	}
 	if callerAgent == nil {
 		return &agenttool.RunResult{ErrMsg: "no agent resolved for the parent session"}
 	}
@@ -651,7 +667,7 @@ func (r *LocalRuntime) RunAgent(ctx context.Context, params agenttool.RunParams)
 	if guardErr != "" {
 		return &agenttool.RunResult{ErrMsg: guardErr}
 	}
-	return r.runCollecting(ctx, params.ParentSession, SubSessionConfig{
+	return r.runCollecting(ctx, params.ParentSession, caller, SubSessionConfig{
 		Task:              params.Task,
 		ExpectedOutput:    params.ExpectedOutput,
 		AgentName:         params.AgentName,
@@ -665,7 +681,7 @@ func (r *LocalRuntime) RunAgent(ctx context.Context, params agenttool.RunParams)
 	}, params.OnContent)
 }
 
-func (r *LocalRuntime) handleTaskTransfer(ctx context.Context, sess *session.Session, toolCall tools.ToolCall, evts EventSink, _ tools.Runtime) (*tools.ToolCallResult, error) {
+func (r *LocalRuntime) handleTaskTransfer(ctx context.Context, sess *session.Session, toolCall tools.ToolCall, evts EventSink, rt tools.Runtime) (*tools.ToolCallResult, error) {
 	var params struct {
 		Agent          string `json:"agent"`
 		Task           string `json:"task"`
@@ -675,10 +691,11 @@ func (r *LocalRuntime) handleTaskTransfer(ctx context.Context, sess *session.Ses
 		return nil, fmt.Errorf("invalid arguments: %w", err)
 	}
 
-	// Resolve the caller session-aware: nested transfer_task from a pinned
-	// background session must attribute the call to the pinned agent, not
-	// the shared current agent (#3886).
-	a := r.resolveSessionAgent(sess)
+	// Resolve the caller from the dispatcher's batch snapshot: a nested
+	// transfer_task from a pinned background session must attribute the call to
+	// the pinned agent (#3886), and a sibling transfer_task in the same parallel
+	// batch may already have swapped the shared current agent (#4156).
+	a := r.callerAgent(rt, sess)
 	if a == nil {
 		return nil, errors.New("no agent resolved for the calling session")
 	}
@@ -735,6 +752,7 @@ func (r *LocalRuntime) handleTaskTransfer(ctx context.Context, sess *session.Ses
 			NonInteractive:    sess.NonInteractive,
 			DelegationLineage: childLineage,
 		},
+		CallerAgent:        a,
 		SwitchCurrentAgent: true,
 	})
 }

pkg/runtime/agent_delegation_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/docker/docker-agent/pkg/team"
 	"github.com/docker/docker-agent/pkg/tools"
 	agenttool "github.com/docker/docker-agent/pkg/tools/builtin/agent"
+	"github.com/docker/docker-agent/pkg/tools/builtin/transfertask"
 )
 
 func TestBuildTaskSystemMessage(t *testing.T) {
@@ -1189,6 +1190,115 @@ func TestTransferTask_ConcurrentPinnedNestedTransfersStayIsolated(t *testing.T)
 	assert.Equal(t, "workerB", completedB.GetAgentName())
 }
 
+// callerRuntime is a [tools.Runtime] reporting a fixed caller agent, standing in
+// for the per-call handle the dispatcher builds from its pre-fan-out snapshot.
+type callerRuntime struct {
+	tools.NopRuntime
+
+	caller *agent.Agent
+}
+
+func (r callerRuntime) CallerAgent() *agent.Agent { return r.caller }
+
+// TestTransferTask_UsesBatchCallerSnapshotNotSharedCurrentAgent pins the exact
+// failure reported in #4156: a sibling transfer_task from the same parallel
+// batch has already swapped the shared current agent to its own target, so
+// re-resolving the caller from it identified the target as its own caller and
+// rejected the transfer with "No agents are configured in this list".
+func TestTransferTask_UsesBatchCallerSnapshotNotSharedCurrentAgent(t *testing.T) {
+	t.Parallel()
+
+	drafter := agent.New("drafter", "Drafter agent", agent.WithModel(&mockProvider{
+		id:     "test/mock-model",
+		stream: newStreamBuilder().AddContent("drafter done").AddStopWithUsage(10, 5).Build(),
+	}))
+	root := agent.New("root", "Root agent",
+		agent.WithModel(&mockProvider{id: "test/mock-model", stream: &mockStream{}}),
+		agent.WithSubAgents(drafter),
+	)
+
+	tm := team.New(team.WithAgents(root, drafter))
+	rt, err := NewLocalRuntime(t.Context(), tm,
+		WithSessionCompaction(false),
+		WithModelStore(mockModelStore{}),
+	)
+	require.NoError(t, err)
+
+	// Stand in for the sibling that already claimed the shared current agent.
+	rt.setCurrentAgent("drafter")
+
+	sess := session.New(session.WithUserMessage("Test"), session.WithToolsApproved(true))
+	result, err := rt.handleTaskTransfer(t.Context(), sess, transferToolCall("drafter"),
+		NewChannelSink(make(chan Event, 128)), callerRuntime{caller: root})
+	require.NoError(t, err)
+	require.NotNil(t, result)
+	require.False(t, result.IsError,
+		"the caller must come from the batch snapshot (root), not the swapped current agent: %s", result.Output)
+	assert.Equal(t, "drafter done", result.Output)
+}
+
+// TestTransferTask_ConcurrentForegroundTransfersStayIsolated reproduces #4156:
+// the model issues three transfer_task calls in a single response and the
+// dispatcher runs them in parallel. Each goroutine used to re-read the shared
+// current agent that a sibling's swapCurrentAgent had already mutated, so the
+// later calls validated the target against the wrong caller and failed with
+// "target agent not in sub-agents list", while a child that did start could run
+// its turns as another sibling's agent. Every transfer must instead run as its
+// own target and leave the shared current agent back at root.
+func TestTransferTask_ConcurrentForegroundTransfersStayIsolated(t *testing.T) {
+	t.Parallel()
+
+	targets := []string{"drafter", "reviewer", "tester"}
+	subAgents := make([]*agent.Agent, 0, len(targets))
+	for _, name := range targets {
+		prov := &mockProvider{id: "test/mock-model", stream: newStreamBuilder().
+			AddContent(name+" done").AddStopWithUsage(10, 5).Build()}
+		subAgents = append(subAgents, agent.New(name, name+" agent", agent.WithModel(prov)))
+	}
+
+	// One assistant response carrying all three calls — the parallel tool use
+	// that triggers the race.
+	batch := newStreamBuilder()
+	for i, name := range targets {
+		id := transferCallID(i)
+		batch.AddToolCallName(id, transfertask.ToolNameTransferTask).
+			AddToolCallArguments(id, fmt.Sprintf(`{"agent":%q,"task":"chunk %d","expected_output":"result"}`, name, i))
+	}
+	rootProv := &queueProvider{id: "test/mock-model", streams: []chat.MessageStream{
+		batch.AddToolCallStopWithUsage(10, 5).Build(),
+		newStreamBuilder().AddContent("all delegated").AddStopWithUsage(10, 5).Build(),
+	}}
+
+	root := agent.New("root", "Root agent",
+		agent.WithModel(rootProv),
+		agent.WithSubAgents(subAgents...),
+		agent.WithToolSets(transfertask.New()),
+	)
+
+	tm := team.New(team.WithAgents(append([]*agent.Agent{root}, subAgents...)...))
+	rt, err := NewLocalRuntime(t.Context(), tm,
+		WithSessionCompaction(false),
+		WithModelStore(mockModelStore{}),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = rt.Close() })
+
+	sess := session.New(session.WithUserMessage("split the diff"), session.WithToolsApproved(true))
+	_, err = rt.Run(t.Context(), sess)
+	require.NoError(t, err)
+
+	for i, name := range targets {
+		out := toolResultContent(t, sess, transferCallID(i))
+		assert.Equal(t, name+" done", out,
+			"transfer to %s must run as %s, not as a sibling's target", name, name)
+	}
+	assert.Equal(t, "root", rt.CurrentAgent().Name(),
+		"the shared current agent must be back at root once the batch drains")
+}
+
+// transferCallID names the nth transfer_task call of a parallel batch.
+func transferCallID(i int) string { return fmt.Sprintf("call_transfer_%d", i) }
+
 func TestTransferTask_DepthBoundary(t *testing.T) {
 	t.Parallel()
 

pkg/runtime/runtime.go

@@ -243,6 +243,7 @@ type LocalRuntime struct {
 	unmanagedOAuthRedirectURI string
 	nonInteractive            bool
 	startupInfoEmitted        atomic.Bool        // Track if startup info has been emitted to avoid unnecessary duplication
+	agentSwitchInFlight       atomic.Bool        // Claims the shared current agent for one foreground delegation (#4156)
 	elicitation               elicitationBridge  // Owns the per-stream events channel for outbound elicitation requests
 	elicitationWaiters        elicitationWaiters // Routes elicitation responses to the request awaiting them, keyed by ID (#3584)
 	elicitationDeclines       elicitationDeclineNotes
@@ -1322,6 +1323,28 @@ func (r *LocalRuntime) resolveSessionAgent(sess *session.Session) *agent.Agent {
 	return r.agents.ResolveSession(sess)
 }
 
+// callerAgentProvider is the optional interface a [tools.Runtime] implements to
+// report the agent that owns the in-flight tool-call batch. [toolexec] snapshots
+// it once before dispatching the batch in parallel.
+type callerAgentProvider interface {
+	CallerAgent() *agent.Agent
+}
+
+// callerAgent returns the agent that issued the in-flight tool call. It prefers
+// rt's batch snapshot over resolving from the session: a sibling transfer_task
+// running in the same parallel batch may already have swapped the shared current
+// agent, which would misidentify this call's caller (#4156). Hosts without a
+// snapshot (NopRuntime in tests, standalone skill invocations) fall back to
+// session resolution.
+func (r *LocalRuntime) callerAgent(rt tools.Runtime, sess *session.Session) *agent.Agent {
+	if p, ok := rt.(callerAgentProvider); ok {
+		if a := p.CallerAgent(); a != nil {
+			return a
+		}
+	}
+	return r.resolveSessionAgent(sess)
+}
+
 // CurrentAgentSkillsToolset returns the skills toolset for the current agent, or nil if not enabled.
 func (r *LocalRuntime) CurrentAgentSkillsToolset() *skills.ToolSet {
 	return agentSkillsToolset(r.CurrentAgent())

pkg/runtime/skill_runner.go

@@ -39,11 +39,13 @@ func (r *LocalRuntime) RunSkillFork(ctx context.Context, sess *session.Session,
 // tool call's runtime handle. Standalone invocations pass nil and skip embedded
 // commands because no tool call exists to own an approval prompt.
 func (r *LocalRuntime) runSkillFork(ctx context.Context, sess *session.Session, args skills.RunSkillArgs, evts EventSink, rt tools.Runtime) (*tools.ToolCallResult, error) {
-	// The caller resolves from the session, not the shared current agent:
-	// a fork skill invoked from a pinned background session must use the
-	// pinned agent's skills, identity, and model override, no matter where
-	// the concurrent foreground loop points (#3886).
-	caller := r.resolveSessionAgent(sess)
+	// The caller never resolves from the shared current agent: a fork skill
+	// invoked from a pinned background session must use the pinned agent's
+	// skills, identity, and model override (#3886), and a sibling transfer_task
+	// in the same parallel batch may already have swapped it — which would run
+	// the skill as the wrong agent (#4156). rt is nil for standalone
+	// invocations; callerAgent falls back to session resolution there.
+	caller := r.callerAgent(rt, sess)
 	if caller == nil {
 		return nil, errors.New("no agent resolved for the calling session")
 	}
@@ -131,6 +133,7 @@ func (r *LocalRuntime) runSkillFork(ctx context.Context, sess *session.Session,
 	// session), pin the child to the same agent so RunStream resolves it
 	// as the pinned caller instead of the shared current agent.
 	return r.runForwarding(ctx, sess, evts, delegationRequest{
+		CallerAgent: caller,
 		SubSessionConfig: SubSessionConfig{
 			Task:                prepared.Task,
 			SystemMessage:       skills.BuildSkillSystemMessage(prepared, sess.AttachedFilesSnapshot()),

pkg/runtime/toolexec/dispatcher.go

@@ -1057,6 +1057,12 @@ func (r callRuntime) EmitOutput(ctx context.Context, output string) {
 	r.c.em.EmitToolCallOutput(r.c.tc.ID, r.c.tool, output, r.c.a.Name())
 }
 
+// CallerAgent reports the agent that owns this tool-call batch, snapshotted
+// by [Dispatcher.Process] before the parallel fan-out. Handlers must prefer it
+// over re-resolving from the session: a sibling transfer_task in the same batch
+// may already have swapped the runtime's shared current agent (#4156).
+func (r callRuntime) CallerAgent() *agent.Agent { return r.c.a }
+
 func (r callRuntime) Recall(ctx context.Context, message string) error {
 	if r.c.d.Recall == nil {
 		return tools.ErrRecallNotSupported