From 811d91dd5fb06924f5bb91cbd27faa41850856ed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Mon, 13 Jul 2026 14:08:16 +0200 Subject: [PATCH] fix(runtime): reliable, correlation-safe, non-blocking elicitation delivery (#3584) --- docs/guides/go-sdk/index.md | 17 + pkg/api/types.go | 7 + pkg/app/app_test.go | 3 +- pkg/chatserver/agent.go | 2 +- pkg/cli/runner.go | 8 +- pkg/cli/runner_test.go | 3 +- pkg/embeddedchat/embeddedchat.go | 4 +- pkg/embeddedchat/embeddedchat_test.go | 4 +- pkg/leantui/update_test.go | 3 +- pkg/runtime/agent_delegation.go | 35 +- pkg/runtime/client.go | 4 +- pkg/runtime/commands_test.go | 3 +- pkg/runtime/contract_test.go | 2 +- pkg/runtime/elicitation.go | 446 +++++++++++- pkg/runtime/elicitation_concurrency_test.go | 641 ++++++++++++++++++ pkg/runtime/elicitation_test.go | 10 +- pkg/runtime/event.go | 60 +- pkg/runtime/remote_client.go | 7 +- pkg/runtime/remote_contract_test.go | 2 +- pkg/runtime/remote_runtime.go | 39 +- pkg/runtime/runtime.go | 48 +- pkg/runtime/runtime_test.go | 33 +- pkg/server/server.go | 2 +- pkg/server/session_manager.go | 7 +- pkg/server/session_manager_test.go | 2 +- pkg/tui/compact_routing_test.go | 3 +- pkg/tui/page/chat/queue_test.go | 3 +- pkg/tui/service/supervisor/supervisor_test.go | 10 +- 28 files changed, 1279 insertions(+), 129 deletions(-) create mode 100644 pkg/runtime/elicitation_concurrency_test.go diff --git a/docs/guides/go-sdk/index.md b/docs/guides/go-sdk/index.md index 6fa025a031..11ae699894 100644 --- a/docs/guides/go-sdk/index.md +++ b/docs/guides/go-sdk/index.md @@ -129,6 +129,23 @@ if err := chat.Restart(); err != nil { For advanced use (custom elicitation, raw event inspection), call `chat.Runtime()` to access the underlying `runtime.Runtime` directly. +> [!WARNING] +> **Breaking change: `Runtime.ResumeElicitation` (#3584)** +> +> `Runtime.ResumeElicitation` gained an `elicitationID` parameter so responses can +> be correlated with a specific concurrent elicitation request (needed once +> multiple background jobs can be eliciting input at the same time). It is +> declared **variadic** (`elicitationID ...string`) specifically so existing +> *callers* of the 3-argument form keep compiling unchanged — `rt.ResumeElicitation(ctx, action, content)` +> still works and falls back to resolving the sole pending request. +> +> If you implement your own `runtime.Runtime` (rather than embedding +> `runtime.LocalRuntime`/`runtime.RemoteRuntime`), you do need to update your +> method's signature to match, and also add an `OnElicitationRequest(handler +> func(runtime.Event))` method (a no-op is fine if your runtime never raises +> elicitations) — both are required interface methods, matching the existing +> no-op-able pattern already used by `OnToolsChanged`/`OnBackgroundEvent`. + ## Optional Provider Build Tags By default docker-agent includes all four cloud providers (OpenAI, Anthropic, Google, Amazon Bedrock). When embedding docker-agent in your own binary you can compile out unneeded providers — together with their transitive SDK dependencies — to reduce binary size. diff --git a/pkg/api/types.go b/pkg/api/types.go index a4b2995474..a0e772031d 100644 --- a/pkg/api/types.go +++ b/pkg/api/types.go @@ -174,6 +174,13 @@ type DesktopTokenResponse struct { type ResumeElicitationRequest struct { Action string `json:"action"` // "accept", "decline", or "cancel" Content map[string]any `json:"content"` // The submitted form data (only present when action is "accept") + // ElicitationID correlates this response with a specific concurrent + // elicitation request (see the elicitation_id field on the + // ElicitationRequestEvent stream event). Optional and additive: when + // empty, the server falls back to resolving the sole pending request, + // for backward compatibility with clients that predate per-request + // correlation (#3584). + ElicitationID string `json:"elicitation_id,omitempty"` } // SteerSessionRequest represents a request to inject user messages into a diff --git a/pkg/app/app_test.go b/pkg/app/app_test.go index 24a965a2e0..a2053d11a0 100644 --- a/pkg/app/app_test.go +++ b/pkg/app/app_test.go @@ -59,7 +59,7 @@ func (m *mockRuntime) Run(ctx context.Context, sess *session.Session) ([]session return nil, nil } func (m *mockRuntime) Resume(ctx context.Context, req runtime.ResumeRequest) {} -func (m *mockRuntime) ResumeElicitation(ctx context.Context, action tools.ElicitationAction, content map[string]any) error { +func (m *mockRuntime) ResumeElicitation(ctx context.Context, action tools.ElicitationAction, content map[string]any, elicitationID ...string) error { return nil } func (m *mockRuntime) SessionStore() session.Store { return m.store } @@ -108,6 +108,7 @@ func (m *mockRuntime) AvailableModels(context.Context) []runtime.ModelChoice { r func (m *mockRuntime) SupportsModelSwitching() bool { return false } func (m *mockRuntime) OnToolsChanged(func(runtime.Event)) {} func (m *mockRuntime) OnBackgroundEvent(func(runtime.Event)) {} +func (m *mockRuntime) OnElicitationRequest(func(runtime.Event)) {} // Verify mockRuntime implements runtime.Runtime var _ runtime.Runtime = (*mockRuntime)(nil) diff --git a/pkg/chatserver/agent.go b/pkg/chatserver/agent.go index f28eda8aa1..4673bccf3a 100644 --- a/pkg/chatserver/agent.go +++ b/pkg/chatserver/agent.go @@ -238,7 +238,7 @@ func runAgentLoop(ctx context.Context, rt runtime.Runtime, sess *session.Session case *runtime.ElicitationRequestEvent: // Required: the runtime blocks until we respond, regardless // of NonInteractive. Decline so the tool call fails fast. - _ = rt.ResumeElicitation(ctx, tools.ElicitationActionDecline, nil) + _ = rt.ResumeElicitation(ctx, tools.ElicitationActionDecline, nil, e.ElicitationID) case *runtime.MaxIterationsReachedEvent: // Defensive: in non-interactive mode the runtime already // stops on its own and this Resume is dropped. diff --git a/pkg/cli/runner.go b/pkg/cli/runner.go index eb842d807c..dcc0223a6c 100644 --- a/pkg/cli/runner.go +++ b/pkg/cli/runner.go @@ -124,7 +124,7 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess // still Reject (the hook said Ask, not Approve). rt.Resume(ctx, runtime.ResumeReject("")) case *runtime.ElicitationRequestEvent: - _ = rt.ResumeElicitation(ctx, "decline", nil) + _ = rt.ResumeElicitation(ctx, "decline", nil, e.ElicitationID) case *runtime.MaxIterationsReachedEvent: switch handleMaxIterationsAutoApprove(cfg.AutoApprove, &autoExtensions, e.MaxIterations) { case maxIterContinue: @@ -242,7 +242,7 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess serverURL, ok := e.Meta["docker-agent/server_url"].(string) if !ok || serverURL == "" { slog.WarnContext(ctx, "Skipping elicitation: missing or invalid server_url (non-interactive session?)") - _ = rt.ResumeElicitation(ctx, "decline", nil) + _ = rt.ResumeElicitation(ctx, "decline", nil, e.ElicitationID) return nil } @@ -254,9 +254,9 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess switch result { case ConfirmationApprove: - _ = rt.ResumeElicitation(ctx, "accept", nil) + _ = rt.ResumeElicitation(ctx, "accept", nil, e.ElicitationID) case ConfirmationReject: - _ = rt.ResumeElicitation(ctx, "decline", nil) + _ = rt.ResumeElicitation(ctx, "decline", nil, e.ElicitationID) return errors.New("OAuth authorization rejected by user") } } diff --git a/pkg/cli/runner_test.go b/pkg/cli/runner_test.go index 91b6bebe0e..a45f1d1371 100644 --- a/pkg/cli/runner_test.go +++ b/pkg/cli/runner_test.go @@ -69,7 +69,7 @@ func (m *mockRuntime) Run(context.Context, *session.Session) ([]session.Message, return nil, nil } -func (m *mockRuntime) ResumeElicitation(_ context.Context, action tools.ElicitationAction, _ map[string]any) error { +func (m *mockRuntime) ResumeElicitation(_ context.Context, action tools.ElicitationAction, _ map[string]any, _ ...string) error { m.mu.Lock() defer m.mu.Unlock() m.elicitationDeclines++ @@ -110,6 +110,7 @@ func (m *mockRuntime) AvailableModels(context.Context) []runtime.ModelChoice func (m *mockRuntime) SupportsModelSwitching() bool { return false } func (m *mockRuntime) OnToolsChanged(func(runtime.Event)) {} func (m *mockRuntime) OnBackgroundEvent(func(runtime.Event)) {} +func (m *mockRuntime) OnElicitationRequest(func(runtime.Event)) {} func (m *mockRuntime) RegenerateTitle(context.Context, *session.Session, chan runtime.Event) {} func (m *mockRuntime) Resume(_ context.Context, req runtime.ResumeRequest) { diff --git a/pkg/embeddedchat/embeddedchat.go b/pkg/embeddedchat/embeddedchat.go index 3f1d07b1f8..42476f1e9d 100644 --- a/pkg/embeddedchat/embeddedchat.go +++ b/pkg/embeddedchat/embeddedchat.go @@ -81,7 +81,7 @@ type ToolActivity struct { type runtimeRunner interface { RunStream(ctx context.Context, sess *session.Session) <-chan dagentruntime.Event Resume(ctx context.Context, req dagentruntime.ResumeRequest) - ResumeElicitation(ctx context.Context, action tools.ElicitationAction, content map[string]any) error + ResumeElicitation(ctx context.Context, action tools.ElicitationAction, content map[string]any, elicitationID ...string) error Close() error } @@ -303,7 +303,7 @@ func (s *Session) forwardEvents(ctx context.Context, events <-chan dagentruntime // This headless wrapper has no built-in elicitation UI. Decline so the // run cannot hang forever; embedders that need elicitation can consume // RuntimeEvent directly by driving the runtime themselves. - _ = s.rt.ResumeElicitation(ctx, tools.ElicitationActionDecline, nil) + _ = s.rt.ResumeElicitation(ctx, tools.ElicitationActionDecline, nil, e.ElicitationID) case *dagentruntime.MaxIterationsReachedEvent: s.rt.Resume(ctx, dagentruntime.ResumeReject("")) case *dagentruntime.ErrorEvent: diff --git a/pkg/embeddedchat/embeddedchat_test.go b/pkg/embeddedchat/embeddedchat_test.go index 95a145e892..b4d2d8fc9b 100644 --- a/pkg/embeddedchat/embeddedchat_test.go +++ b/pkg/embeddedchat/embeddedchat_test.go @@ -62,7 +62,7 @@ func (f *fakeRuntime) Resume(_ context.Context, req dagentruntime.ResumeRequest) f.resumes = append(f.resumes, req) } -func (f *fakeRuntime) ResumeElicitation(_ context.Context, action tools.ElicitationAction, _ map[string]any) error { +func (f *fakeRuntime) ResumeElicitation(_ context.Context, action tools.ElicitationAction, _ map[string]any, _ ...string) error { f.elicitations = append(f.elicitations, action) return nil } @@ -166,7 +166,7 @@ func TestSessionSendDeclinesElicitationAndRejectsMaxIterations(t *testing.T) { out, err := s.Send(t.Context(), "hi") require.NoError(t, err) - rt.events <- dagentruntime.ElicitationRequest("authorize", "url", nil, "https://example.com", "id", nil, "agent") + rt.events <- dagentruntime.ElicitationRequest("authorize", "url", nil, "https://example.com", "id", "", "sess", nil, "agent") rt.events <- dagentruntime.MaxIterationsReached(3) close(rt.events) diff --git a/pkg/leantui/update_test.go b/pkg/leantui/update_test.go index a2c06d52a0..9a589c21c8 100644 --- a/pkg/leantui/update_test.go +++ b/pkg/leantui/update_test.go @@ -56,7 +56,7 @@ func (r *cycleThinkingRuntime) Run(context.Context, *session.Session) ([]session return nil, nil } func (r *cycleThinkingRuntime) Resume(context.Context, runtime.ResumeRequest) {} -func (r *cycleThinkingRuntime) ResumeElicitation(context.Context, tools.ElicitationAction, map[string]any) error { +func (r *cycleThinkingRuntime) ResumeElicitation(context.Context, tools.ElicitationAction, map[string]any, ...string) error { return nil } func (r *cycleThinkingRuntime) SessionStore() session.Store { return nil } @@ -120,6 +120,7 @@ func (r *cycleThinkingRuntime) AvailableModels(context.Context) []runtime.ModelC func (r *cycleThinkingRuntime) SupportsModelSwitching() bool { return r.supports } func (r *cycleThinkingRuntime) OnToolsChanged(func(runtime.Event)) {} func (r *cycleThinkingRuntime) OnBackgroundEvent(func(runtime.Event)) {} +func (r *cycleThinkingRuntime) OnElicitationRequest(func(runtime.Event)) {} var _ runtime.Runtime = (*cycleThinkingRuntime)(nil) diff --git a/pkg/runtime/agent_delegation.go b/pkg/runtime/agent_delegation.go index 01df570672..f1af4c6f38 100644 --- a/pkg/runtime/agent_delegation.go +++ b/pkg/runtime/agent_delegation.go @@ -378,6 +378,15 @@ func (r *LocalRuntime) runCollecting(ctx context.Context, parent *session.Sessio if usage, ok := event.(*TokenUsageEvent); ok { r.emitBackgroundEvent(usage) } + // Elicitation requests are NOT re-forwarded here: elicitationHandler + // already delivered this event to the OnElicitationRequest sink + // directly, synchronously, and exactly once (#3584). Forwarding it + // again here — as the bridge's best-effort copy happens to flow + // through this same channel when this background sub-session + // currently owns the bridge slot — used to cause a second sink + // delivery for the same request, which only a stateful App-side + // dedupe (since removed) papered over. This branch is intentionally + // absent; ElicitationRequestEvents seen here are simply ignored. if errEvt, ok := event.(*ErrorEvent); ok { errMsg = errEvt.Error break @@ -416,15 +425,31 @@ func (r *LocalRuntime) runCollecting(ctx context.Context, parent *session.Sessio // no explanation. Prepend an actionable note so the model (and, through it, // the user) learns the server must be authorized interactively first. if note := backgroundAuthRequiredNote(child); note != "" { - if result != "" { - result = note + "\n\n" + result - } else { - result = note - } + result = prependNote(result, note) + } + // Mid-call elicitations that were auto-declined because this background + // session had no UI to answer them (see elicitationHandler) are recorded + // against this sub-session's ID; surface them the same way (#3584). + for _, note := range r.elicitationDeclines.drain(s.ID) { + result = prependNote(result, note) } return &agenttool.RunResult{Result: result} } +// prependNote prepends note to result, separated by a blank line, handling +// the case where either side is empty. Used to surface model-readable +// context (OAuth-required, elicitation auto-declined) ahead of a background +// sub-session's actual response. +func prependNote(result, note string) string { + if note == "" { + return result + } + if result == "" { + return note + } + return note + "\n\n" + result +} + // backgroundAuthRequiredNote returns a model-readable note naming the child // agent's MCP toolsets that could not start because they require first-time // interactive OAuth authorization, which a background agent cannot complete diff --git a/pkg/runtime/client.go b/pkg/runtime/client.go index 60d159ba6a..57388a8410 100644 --- a/pkg/runtime/client.go +++ b/pkg/runtime/client.go @@ -456,8 +456,8 @@ func (c *Client) DeleteRemoteSession(ctx context.Context, sessionID string) erro return c.doRequest(ctx, http.MethodDelete, "/api/sessions/"+sessionID, nil, nil) } -func (c *Client) ResumeElicitation(ctx context.Context, sessionID string, action tools.ElicitationAction, content map[string]any) error { - req := api.ResumeElicitationRequest{Action: string(action), Content: content} +func (c *Client) ResumeElicitation(ctx context.Context, sessionID string, action tools.ElicitationAction, content map[string]any, elicitationID ...string) error { + req := api.ResumeElicitationRequest{Action: string(action), Content: content, ElicitationID: firstElicitationID(elicitationID)} return c.doRequest(ctx, http.MethodPost, "/api/sessions/"+sessionID+"/elicitation", req, nil) } diff --git a/pkg/runtime/commands_test.go b/pkg/runtime/commands_test.go index 8e44007914..a32aa96c55 100644 --- a/pkg/runtime/commands_test.go +++ b/pkg/runtime/commands_test.go @@ -49,7 +49,7 @@ func (m *mockRuntime) Run(context.Context, *session.Session) ([]session.Message, return nil, nil } func (m *mockRuntime) Resume(context.Context, ResumeRequest) {} -func (m *mockRuntime) ResumeElicitation(context.Context, tools.ElicitationAction, map[string]any) error { +func (m *mockRuntime) ResumeElicitation(context.Context, tools.ElicitationAction, map[string]any, ...string) error { return nil } func (m *mockRuntime) SessionStore() session.Store { return nil } @@ -93,6 +93,7 @@ func (m *mockRuntime) AvailableModels(context.Context) []ModelChoice { return ni func (m *mockRuntime) SupportsModelSwitching() bool { return false } func (m *mockRuntime) OnToolsChanged(func(Event)) {} func (m *mockRuntime) OnBackgroundEvent(func(Event)) {} +func (m *mockRuntime) OnElicitationRequest(func(Event)) {} func (m *mockRuntime) RegenerateTitle(context.Context, *session.Session, chan Event) { } diff --git a/pkg/runtime/contract_test.go b/pkg/runtime/contract_test.go index d3891b6464..47519fff53 100644 --- a/pkg/runtime/contract_test.go +++ b/pkg/runtime/contract_test.go @@ -124,7 +124,7 @@ func runRuntimeContract(t *testing.T, newRT func(t *testing.T) Runtime) { t.Cleanup(func() { _ = rt.Close() }) ctx, cancel := context.WithCancel(t.Context()) cancel() - _ = rt.ResumeElicitation(ctx, tools.ElicitationActionDecline, nil) + _ = rt.ResumeElicitation(ctx, tools.ElicitationActionDecline, nil, "") }) } diff --git a/pkg/runtime/elicitation.go b/pkg/runtime/elicitation.go index e868b4b80a..fdc3468635 100644 --- a/pkg/runtime/elicitation.go +++ b/pkg/runtime/elicitation.go @@ -6,10 +6,14 @@ import ( "fmt" "log/slog" "sync" + "sync/atomic" + "github.com/google/uuid" "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/docker/docker-agent/pkg/telemetry/genai" "github.com/docker/docker-agent/pkg/tools" + mcptools "github.com/docker/docker-agent/pkg/tools/mcp" ) // ElicitationResult represents the result of an elicitation request. @@ -40,6 +44,12 @@ type ElicitationRequestHandler func(ctx context.Context, message string, schema // configured (no RunStream is active). var errNoElicitationChannel = errors.New("no events channel available for elicitation") +// errNoSuchElicitation is returned by ResumeElicitation when the given +// elicitation ID (or, in the empty-ID fallback, "the single pending +// request") no longer has a registered waiter — already answered, timed +// out, or never existed. +var errNoSuchElicitation = errors.New("no elicitation request in progress") + // elicitationBridge owns the events channel that the runtime's MCP // elicitation handler sends requests to. Each RunStream call swaps in its // own channel on entry and the previous one back on exit, so nested @@ -54,6 +64,17 @@ var errNoElicitationChannel = errors.New("no events channel available for elicit // contract testable in isolation (with the race detector) without // spinning up a runtime, and keeps LocalRuntime free of the two raw // fields it used to expose. +// +// Concurrent (non-nested) RunStreams — most notably background jobs +// started via run_background_agent — can swap this single slot out from +// under each other; see elicitationWaiters and OnElicitationRequest for +// the routing/delivery fix (#3584). The bridge itself is kept only as a +// best-effort secondary delivery path for remote/SSE consumers that read +// events directly off a RunStream channel (see remote_runtime.go). It is +// never allowed to hold up the reliable sink or response processing: send +// is bounded by the caller's ctx (see elicitationHandler), and callers +// invoke it from a detached goroutine so a wedged or abandoned channel +// cannot block the request/response path at all (#3584 review item 1). type elicitationBridge struct { mu sync.RWMutex ch chan Event @@ -70,13 +91,18 @@ func (b *elicitationBridge) swap(ch chan Event) chan Event { } // send delivers ev to the current channel, holding the read lock across -// the send. This blocks concurrent teardown until the send completes, -// preserving the invariant that the channel reference held by an -// in-flight sender stays open until that sender finishes. +// the send so a concurrent restoreAndClose cannot close the channel out +// from under an in-flight send without going through recover() below. The +// send itself is bounded by ctx: if ctx is done before the channel accepts +// the event, send returns ctx.Err() instead of blocking forever. Combined +// with callers invoking send from a detached goroutine (see +// elicitationHandler), a full or abandoned channel can no longer delay — +// let alone indefinitely block — the reliable sink delivery or response +// handling that used to be sequenced before this call (#3584 item 1). // -// Returns errNoElicitationChannel when no channel is configured or when -// a defensive recover catches an externally closed channel. -func (b *elicitationBridge) send(ev Event) (err error) { +// Returns errNoElicitationChannel when no channel is configured or when a +// defensive recover catches an externally closed channel. +func (b *elicitationBridge) send(ctx context.Context, ev Event) (err error) { b.mu.RLock() defer b.mu.RUnlock() defer func() { @@ -87,8 +113,12 @@ func (b *elicitationBridge) send(ev Event) (err error) { if b.ch == nil { return errNoElicitationChannel } - b.ch <- ev - return nil + select { + case b.ch <- ev: + return nil + case <-ctx.Done(): + return ctx.Err() + } } // restoreAndClose restores the previous stream channel and closes the current @@ -100,9 +130,11 @@ func (b *elicitationBridge) send(ev Event) (err error) { // lock makes restoreAndClose wait for any in-flight send to finish, because // send holds the read lock across "b.ch <- ev". If the stream consumer has // gone away and current is full (or unbuffered), that parked send never -// drains, so this call blocks on Lock forever and the teardown goroutine -// leaks. A leaked goroutine is the deliberate, accepted alternative to -// crashing the whole process with a send-on-closed-channel panic. +// drains until its own ctx is done, so this call blocks on Lock until then. A +// bounded wait is the deliberate, accepted alternative to crashing the whole +// process with a send-on-closed-channel panic; #3584 bounded the wait (send +// used to have no ctx at all and could block indefinitely) and moved the +// caller onto a detached goroutine so this can never stall the request path. func (b *elicitationBridge) restoreAndClose(current, previous chan Event) { b.mu.Lock() defer b.mu.Unlock() @@ -110,34 +142,306 @@ func (b *elicitationBridge) restoreAndClose(current, previous chan Event) { close(current) } +// waiterState is the terminal-state machine for a single elicitationWaiter. +// Exactly one of resolve/cancel ever wins the transition out of pending, +// closing the #3584 cancellation-vs-response race: a resolve that already +// flipped the state keeps its value in the channel for the ctx.Done() branch +// to drain, instead of the handler discarding a response ResumeElicitation +// already reported as delivered. +type waiterState int32 + +const ( + waiterPending waiterState = iota + waiterResolved + waiterCanceled +) + +// elicitationWaiter is one pending elicitation request's response slot. +type elicitationWaiter struct { + ch chan ElicitationResult + state atomic.Int32 +} + +func newElicitationWaiter() *elicitationWaiter { + return &elicitationWaiter{ch: make(chan ElicitationResult, 1)} +} + +// tryResolve attempts to deliver result, winning the terminal-state race +// only if the waiter is still pending. Returns false without sending when +// the waiter was already resolved or cancelled. +func (w *elicitationWaiter) tryResolve(result ElicitationResult) bool { + if !w.state.CompareAndSwap(int32(waiterPending), int32(waiterResolved)) { + return false + } + w.ch <- result + return true +} + +// tryCancel attempts to mark the waiter cancelled, winning the terminal-state +// race only if it is still pending. Returns false when resolve already won — +// the caller must then receive from ch instead of treating this as a +// cancellation, since a value is already there (or is about to land). +func (w *elicitationWaiter) tryCancel() bool { + return w.state.CompareAndSwap(int32(waiterPending), int32(waiterCanceled)) +} + +// elicitationWaiters routes an elicitation response to the specific request +// that is waiting for it, keyed by a correlation ID that is unique per +// request (see elicitationHandler). This replaces the single shared +// elicitationRequestCh, which could only ever have one request in flight: +// with concurrent (background-job) elicitations, a response arriving on +// that shared channel could be delivered to an arbitrary waiter, and +// ResumeElicitation had no way to tell "no request in flight" from "the +// request hasn't parked on the channel yet" (a TOCTOU race). +// +// Each waiter is registered BEFORE the corresponding request event is +// emitted, so a response that arrives immediately after — even before the +// handler reaches its receive — is never lost. The registry key is always an +// internally-generated ID (never the MCP wire ElicitationID, which two +// different MCP servers can coincidentally reuse): see elicitationHandler. +type elicitationWaiters struct { + mu sync.Mutex + pending map[string]*elicitationWaiter +} + +// register creates a waiter for id and stores it. The channel is buffered +// (capacity 1), so resolve never blocks even if the registrant hasn't +// reached its receive yet. +func (w *elicitationWaiters) register(id string) *elicitationWaiter { + wt := newElicitationWaiter() + w.mu.Lock() + defer w.mu.Unlock() + if w.pending == nil { + w.pending = make(map[string]*elicitationWaiter) + } + w.pending[id] = wt + return wt +} + +// abandon removes id's waiter from the registry, if it is still the one +// registered (defends against a hypothetical ID reuse racing a fresh +// register call), without touching its terminal state. Called once a waiter +// is done being awaited via any path, so a later resolve() for a reused ID +// cannot be confused with this one. +func (w *elicitationWaiters) abandon(id string, wt *elicitationWaiter) { + w.mu.Lock() + defer w.mu.Unlock() + if w.pending[id] == wt { + delete(w.pending, id) + } +} + +// cancel marks wt cancelled if it is still pending and removes it from the +// registry. Returns true when this call won the cancel-vs-resolve race — the +// caller (elicitationHandler's ctx.Done() branch) should then return ctx.Err(). +// Returns false when resolve() already won: the caller must receive from +// wt.ch instead, since a result is already there (or is about to land — the +// buffered send in tryResolve never blocks). +func (w *elicitationWaiters) cancel(id string, wt *elicitationWaiter) bool { + won := wt.tryCancel() + if won { + w.abandon(id, wt) + } + return won +} + +// resolve delivers result to the waiter registered for id and returns true. +// Returns false without side effects when no waiter is currently registered +// for that ID, or when it was already resolved/cancelled — already +// answered, timed out, or unknown. +func (w *elicitationWaiters) resolve(id string, result ElicitationResult) bool { + w.mu.Lock() + wt, ok := w.pending[id] + if ok { + delete(w.pending, id) + } + w.mu.Unlock() + if !ok { + return false + } + return wt.tryResolve(result) +} + +// resolveSingle delivers result to the sole pending waiter. It exists for +// backward compatibility with clients that don't send an elicitation_id +// (the pre-#3584 API contract only ever supported one request in flight). +// It is a deliberate no-op — returning false — when zero or more than one +// request is pending, since there is then no way to disambiguate which one +// the caller meant. +func (w *elicitationWaiters) resolveSingle(result ElicitationResult) bool { + w.mu.Lock() + if len(w.pending) != 1 { + w.mu.Unlock() + return false + } + var id string + var wt *elicitationWaiter + for k, v := range w.pending { + id, wt = k, v + } + delete(w.pending, id) + w.mu.Unlock() + return wt.tryResolve(result) +} + +// count reports the number of elicitations currently awaiting a response. +func (w *elicitationWaiters) count() int { + w.mu.Lock() + defer w.mu.Unlock() + return len(w.pending) +} + +// firstElicitationID extracts the (at most one meaningful) elicitation ID +// from a variadic parameter. The parameter is variadic — rather than a +// required positional string — purely so pre-#3584 callers of +// Runtime.ResumeElicitation (3 args) keep compiling unchanged; see +// docs/guides/go-sdk/index.md for the Go API compatibility note (this +// project's CHANGELOG.md is generated per release from merged PRs, not +// edited alongside the change, so it carries no such note pre-release). +func firstElicitationID(elicitationID []string) string { + if len(elicitationID) == 0 { + return "" + } + return elicitationID[0] +} + // ResumeElicitation sends an elicitation response back to a waiting -// elicitation request. Returns an error if no elicitation is in progress -// or if the context is cancelled before the response can be delivered. -func (r *LocalRuntime) ResumeElicitation(ctx context.Context, action tools.ElicitationAction, content map[string]any) error { - slog.DebugContext(ctx, "Resuming runtime with elicitation response", "agent", r.currentAgentName(), "action", action) +// elicitation request. When elicitationID is non-empty it is routed to that +// specific request; when empty, it falls back to resolving the sole pending +// request for backward compatibility with callers that predate per-request +// correlation. Returns an error if no matching elicitation is in progress. +// +// Unlike the old shared-channel implementation, this never blocks on ctx: +// each waiter channel is buffered (capacity 1) and registered before its +// request event is emitted, so resolving it is always a non-blocking map +// lookup plus a buffered send — there is no TOCTOU window to race. A +// resolve that raced a ctx-cancellation on the handler side is decided by +// an atomic terminal state (see elicitationWaiter), so a true here always +// means the response will reach the caller, never a discarded one. +func (r *LocalRuntime) ResumeElicitation(ctx context.Context, action tools.ElicitationAction, content map[string]any, elicitationID ...string) error { + id := firstElicitationID(elicitationID) + slog.DebugContext(ctx, "Resuming runtime with elicitation response", "agent", r.currentAgentName(), "action", action, "elicitation_id", id) result := ElicitationResult{ Action: action, Content: content, } - select { - case <-ctx.Done(): - slog.DebugContext(ctx, "Context cancelled while sending elicitation response") - return ctx.Err() - case r.elicitationRequestCh <- result: - slog.DebugContext(ctx, "Elicitation response sent successfully", "action", action) - return nil - default: - slog.DebugContext(ctx, "Elicitation channel not ready") - return errors.New("no elicitation request in progress") + var ok bool + if id != "" { + ok = r.elicitationWaiters.resolve(id, result) + } else { + ok = r.elicitationWaiters.resolveSingle(result) + } + if !ok { + slog.DebugContext(ctx, "No matching elicitation request in progress", "elicitation_id", id) + return errNoSuchElicitation + } + slog.DebugContext(ctx, "Elicitation response sent successfully", "action", action) + return nil +} + +// OnElicitationRequest registers a handler invoked whenever an MCP toolset +// raises an elicitation request. This is the reliable route for +// background-job elicitations (run_background_agent): their RunStream runs +// on a detached goroutine and can race concurrent streams for the bridge's +// single channel slot (#3584), so elicitationHandler calls this sink +// directly, synchronously, and unconditionally — before it ever touches the +// best-effort bridge — as the single, exactly-once delivery point. Embedders +// (e.g. the TUI's App, or the API server for session-scoped SSE delivery) +// register a handler here that forwards the event to their UI/transport. +func (r *LocalRuntime) OnElicitationRequest(handler func(Event)) { + r.elicitationSinkMu.Lock() + defer r.elicitationSinkMu.Unlock() + r.onElicitationRequest = handler +} + +// emitElicitationRequest forwards an elicitation request event to the +// registered sink, if any. Besides [LocalRuntime.EmitElicitationRequestForTesting], +// this is the ONLY call site that invokes the sink (see elicitationHandler): +// production callers must not add a second delivery path (e.g. re-forwarding +// an event observed on a RunStream channel), or the exactly-once guarantee +// this type documents no longer holds (#3584 item 5 — dual delivery +// previously required a stateful App-side dedupe to paper over). +func (r *LocalRuntime) emitElicitationRequest(event Event) { + r.elicitationSinkMu.RLock() + handler := r.onElicitationRequest + r.elicitationSinkMu.RUnlock() + if handler != nil { + handler(event) + } +} + +// EmitElicitationRequestForTesting invokes whatever OnElicitationRequest sink +// is currently registered, exactly as elicitationHandler would, but without +// the real MCP elicitation handshake. elicitationHandler is unexported, so +// callers outside this package (e.g. pkg/server) cannot drive it directly to +// prove a *specific* runtime instance has the expected sink wired; this +// gives them a seam to do that instead of reconstructing the sink separately +// and invoking it in isolation, which would pass even if the runtime under +// test was never actually wired up (#3584 re-review should-fix 1). +func (r *LocalRuntime) EmitElicitationRequestForTesting(event Event) { + r.emitElicitationRequest(event) +} + +// hasElicitationSink reports whether an embedder has registered an +// OnElicitationRequest handler. Used to decide whether a background +// session's elicitation has any chance of reaching a user (see +// elicitationHandler's headless fast-decline path). +func (r *LocalRuntime) hasElicitationSink() bool { + r.elicitationSinkMu.RLock() + defer r.elicitationSinkMu.RUnlock() + return r.onElicitationRequest != nil +} + +// elicitationDeclineNotes accumulates model-readable notes for elicitations +// that were auto-declined because a background session had no UI available +// to answer them (see elicitationHandler). runCollecting drains these after +// the sub-session completes and prepends them to the tool result, mirroring +// backgroundAuthRequiredNote's #3200 pattern for OAuth-at-Start failures. +type elicitationDeclineNotes struct { + mu sync.Mutex + bySession map[string][]string +} + +// record appends note under sessionID. No-op when either is empty. +func (n *elicitationDeclineNotes) record(sessionID, note string) { + if sessionID == "" || note == "" { + return } + n.mu.Lock() + defer n.mu.Unlock() + if n.bySession == nil { + n.bySession = make(map[string][]string) + } + n.bySession[sessionID] = append(n.bySession[sessionID], note) +} + +// drain returns and clears the notes recorded for sessionID. +func (n *elicitationDeclineNotes) drain(sessionID string) []string { + n.mu.Lock() + defer n.mu.Unlock() + notes := n.bySession[sessionID] + delete(n.bySession, sessionID) + return notes +} + +// backgroundElicitationDeclinedNote returns a model-readable explanation for +// an elicitation that was auto-declined because it originated from a +// background (non-interactive) session with no UI available to answer it. +func backgroundElicitationDeclinedNote(message string) string { + return fmt.Sprintf( + "Note: a tool requested user input (%q) while running as a background task. "+ + "Background tasks have no interactive UI to answer such requests, so it was "+ + "automatically declined. If the tool truly needs user input, ask the user to "+ + "run this task in the foreground instead.", + message, + ) } // elicitationHandler is the MCP-toolset-side hook that turns an inbound -// elicitation request from a server into an ElicitationRequest event on the -// active stream's events channel and waits for the embedder's response on -// elicitationRequestCh. +// elicitation request from a server into an ElicitationRequest event and +// waits for the embedder's response, correlated by elicitation ID. func (r *LocalRuntime) elicitationHandler(ctx context.Context, req *mcp.ElicitParams) (tools.ElicitationResult, error) { slog.DebugContext(ctx, "Elicitation request received from MCP server", "message", req.Message) @@ -150,30 +454,98 @@ func (r *LocalRuntime) elicitationHandler(ctx context.Context, req *mcp.ElicitPa }, nil } + // A background session (run_background_agent) marks its context so + // toolset Start() OAuth fails fast instead of eliciting (#3200). Mid-call + // elicitations reach here regardless of that marker, so extend the same + // fast-fail idea: if this call is running in such a context AND no + // embedder has registered a sink to surface it (headless use — e.g. the + // --exec CLI path, which never registers OnElicitationRequest), nobody at + // all can answer this request. Decline immediately with a model-readable + // note instead of parking a goroutine forever (#3584). + if !mcptools.InteractivePromptsAllowed(ctx) && !r.hasElicitationSink() { + slog.WarnContext(ctx, "Declining elicitation: background session has no UI to answer it", "message", req.Message) + r.elicitationDeclines.record(genai.ConversationIDFromContext(ctx), backgroundElicitationDeclinedNote(req.Message)) + return tools.ElicitationResult{ + Action: tools.ElicitationActionDecline, + }, nil + } + r.executeOnUserInputHooks(ctx, "", "elicitation") + // The registry key (and the ElicitationID surfaced to clients for + // ResumeElicitation routing) is always a freshly generated, internal + // ID — never the MCP wire req.ElicitationID. The wire value is only + // ever set for URL-mode elicitations and is chosen by the originating + // MCP server; two independent servers (e.g. two background jobs each + // talking to their own MCP process) can legitimately reuse the same + // value. Trusting it as the registry key would let the second + // request's register() silently evict the first request's waiter, + // orphaning it (#3584 review item 2a). The wire ID is preserved + // separately on the event (ServerElicitationID) for callers that want + // to correlate with server-side logs; it is never used for routing. + correlationID := uuid.NewString() + + // Register the waiter BEFORE emitting the request event. This is the + // #3584 TOCTOU fix: previously a response that arrived before the + // handler reached its receive on the shared channel was lost because + // there was nothing to receive it into yet. + wt := r.elicitationWaiters.register(correlationID) + defer r.elicitationWaiters.abandon(correlationID, wt) + slog.DebugContext(ctx, "Sending elicitation request event to client", "message", req.Message, "mode", req.Mode, "requested_schema", req.RequestedSchema, - "url", req.URL) + "url", req.URL, + "elicitation_id", correlationID, + "server_elicitation_id", req.ElicitationID) slog.DebugContext(ctx, "Elicitation request meta", "meta", req.Meta) - if err := r.elicitation.send( - ElicitationRequest(req.Message, req.Mode, req.RequestedSchema, req.URL, req.ElicitationID, req.Meta, r.currentAgentName()), - ); err != nil { - return tools.ElicitationResult{}, err - } + sessionID := genai.ConversationIDFromContext(ctx) + ev := ElicitationRequest(req.Message, req.Mode, req.RequestedSchema, req.URL, correlationID, req.ElicitationID, sessionID, req.Meta, r.currentAgentName()) + + // Reliable delivery: invoked synchronously, unconditionally, and exactly + // once, BEFORE anything that could block (#3584 review item 1). This + // must never be gated behind the best-effort bridge below. + r.emitElicitationRequest(ev) + + // Best-effort secondary delivery on the owning stream's events channel, + // kept for remote/SSE consumers that read directly off RunStream + // (remote_runtime.go depends on it). Dispatched on a detached goroutine, + // bounded by ctx, so a wedged or abandoned bridge channel (concurrent + // RunStreams racing the swap-based single slot, or a dead consumer) can + // never delay — let alone block — sink delivery or the response wait + // below (#3584 review item 1). runCollecting no longer treats a bridge + // delivery as a second source of truth (#3584 review item 5): this send + // exists solely for out-of-process consumers. + go func() { + if err := r.elicitation.send(ctx, ev); err != nil { + slog.DebugContext(ctx, "Elicitation bridge send failed or abandoned; relying on the registered sink", "error", err) + } + }() - // Wait for response from the client. + // Wait for the response addressed to this specific request. The + // ctx.Done() branch cannot simply return ctx.Err(): resolve() may have + // already won the terminal-state race an instant earlier and be about + // to (or have already) delivered into wt.ch, in which case + // ResumeElicitation already reported success to its caller and this + // handler must not silently discard that response (#3584 review item + // 2b). cancel() decides the winner atomically. select { - case result := <-r.elicitationRequestCh: + case result := <-wt.ch: return tools.ElicitationResult{ Action: result.Action, Content: result.Content, }, nil case <-ctx.Done(): slog.DebugContext(ctx, "Context cancelled while waiting for elicitation response") - return tools.ElicitationResult{}, ctx.Err() + if r.elicitationWaiters.cancel(correlationID, wt) { + return tools.ElicitationResult{}, ctx.Err() + } + result := <-wt.ch + return tools.ElicitationResult{ + Action: result.Action, + Content: result.Content, + }, nil } } diff --git a/pkg/runtime/elicitation_concurrency_test.go b/pkg/runtime/elicitation_concurrency_test.go new file mode 100644 index 0000000000..85c8292021 --- /dev/null +++ b/pkg/runtime/elicitation_concurrency_test.go @@ -0,0 +1,641 @@ +package runtime + +import ( + "context" + "fmt" + "maps" + "sync" + "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/chat" + "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" + mcptools "github.com/docker/docker-agent/pkg/tools/mcp" +) + +// --- elicitationWaiters: unit + concurrency regression tests (#3584) --- + +func TestElicitationWaiters_ResolveRoutesToCorrectID(t *testing.T) { + t.Parallel() + + var w elicitationWaiters + wtA := w.register("a") + wtB := w.register("b") + + require.True(t, w.resolve("b", ElicitationResult{Action: tools.ElicitationActionDecline})) + require.True(t, w.resolve("a", ElicitationResult{Action: tools.ElicitationActionAccept})) + + select { + case result := <-wtA.ch: + assert.Equal(t, tools.ElicitationActionAccept, result.Action, "waiter a must receive a's response, not b's") + default: + t.Fatal("waiter a never received its response") + } + select { + case result := <-wtB.ch: + assert.Equal(t, tools.ElicitationActionDecline, result.Action, "waiter b must receive b's response, not a's") + default: + t.Fatal("waiter b never received its response") + } +} + +// TestElicitationWaiters_ResolveBeforeReceiveIsNotLost pins the TOCTOU fix: +// registering the waiter before the request event is emitted means a +// response that arrives before anyone has read from the channel is still +// captured (the channel is buffered), instead of the old shared unbuffered +// elicitationRequestCh's `default:` branch reporting "no elicitation request +// in progress". See TestElicitationHandler_TOCTOU_ResolveImmediatelyAfterRegister +// below for the same pin exercised through the real handler path. +func TestElicitationWaiters_ResolveBeforeReceiveIsNotLost(t *testing.T) { + t.Parallel() + + var w elicitationWaiters + wt := w.register("id-1") + + ok := w.resolve("id-1", ElicitationResult{Action: tools.ElicitationActionAccept, Content: map[string]any{"k": "v"}}) + require.True(t, ok, "resolve must succeed even though nothing has received from the channel yet") + + select { + case result := <-wt.ch: + assert.Equal(t, tools.ElicitationActionAccept, result.Action) + assert.Equal(t, map[string]any{"k": "v"}, result.Content) + default: + t.Fatal("the buffered waiter channel should already hold the resolved result") + } +} + +func TestElicitationWaiters_ResolveUnknownIDReturnsFalse(t *testing.T) { + t.Parallel() + + var w elicitationWaiters + assert.False(t, w.resolve("missing", ElicitationResult{})) +} + +func TestElicitationWaiters_ResolveSingle_FallsBackForEmptyID(t *testing.T) { + t.Parallel() + + var w elicitationWaiters + wt := w.register("only-one") + + require.True(t, w.resolveSingle(ElicitationResult{Action: tools.ElicitationActionAccept})) + select { + case result := <-wt.ch: + assert.Equal(t, tools.ElicitationActionAccept, result.Action) + default: + t.Fatal("resolveSingle should have delivered to the sole pending waiter") + } +} + +// TestElicitationWaiters_ResolveSingle_AmbiguousWithMultiplePending verifies +// resolveSingle refuses to guess when more than one request is in flight — +// the whole point of per-ID correlation is that an empty-ID caller cannot +// safely disambiguate concurrent requests. +func TestElicitationWaiters_ResolveSingle_AmbiguousWithMultiplePending(t *testing.T) { + t.Parallel() + + var w elicitationWaiters + w.register("a") + w.register("b") + + assert.False(t, w.resolveSingle(ElicitationResult{Action: tools.ElicitationActionAccept})) + assert.Equal(t, 2, w.count(), "an ambiguous resolveSingle must not consume either waiter") +} + +func TestElicitationWaiters_Abandon(t *testing.T) { + t.Parallel() + + var w elicitationWaiters + wt := w.register("a") + require.Equal(t, 1, w.count()) + + w.abandon("a", wt) + assert.Equal(t, 0, w.count()) + assert.False(t, w.resolve("a", ElicitationResult{}), "abandoned waiter must not be resolvable") +} + +// TestElicitationWaiters_DuplicateWireIDsDoNotCollide pins #3584 review item +// 2a: registry keys are always internally-generated IDs, never the MCP wire +// ElicitationID, precisely because two independent MCP servers (e.g. two +// concurrent background jobs, each talking to their own server) can +// legitimately reuse the same wire ID. Registering two waiters under +// different (internal) IDs — even when both requests logically share one +// wire ID — must never let the second registration evict the first's +// channel. +func TestElicitationWaiters_DuplicateWireIDsDoNotCollide(t *testing.T) { + t.Parallel() + + var w elicitationWaiters + // Simulates two servers both using wire ID "1": elicitationHandler + // always mints a fresh internal correlation ID regardless, so the two + // registrations land under distinct keys. + wtServerA := w.register("internal-a") + wtServerB := w.register("internal-b") + + require.Equal(t, 2, w.count(), "distinct internal IDs must not collide even if the wire IDs would have") + + require.True(t, w.resolve("internal-a", ElicitationResult{Action: tools.ElicitationActionAccept, Content: map[string]any{"who": "a"}})) + require.True(t, w.resolve("internal-b", ElicitationResult{Action: tools.ElicitationActionAccept, Content: map[string]any{"who": "b"}})) + + select { + case result := <-wtServerA.ch: + assert.Equal(t, map[string]any{"who": "a"}, result.Content, "server A's waiter must not have been evicted or overwritten") + default: + t.Fatal("server A's waiter never received its response") + } + select { + case result := <-wtServerB.ch: + assert.Equal(t, map[string]any{"who": "b"}, result.Content, "server B's waiter must not have been evicted or overwritten") + default: + t.Fatal("server B's waiter never received its response") + } +} + +// TestElicitationWaiter_CancelWinsWhenFirst pins the #3584 review item 2b +// cancellation-vs-response race: when the handler's ctx.Done() branch wins +// the terminal-state CAS first, a subsequent resolve() must be told it lost +// (return false) instead of silently succeeding into a channel nobody will +// ever read again. +func TestElicitationWaiter_CancelWinsWhenFirst(t *testing.T) { + t.Parallel() + + var w elicitationWaiters + wt := w.register("a") + + require.True(t, w.cancel("a", wt), "cancel must win when nothing resolved yet") + assert.False(t, w.resolve("a", ElicitationResult{Action: tools.ElicitationActionAccept}), + "a resolve racing after cancel already won must not report success") + assert.Equal(t, 0, w.count(), "cancel must remove the waiter from the registry") +} + +// TestElicitationWaiter_ResolveWinsWhenFirst is the mirror image: resolve() +// wins the race first, so the handler's cancel() call must lose and report +// false, telling the handler to drain the value from the channel instead of +// discarding a response ResumeElicitation already reported as delivered. +func TestElicitationWaiter_ResolveWinsWhenFirst(t *testing.T) { + t.Parallel() + + var w elicitationWaiters + wt := w.register("a") + + require.True(t, w.resolve("a", ElicitationResult{Action: tools.ElicitationActionAccept, Content: map[string]any{"k": "v"}})) + assert.False(t, w.cancel("a", wt), "cancel must lose once resolve already won the terminal-state race") + + select { + case result := <-wt.ch: + assert.Equal(t, map[string]any{"k": "v"}, result.Content, "the resolved value must still be retrievable by the loser of the race") + default: + t.Fatal("resolve's value must be in the channel even though cancel lost the race") + } +} + +// TestElicitationWaiters_ConcurrentRegisterResolveDeregister runs many +// concurrent request/response pairs through the registry to catch data races +// (run with -race) and confirm no response is ever misdelivered under +// concurrent load, mirroring the concurrent background-job scenario from +// the audit. +func TestElicitationWaiters_ConcurrentRegisterResolveDeregister(t *testing.T) { + t.Parallel() + + var w elicitationWaiters + const n = 200 + + var wg sync.WaitGroup + for i := range n { + wg.Go(func() { + id := fmt.Sprintf("req-%d", i) + wt := w.register(id) + defer w.abandon(id, wt) + + done := make(chan struct{}) + go func() { + defer close(done) + ok := w.resolve(id, ElicitationResult{Action: tools.ElicitationActionAccept, Content: map[string]any{"i": i}}) + assert.True(t, ok) + }() + + select { + case result := <-wt.ch: + assert.Equal(t, map[string]any{"i": i}, result.Content, "response must route back to its own request") + case <-time.After(2 * time.Second): + t.Errorf("waiter %s never received its response", id) + } + <-done + }) + } + wg.Wait() +} + +// TestElicitationWaiters_ConcurrentResolveCancelRace hammers a single waiter +// with concurrent resolve/cancel attempts (run with -race) to confirm the +// terminal-state CAS lets exactly one of them win, never both and never +// neither. +func TestElicitationWaiters_ConcurrentResolveCancelRace(t *testing.T) { + t.Parallel() + + const n = 300 + for i := range n { + var w elicitationWaiters + id := fmt.Sprintf("race-%d", i) + wt := w.register(id) + + var wg sync.WaitGroup + var resolveWon, cancelWon atomicBool + wg.Add(2) + go func() { + defer wg.Done() + if w.resolve(id, ElicitationResult{Action: tools.ElicitationActionAccept}) { + resolveWon.set(true) + } + }() + go func() { + defer wg.Done() + if w.cancel(id, wt) { + cancelWon.set(true) + } + }() + wg.Wait() + + require.NotEqual(t, resolveWon.get(), cancelWon.get(), "exactly one of resolve/cancel must win, never both or neither") + } +} + +// atomicBool is a tiny test-local helper; sync/atomic.Bool is available but +// spelling out set/get keeps the race-loop above terse. +type atomicBool struct { + mu sync.Mutex + v bool +} + +func (a *atomicBool) set(v bool) { + a.mu.Lock() + defer a.mu.Unlock() + a.v = v +} + +func (a *atomicBool) get() bool { + a.mu.Lock() + defer a.mu.Unlock() + return a.v +} + +// --- elicitationBridge: bounded/non-blocking send (#3584 review item 1) --- + +// TestElicitationBridge_SendBlocksUntilCtxDone pins the review-item-1 fix: an +// unbuffered, unconsumed bridge channel used to block send() forever with no +// way out. send() must now be bounded by ctx and release with ctx.Err() +// instead of hanging, and must never panic. +func TestElicitationBridge_SendBlocksUntilCtxDone(t *testing.T) { + t.Parallel() + + var b elicitationBridge + ch := make(chan Event) // unbuffered, nobody ever reads it + b.swap(ch) + + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + err := b.send(ctx, Warning("hello", "agent")) + elapsed := time.Since(start) + + require.ErrorIs(t, err, context.DeadlineExceeded, "send on a full/abandoned channel must release via ctx, not block forever") + assert.Less(t, elapsed, 2*time.Second, "send must not block substantially past the ctx deadline") +} + +// TestElicitationBridge_SendNeverBlocksReliableSink is the end-to-end version +// of the review-item-1 fix: a wedged bridge channel must not delay — let +// alone block — elicitationHandler's reliable OnElicitationRequest sink +// delivery or its subsequent wait for a response. Before the fix, the bridge +// send was awaited synchronously and BEFORE the sink call, so an abandoned +// channel meant the sink (and thus the whole request) never even started. +func TestElicitationBridge_SendNeverBlocksReliableSink(t *testing.T) { + t.Parallel() + + rt := newElicitationTestRuntime(t) + + // Wedge the bridge: swap in an unbuffered channel with no reader, as if + // a concurrent RunStream's swap left a dead consumer behind. + wedged := make(chan Event) + rt.elicitation.swap(wedged) + + sinkCalled := make(chan Event, 1) + rt.OnElicitationRequest(func(ev Event) { sinkCalled <- ev }) + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + type handlerResult struct { + result tools.ElicitationResult + err error + } + done := make(chan handlerResult, 1) + go func() { + result, err := rt.elicitationHandler(ctx, &mcp.ElicitParams{Message: "confirm?"}) + done <- handlerResult{result, err} + }() + + // The sink must fire almost immediately, regardless of the wedged bridge. + var ev *ElicitationRequestEvent + select { + case e := <-sinkCalled: + ev = e.(*ElicitationRequestEvent) + case <-time.After(1 * time.Second): + t.Fatal("the reliable sink must not be blocked by a wedged bridge channel") + } + + require.NoError(t, rt.ResumeElicitation(t.Context(), tools.ElicitationActionAccept, nil, ev.ElicitationID)) + + select { + case got := <-done: + require.NoError(t, got.err) + assert.Equal(t, tools.ElicitationActionAccept, got.result.Action) + case <-time.After(1 * time.Second): + t.Fatal("elicitationHandler must not be blocked by a wedged bridge channel") + } +} + +// --- elicitationHandler: headless fast-decline (#3584 item 5) --- + +func TestElicitationHandler_HeadlessBackgroundFastDeclines(t *testing.T) { + t.Parallel() + + rt := newElicitationTestRuntime(t) + + // No OnElicitationRequest sink registered, and the context is marked + // non-interactive the way runStreamLoop marks a background + // (run_background_agent) session's context (#3200). + ctx := mcptools.WithoutInteractivePrompts(t.Context()) + ctx = genai.WithConversationID(ctx, "bg-sess-1") + + result, err := rt.elicitationHandler(ctx, &mcp.ElicitParams{Message: "need sudo password"}) + require.NoError(t, err) + assert.Equal(t, tools.ElicitationActionDecline, result.Action, + "a background session with no UI sink must fast-decline instead of blocking forever") + + notes := rt.elicitationDeclines.drain("bg-sess-1") + require.Len(t, notes, 1) + assert.Contains(t, notes[0], "need sudo password") +} + +// TestElicitationHandler_BackgroundWithSinkStillWaitsForResponse verifies +// that registering an OnElicitationRequest sink is enough to opt a +// background-session elicitation back into the normal wait-for-response +// path instead of being fast-declined. +func TestElicitationHandler_BackgroundWithSinkStillWaitsForResponse(t *testing.T) { + t.Parallel() + + rt := newElicitationTestRuntime(t) + + received := make(chan Event, 1) + rt.OnElicitationRequest(func(ev Event) { received <- ev }) + + ctx := mcptools.WithoutInteractivePrompts(t.Context()) + ctx = genai.WithConversationID(ctx, "bg-sess-2") + + type handlerResult struct { + result tools.ElicitationResult + err error + } + done := make(chan handlerResult, 1) + go func() { + result, err := rt.elicitationHandler(ctx, &mcp.ElicitParams{Message: "confirm?"}) + done <- handlerResult{result, err} + }() + + var ev *ElicitationRequestEvent + select { + case e := <-received: + ev = e.(*ElicitationRequestEvent) + case <-time.After(2 * time.Second): + t.Fatal("sink never received the elicitation request") + } + assert.Equal(t, "bg-sess-2", ev.SessionID, "the event must carry the originating (sub-)session ID") + + require.NoError(t, rt.ResumeElicitation(t.Context(), tools.ElicitationActionAccept, nil, ev.ElicitationID)) + + select { + case got := <-done: + require.NoError(t, got.err) + assert.Equal(t, tools.ElicitationActionAccept, got.result.Action) + case <-time.After(2 * time.Second): + t.Fatal("elicitationHandler never returned") + } + assert.Empty(t, rt.elicitationDeclines.drain("bg-sess-2"), "must not fast-decline once a sink is registered") +} + +// TestElicitationHandler_TOCTOU_ResolveImmediatelyAfterRegister promotes the +// former helper-level TOCTOU pin (TestElicitationWaiters_ResolveBeforeReceiveIsNotLost) +// to exercise the real elicitationHandler code path end-to-end: the sink +// fires synchronously before the handler ever reaches its response select, +// so a caller that resolves the instant it observes the sink-delivered event +// — beating the handler to its `select` — must still have its response +// delivered rather than racing the old shared-channel `default:` branch. +func TestElicitationHandler_TOCTOU_ResolveImmediatelyAfterRegister(t *testing.T) { + t.Parallel() + + rt := newElicitationTestRuntime(t) + + sinkDone := make(chan struct{}) + rt.OnElicitationRequest(func(ev Event) { + // Resolve from inside the sink callback itself, i.e. before + // elicitationHandler's goroutine has any chance to reach its + // `select` on the waiter channel. This is the tightest possible + // version of the TOCTOU window. + e := ev.(*ElicitationRequestEvent) + ok := rt.elicitationWaiters.resolve(e.ElicitationID, ElicitationResult{Action: tools.ElicitationActionAccept, Content: map[string]any{"answered": true}}) + assert.True(t, ok, "resolve issued synchronously from the sink callback must still find the just-registered waiter") + close(sinkDone) + }) + + result, err := rt.elicitationHandler(t.Context(), &mcp.ElicitParams{Message: "confirm?"}) + require.NoError(t, err) + assert.Equal(t, tools.ElicitationActionAccept, result.Action) + assert.Equal(t, map[string]any{"answered": true}, result.Content) + <-sinkDone +} + +// --- Full-stack regression: concurrent background-job elicitations (#3584) --- + +// elicitingToolSet is a minimal toolset whose one tool blocks on an MCP +// elicitation via the handler ConfigureHandlers wires onto it — mirroring +// how a real MCP toolset elicits mid-tool-call — without needing a real MCP +// server. Used to reproduce the reported bug: elicitations raised from +// multiple concurrent background jobs (run_background_agent) must all +// surface and each response must reach its own waiter. +type elicitingToolSet struct { + mu sync.Mutex + handler tools.ElicitationHandler + message string + // received captures the ElicitationResult the tool call got back, so + // tests can verify the response that reached this specific worker + // without having to scrape it back out of the model transcript. + received tools.ElicitationResult + gotResult bool +} + +var _ tools.Elicitable = (*elicitingToolSet)(nil) + +func (e *elicitingToolSet) SetElicitationHandler(handler tools.ElicitationHandler) { + e.mu.Lock() + defer e.mu.Unlock() + e.handler = handler +} + +func (e *elicitingToolSet) snapshot() (tools.ElicitationResult, bool) { + e.mu.Lock() + defer e.mu.Unlock() + return e.received, e.gotResult +} + +func (e *elicitingToolSet) Tools(context.Context) ([]tools.Tool, error) { + return []tools.Tool{{ + Name: "ask_user", + Handler: func(ctx context.Context, _ tools.ToolCall, _ tools.Runtime) (*tools.ToolCallResult, error) { + e.mu.Lock() + handler := e.handler + e.mu.Unlock() + if handler == nil { + return tools.ResultError("no elicitation handler configured"), nil + } + result, err := handler(ctx, &mcp.ElicitParams{Message: e.message}) + if err != nil { + return nil, err + } + e.mu.Lock() + e.received = result + e.gotResult = true + e.mu.Unlock() + return tools.ResultSuccess(fmt.Sprintf("%s:%v", result.Action, result.Content)), nil + }, + }}, nil +} + +// TestConcurrentBackgroundElicitations_AllSurfaceAndRouteToCorrectWaiter is +// the end-to-end regression test for issue #3584 / audit finding A6: two +// concurrent background jobs (as run_background_agent spawns them) each +// raise an elicitation mid-tool-call. Before the fix, the single-slot +// elicitationBridge and shared elicitationRequestCh meant: (a) runCollecting +// silently dropped the ElicitationRequestEvent, so nothing was ever +// displayed, and (b) even if something had displayed it, a response could be +// delivered to the wrong waiter. This asserts both jobs surface via the new +// OnElicitationRequest sink — exactly once each, with no dedupe map required +// (runCollecting no longer re-forwards a bridge-observed copy; see #3584 +// review item 5) — and each receives its own, non-swapped response. +func TestConcurrentBackgroundElicitations_AllSurfaceAndRouteToCorrectWaiter(t *testing.T) { + t.Parallel() + + newWorker := func(name, question string) (*agent.Agent, *elicitingToolSet) { + ts := &elicitingToolSet{message: question} + toolCallStream := newStreamBuilder().AddToolCallWithStop("call_1", "ask_user", "{}").Build() + followUpStream := newStreamBuilder().AddContent("done").AddStopWithUsage(5, 5).Build() + prov := &queueProvider{id: "test/mock-model", streams: []chat.MessageStream{toolCallStream, followUpStream}} + a := agent.New(name, "worker", agent.WithModel(prov), agent.WithToolSets(ts)) + return a, ts + } + + worker1, ts1 := newWorker("worker1", "worker1 needs input") + worker2, ts2 := newWorker("worker2", "worker2 needs input") + root := agent.New("root", "root", agent.WithModel(&mockProvider{id: "test/mock-model", stream: &mockStream{}})) + agent.WithSubAgents(worker1, worker2)(root) + + tm := team.New(team.WithAgents(root, worker1, worker2)) + rt, err := NewLocalRuntime(t.Context(), tm, WithSessionCompaction(false), WithModelStore(mockModelStore{})) + require.NoError(t, err) + + // deliveries counts sink invocations per ElicitationID. Each ID must be + // delivered exactly once: any count > 1 means the exactly-once guarantee + // (elicitationHandler is the sole caller of emitElicitationRequest) has + // regressed, since nothing else in the App/runtime layer masks + // duplicates any more. + var mu sync.Mutex + deliveries := make(map[string]int) + requests := make(map[string]*ElicitationRequestEvent) + rt.OnElicitationRequest(func(ev Event) { + req := ev.(*ElicitationRequestEvent) + mu.Lock() + defer mu.Unlock() + deliveries[req.ElicitationID]++ + requests[req.ElicitationID] = req + }) + + results := make(chan *agenttool.RunResult, 2) + launch := func(agentName string) { + go func() { + parent := session.New(session.WithUserMessage("go"), session.WithToolsApproved(true)) + res := rt.RunAgent(t.Context(), agenttool.RunParams{ + AgentName: agentName, + Task: "do it", + ParentSession: parent, + }) + results <- res + }() + } + launch("worker1") + launch("worker2") + + // Wait until both elicitations have surfaced, then respond to each by + // its own ID with a distinguishable payload so a swapped response would + // be caught by the assertions below. + var reqs map[string]*ElicitationRequestEvent + require.Eventually(t, func() bool { + mu.Lock() + defer mu.Unlock() + reqs = make(map[string]*ElicitationRequestEvent, len(requests)) + maps.Copy(reqs, requests) + return len(reqs) == 2 + }, 5*time.Second, 10*time.Millisecond, "both concurrent background elicitations must surface via the sink") + + mu.Lock() + for id, count := range deliveries { + assert.Equal(t, 1, count, "elicitation %s must be delivered to the sink exactly once", id) + } + mu.Unlock() + + var worker1ID, worker2ID string + for id, ev := range reqs { + switch ev.Message { + case "worker1 needs input": + worker1ID = id + case "worker2 needs input": + worker2ID = id + } + } + require.NotEmpty(t, worker1ID, "worker1's elicitation must have surfaced") + require.NotEmpty(t, worker2ID, "worker2's elicitation must have surfaced") + require.NotEqual(t, worker1ID, worker2ID, "concurrent elicitations must get distinct correlation IDs") + + require.NoError(t, rt.ResumeElicitation(t.Context(), tools.ElicitationActionAccept, map[string]any{"answer": "1"}, worker1ID)) + require.NoError(t, rt.ResumeElicitation(t.Context(), tools.ElicitationActionAccept, map[string]any{"answer": "2"}, worker2ID)) + + var got []*agenttool.RunResult + for range 2 { + select { + case r := <-results: + got = append(got, r) + case <-time.After(5 * time.Second): + t.Fatal("a background job never completed after its elicitation was resumed") + } + } + for _, r := range got { + require.Empty(t, r.ErrMsg, "background job must not fail") + } + + worker1Result, ok1 := ts1.snapshot() + require.True(t, ok1, "worker1's tool call must have received an elicitation result") + worker2Result, ok2 := ts2.snapshot() + require.True(t, ok2, "worker2's tool call must have received an elicitation result") + + assert.Equal(t, map[string]any{"answer": "1"}, worker1Result.Content, + "worker1 must receive its own response, not worker2's (no swap)") + assert.Equal(t, map[string]any{"answer": "2"}, worker2Result.Content, + "worker2 must receive its own response, not worker1's (no swap)") +} diff --git a/pkg/runtime/elicitation_test.go b/pkg/runtime/elicitation_test.go index 11ab050d9a..17d7d4a195 100644 --- a/pkg/runtime/elicitation_test.go +++ b/pkg/runtime/elicitation_test.go @@ -27,7 +27,7 @@ func TestElicitationBridge_SendBeforeSwapReturnsError(t *testing.T) { t.Parallel() var b elicitationBridge - err := b.send(Error("nothing")) + err := b.send(t.Context(), Error("nothing")) assert.ErrorIs(t, err, errNoElicitationChannel) } @@ -55,7 +55,7 @@ func TestElicitationBridge_SendDeliversToCurrentChannel(t *testing.T) { ch := make(chan Event, 1) b.swap(ch) - require.NoError(t, b.send(Error("hello"))) + require.NoError(t, b.send(t.Context(), Error("hello"))) select { case ev := <-ch: @@ -75,7 +75,7 @@ func TestElicitationBridge_SendRecoversClosedChannel(t *testing.T) { b.swap(ch) close(ch) - err := b.send(Error("closed")) + err := b.send(t.Context(), Error("closed")) assert.ErrorIs(t, err, errNoElicitationChannel) } @@ -97,7 +97,7 @@ func TestElicitationBridge_RestoreAndCloseWaitsForInflightSenders(t *testing.T) sendDone := make(chan error, 1) go func() { - sendDone <- b.send(Error("inflight")) + sendDone <- b.send(t.Context(), Error("inflight")) }() // Wait until the sender holds the read lock: TryLock fails only while @@ -167,7 +167,7 @@ func TestElicitationBridge_ConcurrentSendsAndCloseAreSerializedSafely(t *testing for range 10 { wg.Go(func() { for range 5 { - _ = b.send(Error("x")) + _ = b.send(t.Context(), Error("x")) } }) } diff --git a/pkg/runtime/event.go b/pkg/runtime/event.go index 86cd286b2b..f5c3a257ff 100644 --- a/pkg/runtime/event.go +++ b/pkg/runtime/event.go @@ -583,28 +583,52 @@ func (e *PausedEvent) GetSessionID() string { return e.SessionID } type ElicitationRequestEvent struct { AgentContext - Type string `json:"type"` - Message string `json:"message"` - Mode string `json:"mode,omitempty"` // "form" or "url" - Schema any `json:"schema,omitempty"` - URL string `json:"url,omitempty"` - ElicitationID string `json:"elicitation_id,omitempty"` - Meta map[string]any `json:"meta,omitempty"` -} - -func ElicitationRequest(message, mode string, schema any, url, elicitationID string, meta map[string]any, agentName string) Event { + Type string `json:"type"` + Message string `json:"message"` + Mode string `json:"mode,omitempty"` // "form" or "url" + Schema any `json:"schema,omitempty"` + URL string `json:"url,omitempty"` + // ElicitationID is the internally-generated correlation ID used to route + // a ResumeElicitation response back to this specific request (see + // elicitationHandler). It is always set and always unique, regardless of + // whether the originating MCP server supplied its own wire ID: two + // independent servers can coincidentally reuse the same wire ID, so that + // value is never used for routing (#3584). + ElicitationID string `json:"elicitation_id,omitempty"` + // ServerElicitationID is the MCP wire-protocol elicitationId as supplied + // by the originating server, if any (URL-mode elicitations only; form + // elicitations leave it empty). It is informational only — useful for + // correlating with server-side logs — and must never be used as a + // routing key. + ServerElicitationID string `json:"server_elicitation_id,omitempty"` + // SessionID is the session (or sub-session) on whose behalf this + // elicitation was raised. A detached background job's sub-session ID + // differs from its parent/foreground session, which lets consumers (see + // the TUI supervisor) tell apart a foreground stream's own elicitation + // from a still-live background job's when the foreground stream stops + // (#3584 review item 4). + SessionID string `json:"session_id,omitempty"` + Meta map[string]any `json:"meta,omitempty"` +} + +func ElicitationRequest(message, mode string, schema any, url, elicitationID, serverElicitationID, sessionID string, meta map[string]any, agentName string) Event { return &ElicitationRequestEvent{ - Type: "elicitation_request", - Message: message, - Mode: mode, - Schema: schema, - URL: url, - ElicitationID: elicitationID, - Meta: meta, - AgentContext: newAgentContext(agentName), + Type: "elicitation_request", + Message: message, + Mode: mode, + Schema: schema, + URL: url, + ElicitationID: elicitationID, + ServerElicitationID: serverElicitationID, + SessionID: sessionID, + Meta: meta, + AgentContext: newAgentContext(agentName), } } +// GetSessionID makes ElicitationRequestEvent satisfy [SessionScoped]. +func (e *ElicitationRequestEvent) GetSessionID() string { return e.SessionID } + type AuthorizationEvent struct { AgentContext diff --git a/pkg/runtime/remote_client.go b/pkg/runtime/remote_client.go index 2029869d2e..3cff900765 100644 --- a/pkg/runtime/remote_client.go +++ b/pkg/runtime/remote_client.go @@ -22,8 +22,11 @@ type RemoteClient interface { // ResumeSession resumes a paused session with optional rejection reason or tool name ResumeSession(ctx context.Context, id, confirmation, reason, toolName string) error - // ResumeElicitation sends an elicitation response - ResumeElicitation(ctx context.Context, sessionID string, action tools.ElicitationAction, content map[string]any) error + // ResumeElicitation sends an elicitation response. elicitationID is + // additive: pass "" to fall back to the sole-pending-request behavior. + // Variadic for the same Go-API-compatibility reason as + // [Runtime.ResumeElicitation]: at most one value is meaningful. + ResumeElicitation(ctx context.Context, sessionID string, action tools.ElicitationAction, content map[string]any, elicitationID ...string) error // RunAgent executes an agent and returns a channel of streaming events. // model, when non-empty, is applied as a persistent override on the diff --git a/pkg/runtime/remote_contract_test.go b/pkg/runtime/remote_contract_test.go index 2c2fbb1b14..ec055db91b 100644 --- a/pkg/runtime/remote_contract_test.go +++ b/pkg/runtime/remote_contract_test.go @@ -41,7 +41,7 @@ func (s *stubRemoteClient) ResumeSession(context.Context, string, string, string return nil } -func (s *stubRemoteClient) ResumeElicitation(context.Context, string, tools.ElicitationAction, map[string]any) error { +func (s *stubRemoteClient) ResumeElicitation(context.Context, string, tools.ElicitationAction, map[string]any, ...string) error { return nil } diff --git a/pkg/runtime/remote_runtime.go b/pkg/runtime/remote_runtime.go index b849005e23..acc337d31f 100644 --- a/pkg/runtime/remote_runtime.go +++ b/pkg/runtime/remote_runtime.go @@ -394,15 +394,16 @@ func (r *RemoteRuntime) convertSessionMessages(sess *session.Session) []api.Mess } // ResumeElicitation sends an elicitation response back to a waiting elicitation request -func (r *RemoteRuntime) ResumeElicitation(ctx context.Context, action tools.ElicitationAction, content map[string]any) error { - slog.DebugContext(ctx, "Resuming remote runtime with elicitation response", "agent", r.currentAgent, "action", action, "session_id", r.sessionID) +func (r *RemoteRuntime) ResumeElicitation(ctx context.Context, action tools.ElicitationAction, content map[string]any, elicitationID ...string) error { + id := firstElicitationID(elicitationID) + slog.DebugContext(ctx, "Resuming remote runtime with elicitation response", "agent", r.currentAgent, "action", action, "session_id", r.sessionID, "elicitation_id", id) err := r.handleOAuthElicitation(ctx, r.pendingOAuthElicitation) if err != nil { return err } - if err := r.client.ResumeElicitation(ctx, r.sessionID, action, content); err != nil { + if err := r.client.ResumeElicitation(ctx, r.sessionID, action, content, id); err != nil { return err } @@ -420,7 +421,7 @@ func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, req *Elicita if !ok { err := errors.New("server_url missing from elicitation metadata") slog.ErrorContext(ctx, "Failed to extract server_url", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return err } @@ -428,7 +429,7 @@ func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, req *Elicita if !ok { err := errors.New("auth_server_metadata missing from elicitation metadata") slog.ErrorContext(ctx, "Failed to extract auth_server_metadata", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return err } @@ -436,12 +437,12 @@ func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, req *Elicita metadataBytes, err := json.Marshal(authServerMetadata) if err != nil { slog.ErrorContext(ctx, "Failed to marshal auth_server_metadata", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return fmt.Errorf("failed to marshal auth_server_metadata: %w", err) } if err := json.Unmarshal(metadataBytes, &authMetadata); err != nil { slog.ErrorContext(ctx, "Failed to unmarshal auth_server_metadata", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return fmt.Errorf("failed to unmarshal auth_server_metadata: %w", err) } @@ -461,7 +462,7 @@ func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, req *Elicita callbackServer, err := mcp.NewCallbackServer(ctx) if err != nil { slog.ErrorContext(ctx, "Failed to create callback server", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return fmt.Errorf("failed to create callback server: %w", err) } defer func() { @@ -476,7 +477,7 @@ func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, req *Elicita if err := callbackServer.Start(); err != nil { slog.ErrorContext(ctx, "Failed to start callback server", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return fmt.Errorf("failed to start callback server: %w", err) } @@ -489,21 +490,21 @@ func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, req *Elicita clientID, clientSecret, err = mcp.RegisterClient(oauthCtx, &authMetadata, redirectURI, nil) if err != nil { slog.ErrorContext(ctx, "Dynamic client registration failed", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return fmt.Errorf("failed to register client: %w", err) } slog.DebugContext(ctx, "Client registered successfully", "client_id", clientID) } else { err := errors.New("authorization server does not support dynamic client registration") slog.ErrorContext(ctx, "Client registration not supported", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return err } state, err := mcp.GenerateState() if err != nil { slog.ErrorContext(ctx, "Failed to generate state", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return fmt.Errorf("failed to generate state: %w", err) } @@ -526,14 +527,14 @@ func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, req *Elicita code, receivedState, err := mcp.RequestAuthorizationCode(oauthCtx, authURL, callbackServer, state) if err != nil { slog.ErrorContext(ctx, "Failed to get authorization code", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return fmt.Errorf("failed to get authorization code: %w", err) } if receivedState != state { err := fmt.Errorf("state mismatch: expected %s, got %s", state, receivedState) slog.ErrorContext(ctx, "State mismatch in authorization response", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return err } @@ -551,7 +552,7 @@ func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, req *Elicita ) if err != nil { slog.ErrorContext(ctx, "Failed to exchange code for token", "error", err) - _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil) + _ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil, req.ElicitationID) return fmt.Errorf("failed to exchange code for token: %w", err) } @@ -569,7 +570,7 @@ func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, req *Elicita } slog.DebugContext(ctx, "Sending token to server") - if err := r.client.ResumeElicitation(ctx, r.sessionID, tools.ElicitationActionAccept, tokenData); err != nil { + if err := r.client.ResumeElicitation(ctx, r.sessionID, tools.ElicitationActionAccept, tokenData, req.ElicitationID); err != nil { slog.ErrorContext(ctx, "Failed to send token to server", "error", err) return fmt.Errorf("failed to send token to server: %w", err) } @@ -706,6 +707,12 @@ func (r *RemoteRuntime) OnToolsChanged(func(Event)) {} // run server-side and their events are not forwarded out-of-band. func (r *RemoteRuntime) OnBackgroundEvent(func(Event)) {} +// OnElicitationRequest is a no-op for remote runtimes; elicitation requests +// (including from server-side background jobs) arrive as ElicitationRequestEvent +// values on the RunStream channel itself (see the RunStream forwarding loop +// above), so there is no separate out-of-band sink to register. +func (r *RemoteRuntime) OnElicitationRequest(func(Event)) {} + // Close is a no-op for remote runtimes. func (r *RemoteRuntime) Close() error { return nil diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index 8a68df5da9..d443e9a741 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -82,8 +82,16 @@ type Runtime interface { // Resume allows resuming execution after user confirmation. // The ResumeRequest carries the decision type and an optional reason (for rejections). Resume(ctx context.Context, req ResumeRequest) - // ResumeElicitation sends an elicitation response back to a waiting elicitation request - ResumeElicitation(_ context.Context, action tools.ElicitationAction, content map[string]any) error + // ResumeElicitation sends an elicitation response back to a waiting + // elicitation request. elicitationID correlates the response with a + // specific concurrent request (see ElicitationRequestEvent.ElicitationID); + // omit it (or pass "") to fall back to resolving the sole pending + // request, for backward compatibility with older clients. The parameter + // is variadic — rather than a required 4th positional argument — purely + // so pre-#3584 callers of this interface method keep compiling unchanged; + // implementations should treat more than one value as a caller error and + // use only the first. + ResumeElicitation(_ context.Context, action tools.ElicitationAction, content map[string]any, elicitationID ...string) error // SessionStore returns the session store for browsing/loading past sessions. // Returns nil if no persistent session store is configured. SessionStore() session.Store @@ -176,6 +184,14 @@ type Runtime interface { // background work can implement this as a no-op. OnBackgroundEvent(handler func(Event)) + // OnElicitationRequest registers a handler invoked whenever an MCP + // toolset raises an elicitation request, regardless of which stream + // (foreground or a detached background job) is currently "active" for + // the runtime's internal elicitation bridge. This is the reliable + // delivery route for background-job elicitations (#3584); runtimes + // that don't run local background work can implement this as a no-op. + OnElicitationRequest(handler func(Event)) + // QueueStatus returns the current depth and capacity of message queues QueueStatus() QueueStatus @@ -224,16 +240,23 @@ type LocalRuntime struct { managedOAuth bool unmanagedOAuthRedirectURI string nonInteractive bool - startupInfoEmitted atomic.Bool // Track if startup info has been emitted to avoid unnecessary duplication - elicitationRequestCh chan ElicitationResult // Channel for receiving elicitation responses - elicitation elicitationBridge // Owns the per-stream events channel for outbound elicitation requests - sessionStore session.Store - workingDir string // Working directory for hooks execution - env []string // Environment variables for hooks execution - modelSwitcherCfg *ModelSwitcherConfig - providerRegistry *provider.Registry - gatewayModels gatewayModelsCache - dmrModels dmrModelsCache + startupInfoEmitted atomic.Bool // Track if startup info has been emitted to avoid unnecessary duplication + 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 + + // elicitationSinkMu guards onElicitationRequest; elicitation requests can + // arrive from background-job goroutines (runCollecting) concurrently with + // the sink being (re)registered. + elicitationSinkMu sync.RWMutex + onElicitationRequest func(Event) + sessionStore session.Store + workingDir string // Working directory for hooks execution + env []string // Environment variables for hooks execution + modelSwitcherCfg *ModelSwitcherConfig + providerRegistry *provider.Registry + gatewayModels gatewayModelsCache + dmrModels dmrModelsCache // hooksRegistry is the runtime-private hooks.Registry used to build // every Executor. It carries the runtime-owned builtin hooks @@ -614,7 +637,6 @@ func NewLocalRuntime(ctx context.Context, agents *team.Team, opts ...Opt) (*Loca team: agents, agents: newAgentRouter(agents, defaultAgent.Name()), resumeChan: make(chan ResumeRequest), - elicitationRequestCh: make(chan ElicitationResult), steerQueue: NewInMemoryMessageQueue(defaultSteerQueueCapacity), followUpQueue: NewInMemoryMessageQueue(defaultFollowUpQueueCapacity), sessionCompaction: true, diff --git a/pkg/runtime/runtime_test.go b/pkg/runtime/runtime_test.go index 858b556b33..833a925e46 100644 --- a/pkg/runtime/runtime_test.go +++ b/pkg/runtime/runtime_test.go @@ -4145,6 +4145,12 @@ func TestElicitationHandler_NonInteractive(t *testing.T) { assert.Equal(t, tools.ElicitationActionDecline, result.Action, "non-interactive runtime should decline elicitation") } +// TestElicitationHandler_Interactive_NoChannel verifies that an interactive +// elicitationHandler call succeeds even when the elicitation bridge has no +// channel configured (no RunStream is active). This used to be a hard +// failure because delivery depended entirely on the bridge; with per-request +// waiters (#3584) the bridge send is best-effort only, and ResumeElicitation +// still routes the response correctly via the waiter registry. func TestElicitationHandler_Interactive_NoChannel(t *testing.T) { t.Parallel() @@ -4152,7 +4158,7 @@ func TestElicitationHandler_Interactive_NoChannel(t *testing.T) { root := agent.New("root", "test", agent.WithModel(prov)) tm := team.New(team.WithAgents(root)) - // Default runtime (interactive mode) with no events channel set + // Default runtime (interactive mode) with no events channel set on the bridge. rt, err := NewLocalRuntime(t.Context(), tm) require.NoError(t, err) @@ -4160,10 +4166,29 @@ func TestElicitationHandler_Interactive_NoChannel(t *testing.T) { Message: "Authorize OAuth?", } - _, err = rt.elicitationHandler(t.Context(), params) + type handlerResult struct { + result tools.ElicitationResult + err error + } + done := make(chan handlerResult, 1) + go func() { + result, err := rt.elicitationHandler(t.Context(), params) + done <- handlerResult{result: result, err: err} + }() + + require.Eventually(t, func() bool { return rt.elicitationWaiters.count() == 1 }, time.Second, time.Millisecond, + "elicitationHandler must register a waiter even though the bridge has no channel") - require.Error(t, err, "interactive runtime with no events channel should error") - assert.ErrorIs(t, err, errNoElicitationChannel) + require.NoError(t, rt.ResumeElicitation(t.Context(), tools.ElicitationActionAccept, map[string]any{"ok": true}, "")) + + select { + case got := <-done: + require.NoError(t, got.err) + assert.Equal(t, tools.ElicitationActionAccept, got.result.Action) + assert.Equal(t, map[string]any{"ok": true}, got.result.Content) + case <-time.After(2 * time.Second): + t.Fatal("elicitationHandler did not return after ResumeElicitation") + } } // TestRunAgentPersistsSubSessionToStore is the regression test for the diff --git a/pkg/server/server.go b/pkg/server/server.go index fb1f5c3f7d..6f3e307daf 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -494,7 +494,7 @@ func (s *Server) elicitation(c echo.Context) error { return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("invalid request body: %v", err)) } - if err := s.sm.ResumeElicitation(c.Request().Context(), sessionID, req.Action, req.Content); err != nil { + if err := s.sm.ResumeElicitation(c.Request().Context(), sessionID, req.Action, req.Content, req.ElicitationID); err != nil { return echo.NewHTTPError(http.StatusInternalServerError, fmt.Sprintf("failed to resume elicitation: %v", err)) } diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 34251a399a..dc74ce8f46 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -925,8 +925,9 @@ func (sm *SessionManager) recallSession(ctx context.Context, sessionID string, m return nil } -// ResumeElicitation resumes an elicitation request. -func (sm *SessionManager) ResumeElicitation(ctx context.Context, sessionID, action string, content map[string]any) error { +// ResumeElicitation resumes an elicitation request. elicitationID is +// additive: pass "" to fall back to resolving the sole pending request. +func (sm *SessionManager) ResumeElicitation(ctx context.Context, sessionID, action string, content map[string]any, elicitationID string) error { sm.mux.Lock() defer sm.mux.Unlock() rt, exists := sm.runtimeSessions.Load(sessionID) @@ -934,7 +935,7 @@ func (sm *SessionManager) ResumeElicitation(ctx context.Context, sessionID, acti return errors.New("session not found") } - return rt.runtime.ResumeElicitation(ctx, tools.ElicitationAction(action), content) + return rt.runtime.ResumeElicitation(ctx, tools.ElicitationAction(action), content, elicitationID) } // ToggleToolApproval toggles the tool approval mode for a session. diff --git a/pkg/server/session_manager_test.go b/pkg/server/session_manager_test.go index 5289463f89..b1c412dd84 100644 --- a/pkg/server/session_manager_test.go +++ b/pkg/server/session_manager_test.go @@ -69,7 +69,7 @@ func (f *fakeRuntime) Steer(_ context.Context, _ runtime.QueuedMessage) error { func (f *fakeRuntime) FollowUp(_ context.Context, _ runtime.QueuedMessage) error { return nil } -func (f *fakeRuntime) ResumeElicitation(_ context.Context, _ tools.ElicitationAction, _ map[string]any) error { +func (f *fakeRuntime) ResumeElicitation(_ context.Context, _ tools.ElicitationAction, _ map[string]any, _ ...string) error { return nil } diff --git a/pkg/tui/compact_routing_test.go b/pkg/tui/compact_routing_test.go index 8872280019..a128b1ad51 100644 --- a/pkg/tui/compact_routing_test.go +++ b/pkg/tui/compact_routing_test.go @@ -47,7 +47,7 @@ func (stubRuntime) Run(context.Context, *session.Session) ([]session.Message, er return nil, nil } func (stubRuntime) Resume(context.Context, runtime.ResumeRequest) {} -func (stubRuntime) ResumeElicitation(context.Context, tools.ElicitationAction, map[string]any) error { +func (stubRuntime) ResumeElicitation(context.Context, tools.ElicitationAction, map[string]any, ...string) error { return nil } func (stubRuntime) SessionStore() session.Store { return nil } @@ -84,6 +84,7 @@ func (stubRuntime) AvailableModels(context.Context) []runtime.ModelChoice { retu func (stubRuntime) SupportsModelSwitching() bool { return false } func (stubRuntime) OnToolsChanged(func(runtime.Event)) {} func (stubRuntime) OnBackgroundEvent(func(runtime.Event)) {} +func (stubRuntime) OnElicitationRequest(func(runtime.Event)) {} func (stubRuntime) QueueStatus() runtime.QueueStatus { return runtime.QueueStatus{} } func (stubRuntime) TogglePause(context.Context) (bool, error) { return false, nil } func (stubRuntime) Close() error { return nil } diff --git a/pkg/tui/page/chat/queue_test.go b/pkg/tui/page/chat/queue_test.go index 716b0cdbb3..483c78c5fd 100644 --- a/pkg/tui/page/chat/queue_test.go +++ b/pkg/tui/page/chat/queue_test.go @@ -69,7 +69,7 @@ func (queueTestRuntime) Run(context.Context, *session.Session) ([]session.Messag return nil, nil } func (queueTestRuntime) Resume(context.Context, runtime.ResumeRequest) {} -func (queueTestRuntime) ResumeElicitation(context.Context, tools.ElicitationAction, map[string]any) error { +func (queueTestRuntime) ResumeElicitation(context.Context, tools.ElicitationAction, map[string]any, ...string) error { return nil } func (queueTestRuntime) SessionStore() session.Store { return nil } @@ -110,6 +110,7 @@ func (queueTestRuntime) AvailableModels(context.Context) []runtime.ModelChoice { func (queueTestRuntime) SupportsModelSwitching() bool { return false } func (queueTestRuntime) OnToolsChanged(func(runtime.Event)) {} func (queueTestRuntime) OnBackgroundEvent(func(runtime.Event)) {} +func (queueTestRuntime) OnElicitationRequest(func(runtime.Event)) {} func (queueTestRuntime) QueueStatus() runtime.QueueStatus { return runtime.QueueStatus{} } func (queueTestRuntime) TogglePause(context.Context) (bool, error) { return false, nil } func (queueTestRuntime) Close() error { return nil } diff --git a/pkg/tui/service/supervisor/supervisor_test.go b/pkg/tui/service/supervisor/supervisor_test.go index 40ac7741dc..bc3376859c 100644 --- a/pkg/tui/service/supervisor/supervisor_test.go +++ b/pkg/tui/service/supervisor/supervisor_test.go @@ -191,7 +191,7 @@ func TestStreamStarted_SubSessionDoesNotDropPendingEvent(t *testing.T) { t.Parallel() s := newTestSupervisor([]string{"sess-A", "sess-B"}, "sess-B") // sess-A is background - elicitation := runtime.ElicitationRequest("confirm?", "form", nil, "", "eid-1", nil, "agent") + elicitation := runtime.ElicitationRequest("confirm?", "form", nil, "", "eid-1", "", "", nil, "agent") s.runners["sess-A"].PendingEvents = []tea.Msg{elicitation} s.runners["sess-A"].NeedsAttn = true s.runners["sess-A"].IsRunning = true // already running a top-level turn @@ -217,7 +217,7 @@ func TestStreamStopped_SubSessionDoesNotDropPendingEvent(t *testing.T) { t.Parallel() s := newTestSupervisor([]string{"sess-A", "sess-B"}, "sess-B") - elicitation := runtime.ElicitationRequest("confirm?", "form", nil, "", "eid-2", nil, "agent") + elicitation := runtime.ElicitationRequest("confirm?", "form", nil, "", "eid-2", "", "", nil, "agent") s.runners["sess-A"].PendingEvents = []tea.Msg{elicitation} s.runners["sess-A"].NeedsAttn = true s.runners["sess-A"].IsRunning = true @@ -244,7 +244,7 @@ func TestStreamStarted_TopLevelSupersedesStalePending(t *testing.T) { s := newTestSupervisor([]string{"sess-A"}, "sess-A") s.runners["sess-A"].PendingEvents = []tea.Msg{runtime.ElicitationRequest( - "old?", "form", nil, "", "eid-stale", nil, "agent", + "old?", "form", nil, "", "eid-stale", "", "", nil, "agent", )} s.runners["sess-A"].IsRunning = false @@ -270,7 +270,7 @@ func TestStreamStopped_TopLevelClearsPendingAndNeedsAttn(t *testing.T) { }{ { name: "elicitation pending", - pending: runtime.ElicitationRequest("q?", "form", nil, "", "eid-3", nil, "agent"), + pending: runtime.ElicitationRequest("q?", "form", nil, "", "eid-3", "", "", nil, "agent"), }, { name: "tool confirmation pending", @@ -311,7 +311,7 @@ func TestStreamStarted_EmptySessionID_TreatedAsTopLevel(t *testing.T) { s := newTestSupervisor([]string{"sess-A"}, "sess-A") s.runners["sess-A"].PendingEvents = []tea.Msg{runtime.ElicitationRequest( - "old?", "form", nil, "", "eid-old", nil, "agent", + "old?", "form", nil, "", "eid-old", "", "", nil, "agent", )} // Emitter omits SessionID (empty string).