diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatExecuteActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatExecuteActions.ts index 907a1145b99..14b0cb77d71 100644 --- a/src/vs/workbench/contrib/chat/browser/actions/chatExecuteActions.ts +++ b/src/vs/workbench/contrib/chat/browser/actions/chatExecuteActions.ts @@ -896,7 +896,7 @@ class SendToNewChatAction extends Action2 { // Cancel any in-progress request before clearing if (widget.viewModel) { - chatService.cancelCurrentRequestForSession(widget.viewModel.sessionResource, 'newSessionAction'); + await chatService.cancelCurrentRequestForSession(widget.viewModel.sessionResource, 'newSessionAction'); } if (widget.viewModel?.model) { @@ -956,7 +956,7 @@ export class CancelAction extends Action2 { }); } - run(accessor: ServicesAccessor, ...args: unknown[]) { + async run(accessor: ServicesAccessor, ...args: unknown[]) { const context = args[0] as IChatExecuteActionContext | undefined; const widgetService = accessor.get(IChatWidgetService); const logService = accessor.get(ILogService); @@ -975,7 +975,7 @@ export class CancelAction extends Action2 { const chatService = accessor.get(IChatService); if (widget.viewModel) { - chatService.cancelCurrentRequestForSession(widget.viewModel.sessionResource, 'cancelAction'); + await chatService.cancelCurrentRequestForSession(widget.viewModel.sessionResource, 'cancelAction'); } else { telemetryService.publicLog2(ChatStopCancellationNoopEventName, { source: 'cancelAction', diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatQueueActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatQueueActions.ts index 3eac2679afa..ed1b2f3689a 100644 --- a/src/vs/workbench/contrib/chat/browser/actions/chatQueueActions.ts +++ b/src/vs/workbench/contrib/chat/browser/actions/chatQueueActions.ts @@ -219,7 +219,7 @@ export class ChatSendPendingImmediatelyAction extends Action2 { }); } - override run(accessor: ServicesAccessor, ...args: unknown[]): void { + override async run(accessor: ServicesAccessor, ...args: unknown[]): Promise { const chatService = accessor.get(IChatService); const widgetService = accessor.get(IChatWidgetService); const [context] = args; @@ -250,7 +250,7 @@ export class ChatSendPendingImmediatelyAction extends Action2 { ]; chatService.setPendingRequests(context.sessionResource, reordered); - chatService.cancelCurrentRequestForSession(context.sessionResource, 'queueRunNext'); + await chatService.cancelCurrentRequestForSession(context.sessionResource, 'queueRunNext'); chatService.processPendingRequests(context.sessionResource); } } diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts b/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts index 81716d9c338..266efc7b245 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts @@ -2246,7 +2246,7 @@ export class ChatWidget extends Disposable implements IChatWidget { this.chatService.removePendingRequest(this.viewModel.sessionResource, editingRequestId); options.queue ??= editingPendingRequest; } else { - this.chatService.cancelCurrentRequestForSession(this.viewModel.sessionResource, 'acceptInput-editing'); + await this.chatService.cancelCurrentRequestForSession(this.viewModel.sessionResource, 'acceptInput-editing'); options.queue = undefined; } @@ -2266,7 +2266,7 @@ export class ChatWidget extends Disposable implements IChatWidget { options.queue ??= ChatRequestQueueKind.Queued; } if (model.requestNeedsInput.get() && !model.getPendingRequests().length) { - this.chatService.cancelCurrentRequestForSession(this.viewModel.sessionResource, 'acceptInput-needsInput'); + await this.chatService.cancelCurrentRequestForSession(this.viewModel.sessionResource, 'acceptInput-needsInput'); options.queue ??= ChatRequestQueueKind.Queued; } if (requestInProgress) { diff --git a/src/vs/workbench/contrib/chat/common/chatService/chatService.ts b/src/vs/workbench/contrib/chat/common/chatService/chatService.ts index 661a408290a..51cea27eb77 100644 --- a/src/vs/workbench/contrib/chat/common/chatService/chatService.ts +++ b/src/vs/workbench/contrib/chat/common/chatService/chatService.ts @@ -1410,7 +1410,7 @@ export interface IChatService { resendRequest(request: IChatRequestModel, options?: IChatSendRequestOptions): Promise; adoptRequest(sessionResource: URI, request: IChatRequestModel): Promise; removeRequest(sessionResource: URI, requestId: string): Promise; - cancelCurrentRequestForSession(sessionResource: URI, source?: string): void; + cancelCurrentRequestForSession(sessionResource: URI, source?: string): Promise; /** * Sets yieldRequested on the active request for the given session. */ diff --git a/src/vs/workbench/contrib/chat/common/chatService/chatServiceImpl.ts b/src/vs/workbench/contrib/chat/common/chatService/chatServiceImpl.ts index 854edefa586..e91a75cd468 100644 --- a/src/vs/workbench/contrib/chat/common/chatService/chatServiceImpl.ts +++ b/src/vs/workbench/contrib/chat/common/chatService/chatServiceImpl.ts @@ -3,7 +3,7 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { DeferredPromise } from '../../../../../base/common/async.js'; +import { DeferredPromise, raceTimeout } from '../../../../../base/common/async.js'; import { CancellationToken, CancellationTokenSource } from '../../../../../base/common/cancellation.js'; import { toErrorMessage } from '../../../../../base/common/errorMessage.js'; import { BugIndicatingError, ErrorNoTelemetry } from '../../../../../base/common/errors.js'; @@ -67,6 +67,7 @@ class CancellableRequest implements IDisposable { constructor( public readonly cancellationTokenSource: CancellationTokenSource, public requestId: string | undefined, + public readonly responseCompletePromise: Promise | undefined, @ILanguageModelToolsService private readonly toolsService: ILanguageModelToolsService ) { } @@ -679,7 +680,7 @@ export class ChatService extends Disposable implements IChatService { } if (providedSession.progressObs && lastRequest && providedSession.interruptActiveResponseCallback) { - const initialCancellationRequest = this.instantiationService.createInstance(CancellableRequest, new CancellationTokenSource(), undefined); + const initialCancellationRequest = this.instantiationService.createInstance(CancellableRequest, new CancellationTokenSource(), undefined, undefined); this._pendingRequests.set(model.sessionResource, initialCancellationRequest); this.telemetryService.publicLog2(ChatPendingRequestChangeEventName, { action: 'add', source: 'remoteSession', chatSessionId: chatSessionResourceToId(model.sessionResource) }); const cancellationListener = disposables.add(new MutableDisposable()); @@ -689,7 +690,7 @@ export class ChatService extends Disposable implements IChatService { providedSession.interruptActiveResponseCallback?.().then(userConfirmedInterruption => { if (!userConfirmedInterruption) { // User cancelled the interruption - const newCancellationRequest = this.instantiationService.createInstance(CancellableRequest, new CancellationTokenSource(), undefined); + const newCancellationRequest = this.instantiationService.createInstance(CancellableRequest, new CancellationTokenSource(), undefined, undefined); this._pendingRequests.set(model.sessionResource, newCancellationRequest); this.telemetryService.publicLog2(ChatPendingRequestChangeEventName, { action: 'add', source: 'remoteSession', chatSessionId: chatSessionResourceToId(model.sessionResource) }); cancellationListener.value = createCancellationListener(newCancellationRequest.cancellationTokenSource.token); @@ -1230,7 +1231,7 @@ export class ChatService extends Disposable implements IChatService { let shouldProcessPending = false; const rawResponsePromise = sendRequestInternal(); // Note- requestId is not known at this point, assigned later - const cancellableRequest = this.instantiationService.createInstance(CancellableRequest, source, undefined); + const cancellableRequest = this.instantiationService.createInstance(CancellableRequest, source, undefined, rawResponsePromise); this._pendingRequests.set(model.sessionResource, cancellableRequest); this.telemetryService.publicLog2(ChatPendingRequestChangeEventName, { action: 'add', source: 'sendRequest', chatSessionId: chatSessionResourceToId(model.sessionResource) }); rawResponsePromise.finally(() => { @@ -1493,7 +1494,7 @@ export class ChatService extends Disposable implements IChatService { request.response?.complete(); } - cancelCurrentRequestForSession(sessionResource: URI, source?: string): void { + async cancelCurrentRequestForSession(sessionResource: URI, source?: string): Promise { this.trace('cancelCurrentRequestForSession', `session: ${sessionResource}`); const pendingRequest = this._pendingRequests.get(sessionResource); if (!pendingRequest) { @@ -1514,9 +1515,14 @@ export class ChatService extends Disposable implements IChatService { return; } + const responseCompletePromise = pendingRequest.responseCompletePromise; pendingRequest.cancel(); this._pendingRequests.deleteAndDispose(sessionResource); this.telemetryService.publicLog2(ChatPendingRequestChangeEventName, { action: 'remove', source: source ?? 'cancelRequest', requestId: pendingRequest.requestId, chatSessionId: chatSessionResourceToId(sessionResource) }); + + if (responseCompletePromise) { + await raceTimeout(responseCompletePromise, 1000); + } } setYieldRequested(sessionResource: URI): void { diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/localAgentSessionsController.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/localAgentSessionsController.test.ts index c1c62a1246e..ae6d57e0a80 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/localAgentSessionsController.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/localAgentSessionsController.test.ts @@ -149,7 +149,7 @@ class MockChatService implements IChatService { throw new Error('Method not implemented.'); } - cancelCurrentRequestForSession(_sessionResource: URI, _source?: string): void { } + async cancelCurrentRequestForSession(_sessionResource: URI, _source?: string): Promise { } setYieldRequested(_sessionResource: URI): void { } diff --git a/src/vs/workbench/contrib/chat/test/common/chatService/chatService.test.ts b/src/vs/workbench/contrib/chat/test/common/chatService/chatService.test.ts index f0bd3984aae..3ca952c05d5 100644 --- a/src/vs/workbench/contrib/chat/test/common/chatService/chatService.test.ts +++ b/src/vs/workbench/contrib/chat/test/common/chatService/chatService.test.ts @@ -12,6 +12,7 @@ import { constObservable, observableValue } from '../../../../../../base/common/ import { URI } from '../../../../../../base/common/uri.js'; import { assertSnapshot } from '../../../../../../base/test/common/snapshot.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; +import { runWithFakedTimers } from '../../../../../../base/test/common/timeTravelScheduler.js'; import { Range } from '../../../../../../editor/common/core/range.js'; import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js'; import { TestConfigurationService } from '../../../../../../platform/configuration/test/common/testConfigurationService.js'; @@ -40,7 +41,7 @@ import { TestMcpService } from '../../../../mcp/test/common/testMcpService.js'; import { ChatAgentService, IChatAgent, IChatAgentData, IChatAgentImplementation, IChatAgentService } from '../../../common/participants/chatAgents.js'; import { IChatEditingService, IChatEditingSession } from '../../../common/editing/chatEditingService.js'; import { ChatModel, IChatModel, ISerializableChatData } from '../../../common/model/chatModel.js'; -import { ChatRequestQueueKind, ChatSendResult, IChatFollowup, IChatModelReference, IChatService } from '../../../common/chatService/chatService.js'; +import { ChatRequestQueueKind, ChatSendResult, IChatFollowup, IChatModelReference, IChatService, ResponseModelState } from '../../../common/chatService/chatService.js'; import { ChatService } from '../../../common/chatService/chatServiceImpl.js'; import { ChatSlashCommandService, IChatSlashCommandService } from '../../../common/participants/chatSlashCommands.js'; import { IChatVariablesService } from '../../../common/attachments/chatVariables.js'; @@ -535,6 +536,79 @@ suite('ChatService', () => { assert.ok(invokedRequests[1].includes('steering3'), 'Combined message should include steering3'); assert.ok(invokedRequests[1].includes('\n\n'), 'Combined message should use \\n\\n as separator'); }); + test('cancelCurrentRequestForSession waits for response completion', async () => { + const requestStarted = new DeferredPromise(); + const completeRequest = new DeferredPromise(); + + const slowAgent: IChatAgentImplementation = { + async invoke(request, progress, history, token) { + requestStarted.complete(); + const listener = token.onCancellationRequested(() => { + listener.dispose(); + // Simulate some cleanup delay before completing + setTimeout(() => completeRequest.complete(), 10); + }); + await completeRequest.p; + return {}; + }, + }; + + testDisposables.add(chatAgentService.registerAgent('slowAgent', { ...getAgentData('slowAgent'), isDefault: true })); + testDisposables.add(chatAgentService.registerAgentImplementation('slowAgent', slowAgent)); + + const testService = createChatService(); + const modelRef = testDisposables.add(startSessionModel(testService)); + const model = modelRef.object; + + const response = await testService.sendRequest(model.sessionResource, 'test request', { agentId: 'slowAgent' }); + ChatSendResult.assertSent(response); + + await requestStarted.p; + + // Cancel and await - should wait for the response to complete + await testService.cancelCurrentRequestForSession(model.sessionResource, 'test'); + + // After cancel resolves, the response model should have a result + const lastRequest = model.getRequests()[0]; + assert.ok(lastRequest.response, 'Response should exist after cancellation completes'); + assert.strictEqual(lastRequest.response.state, ResponseModelState.Cancelled, 'Response should be in Cancelled state'); + }); + + test('cancelCurrentRequestForSession returns after timeout if response does not complete', async () => { + const requestStarted = new DeferredPromise(); + const completeRequest = new DeferredPromise(); + + const hangingAgent: IChatAgentImplementation = { + async invoke(request, progress, history, token) { + requestStarted.complete(); + // Wait for external signal, ignoring cancellation to simulate a hung agent + await completeRequest.p; + return {}; + }, + }; + + testDisposables.add(chatAgentService.registerAgent('hangingAgent', { ...getAgentData('hangingAgent'), isDefault: true })); + testDisposables.add(chatAgentService.registerAgentImplementation('hangingAgent', hangingAgent)); + + const testService = createChatService(); + const modelRef = testDisposables.add(startSessionModel(testService)); + const model = modelRef.object; + + const response = await testService.sendRequest(model.sessionResource, 'test request', { agentId: 'hangingAgent' }); + ChatSendResult.assertSent(response); + + await requestStarted.p; + + // Cancel should return after timeout even though the agent has not completed. + // Use faked timers so raceTimeout's 1s setTimeout fires instantly. + await runWithFakedTimers({ useFakeTimers: true }, async () => { + await testService.cancelCurrentRequestForSession(model.sessionResource, 'test'); + }); + + // Let the agent finish so the test cleans up properly + completeRequest.complete(); + await response.data.responseCompletePromise; + }); }); diff --git a/src/vs/workbench/contrib/chat/test/common/chatService/mockChatService.ts b/src/vs/workbench/contrib/chat/test/common/chatService/mockChatService.ts index f10ffa492b6..4911cfe126c 100644 --- a/src/vs/workbench/contrib/chat/test/common/chatService/mockChatService.ts +++ b/src/vs/workbench/contrib/chat/test/common/chatService/mockChatService.ts @@ -88,7 +88,7 @@ export class MockChatService implements IChatService { removeRequest(sessionResource: URI, requestId: string): Promise { throw new Error('Method not implemented.'); } - cancelCurrentRequestForSession(sessionResource: URI, source?: string): void { + async cancelCurrentRequestForSession(sessionResource: URI, source?: string): Promise { throw new Error('Method not implemented.'); } setYieldRequested(sessionResource: URI): void { diff --git a/src/vs/workbench/contrib/inlineChat/browser/inlineChatController.ts b/src/vs/workbench/contrib/inlineChat/browser/inlineChatController.ts index 76aa39da942..d26f9270bdd 100644 --- a/src/vs/workbench/contrib/inlineChat/browser/inlineChatController.ts +++ b/src/vs/workbench/contrib/inlineChat/browser/inlineChatController.ts @@ -605,7 +605,7 @@ export class InlineChatController implements IEditorContribution { if (!session) { return; } - this._chatService.cancelCurrentRequestForSession(session.chatModel.sessionResource, 'inlineChatReject'); + await this._chatService.cancelCurrentRequestForSession(session.chatModel.sessionResource, 'inlineChatReject'); await session.editingSession.reject(); session.dispose(); } diff --git a/src/vs/workbench/contrib/inlineChat/browser/inlineChatSessionServiceImpl.ts b/src/vs/workbench/contrib/inlineChat/browser/inlineChatSessionServiceImpl.ts index 4a008a59b0f..6879c6b4391 100644 --- a/src/vs/workbench/contrib/inlineChat/browser/inlineChatSessionServiceImpl.ts +++ b/src/vs/workbench/contrib/inlineChat/browser/inlineChatSessionServiceImpl.ts @@ -86,7 +86,7 @@ export class InlineChatSessionServiceImpl implements IInlineChatSessionService { const store = new DisposableStore(); store.add(toDisposable(() => { - this._chatService.cancelCurrentRequestForSession(chatModel.sessionResource, 'inlineChatSession'); + void this._chatService.cancelCurrentRequestForSession(chatModel.sessionResource, 'inlineChatSession'); chatModel.editingSession?.reject(); this._sessions.delete(uri); this._onDidChangeSessions.fire(this); diff --git a/src/vs/workbench/contrib/terminalContrib/chat/browser/terminalChatWidget.ts b/src/vs/workbench/contrib/terminalContrib/chat/browser/terminalChatWidget.ts index 2be6c5d9e96..be4f138c4ea 100644 --- a/src/vs/workbench/contrib/terminalContrib/chat/browser/terminalChatWidget.ts +++ b/src/vs/workbench/contrib/terminalContrib/chat/browser/terminalChatWidget.ts @@ -412,7 +412,7 @@ export class TerminalChatWidget extends Disposable { if (!model?.sessionResource) { return; } - this._chatService.cancelCurrentRequestForSession(model?.sessionResource, 'terminalChat'); + void this._chatService.cancelCurrentRequestForSession(model?.sessionResource, 'terminalChat'); } async viewInChat(): Promise {