fix(runtime): avoid false on_user_input notifications

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

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
 	}

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)
 

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:

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")
+}

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) {

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 {

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()