diff --git a/pkg/runtime/elicitation.go b/pkg/runtime/elicitation.go index 10893823d..85bb1ee0e 100644 --- a/pkg/runtime/elicitation.go +++ b/pkg/runtime/elicitation.go @@ -486,7 +486,11 @@ func (r *LocalRuntime) elicitationHandler(ctx context.Context, req *mcp.ElicitPa }, nil } - r.executeOnUserInputHooks(ctx, "", "elicitation") + // No *agent.Agent is threaded into MCP handler callbacks, so fall back + // to the current agent here. The session ID is the conversation ID the + // run loop seeded into ctx (empty for elicitations outside a run, e.g. + // startup OAuth probes). + r.executeOnUserInputHooks(ctx, r.CurrentAgent(), genai.ConversationIDFromContext(ctx), "elicitation") // The registry key (and the ElicitationID surfaced to clients for // ResumeElicitation routing) is always a freshly generated, internal diff --git a/pkg/runtime/hooks.go b/pkg/runtime/hooks.go index b1cc35880..00d6eb892 100644 --- a/pkg/runtime/hooks.go +++ b/pkg/runtime/hooks.go @@ -570,10 +570,12 @@ func (r *LocalRuntime) executeAfterLLMCallHooks(ctx context.Context, sess *sessi // executeOnUserInputHooks fires on_user_input when the runtime is about // to wait for the user (tool confirmation, elicitation, max iterations, -// stream stopped). Resolves the agent itself so callsites in code paths -// without an agent handle (like the elicitation handler) stay short. -func (r *LocalRuntime) executeOnUserInputHooks(ctx context.Context, sessionID, logContext string) { - a := r.CurrentAgent() +// stream stopped). The agent is passed explicitly so callsites that +// already resolved the session's agent (pinned sessions included) +// attribute the event to that agent; paths without an agent handle +// (like the elicitation handler) fall back to CurrentAgent at the +// callsite. A nil agent is a no-op. +func (r *LocalRuntime) executeOnUserInputHooks(ctx context.Context, a *agent.Agent, sessionID, logContext string) { if a == nil { return } diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 41878ffcc..8b7850be0 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -214,7 +214,13 @@ func (r *LocalRuntime) finalizeEventChannel(ctx context.Context, sess *session.S // cleanup hooks run even when the stream was interrupted (e.g. Ctrl+C). r.executeSessionEndHooks(context.WithoutCancel(ctx), sess, a) - r.executeOnUserInputHooks(ctx, sess.ID, "stream stopped") + // on_user_input means "the agent is now waiting for the user". Only a + // root interactive stream hands control back to a user when it ends; + // sub-session and non-interactive teardowns (background agents, MCP + // serve, A2A, evals) have nobody to wait for and must not fire it (#4004). + if !sess.IsSubSession() && !sess.NonInteractive { + r.executeOnUserInputHooks(ctx, a, sess.ID, "stream stopped") + } r.telemetry.RecordSessionEnd(ctx) diff --git a/pkg/runtime/loop_steps.go b/pkg/runtime/loop_steps.go index 00c72197a..605bd6bdb 100644 --- a/pkg/runtime/loop_steps.go +++ b/pkg/runtime/loop_steps.go @@ -26,8 +26,8 @@ const ( // returned newMax is what the loop should use going forward. iterationContinue iterationDecision = iota // iterationStop means the loop should exit (limit reached and the - // user/non-interactive policy declined to continue, or context was - // cancelled while waiting for a resume decision). + // user/non-interactive policy declined to continue, or the context + // was cancelled before or while waiting for a resume decision). iterationStop ) @@ -67,7 +67,6 @@ func (r *LocalRuntime) enforceMaxIterations( maxIterMsg := fmt.Sprintf("Maximum iterations reached (%d)", runtimeMaxIterations) r.notifyMaxIterations(ctx, a, sess.ID, maxIterMsg) - r.executeOnUserInputHooks(ctx, sess.ID, "max iterations reached") stopMsg := fmt.Sprintf( "Execution stopped after reaching the configured max_iterations limit (%d).", @@ -89,6 +88,22 @@ func (r *LocalRuntime) enforceMaxIterations( return runtimeMaxIterations, iterationStop } + // A cancelled run can never deliver a resume decision. Bail out + // before signalling a bogus "waiting for user input" (#4004); + // mirrors the ctx.Done() branch below, which also skips the stop + // message. + if ctx.Err() != nil { + slog.DebugContext(ctx, "Context already cancelled at max iterations; stopping", + "agent", a.Name(), + "session_id", sess.ID, + ) + return runtimeMaxIterations, iterationStop + } + + // Only now is the runtime actually waiting for the user; the + // non-interactive auto-stop above never is (#4004). + r.executeOnUserInputHooks(ctx, a, sess.ID, "max iterations reached") + // Wait for user decision (resume / reject) select { case req := <-r.resumeChan: diff --git a/pkg/runtime/on_user_input_test.go b/pkg/runtime/on_user_input_test.go new file mode 100644 index 000000000..7d75e781f --- /dev/null +++ b/pkg/runtime/on_user_input_test.go @@ -0,0 +1,272 @@ +package runtime + +import ( + "context" + "testing" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/hooks" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/team" + "github.com/docker/docker-agent/pkg/telemetry/genai" + "github.com/docker/docker-agent/pkg/tools" + agenttool "github.com/docker/docker-agent/pkg/tools/builtin/agent" +) + +// runtimeWithRecordedUserInput mirrors runtimeWithRecordedSessionResume +// for the on_user_input event: a recording builtin is registered on the +// runtime's private registry post-construction and wired onto the root +// agent. extraAgents become sub-agents of root so delegation paths +// (run_background_agent) can be exercised; only root gets the hook wired +// here, but an extra agent constructed with its own on_user_input config +// shares the same recording builtin. +func runtimeWithRecordedUserInput(t *testing.T, extraAgents ...*agent.Agent) (*LocalRuntime, *recordingBuiltin) { + t.Helper() + + rb := &recordingBuiltin{} + prov := &mockProvider{ + id: "test/mock-model", + stream: newStreamBuilder().AddContent("done").AddStopWithUsage(1, 1).Build(), + } + root := agent.New("root", "instructions", + agent.WithModel(prov), + agent.WithHooks(&hooks.Config{ + OnUserInput: []hooks.Hook{{ + Type: hooks.HookTypeBuiltin, + Command: "test_record_on_user_input", + }}, + }), + ) + if len(extraAgents) > 0 { + agent.WithSubAgents(extraAgents...)(root) + } + tm := team.New(team.WithAgents(append([]*agent.Agent{root}, extraAgents...)...)) + + r, err := NewLocalRuntime(t.Context(), tm, WithSessionCompaction(false), WithModelStore(mockModelStore{})) + require.NoError(t, err) + + require.NoError(t, r.hooksRegistry.RegisterBuiltin("test_record_on_user_input", rb.hook)) + r.buildHooksExecutors() + + return r, rb +} + +// TestOnUserInputHooks_RootInteractiveStreamEnd_FiresOnce pins the +// PR #1847 semantics: when a root interactive stream finishes, the +// agent is waiting for the user's next message, so on_user_input fires +// exactly once with the session id. +func TestOnUserInputHooks_RootInteractiveStreamEnd_FiresOnce(t *testing.T) { + t.Parallel() + + r, rb := runtimeWithRecordedUserInput(t) + sess := session.New(session.WithUserMessage("hi")) + + for range r.RunStream(t.Context(), sess) { + } + + got := rb.snapshot() + require.Len(t, got, 1, "root interactive teardown must fire on_user_input exactly once") + assert.Equal(t, sess.ID, got[0].SessionID) +} + +// TestOnUserInputHooks_NonInteractiveStreamEnd_DoesNotFire is the #4004 +// regression guard: a non-interactive session (MCP serve, A2A, evals) +// has no user to wait for, so its teardown must not fire on_user_input. +func TestOnUserInputHooks_NonInteractiveStreamEnd_DoesNotFire(t *testing.T) { + t.Parallel() + + r, rb := runtimeWithRecordedUserInput(t) + sess := session.New(session.WithUserMessage("hi"), session.WithNonInteractive(true)) + + for range r.RunStream(t.Context(), sess) { + } + + assert.Empty(t, rb.snapshot(), "non-interactive teardown must not fire on_user_input") +} + +// TestOnUserInputHooks_SubSessionStreamEnd_DoesNotFire covers the +// interactive sub-session teardown (transfer_task): the child stream +// ending hands control back to the parent loop, not to the user, so +// on_user_input must not fire (#4004). +func TestOnUserInputHooks_SubSessionStreamEnd_DoesNotFire(t *testing.T) { + t.Parallel() + + r, rb := runtimeWithRecordedUserInput(t) + sub := session.New(session.WithUserMessage("hi"), session.WithParentID("parent-session")) + + for range r.RunStream(t.Context(), sub) { + } + + assert.Empty(t, rb.snapshot(), "sub-session teardown must not fire on_user_input") +} + +// TestOnUserInputHooks_BackgroundAgentTeardown_DoesNotFire reproduces +// the #4004 report: run_background_agent (runCollecting) tears down its +// child stream on normal completion, which used to reach the current +// (root) agent's on_user_input hook and signal a bogus "needs input". +func TestOnUserInputHooks_BackgroundAgentTeardown_DoesNotFire(t *testing.T) { + t.Parallel() + + worker := agent.New("worker", "worker instructions", + agent.WithModel(&mockProvider{ + id: "test/mock-model", + stream: newStreamBuilder().AddContent("worker done").AddStopWithUsage(1, 1).Build(), + })) + r, rb := runtimeWithRecordedUserInput(t, worker) + + parent := session.New(session.WithUserMessage("dispatch"), session.WithToolsApproved(true)) + res := r.RunAgent(t.Context(), agenttool.RunParams{ + AgentName: "worker", + Task: "background work", + ParentSession: parent, + }) + require.Empty(t, res.ErrMsg) + + assert.Empty(t, rb.snapshot(), "background sub-session teardown must not fire on_user_input") +} + +// TestEnforceMaxIterations_Interactive_FiresOnUserInput pins the one +// max-iterations path that genuinely waits for the user: the hook fires +// before blocking on the resume decision. +func TestEnforceMaxIterations_Interactive_FiresOnUserInput(t *testing.T) { + t.Parallel() + + r, rb := runtimeWithRecordedUserInput(t) + a := r.CurrentAgent() + require.NotNil(t, a) + sess := session.New() + events := make(chan Event, 8) + + go func() { r.resumeChan <- ResumeReject("stop") }() + _, decision := r.enforceMaxIterations(t.Context(), sess, a, 10, 10, NewChannelSink(events)) + + assert.Equal(t, iterationStop, decision) + got := rb.snapshot() + require.Len(t, got, 1, "interactive max-iterations wait must fire on_user_input once") + assert.Equal(t, sess.ID, got[0].SessionID) +} + +// TestEnforceMaxIterations_NonInteractive_DoesNotFireOnUserInput is the +// #4004 counterpart: the non-interactive auto-stop never waits for the +// user, so on_user_input must not fire. +func TestEnforceMaxIterations_NonInteractive_DoesNotFireOnUserInput(t *testing.T) { + t.Parallel() + + r, rb := runtimeWithRecordedUserInput(t) + a := r.CurrentAgent() + require.NotNil(t, a) + sess := session.New() + sess.NonInteractive = true + events := make(chan Event, 8) + + _, decision := r.enforceMaxIterations(t.Context(), sess, a, 10, 10, NewChannelSink(events)) + + assert.Equal(t, iterationStop, decision) + assert.Empty(t, rb.snapshot(), "non-interactive auto-stop must not fire on_user_input") +} + +// TestEnforceMaxIterations_ContextCancelled_DoesNotFireOnUserInput pins +// the cancelled-run guard: a run whose context is already done can never +// deliver a resume decision, so the runtime must stop without signalling +// a bogus "waiting for user input" (#4004). The limit event itself still +// fires — only the input wait is skipped. +func TestEnforceMaxIterations_ContextCancelled_DoesNotFireOnUserInput(t *testing.T) { + t.Parallel() + + r, rb := runtimeWithRecordedUserInput(t) + a := r.CurrentAgent() + require.NotNil(t, a) + sess := session.New() + events := make(chan Event, 8) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, decision := r.enforceMaxIterations(ctx, sess, a, 10, 10, NewChannelSink(events)) + + assert.Equal(t, iterationStop, decision) + assert.Empty(t, rb.snapshot(), "a cancelled context must not fire on_user_input") + + select { + case ev := <-events: + assert.IsType(t, &MaxIterationsReachedEvent{}, ev, + "MaxIterationsReached must still be emitted before the cancelled-context stop") + default: + t.Fatal("expected MaxIterationsReached to be emitted") + } +} + +// TestEnforceMaxIterations_PinnedSessionAgent_AttributesHookToThatAgent +// pins the attribution fix: the hook fires on the agent the callsite +// resolved for the session (here a pinned worker), not on whatever +// CurrentAgent happens to be. +func TestEnforceMaxIterations_PinnedSessionAgent_AttributesHookToThatAgent(t *testing.T) { + t.Parallel() + + worker := agent.New("worker", "worker instructions", + agent.WithModel(&mockProvider{ + id: "test/mock-model", + stream: newStreamBuilder().AddContent("worker done").AddStopWithUsage(1, 1).Build(), + }), + agent.WithHooks(&hooks.Config{ + OnUserInput: []hooks.Hook{{ + Type: hooks.HookTypeBuiltin, + Command: "test_record_on_user_input", + }}, + }), + ) + r, rb := runtimeWithRecordedUserInput(t, worker) + sess := session.New() + events := make(chan Event, 8) + + go func() { r.resumeChan <- ResumeReject("stop") }() + _, decision := r.enforceMaxIterations(t.Context(), sess, worker, 10, 10, NewChannelSink(events)) + + assert.Equal(t, iterationStop, decision) + got := rb.snapshot() + require.Len(t, got, 1, "the pinned agent's on_user_input hook must fire once") + assert.Equal(t, "worker", got[0].AgentName, + "the hook must be attributed to the session's agent, not CurrentAgent") +} + +// TestOnUserInputHooks_Elicitation_CarriesConversationID verifies the +// elicitation wait fires on_user_input with the conversation id seeded +// into ctx by the run loop, instead of an empty session id. +func TestOnUserInputHooks_Elicitation_CarriesConversationID(t *testing.T) { + t.Parallel() + + r, rb := runtimeWithRecordedUserInput(t) + + sinkCalled := make(chan Event, 1) + r.OnElicitationRequest(func(ev Event) { sinkCalled <- ev }) + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + ctx = genai.WithConversationID(ctx, "elicit-sess-1") + + done := make(chan error, 1) + go func() { + _, err := r.elicitationHandler(ctx, &mcp.ElicitParams{Message: "confirm?"}) + done <- err + }() + + var ev *ElicitationRequestEvent + select { + case e := <-sinkCalled: + ev = e.(*ElicitationRequestEvent) + case <-time.After(time.Second): + t.Fatal("elicitation request event not delivered to the sink") + } + require.NoError(t, r.ResumeElicitation(t.Context(), tools.ElicitationActionAccept, nil, ev.ElicitationID)) + require.NoError(t, <-done) + + got := rb.snapshot() + require.Len(t, got, 1, "elicitation wait must fire on_user_input once") + assert.Equal(t, "elicit-sess-1", got[0].SessionID, + "on_user_input must carry the conversation id from ctx, not an empty session id") +} diff --git a/pkg/runtime/tool_dispatch.go b/pkg/runtime/tool_dispatch.go index 6049b656c..72e9a9a9d 100644 --- a/pkg/runtime/tool_dispatch.go +++ b/pkg/runtime/tool_dispatch.go @@ -126,8 +126,8 @@ func (h *hookDispatcher) Dispatch(ctx context.Context, a *agent.Agent, event hoo return h.r.dispatchHook(ctx, a, event, in, h.events) } -func (h *hookDispatcher) NotifyUserInput(ctx context.Context, sessionID, label string) { - h.r.executeOnUserInputHooks(ctx, sessionID, label) +func (h *hookDispatcher) NotifyUserInput(ctx context.Context, a *agent.Agent, sessionID, label string) { + h.r.executeOnUserInputHooks(ctx, a, sessionID, label) } func (h *hookDispatcher) NotifyApprovalDecision(ctx context.Context, sess *session.Session, a *agent.Agent, tc tools.ToolCall, decision, source, safetyLabel string) { diff --git a/pkg/runtime/toolexec/dispatcher.go b/pkg/runtime/toolexec/dispatcher.go index fb08e4d1a..6f4f105f0 100644 --- a/pkg/runtime/toolexec/dispatcher.go +++ b/pkg/runtime/toolexec/dispatcher.go @@ -117,8 +117,8 @@ type HookDispatcher interface { // NotifyUserInput is invoked just before the dispatcher blocks waiting // for the user (tool confirmation). Implementations typically fire - // [hooks.EventOnUserInput]. - NotifyUserInput(ctx context.Context, sessionID, label string) + // [hooks.EventOnUserInput] attributed to the supplied agent. + NotifyUserInput(ctx context.Context, a *agent.Agent, sessionID, label string) // NotifyApprovalDecision is invoked once per tool call after the // approval pipeline (auto-allow, deny, user confirmation, ...) has @@ -829,7 +829,7 @@ func (c *call) askUser(ctx context.Context, runTool func() CallOutcome) CallOutc c.em.EmitToolCallConfirmation(c.tc, c.tool, c.a.Name(), c.confirmationMetadata(hookMeta)) if c.d.Hooks != nil { - c.d.Hooks.NotifyUserInput(ctx, c.sess.ID, "tool confirmation") + c.d.Hooks.NotifyUserInput(ctx, c.a, c.sess.ID, "tool confirmation") } select { diff --git a/pkg/runtime/toolexec/helpers_test.go b/pkg/runtime/toolexec/helpers_test.go index a953d3931..e8459276b 100644 --- a/pkg/runtime/toolexec/helpers_test.go +++ b/pkg/runtime/toolexec/helpers_test.go @@ -59,7 +59,7 @@ func (s *stubHookDispatcher) Dispatch(_ context.Context, _ *agent.Agent, event h return s.on[event] } -func (s *stubHookDispatcher) NotifyUserInput(context.Context, string, string) {} +func (s *stubHookDispatcher) NotifyUserInput(context.Context, *agent.Agent, string, string) {} func (s *stubHookDispatcher) NotifyApprovalDecision(_ context.Context, _ *session.Session, _ *agent.Agent, _ tools.ToolCall, decision, source, _ string) { s.mu.Lock() defer s.mu.Unlock()