From 47b0f08738343b068d03fcd01f2aa0db0f14b9c6 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Tue, 4 Aug 2026 08:58:50 -0700 Subject: [PATCH 1/2] agentHost: drive tool execution from the session input queue Subagent tool calls could stall indefinitely. A user reported 16 subagents running overnight that "keep stalling and dying for no apparent reason", needing the main agent to repeatedly repair them. Log analysis found 16 permission requests that were never answered, and 80 subagent chat channels unsubscribed ~12ms after a single provider error. The cause is structural rather than a single bug. Answering a tool call was owned by the per-turn chat observer: it rendered the call AND invoked the tool AND dispatched the outcome. So anything that tore down an observer -- a provider error disposing the parent turn's store, a turn ending, a reconnect, or simply never observing a subagent chat -- left the agent blocked on an obligation nobody was left to answer. Invert the relationship. The protocol already maintains SessionState.inputNeeded: a session-level queue of every outstanding blocker, each entry self-sufficient so a client can answer it without subscribing to the owning chat. It is a derived projection recomputed from tool-call status, so it is a set that can be re-read rather than a stream that can be missed. Make that queue the driver: - A session-level watcher owns all four blocker kinds and is the single caller of invokeTool. Chat observers only render. - One shared ChatToolInvocation per call, created by whichever side arrives first, so the card an observer renders in its subagent group is the same object the watcher executes. - Claimed calls run with chat context so confirmations render inline. Unclaimed non-confirmable calls run headlessly. Unclaimed confirmable calls wait for an observer, then deny rather than surface a modal nobody can see. - Chat input requests and MCP authentication get the same treatment; both could previously stall with no surface at all. This removes the class rather than the instances: an obligation is now answered because the session says it is outstanding, not because some particular observer happened to still be alive. Also stop counting toolClientExecution entries as user-blocking. That entry means a client is running the tool, not that a user was asked, so it must not raise InputNeeded -- otherwise every client tool call flags the session as needing input for its whole duration, and an approved call keeps presenting as blocked. Mirrors microsoft/agent-host-protocol#380. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../protocol/channels-session/reducer.ts | 16 +- .../state/protocol/channels-session/state.ts | 9 +- .../agentHost/node/agentSideEffects.ts | 13 +- .../test/node/agentSideEffects.test.ts | 16 +- .../agentHost/agentHostSessionHandler.ts | 661 ++++++++++----- .../agentHostClientTools.test.ts | 769 +++++++++++++++++- 6 files changed, 1248 insertions(+), 236 deletions(-) diff --git a/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts b/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts index 0de727342d802..10e2bf4db2d64 100644 --- a/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts +++ b/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts @@ -7,7 +7,7 @@ // DO NOT EDIT -- auto-generated by scripts/sync-agent-host-protocol.ts import { ActionType } from '../common/actions.js'; -import { SessionLifecycle, SessionStatus, CustomizationType, McpServerStatus, type SessionState, type SessionInputRequest, type McpServerCustomization } from './state.js'; +import { SessionLifecycle, SessionStatus, SessionInputRequestKind, CustomizationType, McpServerStatus, type SessionState, type SessionInputRequest, type McpServerCustomization } from './state.js'; import type { SessionAction } from '../action-origin.generated.js'; import { softAssertNever } from '../common/reducer-helpers.js'; @@ -30,8 +30,20 @@ function withStatusFlag(status: SessionStatus, flag: SessionStatus, set: boolean * `InProgress` while an already-idle session stays idle. Orthogonal flags * (`IsRead` / `IsArchived`) are preserved. */ +/** + * Whether an entry blocks on the *user*. + * + * {@link SessionInputRequestKind.ToolClientExecution} is work delegated to a + * client, not a prompt: the call has already cleared its confirmation gate and + * is simply running somewhere else. Counting it would report a session as + * awaiting the user for the entire duration of every client tool call. + */ +function awaitsUser(request: SessionInputRequest): boolean { + return request.kind !== SessionInputRequestKind.ToolClientExecution; +} + function withInputNeededStatus(status: SessionStatus, inputNeeded: readonly SessionInputRequest[]): SessionStatus { - if (inputNeeded.length > 0) { + if (inputNeeded.some(awaitsUser)) { return (status & ~STATUS_ACTIVITY_MASK) | SessionStatus.InputNeeded; } return status & ~(SessionStatus.InputNeeded & ~SessionStatus.InProgress); diff --git a/src/vs/platform/agentHost/common/state/protocol/channels-session/state.ts b/src/vs/platform/agentHost/common/state/protocol/channels-session/state.ts index f31cbd9a95630..2ae223ba4d1bd 100644 --- a/src/vs/platform/agentHost/common/state/protocol/channels-session/state.ts +++ b/src/vs/platform/agentHost/common/state/protocol/channels-session/state.ts @@ -172,9 +172,12 @@ export interface SessionState extends SessionMetadata { * Each entry is self-sufficient: it carries the owning chat's URI plus every * identifier the client needs to respond. A client answers by dispatching the * ordinary `chat/*` action to that chat's channel — see - * {@link SessionInputRequest} for the per-variant response path. A present, - * non-empty list implies {@link SessionStatus.InputNeeded} on - * {@link SessionSummary.status}. + * {@link SessionInputRequest} for the per-variant response path. A list + * holding any entry other than + * {@link SessionInputRequestKind.ToolClientExecution} implies + * {@link SessionStatus.InputNeeded} on {@link SessionSummary.status}; + * client-execution entries are work delegated to a client rather than a + * prompt, so they leave the session's activity unchanged. * * Host-managed: the host upserts entries with `session/inputNeededSet` as * chats raise requests and removes them with `session/inputNeededRemoved` diff --git a/src/vs/platform/agentHost/node/agentSideEffects.ts b/src/vs/platform/agentHost/node/agentSideEffects.ts index c7bf2b90b4c09..c5c5fd4fa1248 100644 --- a/src/vs/platform/agentHost/node/agentSideEffects.ts +++ b/src/vs/platform/agentHost/node/agentSideEffects.ts @@ -485,12 +485,11 @@ export class AgentSideEffects extends Disposable { const authenticationId = this._toolAuthenticationNeededId(chatUri, turnId, toolCallId); const toolCall = this._findToolCall(chatUri, turnId, toolCallId); - // A call auto-approved by the session's bypass setting is run - // automatically by the owning client and never blocks on the user, so - // keep it out of the session `inputNeeded` queue (which would flash - // "input needed" in the sessions list). `autoApproveBySetting` covers - // only the parameter gate; a `PendingResultConfirmation` is a genuine - // prompt and is still surfaced. + // A parameter gate auto-approved by the session's bypass setting never + // blocks on the user, so keep it out of the session `inputNeeded` queue + // (which would flash "input needed" in the sessions list). + // `autoApproveBySetting` covers only the parameter gate; a + // `PendingResultConfirmation` is a genuine prompt and is still surfaced. const autoApproved = !!toolCall && readToolCallMeta(toolCall).autoApproveBySetting === true; const suppressAutoApprovedConfirmation = autoApproved && toolCall?.status === ToolCallStatus.PendingConfirmation; @@ -508,7 +507,7 @@ export class AgentSideEffects extends Disposable { } const contributor = toolCall?.contributor; - if (!autoApproved && toolCall?.status === ToolCallStatus.Running && contributor?.kind === ToolCallContributorKind.Client) { + if (toolCall?.status === ToolCallStatus.Running && contributor?.kind === ToolCallContributorKind.Client) { this._setSessionInputNeeded(chatUri, { id: clientExecutionId, kind: SessionInputRequestKind.ToolClientExecution, diff --git a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts index d35c8769e6107..1ae60391b2ceb 100644 --- a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts +++ b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts @@ -4975,6 +4975,10 @@ suite('AgentSideEffects', () => { return stateManager.getSessionState(sessionUri.toString())?.inputNeeded ?? []; } + function sessionStatus() { + return stateManager.getSessionState(sessionUri.toString())?.status; + } + test('chat input request mirrors its unresolved response part and is removed on completion', () => { setupSession(); startTurn('turn-1'); @@ -5178,7 +5182,7 @@ suite('AgentSideEffects', () => { assert.deepStrictEqual(sessionInputNeeded(), []); }); - test('auto-approved tool call is kept out of the session inputNeeded queue', () => { + test('auto-approved tool call still surfaces its client execution without flagging input needed', () => { setupSession(); startTurn('turn-1'); @@ -5200,7 +5204,15 @@ suite('AgentSideEffects', () => { type: ActionType.ChatToolCallConfirmed, turnId: 'turn-1', toolCallId: 'tc-auto', approved: true, confirmed: ToolCallConfirmationReason.Setting, }); - assert.deepStrictEqual(sessionInputNeeded(), [], 'no client-execution entry while Running'); + + // The client still has to run the call, so it must be discoverable + // from the session channel — but it is not a user prompt, so the + // session must not present as "input needed". + assert.deepStrictEqual( + sessionInputNeeded().map(r => ({ kind: r.kind, clientId: r.kind === SessionInputRequestKind.ToolClientExecution ? r.clientId : undefined })), + [{ kind: SessionInputRequestKind.ToolClientExecution, clientId: 'client-1' }], + ); + assert.strictEqual(sessionStatus(), SessionStatus.InProgress, 'auto-approved client execution must not present as input needed'); }); test('auto-approved tool still surfaces a genuine result confirmation', () => { diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts index ce9b54d40ec59..be1cf2689f73b 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -41,7 +41,7 @@ import type { ChatInputRequestWithPlanReview, IAgentHostPlanReview } from '../.. import { IAgentSubscription, observableFromSubscription } from '../../../../../../platform/agentHost/common/state/agentSubscription.js'; import { ChatTruncatedAction } from '../../../../../../platform/agentHost/common/state/protocol/actions.js'; import { CompletionItemKind as AhpCompletionItemKind, type CompletionItem as AhpCompletionItem } from '../../../../../../platform/agentHost/common/state/protocol/commands.js'; -import { ConfirmationOptionKind, CustomizationType, JsonPrimitive, McpServerAuthRequiredState, McpServerStatus, SessionInputRequestKind, TerminalClaimKind, ToolCallContributorKind, ToolResultContentType, type ConfirmationOption, type ProtectedResourceMetadata, type SessionActiveClient } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { ConfirmationOptionKind, CustomizationType, JsonPrimitive, McpServerAuthRequiredState, McpServerStatus, SessionInputRequestKind, TerminalClaimKind, ToolCallContributorKind, ToolResultContentType, type ConfirmationOption, type ProtectedResourceMetadata, type SessionActiveClient, type SessionInputRequest, type SessionToolClientExecutionRequest } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; import { ActionType, ChatTurnStartedAction, isChatAction, type ClientChatAction, type ClientSessionAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; import { AHP_AUTH_REQUIRED, ProtocolError } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { buildSubagentChatUri, ChatOriginKind, getToolSubagentContent, isChatReadOnly, MessageAttachmentKind, MessageKind, PendingMessageKind, ResponsePartKind, ChatInputAnswerState, ChatInputAnswerValueKind, ChatInputQuestionKind, ChatInputResponseKind, SessionStatus, StateComponents, ToolCallCancellationReason, ToolCallConfirmationReason, ToolCallStatus, TurnState, parseChatUri, mergeSessionWithDefaultChat, readUsageInfoMeta, type ChatState, type ISessionWithDefaultChat, type ClientPluginCustomization, type ICompletedToolCall, type InputRequestResponsePart, type MarkdownResponsePart, type Message, type MessageAttachment, type MessageAnnotationsAttachment, type MessageChatAttachment, type MessageResourceAttachment, type MessageEmbeddedResourceAttachment, type ModelSelection, type PendingMessage, type ReasoningResponsePart, type RootState, type ChatInputAnswer, type ChatInputQuestion, type ChatInputRequest, type SessionState, type StringOrMarkdown, type ToolCallResponsePart, type ToolCallState, type Turn } from '../../../../../../platform/agentHost/common/state/sessionState.js'; @@ -84,7 +84,7 @@ import { ChatElicitationRequestPart } from '../../../common/model/chatProgressTy import { ChatToolInvocation } from '../../../common/model/chatProgressTypes/chatToolInvocation.js'; import { getChatSessionType, isUntitledChatSession } from '../../../common/model/chatUri.js'; import { IChatAgentData, IChatAgentImplementation, IChatAgentRequest, IChatAgentResult, IChatAgentService } from '../../../common/participants/chatAgents.js'; -import { ILanguageModelToolsService, IToolInvocation, IToolResult, stringifyPromptTsxPart, ToolInvocationPresentation } from '../../../common/tools/languageModelToolsService.js'; +import { ILanguageModelToolsService, IToolResult, stringifyPromptTsxPart, ToolInvocationPresentation } from '../../../common/tools/languageModelToolsService.js'; import { IChatWidgetService } from '../../chat.js'; import { getAgentSessionProviderIcon } from '../agentSessions.js'; import { IAgentHostActiveClientService } from './agentHostActiveClientService.js'; @@ -113,6 +113,7 @@ const MAX_INLINED_UNSAVED_EDITOR_BYTES = 1024 * 1024; /** Stable id of the progress row mirroring the host's chat activity, so updates replace it in place. */ const CHAT_ACTIVITY_PROGRESS_ID = 'agentHost.chatActivity'; +export const UNOBSERVED_CLIENT_TOOL_GRACE_MS = 5000; type AgentHostInvocationFailureStage = 'resolveSession' | 'provisionalSession' | 'sessionState' | 'authentication' | 'createSession' | 'subscribeSession' | 'prepareTurn' | 'dispatchTurn' | 'observeTurn'; interface IRestoredSubagentState extends IDisposable { @@ -762,8 +763,36 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC private readonly _serverTurnWatchers = this._register(new DisposableResourceMap()); /** Per-session subscription silently resolving existing MCP authentication grants. */ private readonly _mcpAuthWatchers = this._register(new DisposableResourceMap()); + /** Per-session ownership of actionable protocol requests. */ + private readonly _inputNeededWatchers = this._register(new DisposableResourceMap()); /** Historical turns with file edits, pending hydration into the editing session. */ private readonly _pendingHistoryTurns = new ResourceMap(); + /** + * Requests a turn observer is currently rendering, keyed by + * {@link _toolCallKey} for tool calls and {@link _inputRequestKey} for chat + * input requests (the two key shapes differ in arity, so they cannot + * collide). The session-level responder defers to those observers so the + * inline UI stays in charge of answering. + */ + private readonly _renderedRequests = observableValue>(this, new Set()); + /** Tool calls whose protocol outcome has already been dispatched. */ + private readonly _resolvedToolCalls = new Set(); + /** + * A single {@link ChatToolInvocation} per client tool call, keyed by + * {@link _toolCallKey}. Created lazily by whichever of the session-level + * watcher or the turn observer arrives first, so both act on one object: + * the observer renders it while the watcher executes it. Entries are + * dropped once the call resolves so a later call with the same ids is not + * mistaken for it. + */ + private readonly _clientToolInvocations = new Map(); + /** + * Live `inputNeeded` requests per tool call, keyed by {@link _toolCallKey}. + * One tool call is represented by a succession of requests — a confirmation + * is replaced by a client execution once approved — so the shared state + * above is only released when the last of them goes away. + */ + private readonly _clientToolRetainCounts = new Map(); /** * Per-session set of MCP server ids that already had an authentication * prompt surfaced in the current conversation. A server is removed from the @@ -1249,6 +1278,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC this._draftSyncSubscriptions.deleteAndDispose(sessionResource); this._serverTurnWatchers.deleteAndDispose(sessionResource); this._mcpAuthWatchers.deleteAndDispose(sessionResource); + this._inputNeededWatchers.deleteAndDispose(sessionResource); this._pendingHistoryTurns.delete(sessionResource); this._surfacedMcpAuthServers.delete(sessionResource); const chatURI = this._chatURIsBySessionResource.get(sessionResource); @@ -1826,6 +1856,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC const sessionStr = backendSession.toString(); const chatURI = this._getChatURI(sessionResource); this._watchForMcpAuthentication(backendSession, sessionResource, chatURI); + this._watchForSessionInputNeeded(backendSession, sessionResource); // Seed from the current state so we don't treat any pre-existing active // turn (e.g. one being handled by _reconnectToActiveTurn) as new. @@ -1941,6 +1972,322 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC this._mcpAuthWatchers.set(sessionResource, disposables); } + private _watchForSessionInputNeeded(backendSession: URI, sessionResource: URI): void { + const sessionSub = this._ensureSessionSubscription(backendSession.toString()); + const state = observableFromSubscription(this, sessionSub); + const store = new DisposableStore(); + + const requests = derivedOpts({ equalsFn: equals }, reader => + (state.read(reader)?.inputNeeded ?? []).filter((request): request is SessionInputRequest => + request.kind === SessionInputRequestKind.ChatInput + || request.kind === SessionInputRequestKind.ToolConfirmation + || request.kind === SessionInputRequestKind.ToolClientExecution + || request.kind === SessionInputRequestKind.ToolAuthentication)); + + // This watcher is the single point of truth for how client tools + // execute. A turn observer only ever renders the shared invocation; it + // never invokes the tool. Each outstanding blocker is handled here + // exactly once, keyed by its request id. + store.add(autorunPerKeyedItem(requests, request => request.id, (_requestId, request$, itemStore) => { + const initial = request$.get(); + const chatURI = initial.chat.toString(); + + if (initial.kind === SessionInputRequestKind.ChatInput) { + // A user-facing elicitation with no tool call. If no turn + // observer renders it within the grace window, nobody could + // answer it, so cancel it (the agent asked; nobody was there). + const inputKey = this._inputRequestKey(chatURI, initial.request.id); + let cancelled = false; + itemStore.add(disposableTimeout(() => { + if (cancelled || this._renderedRequests.get().has(inputKey)) { + return; + } + cancelled = true; + this._logService.warn(`[AgentHost] Cancelling chat input request ${initial.request.id}: no session claimed it within ${UNOBSERVED_CLIENT_TOOL_GRACE_MS}ms`); + this._dispatchAction(backendSession, { + type: ActionType.ChatInputCompleted, + requestId: initial.request.id, + response: ChatInputResponseKind.Cancel, + }, chatURI); + }, UNOBSERVED_CLIENT_TOOL_GRACE_MS)); + return; + } + + const key = this._toolCallKey(chatURI, initial.turnId, initial.toolCall.toolCallId); + const cts = new CancellationTokenSource(); + itemStore.add(toDisposable(() => cts.dispose(true))); + itemStore.add(this._retainToolCall(key)); + + if (initial.kind === SessionInputRequestKind.ToolClientExecution) { + if (initial.clientId !== this._config.connection.clientId) { + return; // A different client owns this call. + } + let handled = false; + const execute = (withContext: boolean) => { + if (handled) { + return; + } + handled = true; + void this._executeClientTool(request$.get() as SessionToolClientExecutionRequest, sessionResource, withContext, cts.token); + }; + if (this._renderedRequests.get().has(key)) { + // A turn observer is rendering it, so a live chat request + // exists: run with context so confirmation renders in the + // tool part and any pre-approval is honored. + execute(true); + } else if (!this._clientToolRequiresConfirmation(initial.toolCall)) { + // Unclaimed and cannot pop a confirmation: run headlessly so + // it does not depend on the owning turn still being live. + execute(false); + } else { + // Unclaimed and might pop a confirmation: a headless run + // would surface a modal nobody could see. Wait for an + // observer to claim it; if none does within the grace + // window, deny it. + itemStore.add(autorun(reader => { + if (!handled && this._renderedRequests.read(reader).has(key)) { + execute(true); + } + })); + itemStore.add(disposableTimeout(() => { + if (!handled) { + handled = true; + this._denyClientTool(request$.get() as SessionToolClientExecutionRequest); + } + }, UNOBSERVED_CLIENT_TOOL_GRACE_MS)); + } + } else if (initial.kind === SessionInputRequestKind.ToolAuthentication) { + // An MCP tool call blocked on authentication. The token is + // pushed out-of-band via the `authenticate` command, so this + // watcher does not resolve it — but if no observer renders the + // call within the grace window nobody can drive that flow, so + // cancel the call rather than leave the agent blocked forever. + itemStore.add(disposableTimeout(() => { + if (!this._renderedRequests.get().has(key)) { + this._logService.warn(`[AgentHost] Cancelling MCP authentication for ${initial.toolCall.toolName} (callId=${initial.toolCall.toolCallId}): no session claimed it within ${UNOBSERVED_CLIENT_TOOL_GRACE_MS}ms`); + this._resolveToolCall(chatURI, initial.turnId, initial.toolCall.toolCallId, { + type: ActionType.ChatToolCallComplete, + turnId: initial.turnId, + toolCallId: initial.toolCall.toolCallId, + result: { + success: false, + pastTenseMessage: localize('agentHost.mcpToolAuthentication.cancelled', "Cancelled tool call"), + error: { message: localize('agentHost.mcpToolAuthentication.cancelledError', "MCP authentication was cancelled"), code: 'cancelled' }, + }, + }); + } + }, UNOBSERVED_CLIENT_TOOL_GRACE_MS)); + } else { + // A confirmation that no sub/agent observer claims within the + // grace window is auto-denied so the agent is not left blocked + // on a surface that never renders. + itemStore.add(disposableTimeout(() => { + if (!this._renderedRequests.get().has(key)) { + this._logService.warn(`[AgentHost] Denying confirmation for ${initial.toolCall.toolName} (callId=${initial.toolCall.toolCallId}): no session claimed it within ${UNOBSERVED_CLIENT_TOOL_GRACE_MS}ms`); + this._resolveToolCall(chatURI, initial.turnId, initial.toolCall.toolCallId, { + type: ActionType.ChatToolCallConfirmed, + turnId: initial.turnId, + toolCallId: initial.toolCall.toolCallId, + approved: false, + reason: ToolCallCancellationReason.Denied, + }); + } + }, UNOBSERVED_CLIENT_TOOL_GRACE_MS)); + } + })); + + this._inputNeededWatchers.set(sessionResource, store); + } + + /** + * Holds the shared state for a tool call while an `inputNeeded` request + * references it. Once the host stops asking — the request disappears, or the + * watcher is disposed — the outcome is settled, so the dispatch-funnel entry + * and the shared invocation are dropped and a later call with the same ids + * is never mistaken for this one. + */ + private _retainToolCall(key: string): IDisposable { + this._clientToolRetainCounts.set(key, (this._clientToolRetainCounts.get(key) ?? 0) + 1); + return toDisposable(() => { + const remaining = (this._clientToolRetainCounts.get(key) ?? 1) - 1; + if (remaining > 0) { + this._clientToolRetainCounts.set(key, remaining); + return; + } + this._clientToolRetainCounts.delete(key); + this._forgetResolvedToolCall(key); + this._clientToolInvocations.delete(key); + }); + } + + /** + * Returns the shared {@link ChatToolInvocation} for a client tool call, + * creating it on first use via {@link ILanguageModelToolsService.beginToolCall}. + * `sessionResource` is deliberately omitted so `beginToolCall` does not + * append progress into a chat model (which throws once the owning request + * is complete); it still registers the invocation, so a later `invokeTool` + * with a matching `chatStreamToolCallId` attaches to this same object. The + * observer that renders the call and the watcher that executes it therefore + * act on one invocation. + */ + private _ensureClientToolInvocation(chatURI: string, turnId: string, toolCallId: string, toolId: string, subagentInvocationId: string | undefined): ChatToolInvocation | undefined { + const key = this._toolCallKey(chatURI, turnId, toolCallId); + const existing = this._clientToolInvocations.get(key); + if (existing) { + return existing; + } + const invocation = this._toolsService.beginToolCall({ + toolCallId, + toolId, + subagentInvocationId, + sessionResource: undefined, + force: true, + }) as ChatToolInvocation | undefined; + if (invocation) { + this._clientToolInvocations.set(key, invocation); + } + return invocation; + } + + /** + * Whether an unclaimed client tool must wait for a rendering observer + * before running. There is no protocol field for this, so we use the tool's + * static {@link IToolData.canRequestPreApproval} signal: a tool that might + * ask for pre-approval could pop a confirmation, which only makes sense + * inside a live chat request. Limitation: this is a "might" signal — a tool + * may set it yet auto-approve at runtime — so an unclaimed such tool is + * conservatively made to wait (and denied on timeout) rather than risk a + * headless modal nobody can answer. Only consulted for the unclaimed case; + * a claimed call always runs with context regardless. + */ + private _clientToolRequiresConfirmation(toolCall: ToolCallState): boolean { + const clientToolName = toolCall.toolName === RUNTIME_TOOL_SEARCH_TOOL_NAME ? CLIENT_TOOL_SEARCH_REFERENCE_NAME : toolCall.toolName; + return this._toolsService.getToolByName(clientToolName)?.canRequestPreApproval === true; + } + + /** + * The one place a client tool is actually invoked. Ensures the shared + * invocation exists, parses the protocol input (preserving the tool-search + * candidate handling), invokes the tool, and dispatches the protocol + * completion. `withContext` is set when a turn observer is rendering the + * call: a live chat request then exists, so confirmation renders in the + * tool part and any pre-approval is honored. Without it the tool runs + * headlessly, independent of whether the owning turn is live. + */ + private async _executeClientTool(request: SessionToolClientExecutionRequest, sessionResource: URI, withContext: boolean, token: CancellationToken): Promise { + const chatURI = request.chat.toString(); + const toolCall = request.toolCall; + const toolName = toolCall.toolName; + const isToolSearch = toolName === RUNTIME_TOOL_SEARCH_TOOL_NAME; + const clientToolName = isToolSearch ? CLIENT_TOOL_SEARCH_REFERENCE_NAME : toolName; + const toolData = this._toolsService.getToolByName(clientToolName); + + // A tool-search completion (success or failure) must drop the transient + // candidate corpus from `_meta` while preserving any other metadata. + const completionMeta = isToolSearch ? { _meta: metaWithoutToolSearchCandidates(toolCall) } : {}; + + const fail = (message: string, code: string) => this._resolveToolCall(chatURI, request.turnId, toolCall.toolCallId, { + type: ActionType.ChatToolCallComplete, + turnId: request.turnId, + toolCallId: toolCall.toolCallId, + result: { + success: false, + pastTenseMessage: localize('agentHost.clientTool.pastTense', "Couldn't run {0}", toolCall.displayName), + error: { message, code }, + }, + ...completionMeta, + }); + + if (!toolData) { + fail(localize('agentHost.clientTool.unknown', "Tool \"{0}\" is not available on this client.", toolName), 'toolUnavailable'); + return; + } + + // eslint-disable-next-line local/code-no-in-operator + const rawInput = 'toolInput' in toolCall ? toolCall.toolInput : undefined; + let parameters: Record; + try { + const parsed: unknown = JSON.parse(rawInput ?? '{}'); + if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) { + throw new Error('expected JSON object'); + } + parameters = parsed as Record; + } catch { + fail(localize('agentHost.clientTool.badInput', "Invalid tool input for \"{0}\": expected JSON object parameters.", toolName), 'invalidInput'); + return; + } + + const toolSearchCandidates = isToolSearch ? readToolCallMeta(toolCall).toolSearchCandidates : undefined; + if (toolSearchCandidates !== undefined) { + parameters = { ...parameters, candidateTools: toolSearchCandidates }; + } + + const invocation = this._ensureClientToolInvocation(chatURI, request.turnId, toolCall.toolCallId, toolData.id, undefined); + if (!invocation) { + fail(localize('agentHost.clientTool.beginFailed', "Could not create invocation for client tool \"{0}\".", toolName), 'invocationFailed'); + return; + } + + this._logService.info(`[AgentHost] Running client tool: ${toolName} (callId=${toolCall.toolCallId}, withContext=${withContext})`); + let result: IToolResult | undefined; + let error: unknown; + try { + result = await this._toolsService.invokeTool({ + callId: toolCall.toolCallId, + toolId: toolData.id, + parameters, + context: withContext ? { sessionResource } : undefined, + chatStreamToolCallId: toolCall.toolCallId, + preApproved: getClientToolPreApproval(toolCall), + }, async () => 0, token); + } catch (err) { + error = err; + } + + if (token.isCancellationRequested) { + return; + } + if (error !== undefined) { + if (!isCancellationError(error)) { + this._logService.warn(`[AgentHost] Client tool failed: ${toolName}`, error); + } + result = { content: [], toolResultError: error instanceof Error ? error.message : String(error) }; + } + + this._resolveToolCall(chatURI, request.turnId, toolCall.toolCallId, { + type: ActionType.ChatToolCallComplete, + turnId: request.turnId, + toolCallId: toolCall.toolCallId, + result: toolResultToProtocol(result ?? { content: [] }, toolName), + ...completionMeta, + }); + } + + /** + * Denies a client tool call that needs confirmation but that no sub/agent + * observer claimed within the grace window: there is no live surface to + * answer it, so report a failed completion rather than pop a headless + * modal. + */ + private _denyClientTool(request: SessionToolClientExecutionRequest): void { + const toolCall = request.toolCall; + this._logService.warn(`[AgentHost] Denying client tool ${toolCall.toolName} (callId=${toolCall.toolCallId}): it can request confirmation but no session claimed it within ${UNOBSERVED_CLIENT_TOOL_GRACE_MS}ms`); + this._resolveToolCall(request.chat.toString(), request.turnId, toolCall.toolCallId, { + type: ActionType.ChatToolCallComplete, + turnId: request.turnId, + toolCallId: toolCall.toolCallId, + result: { + success: false, + pastTenseMessage: localize('agentHost.clientTool.unclaimed', "Couldn't run {0}", toolCall.displayName), + error: { + message: localize('agentHost.clientTool.unclaimedError', "{0} needs confirmation but no session was available to answer it.", toolCall.displayName), + code: 'clientUnavailable', + }, + }, + }); + this._clientToolInvocations.delete(this._toolCallKey(request.chat.toString(), request.turnId, toolCall.toolCallId)); + } + /** * Tracks protocol state changes for a specific server-initiated turn and * pushes `IChatProgress[]` items into the session's `progressObs`. @@ -2147,17 +2494,16 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC this._logService.info(`[AgentHost] Tool confirmation: toolCallId=${toolCallId}, approved=${approved}, selectedOptionId=${selectedOption?.id}`); const target = this._requireChatURI(chatURI, ActionType.ChatToolCallConfirmed); - if (approved) { - this._config.connection.dispatch(target, { + this._resolveToolCall(target, turnId, toolCallId, approved + ? { type: ActionType.ChatToolCallConfirmed, turnId, toolCallId, approved: true, confirmed: ToolCallConfirmationReason.UserAction, ...(selectedOption ? { selectedOptionId: selectedOption.id } : {}), - }); - } else { - this._config.connection.dispatch(target, { + } + : { type: ActionType.ChatToolCallConfirmed, turnId, toolCallId, @@ -2165,7 +2511,6 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC reason: ToolCallCancellationReason.Denied, ...(selectedOption ? { selectedOptionId: selectedOption.id } : {}), }); - } }).catch(err => { this._logService.warn(`[AgentHost] Tool confirmation failed for toolCallId=${toolCallId}`, err); }); @@ -2751,14 +3096,84 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC const initial = part$.get().toolCall; const contributor = initial.contributor; if (contributor?.kind === ToolCallContributorKind.Client && contributor.clientId === this._config.connection.clientId) { + // Set up before claiming: the claim is what tells the session-level + // watcher it may execute this call, and it must find the shared + // invocation already created when it does. this._setupClientToolCall(initial, part$, store, opts, subagentContext); + store.add(this._markToolCallRendered(opts.chatURI, opts.turnId, initial.toolCallId)); } else if (contributor?.kind === ToolCallContributorKind.Client) { this._setupOtherClientToolCall(initial, part$, store, opts); } else { + store.add(this._markToolCallRendered(opts.chatURI, opts.turnId, initial.toolCallId)); this._setupServerToolCall(initial, part$, store, opts, subagentContext); } } + private _toolCallKey(chatURI: string, turnId: string, toolCallId: string): string { + return `${chatURI}\0${turnId}\0${toolCallId}`; + } + + private _inputRequestKey(chatURI: string, requestId: string): string { + return `${chatURI}\0${requestId}`; + } + + /** Claims a request as rendered until the returned disposable is disposed. */ + private _markRendered(key: string): IDisposable { + this._renderedRequests.set(new Set(this._renderedRequests.get()).add(key), undefined); + return toDisposable(() => { + const next = new Set(this._renderedRequests.get()); + next.delete(key); + this._renderedRequests.set(next, undefined); + }); + } + + /** + * Records that a turn observer is rendering this chat input request, so the + * session-level responder leaves its inline elicitation UI in charge. + */ + private _markInputRequestRendered(chatURI: string, requestId: string): IDisposable { + return this._markRendered(this._inputRequestKey(chatURI, requestId)); + } + + /** + * Records that a turn observer is rendering this tool call, so the + * session-level responder leaves its inline UI in charge. Releasing the + * claim also forgets the funnel entries, which is the only cleanup a tool + * call that never reached `inputNeeded` ever gets. + */ + private _markToolCallRendered(chatURI: string, turnId: string, toolCallId: string): IDisposable { + const key = this._toolCallKey(chatURI, turnId, toolCallId); + const rendered = this._markRendered(key); + return toDisposable(() => { + rendered.dispose(); + this._forgetResolvedToolCall(key); + }); + } + + /** + * Single funnel for tool-call outcomes, so an inline invocation and the + * session-level responder can both offer the action while the protocol + * only ever sees the first answer. Confirming and completing are distinct + * outcomes, so each is tracked separately. + */ + private _resolveToolCall(chatURI: string, turnId: string, toolCallId: string, action: ClientChatAction): void { + const key = `${this._toolCallKey(chatURI, turnId, toolCallId)}\0${action.type}`; + if (this._resolvedToolCalls.has(key)) { + this._logService.trace(`[AgentHost] Tool call outcome was already dispatched: ${toolCallId} (${action.type})`); + return; + } + this._resolvedToolCalls.add(key); + this._config.connection.dispatch(chatURI, action); + } + + private _forgetResolvedToolCall(toolCallKey: string): void { + for (const key of this._resolvedToolCalls) { + if (key.startsWith(`${toolCallKey}\0`)) { + this._resolvedToolCalls.delete(key); + } + } + } + private _setupOtherClientToolCall( initial: ToolCallState, part$: IObservable, @@ -2876,6 +3291,9 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC if (status === ToolCallStatus.Streaming) { updateStreamingToolInvocation(invocation, tc, this._config.connectionAuthority); } else if (enteringConfirmation) { + // A re-ask is a fresh obligation, so a previous answer must not + // suppress this one. + this._forgetResolvedToolCall(this._toolCallKey(opts.chatURI, opts.turnId, toolCallId)); if (!IChatToolInvocation.isComplete(invocation)) { const prepared = toolCallStateToPreparedInvocation(tc, opts.backendSession, this._config.connectionAuthority, opts.sessionResource.authority); invocation.requestConfirmation(prepared); @@ -3035,12 +3453,13 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } /** - * Per-call setup for a client-provided tool. Eagerly creates a streaming - * {@link ChatToolInvocation} so the UI has a handle, then invokes the - * tool once parameters are available. The inner autorun on `part$` is - * idempotent: `invoked` ensures `invokeTool` runs at most once, - * `confirmationDispatched` ensures `ChatToolCallConfirmed` is sent at - * most once. + * Per-call setup for a client-provided tool. The observer only renders: it + * obtains the shared {@link ChatToolInvocation} (created by whichever of + * this observer or the session-level watcher arrives first), emits it into + * this chat so it renders in the correct group, drives subagent + * presentation, and dispatches `ChatToolCallConfirmed` from the + * invocation's confirmation gate. It never invokes the tool — the + * session-level watcher owns execution. */ private _setupClientToolCall( initial: ToolCallState, @@ -3052,9 +3471,9 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC const toolCallId = initial.toolCallId; const toolName = initial.toolName; - // Reconnect adoption: settle any snapshot invocation so the new - // streaming one created by `beginToolCall` can take over the UI - // slot rather than leaving the old instance orphaned. + // Reconnect adoption: settle any snapshot invocation so the shared + // invocation can take over the UI slot rather than leaving the old + // instance orphaned. const adopted = opts.adoptInvocations?.get(toolCallId); if (adopted && !IChatToolInvocation.isComplete(adopted)) { adopted.didExecuteTool(undefined); @@ -3076,14 +3495,8 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC }, opts.chatURI); return; } - const invocation = this._toolsService.beginToolCall({ - toolCallId, - toolId: toolData.id, - subagentInvocationId: opts.subAgentInvocationId, - sessionResource: opts.sessionResource, - force: true, - }) as ChatToolInvocation | undefined; + const invocation = this._ensureClientToolInvocation(opts.chatURI, opts.turnId, toolCallId, toolData.id, opts.subAgentInvocationId); if (!invocation) { this._logService.warn(`[AgentHost] Failed to begin client tool invocation: ${toolName}`); this._dispatchAction(opts.backendSession, { @@ -3107,205 +3520,65 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } this._tryObserveSubagentToolCall(initial, invocation, store, opts, subagentContext); - const cts = new CancellationTokenSource(); - store.add(toDisposable(() => cts.dispose(true))); + // The shared invocation is created with no `sessionResource`, so it + // does not `appendProgress` into a chat model. Emit it explicitly so it + // renders in this chat / subagent group (mirrors `_setupServerToolCall`). + opts.sink([invocation]); - let invoked = false; - let approvedDispatched = false; let confirmationDispatched = false; // Drive `ChatToolCallConfirmed` from the invocation's confirmation - // gate. The autorun runs synchronously many times; the guards keep it - // idempotent. + // gate. The watcher's `invokeTool` transitions the shared invocation; + // this reports the outcome to the protocol. The autorun runs + // synchronously many times; the guard keeps it idempotent. store.add(autorun(reader => { const state = invocation.state.read(reader); - const tc = part$.read(reader).toolCall; - const preApproval = getClientToolPreApproval(tc); - if (state.type === IChatToolInvocation.StateKind.WaitingForConfirmation && preApproval) { - state.confirm(preApproval); - return; - } if (confirmationDispatched) { return; } if (state.type === IChatToolInvocation.StateKind.Executing) { confirmationDispatched = true; - if (cts.token.isCancellationRequested) { - return; - } - approvedDispatched = true; - this._dispatchAction(opts.backendSession, { + this._resolveToolCall(opts.chatURI, opts.turnId, toolCallId, { type: ActionType.ChatToolCallConfirmed, turnId: opts.turnId, toolCallId, approved: true, confirmed: confirmedReasonToProtocol(state.confirmed), - }, opts.chatURI); + }); } else if (state.type === IChatToolInvocation.StateKind.Cancelled) { - // Pre-execution cancellation. If the server already knows - // (cts cancelled), suppress the dispatch — the server - // transitioned the call itself. + // Pre-execution cancellation (a denied confirmation). If the + // protocol call already reached a terminal state the server + // drove it, so suppress the dispatch. confirmationDispatched = true; - if (cts.token.isCancellationRequested) { + const status = part$.read(undefined).toolCall.status; + if (status === ToolCallStatus.Cancelled || status === ToolCallStatus.Completed) { return; } - this._dispatchAction(opts.backendSession, { + this._resolveToolCall(opts.chatURI, opts.turnId, toolCallId, { type: ActionType.ChatToolCallConfirmed, turnId: opts.turnId, toolCallId, approved: false, reason: ToolCallCancellationReason.Denied, - }, opts.chatURI); + }); } })); - const handleSettled = (result: IToolResult | undefined, err: unknown) => { - if (cts.token.isCancellationRequested) { - return; - } - - if (err !== undefined) { - if (!isCancellationError(err)) { - if (!approvedDispatched) { - this._logService.warn(`[AgentHost] Client tool rejected pre-execution: ${toolName}`, err); - } else { - this._logService.warn(`[AgentHost] Client tool invocation failed: ${toolName}`, err); - } - } - - result = { content: [], toolResultError: err instanceof Error ? err.message : String(err) }; - } - - const protocolToolCall = part$.get().toolCall; - const isProtocolToolCallComplete = protocolToolCall.status === ToolCallStatus.Completed || protocolToolCall.status === ToolCallStatus.Cancelled; - if (!isProtocolToolCallComplete) { - // The tool-search ready action stashes the (potentially large) - // deferred-tool corpus in `_meta.toolSearchCandidates` purely to - // seed this invocation. The completion reducer keeps the prior - // `_meta` when the action omits one, so without an explicit - // replacement the corpus would persist on the completed call and - // across reconnects. Carry a candidate-stripped `_meta` on the - // tool-search completion to drop it once the search has run. - const clearedMeta = toolName === RUNTIME_TOOL_SEARCH_TOOL_NAME - ? metaWithoutToolSearchCandidates(protocolToolCall) - : undefined; - this._dispatchAction(opts.backendSession, { - type: ActionType.ChatToolCallComplete, - turnId: opts.turnId, - toolCallId, - result: toolResultToProtocol(result ?? { content: [] }, toolName), - ...(clearedMeta !== undefined ? { _meta: clearedMeta } : {}), - }, opts.chatURI); - } - }; - - // React to part$ updates: route external cancellation, and try to - // invoke once parameters are present. Idempotent via `invoked` and - // `cts.token.isCancellationRequested`. + // Presentational: keep subagent observation current, and if the + // protocol call reaches a terminal state while the shared invocation is + // still streaming (the watcher never ran it), settle the card so the UI + // is not stuck. store.add(autorun(reader => { const tc = part$.read(reader).toolCall; - const state = invocation.state.read(reader); this._tryObserveSubagentToolCall(tc, invocation, store, opts, subagentContext); - const preApproval = getClientToolPreApproval(tc); - if (state.type === IChatToolInvocation.StateKind.WaitingForConfirmation && preApproval) { - state.confirm(preApproval); - } - if (tc.status === ToolCallStatus.Cancelled || tc.status === ToolCallStatus.Completed) { - // The protocol tool call reached a terminal state. If this was - // driven by the server (e.g. the client-tool bridge abandoned the - // call because the client was considered disconnected, the turn was - // superseded, or a reconnect occurred) while our local `invokeTool` - // is still running, cancel it so the tool cleans up (e.g. dismisses a - // pending question carousel) instead of blocking forever on an answer - // nobody will consume. In the normal path we complete the call - // ourselves first, so `invokeTool` has already settled and this - // cancellation is a harmless no-op. - if (state.type === IChatToolInvocation.StateKind.Streaming) { - const fileEdits = finalizeToolInvocation(invocation, tc, opts.backendSession, this._config.connectionAuthority); - if (fileEdits.length > 0) { - opts.onFileEdits?.(tc, fileEdits); - } - } - if (cts.token.isCancellationRequested) { - return; - } - cts.cancel(); - if (!invoked && tc.status === ToolCallStatus.Cancelled && state.type !== IChatToolInvocation.StateKind.Streaming) { - // No `invokeTool` is listening to the CTS — transition - // the invocation to `Cancelled` ourselves. - invocation.cancelFromStreaming(ToolConfirmKind.Skipped); - } - return; - } - if (invoked || cts.token.isCancellationRequested) { - return; - } - const toolSearchCandidates = toolName === RUNTIME_TOOL_SEARCH_TOOL_NAME - ? readToolCallMeta(tc).toolSearchCandidates - : undefined; - if (toolName === RUNTIME_TOOL_SEARCH_TOOL_NAME && toolSearchCandidates === undefined) { - return; - } - // eslint-disable-next-line local/code-no-in-operator - let toolInput = 'toolInput' in tc ? tc.toolInput : undefined; - if (toolInput === undefined) { - // Still streaming — parameters may still be arriving. Once - // we move past Streaming, treat a missing toolInput as `{}` - // so zero-argument tools are not stuck. - if (tc.status === ToolCallStatus.Streaming) { - return; - } - toolInput = '{}'; - } - invoked = true; - - let parameters: Record = {}; - try { - const parsed: unknown = JSON.parse(toolInput); - if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) { - throw new Error('expected JSON object'); + if ((tc.status === ToolCallStatus.Cancelled || tc.status === ToolCallStatus.Completed) + && invocation.state.read(reader).type === IChatToolInvocation.StateKind.Streaming) { + const fileEdits = finalizeToolInvocation(invocation, tc, opts.backendSession, this._config.connectionAuthority); + if (fileEdits.length > 0) { + opts.onFileEdits?.(tc, fileEdits); } - parameters = parsed as Record; - } catch { - this._logService.warn(`[AgentHost] Failed to parse tool input for ${toolName}`); - const clearedMeta = toolName === RUNTIME_TOOL_SEARCH_TOOL_NAME - ? metaWithoutToolSearchCandidates(tc) - : undefined; - this._dispatchAction(opts.backendSession, { - type: ActionType.ChatToolCallComplete, - turnId: opts.turnId, - toolCallId, - result: { - success: false, - pastTenseMessage: `Failed to execute ${toolName}`, - error: { message: `Invalid tool input for "${toolName}": expected JSON object parameters` }, - }, - ...(clearedMeta !== undefined ? { _meta: clearedMeta } : {}), - }, opts.chatURI); - return; - } - if (toolSearchCandidates !== undefined) { - parameters = { ...parameters, candidateTools: toolSearchCandidates }; + invocation.cancelFromStreaming(ToolConfirmKind.Skipped); } - - const inv: IToolInvocation = { - callId: toolCallId, - toolId: invocation.toolId, - parameters, - context: { sessionResource: opts.sessionResource }, - chatStreamToolCallId: toolCallId, - // If the agent host already resolved auto-approval for this call, - // pass it through so the invocation transitions straight to - // executing instead of briefly flashing a confirmation prompt - // (which would flicker "needs input" in the sessions list). - preApproved: getClientToolPreApproval(tc), - }; - const noOpCountTokens = async () => 0; - this._logService.info(`[AgentHost] Invoking client tool: ${toolName} (callId=${toolCallId})`); - this._toolsService.invokeTool(inv, noOpCountTokens, cts.token).then( - result => handleSettled(result, undefined), - err => handleSettled(undefined, err), - ); })); } @@ -3315,6 +3588,10 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC opts: IObserveTurnOptions, ): void { const inputReq = part$.get().request; + // Claim the elicitation so the session-level responder does not cancel + // it while an observer is rendering it. This covers all three render + // paths below, since each is reached only through this method. + store.add(this._markInputRequestRendered(opts.chatURI, inputReq.id)); const planReview = (inputReq as ChatInputRequestWithPlanReview).planReview; if (planReview) { this._setupPlanReviewInputRequest(part$, planReview, store, opts); @@ -3575,11 +3852,11 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } const terminalContent = getTerminalContent(tc.content); const terminalUri = terminalContent?.resource; - if (!terminalContent || !terminalUri || !tc.toolInput) { + const toolInput = tc.toolInput; + if (!terminalContent || !terminalUri || !toolInput) { return; } invocation.presentation = undefined; - const toolInput = tc.toolInput; const sessionId = makeAhpTerminalToolSessionId(terminalUri, backendSession); const terminalCommandUri = URI.parse(terminalUri); const isPty = terminalContent.isPty !== false; diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts index c6402bb9032a5..9ff6aca18394c 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts @@ -15,15 +15,16 @@ import { URI } from '../../../../../../base/common/uri.js'; import { constObservable, observableValue, autorun } from '../../../../../../base/common/observable.js'; import { mock } from '../../../../../../base/test/common/mock.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; +import { runWithFakedTimers } from '../../../../../../base/test/common/timeTravelScheduler.js'; import { ILogService, NullLogService } from '../../../../../../platform/log/common/log.js'; import { IConfigurationChangeEvent, IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js'; import { AgentSession, IAgentHostService } from '../../../../../../platform/agentHost/common/agentService.js'; import { CLIENT_TOOL_SEARCH_REFERENCE_NAME, RUNTIME_TOOL_SEARCH_TOOL_NAME } from '../../../../../../platform/agentHost/common/toolSearchConstants.js'; import { isChatAction, isSessionAction, type ActionEnvelope, type ChatAction, type IRootConfigChangedAction, type SessionAction, type TerminalAction, type INotification, type ClientAnnotationsAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; -import { buildDefaultChatUri, buildSubagentChatUri, createChatState, createDefaultChatSummary, MessageKind, SessionLifecycle, SessionStatus, createSessionState, StateComponents, parseDefaultChatUri, type ChatState, type SessionState, type SessionSummary, type RootState } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { buildDefaultChatUri, buildSubagentChatUri, createChatState, createDefaultChatSummary, ChatInputResponseKind, MessageKind, SessionLifecycle, SessionStatus, createSessionState, StateComponents, parseDefaultChatUri, ToolCallCancellationReason, type ChatState, type SessionState, type SessionSummary, type RootState } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { chatReducer, sessionReducer } from '../../../../../../platform/agentHost/common/state/sessionReducers.js'; import { ActionType } from '../../../../../../platform/agentHost/common/state/protocol/actions.js'; -import { ToolCallConfirmationReason, ToolCallContributorKind, ToolResultContentType } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { McpAuthRequiredReason, SessionInputRequestKind, ToolCallConfirmationReason, ToolCallContributorKind, ToolCallStatus, ToolResultContentType } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; import { IChatAgentService } from '../../../common/participants/chatAgents.js'; import { IChatProgress, IChatService, IChatToolInvocation, ToolConfirmKind } from '../../../common/chatService/chatService.js'; import { IChatEditingService } from '../../../common/editing/chatEditingService.js'; @@ -34,7 +35,7 @@ import { PieceCtorKind, PromptNodeType } from '../../../common/tools/promptTsxTy import { IProductService } from '../../../../../../platform/product/common/productService.js'; import { TestInstantiationService } from '../../../../../../platform/instantiation/test/common/instantiationServiceMock.js'; import { IWorkspaceContextService } from '../../../../../../platform/workspace/common/workspace.js'; -import { AgentHostSessionHandler, toolDataToDefinition, toolResultToProtocol } from '../../../browser/agentSessions/agentHost/agentHostSessionHandler.js'; +import { AgentHostSessionHandler, toolDataToDefinition, toolResultToProtocol, UNOBSERVED_CLIENT_TOOL_GRACE_MS } from '../../../browser/agentSessions/agentHost/agentHostSessionHandler.js'; import { AgentHostActiveClientService, IAgentHostActiveClientService } from '../../../browser/agentSessions/agentHost/agentHostActiveClientService.js'; import { IAgentHostCustomizationService, NullAgentHostCustomizationService } from '../../../browser/agentSessions/agentHost/agentHostCustomizationService.js'; import { IAgentHostToolSetEnablementService, IToolEnablementState } from '../../../browser/agentSessions/agentHost/agentHostToolSetEnablementService.js'; @@ -52,6 +53,7 @@ import { IAgentHostSessionWorkingDirectoryResolver } from '../../../browser/agen import { IAgentHostUntitledProvisionalSessionService } from '../../../browser/agentSessions/agentHost/agentHostUntitledProvisionalSessionService.js'; import { ILanguageModelToolsService, IToolData, IToolInvocation, IToolResult, ToolAndToolSetEnablementMap, ToolDataSource } from '../../../common/tools/languageModelToolsService.js'; import { IChatSessionsService } from '../../../common/chatSessionsService.js'; +import { IChatWidgetService } from '../../../browser/chat.js'; import { ICustomizationHarnessService } from '../../../common/customizationHarnessService.js'; import { IAgentPluginService } from '../../../common/plugins/agentPluginService.js'; import { IOutputService } from '../../../../../services/output/common/output.js'; @@ -531,6 +533,9 @@ suite('AgentHostClientTools', () => { registerChatSessionContentProvider: () => toDisposable(() => { }), registerChatSessionContribution: () => toDisposable(() => { }), }); + instantiationService.stub(IChatWidgetService, { + getWidgetBySessionResource: () => undefined, + }); instantiationService.stub(IDefaultAccountService, { onDidChangeDefaultAccount: Event.None, getDefaultAccount: async () => null }); instantiationService.stub(IAuthenticationService, { onDidChangeSessions: Event.None }); instantiationService.stub(ILanguageModelsService, { @@ -678,6 +683,18 @@ suite('AgentHostClientTools', () => { inputSchema: { type: 'object', properties: { query: { type: 'string' } } }, }; + // A tool that might ask for pre-approval: the handler treats it as + // requiring confirmation, so an unclaimed call waits for an observer. + const testConfirmTool: IToolData = { + id: 'vscode.deleteAll', + toolReferenceName: 'deleteAll', + displayName: 'Delete Everything', + modelDescription: 'A destructive action that needs confirmation', + source: ToolDataSource.Internal, + canRequestPreApproval: true, + inputSchema: { type: 'object', properties: {} }, + }; + async function provideSessionWithReadyRunTaskTool(handler: AgentHostSessionHandler, connection: MockAgentHostConnection): Promise { const sessionResource = URI.parse('agent-host-copilot:/session-1'); const backendSession = AgentSession.uri('copilot', 'session-1').toString(); @@ -706,6 +723,13 @@ suite('AgentHostClientTools', () => { } as ChatAction); await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + applyRunningClientExecution(connection, buildDefaultChatUri(backendSession), 'turn-1', { + toolCallId: 'tool-call-1', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + }); await timeout(0); await timeout(0); } @@ -736,6 +760,47 @@ suite('AgentHostClientTools', () => { }); } + // The watcher is the single point of truth for client-tool execution: + // it only acts on a `ToolClientExecution` blocker. Tests that drive a + // client tool through a chat turn must therefore also surface the + // matching running record so the tool actually runs. + function applyRunningClientExecution( + connection: MockAgentHostConnection, + chat: string, + turnId: string, + toolCall: { + toolCallId: string; + toolName: string; + displayName: string; + invocationMessage: string; + toolInput: string; + confirmed?: ToolCallConfirmationReason; + _meta?: Record; + }, + ): void { + connection.applySessionAction(URI.parse(AgentSession.uri('copilot', 'session-1').toString()), { + type: ActionType.SessionInputNeededSet, + request: { + id: `exec-${toolCall.toolCallId}`, + kind: SessionInputRequestKind.ToolClientExecution, + clientId: connection.clientId, + chat, + turnId, + toolCall: { + status: ToolCallStatus.Running, + toolCallId: toolCall.toolCallId, + toolName: toolCall.toolName, + displayName: toolCall.displayName, + invocationMessage: toolCall.invocationMessage, + toolInput: toolCall.toolInput, + confirmed: toolCall.confirmed ?? ToolCallConfirmationReason.NotNeeded, + contributor: { kind: ToolCallContributorKind.Client, clientId: connection.clientId }, + ...(toolCall._meta ? { _meta: toolCall._meta } : {}), + }, + }, + }); + } + test('maps tool data to protocol definitions', async () => { const { connection } = createHandlerWithMocks(disposables, [testRunTestsTool, testRunTaskTool, testUnlistedTool]); @@ -791,6 +856,13 @@ suite('AgentHostClientTools', () => { } as ChatAction); await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + applyRunningClientExecution(connection, buildDefaultChatUri(backendSession), 'turn-1', { + toolCallId: 'tool-call-1', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + }); await timeout(0); await timeout(0); @@ -844,6 +916,17 @@ suite('AgentHostClientTools', () => { } as ChatAction); await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + applyRunningClientExecution(connection, chatURI.toString(), 'turn-1', { + toolCallId: 'tool-search-call-1', + toolName: RUNTIME_TOOL_SEARCH_TOOL_NAME, + displayName: 'Search Tools', + invocationMessage: 'Search Tools', + toolInput: '{"query":"calculator"}', + _meta: { + toolSearchCandidates: [{ name: 'calculator', description: 'Adds numbers' }], + futureMetadata: { preserve: true }, + }, + }); await timeout(0); await timeout(0); @@ -897,6 +980,17 @@ suite('AgentHostClientTools', () => { } as ChatAction); await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + applyRunningClientExecution(connection, chatURI.toString(), 'turn-1', { + toolCallId: 'tool-search-call-invalid', + toolName: RUNTIME_TOOL_SEARCH_TOOL_NAME, + displayName: 'Search Tools', + invocationMessage: 'Search Tools', + toolInput: '{invalid', + _meta: { + toolSearchCandidates: [{ name: 'calculator', description: 'Adds numbers' }], + futureMetadata: { preserve: true }, + }, + }); await timeout(0); await timeout(0); @@ -1038,6 +1132,15 @@ suite('AgentHostClientTools', () => { } as ChatAction); await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + applyRunningClientExecution(connection, buildDefaultChatUri(backendSession), 'turn-1', { + toolCallId: 'tool-call-1', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + confirmed: ToolCallConfirmationReason.Setting, + _meta: { autoApproveBySetting: true }, + }); await timeout(0); await timeout(0); await timeout(0); @@ -1109,6 +1212,14 @@ suite('AgentHostClientTools', () => { } as ChatAction); await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + applyRunningClientExecution(connection, buildDefaultChatUri(backendSession), 'turn-1', { + toolCallId: 'tool-call-1', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); await timeout(0); await timeout(0); await timeout(0); @@ -1128,7 +1239,7 @@ suite('AgentHostClientTools', () => { ); }); - async function reachLocalWaitingForConfirmation(handler: AgentHostSessionHandler, connection: MockAgentHostConnection): Promise { + async function provideSessionWithPendingConfirmationClientTool(handler: AgentHostSessionHandler, connection: MockAgentHostConnection): Promise { const sessionResource = URI.parse('agent-host-copilot:/session-1'); const backendSession = AgentSession.uri('copilot', 'session-1').toString(); const chatURI = URI.parse(buildDefaultChatUri(backendSession)); @@ -1148,8 +1259,9 @@ suite('AgentHostClientTools', () => { contributor: { kind: ToolCallContributorKind.Client, clientId: connection.clientId }, } as ChatAction); // No `confirmed` and no auto-approve metadata: the protocol call - // stays `PendingConfirmation`, so the local invocation must reach - // `WaitingForConfirmation` and block on the confirmation gate. + // stays `PendingConfirmation`. Under the single-watcher model the + // client never drives a local confirmation gate, so nothing runs + // until the host surfaces a running client-execution record. connection.applySessionAction(chatURI, { type: ActionType.ChatToolCallReady, turnId: 'turn-1', @@ -1165,7 +1277,7 @@ suite('AgentHostClientTools', () => { return chatURI; } - test('resolves a waiting client tool confirmation when the agent host approves it late, preserving the reason', async () => { + test('confirms and completes a client tool once the agent host surfaces it as running, preserving the reason', async () => { const reasons = [ ToolCallConfirmationReason.NotNeeded, ToolCallConfirmationReason.Setting, @@ -1175,21 +1287,20 @@ suite('AgentHostClientTools', () => { const results: unknown[] = []; for (const reason of reasons) { const local = disposables.add(new DisposableStore()); - const { handler, connection, toolsService } = createHandlerWithMocks(local, [testRunTaskTool], { requireConfirmation: true }); - const chatURI = await reachLocalWaitingForConfirmation(handler, connection); - - const sawWaitingForConfirmation = (toolsService.recordedStateKinds.get('tool-call-1') ?? []).includes(IChatToolInvocation.StateKind.WaitingForConfirmation); + const { handler, connection } = createHandlerWithMocks(local, [testRunTaskTool], { requireConfirmation: true }); + const chatURI = await provideSessionWithPendingConfirmationClientTool(handler, connection); - // The agent host approves the call after the fact, transitioning - // the protocol tool call to `Running` with the resolved reason. - connection.applySessionAction(chatURI, { - type: ActionType.ChatToolCallReady, - turnId: 'turn-1', + // The agent host confirms the call by surfacing it as a running + // client execution with the resolved reason; the watcher then + // runs it pre-approved, so it never re-prompts locally. + applyRunningClientExecution(connection, chatURI.toString(), 'turn-1', { toolCallId: 'tool-call-1', + toolName: 'runTask', + displayName: 'Run Task', invocationMessage: 'Run Task', toolInput: '{"task":"build"}', confirmed: reason, - } as ChatAction); + }); await timeout(0); await timeout(0); @@ -1198,7 +1309,6 @@ suite('AgentHostClientTools', () => { && entry.action.toolCallId === 'tool-call-1'); results.push({ reason, - sawWaitingForConfirmation, dispatchedConfirmed: confirmedAction && confirmedAction.action.type === ActionType.ChatToolCallConfirmed && confirmedAction.action.approved ? confirmedAction.action.confirmed : undefined, @@ -1213,22 +1323,18 @@ suite('AgentHostClientTools', () => { assert.deepStrictEqual(results, reasons.map(reason => ({ reason, - sawWaitingForConfirmation: true, dispatchedConfirmed: reason, completed: true, }))); }); - test('does not confirm or execute a waiting client tool when the protocol call completes while still pending', async () => { + test('does not confirm or execute a pending client tool that completes without ever running', async () => { const { handler, connection, toolsService } = createHandlerWithMocks(disposables, [testRunTaskTool], { requireConfirmation: true }); - const chatURI = await reachLocalWaitingForConfirmation(handler, connection); + const chatURI = await provideSessionWithPendingConfirmationClientTool(handler, connection); - const sawWaitingForConfirmation = (toolsService.recordedStateKinds.get('tool-call-1') ?? []).includes(IChatToolInvocation.StateKind.WaitingForConfirmation); - - // The reducer synthesizes `confirmed: NotNeeded` when a completion - // arrives during `PendingConfirmation`. That is not evidence of a - // genuine approval, so the still-waiting local invocation must not - // be confirmed or driven through execution. + // The call completes while still `PendingConfirmation`, with no + // running client-execution record. The watcher never runs it, so it + // is never confirmed or driven through execution. connection.applySessionAction(chatURI, { type: ActionType.ChatToolCallComplete, turnId: 'turn-1', @@ -1239,14 +1345,14 @@ suite('AgentHostClientTools', () => { await timeout(0); assert.deepStrictEqual({ - sawWaitingForConfirmation, + invoked: toolsService.invokedToolCalls.filter(invocation => invocation.chatStreamToolCallId === 'tool-call-1').length, sawExecuting: (toolsService.recordedStateKinds.get('tool-call-1') ?? []).includes(IChatToolInvocation.StateKind.Executing), dispatchedApproval: connection.dispatchedActions.some(entry => isChatAction(entry.action) && entry.action.type === ActionType.ChatToolCallConfirmed && entry.action.toolCallId === 'tool-call-1' && entry.action.approved === true), }, { - sawWaitingForConfirmation: true, + invoked: 0, sawExecuting: false, dispatchedApproval: false, }); @@ -1301,6 +1407,573 @@ suite('AgentHostClientTools', () => { 'the initial snapshot invocation should be completed, not orphaned'); }); + test('auto-denies an unclaimed session confirmation after the grace period', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection } = createHandlerWithMocks(disposables, []); + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const subagentChat = buildSubagentChatUri(backendSession, 'task-call-1'); + await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + + // No turn observer ever renders this confirmation, so nothing can + // answer it; the watcher denies it once the grace window expires. + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request: { + id: 'approval-1', + kind: SessionInputRequestKind.ToolConfirmation, + chat: subagentChat, + turnId: 'subagent-turn-1', + toolCall: { + status: ToolCallStatus.PendingConfirmation, + toolCallId: 'powershell-call-1', + toolName: 'powershell', + displayName: 'PowerShell', + invocationMessage: 'Run PowerShell', + }, + }, + }); + await timeout(UNOBSERVED_CLIENT_TOOL_GRACE_MS + 1); + + assert.deepStrictEqual( + connection.dispatchedActions + .filter(entry => entry.action.type === ActionType.ChatToolCallConfirmed && entry.action.toolCallId === 'powershell-call-1') + .map(entry => ({ channel: entry.channel, action: entry.action })), + [{ + channel: subagentChat, + action: { + type: ActionType.ChatToolCallConfirmed, + turnId: 'subagent-turn-1', + toolCallId: 'powershell-call-1', + approved: false, + reason: ToolCallCancellationReason.Denied, + }, + }], + ); + })); + + test('cancels an unclaimed chat input request after the grace period', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection } = createHandlerWithMocks(disposables, []); + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const subagentChat = buildSubagentChatUri(backendSession, 'task-call-1'); + await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + + // No turn observer renders this elicitation, so nothing can answer + // it; the watcher cancels it once the grace window expires. + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request: { + id: 'input-1', + kind: SessionInputRequestKind.ChatInput, + chat: subagentChat, + request: { id: 'elicit-1', message: 'Pick one', questions: [] }, + }, + }); + await timeout(5001); + + assert.deepStrictEqual( + connection.dispatchedActions + .filter(entry => entry.action.type === ActionType.ChatInputCompleted) + .map(entry => ({ channel: entry.channel, action: entry.action })), + [{ + channel: subagentChat, + action: { + type: ActionType.ChatInputCompleted, + requestId: 'elicit-1', + response: ChatInputResponseKind.Cancel, + }, + }], + ); + })); + + test('does not cancel a chat input request a turn observer is rendering', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection } = createHandlerWithMocks(disposables, []); + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const chatURI = buildDefaultChatUri(backendSession); + + // The default-chat turn observer renders the elicitation, so it + // claims the request and the watcher must leave it alone. + connection.applySessionAction(URI.parse(chatURI), { + type: ActionType.ChatTurnStarted, + turnId: 'turn-1', + startedAt: '2025-01-01T00:00:00.000Z', + message: { text: 'ask me', origin: { kind: MessageKind.User } }, + } as ChatAction); + connection.applySessionAction(URI.parse(chatURI), { + type: ActionType.ChatInputRequested, + request: { id: 'elicit-1', message: 'Pick one', questions: [] }, + } as ChatAction); + await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + await timeout(0); + + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request: { + id: 'input-1', + kind: SessionInputRequestKind.ChatInput, + chat: chatURI, + request: { id: 'elicit-1', message: 'Pick one', questions: [] }, + }, + }); + await timeout(5001); + + assert.strictEqual(connection.dispatchedActions.some(entry => entry.action.type === ActionType.ChatInputCompleted), false); + + // Settle the elicitation so the rendered carousel's cancellation + // listener is disposed before teardown. + connection.applySessionAction(URI.parse(chatURI), { + type: ActionType.ChatInputCompleted, + requestId: 'elicit-1', + response: ChatInputResponseKind.Cancel, + } as ChatAction); + await timeout(0); + })); + + test('cancels an unclaimed MCP authentication tool call after the grace period', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection } = createHandlerWithMocks(disposables, []); + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const subagentChat = buildSubagentChatUri(backendSession, 'task-call-1'); + await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + + // No turn observer renders this auth-required MCP tool call, so + // nobody can drive authentication; the watcher cancels it once the + // grace window expires. + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request: { + id: 'auth-1', + kind: SessionInputRequestKind.ToolAuthentication, + chat: subagentChat, + turnId: 'subagent-turn-1', + toolCall: { + status: ToolCallStatus.AuthRequired, + toolCallId: 'mcp-call-1', + toolName: 'notionSearch', + displayName: 'Notion Search', + invocationMessage: 'Search Notion', + confirmed: ToolCallConfirmationReason.UserAction, + contributor: { kind: ToolCallContributorKind.MCP, customizationId: 'notion-mcp' }, + auth: { reason: McpAuthRequiredReason.Required, resource: { resource: 'https://mcp.notion.com/mcp', authorization_servers: [] } }, + }, + }, + }); + await timeout(5001); + + assert.deepStrictEqual( + connection.dispatchedActions + .filter(entry => entry.action.type === ActionType.ChatToolCallComplete && entry.action.toolCallId === 'mcp-call-1') + .map(entry => ({ channel: entry.channel, action: entry.action })), + [{ + channel: subagentChat, + action: { + type: ActionType.ChatToolCallComplete, + turnId: 'subagent-turn-1', + toolCallId: 'mcp-call-1', + result: { + success: false, + pastTenseMessage: 'Cancelled tool call', + error: { message: 'MCP authentication was cancelled', code: 'cancelled' }, + }, + }, + }], + ); + })); + + test('does not cancel an MCP authentication tool call a turn observer is rendering', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection } = createHandlerWithMocks(disposables, []); + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const chatURI = buildDefaultChatUri(backendSession); + + // The default-chat observer renders the MCP tool call as it pauses + // for authentication, so it claims the call and the watcher must + // leave it alone. + connection.applySessionAction(URI.parse(chatURI), { + type: ActionType.ChatTurnStarted, + turnId: 'turn-1', + startedAt: '2025-01-01T00:00:00.000Z', + message: { text: 'search notion', origin: { kind: MessageKind.User } }, + } as ChatAction); + connection.applySessionAction(URI.parse(chatURI), { + type: ActionType.ChatToolCallStart, + turnId: 'turn-1', + toolCallId: 'mcp-call-1', + toolName: 'notionSearch', + displayName: 'Notion Search', + contributor: { kind: ToolCallContributorKind.MCP, customizationId: 'notion-mcp' }, + } as ChatAction); + connection.applySessionAction(URI.parse(chatURI), { + type: ActionType.ChatToolCallReady, + turnId: 'turn-1', + toolCallId: 'mcp-call-1', + invocationMessage: 'Search Notion', + toolInput: '{}', + confirmed: ToolCallConfirmationReason.NotNeeded, + } as ChatAction); + await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + await timeout(0); + connection.applySessionAction(URI.parse(chatURI), { + type: ActionType.ChatToolCallAuthRequired, + turnId: 'turn-1', + toolCallId: 'mcp-call-1', + auth: { reason: McpAuthRequiredReason.Required, resource: { resource: 'https://mcp.notion.com/mcp', authorization_servers: [] } }, + } as ChatAction); + await timeout(0); + + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request: { + id: 'auth-1', + kind: SessionInputRequestKind.ToolAuthentication, + chat: chatURI, + turnId: 'turn-1', + toolCall: { + status: ToolCallStatus.AuthRequired, + toolCallId: 'mcp-call-1', + toolName: 'notionSearch', + displayName: 'Notion Search', + invocationMessage: 'Search Notion', + confirmed: ToolCallConfirmationReason.UserAction, + contributor: { kind: ToolCallContributorKind.MCP, customizationId: 'notion-mcp' }, + auth: { reason: McpAuthRequiredReason.Required, resource: { resource: 'https://mcp.notion.com/mcp', authorization_servers: [] } }, + }, + }, + }); + await timeout(5001); + + assert.strictEqual(connection.dispatchedActions.some(entry => entry.action.type === ActionType.ChatToolCallComplete && entry.action.toolCallId === 'mcp-call-1'), false); + })); + + test('renders a subagent client tool as the same invocation the watcher executes', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + // The subagent observer renders the shared invocation and the + // watcher executes it: both act on one object, invoked exactly once, + // and the card renders in the subagent's own group. + const { handler, connection, toolsService } = createHandlerWithMocks(disposables, [testSubagentTool, testRunTaskTool]); + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const parentToolCallId = 'client-task-1'; + const subagentChat = buildSubagentChatUri(backendSession, parentToolCallId); + const parentChat = URI.parse(buildDefaultChatUri(backendSession)); + + connection.applySessionAction(parentChat, { + type: ActionType.ChatTurnStarted, + turnId: 'turn-1', + startedAt: '2025-01-01T00:00:00.000Z', + message: { text: 'delegate work', origin: { kind: MessageKind.User } }, + }); + connection.applySessionAction(parentChat, { + type: ActionType.ChatToolCallStart, + turnId: 'turn-1', + toolCallId: parentToolCallId, + toolName: 'task', + displayName: 'Delegated Task', + contributor: { kind: ToolCallContributorKind.Client, clientId: connection.clientId }, + _meta: { toolKind: 'subagent', subagentChatUri: subagentChat }, + }); + const session = await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + await timeout(0); + connection.applySessionAction(parentChat, { + type: ActionType.ChatToolCallReady, + turnId: 'turn-1', + toolCallId: parentToolCallId, + invocationMessage: 'Delegating task', + toolInput: '{}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); + + // The subagent runs a client tool. + connection.applySessionAction(URI.parse(subagentChat), { + type: ActionType.ChatTurnStarted, + turnId: 'sub-turn-1', + startedAt: '2025-01-01T00:00:00.000Z', + message: { text: '', origin: { kind: MessageKind.User } }, + }); + connection.applySessionAction(URI.parse(subagentChat), { + type: ActionType.ChatToolCallStart, + turnId: 'sub-turn-1', + toolCallId: 'runTask-call-1', + toolName: 'runTask', + displayName: 'Run Task', + contributor: { kind: ToolCallContributorKind.Client, clientId: connection.clientId }, + }); + connection.applySessionAction(URI.parse(subagentChat), { + type: ActionType.ChatToolCallReady, + turnId: 'sub-turn-1', + toolCallId: 'runTask-call-1', + invocationMessage: 'Run Task', + toolInput: '{}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); + await timeout(0); + + // The host reports it as a running client-execution obligation. + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request: { + id: 'exec-1', + kind: SessionInputRequestKind.ToolClientExecution, + clientId: connection.clientId, + chat: subagentChat, + turnId: 'sub-turn-1', + toolCall: { + status: ToolCallStatus.Running, + toolCallId: 'runTask-call-1', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }, + }, + }); + await timeout(0); + + const rendered = (session as unknown as { progressObs: { get(): IChatProgress[] } }).progressObs.get() + .find((part): part is ChatToolInvocation => part instanceof ChatToolInvocation && part.toolCallId === 'runTask-call-1'); + + assert.deepStrictEqual({ + renderedInSubagentGroup: rendered?.subAgentInvocationId, + renderedIsTheBegunInvocation: rendered === toolsService.begunToolCalls.find(inv => inv.toolCallId === 'runTask-call-1'), + begun: toolsService.begunToolCalls.filter(inv => inv.toolCallId === 'runTask-call-1').length, + invoked: toolsService.invokedToolCalls.filter(inv => inv.chatStreamToolCallId === 'runTask-call-1').length, + }, { + renderedInSubagentGroup: parentToolCallId, + renderedIsTheBegunInvocation: true, + begun: 1, + invoked: 1, + }); + })); + + test('runs an unclaimed non-confirmable client tool headlessly without waiting for the grace window', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection, toolsService } = createHandlerWithMocks(disposables, [testRunTaskTool]); + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const subagentChat = buildSubagentChatUri(backendSession, 'task-call-1'); + await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request: { + id: 'execution-1', + kind: SessionInputRequestKind.ToolClientExecution, + chat: subagentChat, + turnId: 'subagent-turn-1', + clientId: connection.clientId, + toolCall: { + status: ToolCallStatus.Running, + toolCallId: 'client-tool-1', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + confirmed: ToolCallConfirmationReason.NotNeeded, + contributor: { kind: ToolCallContributorKind.Client, clientId: connection.clientId }, + }, + }, + }); + // No grace wait: a non-confirmable tool that nobody is rendering + // runs immediately and headlessly. + await timeout(0); + + assert.deepStrictEqual({ + // Executed headlessly: no chat `context`, so the invocation does + // not depend on the owning turn still being live. + invocation: toolsService.invokedToolCalls.map(call => ({ + callId: call.callId, + parameters: call.parameters, + hasContext: call.context !== undefined, + preApprovedKind: call.preApproved?.type, + })), + completion: connection.dispatchedActions.find(entry => + entry.channel === subagentChat + && entry.action.type === ActionType.ChatToolCallComplete), + }, { + invocation: [{ + callId: 'client-tool-1', + parameters: { task: 'build' }, + hasContext: false, + preApprovedKind: ToolConfirmKind.ConfirmationNotNeeded, + }], + completion: { + channel: subagentChat, + action: { + type: ActionType.ChatToolCallComplete, + turnId: 'subagent-turn-1', + toolCallId: 'client-tool-1', + result: { + success: true, + pastTenseMessage: 'Ran runTask', + content: [{ type: 'text', text: 'done' }], + error: undefined, + }, + }, + }, + }); + })); + + test('executes a claimed client tool exactly once, with chat context', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection, toolsService } = createHandlerWithMocks(disposables, [testRunTaskTool]); + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const chat = buildDefaultChatUri(backendSession); + connection.applySessionAction(URI.parse(chat), { + type: ActionType.ChatTurnStarted, + turnId: 'turn-1', + startedAt: '2025-01-01T00:00:00.000Z', + message: { text: 'run the task', origin: { kind: MessageKind.User } }, + }); + connection.applySessionAction(URI.parse(chat), { + type: ActionType.ChatToolCallStart, + turnId: 'turn-1', + toolCallId: 'client-tool-1', + toolName: 'runTask', + displayName: 'Run Task', + contributor: { kind: ToolCallContributorKind.Client, clientId: connection.clientId }, + }); + connection.applySessionAction(URI.parse(chat), { + type: ActionType.ChatToolCallReady, + turnId: 'turn-1', + toolCallId: 'client-tool-1', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); + await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request: { + id: 'execution-1', + kind: SessionInputRequestKind.ToolClientExecution, + chat, + turnId: 'turn-1', + clientId: connection.clientId, + toolCall: { + status: ToolCallStatus.Running, + toolCallId: 'client-tool-1', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + confirmed: ToolCallConfirmationReason.NotNeeded, + contributor: { kind: ToolCallContributorKind.Client, clientId: connection.clientId }, + }, + }, + }); + await timeout(5001); + + assert.deepStrictEqual({ + // A live turn observer renders the call, so the watcher runs it + // once with chat context (not per-observer, not headless). + invocations: toolsService.invokedToolCalls + .filter(invocation => invocation.chatStreamToolCallId === 'client-tool-1') + .map(invocation => invocation.context !== undefined), + declines: connection.dispatchedActions.filter(entry => + entry.action.type === ActionType.ChatToolCallComplete + && entry.action.result.error?.code === 'clientUnavailable').length, + }, { + invocations: [true], + declines: 0, + }); + })); + + test('denies an unclaimed confirmable client tool after the grace window without executing it', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection, toolsService } = createHandlerWithMocks(disposables, [testConfirmTool]); + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const subagentChat = buildSubagentChatUri(backendSession, 'task-call-1'); + await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + + // A tool that might ask for confirmation, with no observer to render + // it: running headlessly would pop a modal nobody could answer, so + // the watcher waits and then denies once the grace window expires. + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request: { + id: 'execution-1', + kind: SessionInputRequestKind.ToolClientExecution, + chat: subagentChat, + turnId: 'subagent-turn-1', + clientId: connection.clientId, + toolCall: { + status: ToolCallStatus.Running, + toolCallId: 'client-tool-1', + toolName: 'deleteAll', + displayName: 'Delete Everything', + invocationMessage: 'Delete everything', + toolInput: '{}', + confirmed: ToolCallConfirmationReason.UserAction, + contributor: { kind: ToolCallContributorKind.Client, clientId: connection.clientId }, + }, + }, + }); + await timeout(5001); + + assert.deepStrictEqual({ + invocations: toolsService.invokedToolCalls.filter(invocation => invocation.chatStreamToolCallId === 'client-tool-1').length, + denial: connection.dispatchedActions.find(entry => + entry.channel === subagentChat + && entry.action.type === ActionType.ChatToolCallComplete + && entry.action.toolCallId === 'client-tool-1')?.action, + }, { + invocations: 0, + denial: { + type: ActionType.ChatToolCallComplete, + turnId: 'subagent-turn-1', + toolCallId: 'client-tool-1', + result: { + success: false, + pastTenseMessage: 'Couldn\'t run Delete Everything', + error: { + message: 'Delete Everything needs confirmation but no session was available to answer it.', + code: 'clientUnavailable', + }, + }, + }, + }); + })); + + test('does not run foreign or already-resolved client tools', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection } = createHandlerWithMocks(disposables, [testRunTaskTool]); + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const subagentChat = buildSubagentChatUri(backendSession, 'task-call-1'); + await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + const request = { + id: 'execution-1', + kind: SessionInputRequestKind.ToolClientExecution, + chat: subagentChat, + turnId: 'subagent-turn-1', + clientId: 'other-client', + toolCall: { + status: ToolCallStatus.Running, + toolCallId: 'client-tool-1', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + confirmed: ToolCallConfirmationReason.NotNeeded, + contributor: { kind: ToolCallContributorKind.Client, clientId: 'other-client' }, + }, + } as const; + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request, + }); + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededSet, + request: { ...request, id: 'execution-2', clientId: connection.clientId }, + }); + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionInputNeededRemoved, + id: 'execution-2', + }); + await timeout(5001); + + assert.strictEqual(connection.dispatchedActions.some(entry => entry.action.type === ActionType.ChatToolCallComplete), false); + })); + test('invokes a client tool inside a subagent session and dispatches completion against the subagent URI', async () => { // Regression: a client-provided tool running inside a subagent // must be invoked locally (the renderer owns the tool @@ -1373,6 +2046,15 @@ suite('AgentHostClientTools', () => { }); await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + applyRunningClientExecution(connection, subagentChat, 'sub-turn-1', { + toolCallId: 'inner-tool-call-1', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); + await timeout(0); await timeout(0); // The inner client tool must have been invoked locally — without @@ -1435,7 +2117,18 @@ suite('AgentHostClientTools', () => { toolInput: '{}', confirmed: ToolCallConfirmationReason.NotNeeded, }); - + // The delegated `task` tool is client-contributed, so the watcher + // runs it locally; invoking it is what prepares the subagent + // container (mock sets the `Prepared delegated task` description). + applyRunningClientExecution(connection, buildDefaultChatUri(backendSession), 'turn-1', { + toolCallId: parentToolCallId, + toolName: 'task', + displayName: 'Delegated Task', + invocationMessage: 'Delegating task', + toolInput: '{}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); + await timeout(0); connection.applySessionAction(URI.parse(subagentChat), { type: ActionType.ChatTurnStarted, turnId: 'sub-turn-1', @@ -1551,6 +2244,14 @@ suite('AgentHostClientTools', () => { }); await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + applyRunningClientExecution(connection, subagentChat2, 'sub-turn-2', { + toolCallId: 'deep-tool-call', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); for (let i = 0; i < 200 && !connection.dispatchedActions.some(e => isChatAction(e.action) && e.action.type === ActionType.ChatToolCallComplete && e.action.toolCallId === 'deep-tool-call'); i++) { await timeout(1); } @@ -1633,6 +2334,14 @@ suite('AgentHostClientTools', () => { }); await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + applyRunningClientExecution(connection, subagentChat2, 'sub-turn-2', { + toolCallId: 'deep-tool-call', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); for (let i = 0; i < 200 && !connection.dispatchedActions.some(e => isChatAction(e.action) && e.action.type === ActionType.ChatToolCallComplete && e.action.toolCallId === 'deep-tool-call'); i++) { await timeout(1); } From a4f501e79363cbd9446244b982ce309ff8253b3c Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Tue, 4 Aug 2026 09:35:06 -0700 Subject: [PATCH 2/2] agentHost: share one input-needed watcher per backend session Sibling resources (default, peer and subagent chats) can be open against the same backend session at once, and each installed its own session-level watcher over the same inputNeeded queue. Each had independent per-request state, so one client-tool request executed the tool once per open resource; _resolveToolCall only deduplicates the eventual dispatch, long after the tool's side effects have already run N times. Ref-count a single watcher per backend session instead, keeping it alive while any sibling holds a reference. The resource-to-backend mapping is recorded at install time rather than resolved during teardown, when provisional session state may already be gone. The claim registry now records which observer is rendering a request, so a claimed tool executes with that observer's chat context instead of whichever sibling happened to install the watcher. Also reattach the withInputNeededStatus documentation, which described the old "any non-empty queue" rule and had come loose from its function. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../protocol/channels-session/reducer.ts | 18 +-- .../agentHost/agentHostSessionHandler.ts | 137 +++++++++++++----- .../agentHostClientTools.test.ts | 96 +++++++++++- 3 files changed, 207 insertions(+), 44 deletions(-) diff --git a/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts b/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts index 10e2bf4db2d64..04890581f8d47 100644 --- a/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts +++ b/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts @@ -21,15 +21,6 @@ function withStatusFlag(status: SessionStatus, flag: SessionStatus, set: boolean return set ? status | flag : status & ~flag; } -/** - * Reflects the session-level {@link SessionState.inputNeeded | input queue} - * into the activity bits of `status`. A non-empty queue promotes the activity - * to {@link SessionStatus.InputNeeded}; emptying it clears the - * input-needed-specific bit. Since `InputNeeded` implies - * {@link SessionStatus.InProgress}, an unblocked turn falls back to - * `InProgress` while an already-idle session stays idle. Orthogonal flags - * (`IsRead` / `IsArchived`) are preserved. - */ /** * Whether an entry blocks on the *user*. * @@ -42,6 +33,15 @@ function awaitsUser(request: SessionInputRequest): boolean { return request.kind !== SessionInputRequestKind.ToolClientExecution; } +/** + * Reflects the session-level {@link SessionState.inputNeeded | input queue} + * into the activity bits of `status`. A queue holding any user-blocking entry + * promotes the activity to {@link SessionStatus.InputNeeded}; draining those + * entries clears the input-needed-specific bit. Since `InputNeeded` implies + * {@link SessionStatus.InProgress}, an unblocked turn falls back to + * `InProgress` while an already-idle session stays idle. Orthogonal flags + * (`IsRead` / `IsArchived`) are preserved. + */ function withInputNeededStatus(status: SessionStatus, inputNeeded: readonly SessionInputRequest[]): SessionStatus { if (inputNeeded.some(awaitsUser)) { return (status & ~STATUS_ACTIVITY_MASK) | SessionStatus.InputNeeded; diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts index be1cf2689f73b..74b668a2274ee 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -763,18 +763,37 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC private readonly _serverTurnWatchers = this._register(new DisposableResourceMap()); /** Per-session subscription silently resolving existing MCP authentication grants. */ private readonly _mcpAuthWatchers = this._register(new DisposableResourceMap()); - /** Per-session ownership of actionable protocol requests. */ - private readonly _inputNeededWatchers = this._register(new DisposableResourceMap()); + /** + * Ownership of actionable protocol requests, keyed by backend session URI + * string. `inputNeeded` is a session-level queue and the single caller of + * {@link invokeTool} for client tools, so it must be handled exactly once + * per backend session no matter how many sibling chat resources (default + * chat, peer chats, subagent chats) are open against it. Each such resource + * holds a reference; the shared watcher stays alive while any reference + * remains and is disposed only when the last one is released. + */ + private readonly _inputNeededWatchers = new Map }>(); + /** + * Backend session each open resource's {@link _inputNeededWatchers} + * reference belongs to, recorded when the reference is installed. Teardown + * uses this to release the right reference without re-deriving the backend + * session via {@link _resolveSessionUri}, whose provisional mapping may + * already be cleared by then. + */ + private readonly _inputNeededWatcherBackends = new ResourceMap(); /** Historical turns with file edits, pending hydration into the editing session. */ private readonly _pendingHistoryTurns = new ResourceMap(); /** * Requests a turn observer is currently rendering, keyed by * {@link _toolCallKey} for tool calls and {@link _inputRequestKey} for chat * input requests (the two key shapes differ in arity, so they cannot - * collide). The session-level responder defers to those observers so the - * inline UI stays in charge of answering. + * collide). The value is the claiming observer's session resource, which + * the session-level responder uses as the chat context when it executes a + * client tool so the tool runs against the chat that is actually rendering + * it. The session-level responder defers to those observers so the inline + * UI stays in charge of answering. */ - private readonly _renderedRequests = observableValue>(this, new Set()); + private readonly _renderedRequests = observableValue>(this, new Map()); /** Tool calls whose protocol outcome has already been dispatched. */ private readonly _resolvedToolCalls = new Set(); /** @@ -869,6 +888,16 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC super(); this._config = config; + // The `inputNeeded` watchers live in a plain map (they are shared and + // ref-counted across sibling resources), so dispose any that survive + // when the handler goes away. + this._register(toDisposable(() => { + for (const { store } of this._inputNeededWatchers.values()) { + store.dispose(); + } + this._inputNeededWatchers.clear(); + this._inputNeededWatcherBackends.clear(); + })); // Drop MCP servers from the per-session surfaced set once they reach the // running state so a later auth requirement for the same server prompts // again. @@ -1278,7 +1307,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC this._draftSyncSubscriptions.deleteAndDispose(sessionResource); this._serverTurnWatchers.deleteAndDispose(sessionResource); this._mcpAuthWatchers.deleteAndDispose(sessionResource); - this._inputNeededWatchers.deleteAndDispose(sessionResource); + this._releaseSessionInputNeeded(sessionResource); this._pendingHistoryTurns.delete(sessionResource); this._surfacedMcpAuthServers.delete(sessionResource); const chatURI = this._chatURIsBySessionResource.get(sessionResource); @@ -1973,9 +2002,24 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } private _watchForSessionInputNeeded(backendSession: URI, sessionResource: URI): void { - const sessionSub = this._ensureSessionSubscription(backendSession.toString()); + // Record which backend session this resource's reference belongs to so + // teardown can release it even after provisional state is cleared. + this._inputNeededWatcherBackends.set(sessionResource, backendSession); + + const sessionKey = backendSession.toString(); + const existing = this._inputNeededWatchers.get(sessionKey); + if (existing) { + // Sibling resources against the same backend session share the one + // watcher: only add a reference so the single session-level queue + // is not handled — and client tools not executed — more than once. + existing.refs.add(sessionResource.toString()); + return; + } + + const sessionSub = this._ensureSessionSubscription(sessionKey); const state = observableFromSubscription(this, sessionSub); const store = new DisposableStore(); + this._inputNeededWatchers.set(sessionKey, { store, refs: new Set([sessionResource.toString()]) }); const requests = derivedOpts({ equalsFn: equals }, reader => (state.read(reader)?.inputNeeded ?? []).filter((request): request is SessionInputRequest => @@ -2023,30 +2067,33 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC return; // A different client owns this call. } let handled = false; - const execute = (withContext: boolean) => { + const execute = (contextSessionResource: URI | undefined) => { if (handled) { return; } handled = true; - void this._executeClientTool(request$.get() as SessionToolClientExecutionRequest, sessionResource, withContext, cts.token); + void this._executeClientTool(request$.get() as SessionToolClientExecutionRequest, contextSessionResource, cts.token); }; - if (this._renderedRequests.get().has(key)) { + const claimant = this._renderedRequests.get().get(key); + if (claimant) { // A turn observer is rendering it, so a live chat request - // exists: run with context so confirmation renders in the - // tool part and any pre-approval is honored. - execute(true); + // exists: run with that observer's session as context so + // confirmation renders in the tool part, any pre-approval + // is honored, and side effects attribute to the right chat. + execute(claimant); } else if (!this._clientToolRequiresConfirmation(initial.toolCall)) { // Unclaimed and cannot pop a confirmation: run headlessly so // it does not depend on the owning turn still being live. - execute(false); + execute(undefined); } else { // Unclaimed and might pop a confirmation: a headless run // would surface a modal nobody could see. Wait for an // observer to claim it; if none does within the grace // window, deny it. itemStore.add(autorun(reader => { - if (!handled && this._renderedRequests.read(reader).has(key)) { - execute(true); + const claimed = this._renderedRequests.read(reader).get(key); + if (!handled && claimed) { + execute(claimed); } })); itemStore.add(disposableTimeout(() => { @@ -2095,8 +2142,29 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC }, UNOBSERVED_CLIENT_TOOL_GRACE_MS)); } })); + } - this._inputNeededWatchers.set(sessionResource, store); + /** + * Releases this resource's reference to the shared per-backend-session + * {@link _watchForSessionInputNeeded} watcher, disposing it only once the + * last sibling resource has let go. + */ + private _releaseSessionInputNeeded(sessionResource: URI): void { + const backendSession = this._inputNeededWatcherBackends.get(sessionResource); + this._inputNeededWatcherBackends.delete(sessionResource); + if (!backendSession) { + return; + } + const sessionKey = backendSession.toString(); + const entry = this._inputNeededWatchers.get(sessionKey); + if (!entry) { + return; + } + entry.refs.delete(sessionResource.toString()); + if (entry.refs.size === 0) { + this._inputNeededWatchers.delete(sessionKey); + entry.store.dispose(); + } } /** @@ -2169,12 +2237,13 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC * The one place a client tool is actually invoked. Ensures the shared * invocation exists, parses the protocol input (preserving the tool-search * candidate handling), invokes the tool, and dispatches the protocol - * completion. `withContext` is set when a turn observer is rendering the - * call: a live chat request then exists, so confirmation renders in the - * tool part and any pre-approval is honored. Without it the tool runs - * headlessly, independent of whether the owning turn is live. + * completion. `contextSessionResource` is set when a turn observer is + * rendering the call: a live chat request then exists, so confirmation + * renders in the tool part, any pre-approval is honored, and side effects + * attribute to that observer's chat. Without it the tool runs headlessly, + * independent of whether the owning turn is live. */ - private async _executeClientTool(request: SessionToolClientExecutionRequest, sessionResource: URI, withContext: boolean, token: CancellationToken): Promise { + private async _executeClientTool(request: SessionToolClientExecutionRequest, contextSessionResource: URI | undefined, token: CancellationToken): Promise { const chatURI = request.chat.toString(); const toolCall = request.toolCall; const toolName = toolCall.toolName; @@ -2228,7 +2297,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC return; } - this._logService.info(`[AgentHost] Running client tool: ${toolName} (callId=${toolCall.toolCallId}, withContext=${withContext})`); + this._logService.info(`[AgentHost] Running client tool: ${toolName} (callId=${toolCall.toolCallId}, withContext=${contextSessionResource !== undefined})`); let result: IToolResult | undefined; let error: unknown; try { @@ -2236,7 +2305,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC callId: toolCall.toolCallId, toolId: toolData.id, parameters, - context: withContext ? { sessionResource } : undefined, + context: contextSessionResource ? { sessionResource: contextSessionResource } : undefined, chatStreamToolCallId: toolCall.toolCallId, preApproved: getClientToolPreApproval(toolCall), }, async () => 0, token); @@ -3100,11 +3169,11 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC // watcher it may execute this call, and it must find the shared // invocation already created when it does. this._setupClientToolCall(initial, part$, store, opts, subagentContext); - store.add(this._markToolCallRendered(opts.chatURI, opts.turnId, initial.toolCallId)); + store.add(this._markToolCallRendered(opts.chatURI, opts.turnId, initial.toolCallId, opts.sessionResource)); } else if (contributor?.kind === ToolCallContributorKind.Client) { this._setupOtherClientToolCall(initial, part$, store, opts); } else { - store.add(this._markToolCallRendered(opts.chatURI, opts.turnId, initial.toolCallId)); + store.add(this._markToolCallRendered(opts.chatURI, opts.turnId, initial.toolCallId, opts.sessionResource)); this._setupServerToolCall(initial, part$, store, opts, subagentContext); } } @@ -3118,10 +3187,10 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } /** Claims a request as rendered until the returned disposable is disposed. */ - private _markRendered(key: string): IDisposable { - this._renderedRequests.set(new Set(this._renderedRequests.get()).add(key), undefined); + private _markRendered(key: string, sessionResource: URI): IDisposable { + this._renderedRequests.set(new Map(this._renderedRequests.get()).set(key, sessionResource), undefined); return toDisposable(() => { - const next = new Set(this._renderedRequests.get()); + const next = new Map(this._renderedRequests.get()); next.delete(key); this._renderedRequests.set(next, undefined); }); @@ -3131,8 +3200,8 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC * Records that a turn observer is rendering this chat input request, so the * session-level responder leaves its inline elicitation UI in charge. */ - private _markInputRequestRendered(chatURI: string, requestId: string): IDisposable { - return this._markRendered(this._inputRequestKey(chatURI, requestId)); + private _markInputRequestRendered(chatURI: string, requestId: string, sessionResource: URI): IDisposable { + return this._markRendered(this._inputRequestKey(chatURI, requestId), sessionResource); } /** @@ -3141,9 +3210,9 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC * claim also forgets the funnel entries, which is the only cleanup a tool * call that never reached `inputNeeded` ever gets. */ - private _markToolCallRendered(chatURI: string, turnId: string, toolCallId: string): IDisposable { + private _markToolCallRendered(chatURI: string, turnId: string, toolCallId: string, sessionResource: URI): IDisposable { const key = this._toolCallKey(chatURI, turnId, toolCallId); - const rendered = this._markRendered(key); + const rendered = this._markRendered(key, sessionResource); return toDisposable(() => { rendered.dispose(); this._forgetResolvedToolCall(key); @@ -3591,7 +3660,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC // Claim the elicitation so the session-level responder does not cancel // it while an observer is rendering it. This covers all three render // paths below, since each is reached only through this method. - store.add(this._markInputRequestRendered(opts.chatURI, inputReq.id)); + store.add(this._markInputRequestRendered(opts.chatURI, inputReq.id, opts.sessionResource)); const planReview = (inputReq as ChatInputRequestWithPlanReview).planReview; if (planReview) { this._setupPlanReviewInputRequest(part$, planReview, store, opts); diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts index 9ff6aca18394c..c28886cd50817 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts @@ -21,7 +21,7 @@ import { IConfigurationChangeEvent, IConfigurationService } from '../../../../.. import { AgentSession, IAgentHostService } from '../../../../../../platform/agentHost/common/agentService.js'; import { CLIENT_TOOL_SEARCH_REFERENCE_NAME, RUNTIME_TOOL_SEARCH_TOOL_NAME } from '../../../../../../platform/agentHost/common/toolSearchConstants.js'; import { isChatAction, isSessionAction, type ActionEnvelope, type ChatAction, type IRootConfigChangedAction, type SessionAction, type TerminalAction, type INotification, type ClientAnnotationsAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; -import { buildDefaultChatUri, buildSubagentChatUri, createChatState, createDefaultChatSummary, ChatInputResponseKind, MessageKind, SessionLifecycle, SessionStatus, createSessionState, StateComponents, parseDefaultChatUri, ToolCallCancellationReason, type ChatState, type SessionState, type SessionSummary, type RootState } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { buildChatUri, buildDefaultChatUri, buildSubagentChatUri, createChatState, createDefaultChatSummary, ChatInputResponseKind, MessageKind, SessionLifecycle, SessionStatus, createSessionState, StateComponents, parseDefaultChatUri, ToolCallCancellationReason, type ChatState, type SessionState, type SessionSummary, type RootState } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { chatReducer, sessionReducer } from '../../../../../../platform/agentHost/common/state/sessionReducers.js'; import { ActionType } from '../../../../../../platform/agentHost/common/state/protocol/actions.js'; import { McpAuthRequiredReason, SessionInputRequestKind, ToolCallConfirmationReason, ToolCallContributorKind, ToolCallStatus, ToolResultContentType } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; @@ -1878,6 +1878,100 @@ suite('AgentHostClientTools', () => { }); })); + // Two sibling resources (default chat + peer chat) share one backend + // session and therefore one session-level `inputNeeded` queue. Opening + // each used to install its own watcher, so a single client-tool request + // was invoked once per open resource — running real side effects N + // times. The watcher is now ref-counted per backend session, so it + // executes exactly once no matter how many siblings are open. + async function openSiblingResourcesWithClaimedClientTool( + handler: AgentHostSessionHandler, + connection: MockAgentHostConnection, + ): Promise<{ sessionResource: URI; peerResource: URI; chat: string }> { + const sessionResource = URI.parse('agent-host-copilot:/session-1'); + const peerResource = URI.from({ scheme: 'agent-host-copilot', path: '/session-1', fragment: 'peer-1' }); + const backendSession = AgentSession.uri('copilot', 'session-1').toString(); + const chat = buildDefaultChatUri(backendSession); + const peerChat = buildChatUri(backendSession, 'peer-1'); + const summary: SessionSummary = { + resource: backendSession, + provider: 'copilot', + title: 'Test', + status: SessionStatus.Idle, + createdAt: '2025-01-01T00:00:00.000Z', + modifiedAt: '2025-01-01T00:00:00.000Z', + }; + + // Advertise the peer chat so the sibling resource resolves and + // installs its own turn/inputNeeded watchers against the shared + // backend session. + connection.applySessionAction(URI.parse(backendSession), { + type: ActionType.SessionChatAdded, + summary: createDefaultChatSummary(summary, peerChat), + } as SessionAction); + + connection.applySessionAction(URI.parse(chat), { + type: ActionType.ChatTurnStarted, + turnId: 'turn-1', + startedAt: '2025-01-01T00:00:00.000Z', + message: { text: 'run the task', origin: { kind: MessageKind.User } }, + }); + connection.applySessionAction(URI.parse(chat), { + type: ActionType.ChatToolCallStart, + turnId: 'turn-1', + toolCallId: 'client-tool-1', + toolName: 'runTask', + displayName: 'Run Task', + contributor: { kind: ToolCallContributorKind.Client, clientId: connection.clientId }, + }); + connection.applySessionAction(URI.parse(chat), { + type: ActionType.ChatToolCallReady, + turnId: 'turn-1', + toolCallId: 'client-tool-1', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); + + // Only the default chat carries the tool call, so only its observer + // claims it — the peer observer renders an empty chat. + await handler.provideChatSessionContent(sessionResource, CancellationToken.None); + await handler.provideChatSessionContent(peerResource, CancellationToken.None); + + applyRunningClientExecution(connection, chat, 'turn-1', { + toolCallId: 'client-tool-1', + toolName: 'runTask', + displayName: 'Run Task', + invocationMessage: 'Run Task', + toolInput: '{"task":"build"}', + }); + await timeout(5001); + return { sessionResource, peerResource, chat }; + } + + test('two sibling resources on one backend session execute a client tool exactly once', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection, toolsService } = createHandlerWithMocks(disposables, [testRunTaskTool]); + await openSiblingResourcesWithClaimedClientTool(handler, connection); + + assert.deepStrictEqual({ + invocations: toolsService.invokedToolCalls.filter(invocation => invocation.chatStreamToolCallId === 'client-tool-1').length, + }, { + invocations: 1, + }); + })); + + test('a claimed client tool executes with the claiming observer\'s session resource as context', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { handler, connection, toolsService } = createHandlerWithMocks(disposables, [testRunTaskTool]); + const { sessionResource } = await openSiblingResourcesWithClaimedClientTool(handler, connection); + + assert.deepStrictEqual( + toolsService.invokedToolCalls + .filter(invocation => invocation.chatStreamToolCallId === 'client-tool-1') + .map(invocation => invocation.context?.sessionResource.toString()), + [sessionResource.toString()], + ); + })); + test('denies an unclaimed confirmable client tool after the grace window without executing it', () => runWithFakedTimers({ useFakeTimers: true }, async () => { const { handler, connection, toolsService } = createHandlerWithMocks(disposables, [testConfirmTool]); const sessionResource = URI.parse('agent-host-copilot:/session-1');