From b6fdcaf71b5fa3dd00fd5c443eed510b0571872b Mon Sep 17 00:00:00 2001 From: bernard-code-lab Date: Sun, 6 Sep 2026 17:48:26 -0300 Subject: [PATCH] fix: make transfer_task concurrency-safe under parallel tool calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a model issues several transfer_task calls in one response, the dispatcher runs them in parallel goroutines. Each one re-resolved its caller from 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", and a child that did start could run its turns as another sibling's agent. Resolve the caller from the snapshot Dispatcher.Process already takes before the fan-out, exposed on the per-call tools.Runtime handle, and let only one delegation at a time own the shared current agent — siblings that lose the claim pin their child session instead, the isolation mode background delegations already use. run_skill shares runForwarding and had the same latent defect when it landed in a batch alongside transfer_task; it is fixed by the same change. run_background_agent resolves its caller the same way, but agenttool's HandleRun does not receive a tools.Runtime, so that path is left for a follow-up. Its children are already pinned, which bounds the impact. Fixes #4156 Signed-off-by: bernard-code-lab --- pkg/runtime/agent_delegation.go | 76 +++++++++++------- pkg/runtime/agent_delegation_test.go | 110 +++++++++++++++++++++++++++ pkg/runtime/runtime.go | 23 ++++++ pkg/runtime/skill_runner.go | 13 ++-- pkg/runtime/toolexec/dispatcher.go | 6 ++ 5 files changed, 194 insertions(+), 34 deletions(-) diff --git a/pkg/runtime/agent_delegation.go b/pkg/runtime/agent_delegation.go index 2bfa134ee..c4d8372a6 100644 --- a/pkg/runtime/agent_delegation.go +++ b/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, }) } diff --git a/pkg/runtime/agent_delegation_test.go b/pkg/runtime/agent_delegation_test.go index 2ff78f75f..a64548fab 100644 --- a/pkg/runtime/agent_delegation_test.go +++ b/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() diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index 82439c23a..d6ad08e16 100644 --- a/pkg/runtime/runtime.go +++ b/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()) diff --git a/pkg/runtime/skill_runner.go b/pkg/runtime/skill_runner.go index 40b9ceda8..b2903d191 100644 --- a/pkg/runtime/skill_runner.go +++ b/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()), diff --git a/pkg/runtime/toolexec/dispatcher.go b/pkg/runtime/toolexec/dispatcher.go index 877161314..c3823845f 100644 --- a/pkg/runtime/toolexec/dispatcher.go +++ b/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