From ec450a630ecc504472890ac0ef949ae0274b17e0 Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Sat, 27 Jun 2026 00:11:46 +1000 Subject: [PATCH] fix: removed fallback rest api call in ACP mode for edit in place (#10034) --- .../__tests__/chatSessionController.test.ts | 88 +++++++++++++- .../acp/__tests__/chatSessionStore.test.ts | 89 ++++++++++++++ .../acp/__tests__/elicitationRequests.test.ts | 20 ++- .../acp/__tests__/permissionRequests.test.ts | 8 ++ ui/desktop/src/acp/chatSessionController.ts | 102 +++++++++++----- ui/desktop/src/acp/chatSessionStore.ts | 114 ++++++++++++++++++ ui/desktop/src/acp/elicitationRequests.ts | 10 +- ui/desktop/src/acp/permissionRequests.ts | 6 +- .../components/ToolApprovalButtons.test.tsx | 40 +++++- .../src/components/ToolApprovalButtons.tsx | 63 ++++++---- ui/desktop/src/i18n/messages/en.json | 3 + ui/desktop/src/i18n/messages/es.json | 3 + ui/desktop/src/i18n/messages/hi.json | 3 + ui/desktop/src/i18n/messages/ja.json | 3 + ui/desktop/src/i18n/messages/ko.json | 3 + ui/desktop/src/i18n/messages/ru.json | 3 + ui/desktop/src/i18n/messages/tr.json | 3 + ui/desktop/src/i18n/messages/zh-CN.json | 3 + 18 files changed, 496 insertions(+), 68 deletions(-) diff --git a/ui/desktop/src/acp/__tests__/chatSessionController.test.ts b/ui/desktop/src/acp/__tests__/chatSessionController.test.ts index be0f0b9a7..9f9ad0c7f 100644 --- a/ui/desktop/src/acp/__tests__/chatSessionController.test.ts +++ b/ui/desktop/src/acp/__tests__/chatSessionController.test.ts @@ -35,6 +35,8 @@ vi.mock('../chatSessionStore', () => ({ clearActivePromptAttempt: vi.fn(), startPromptCancellation: vi.fn(), clearPromptCancellation: vi.fn(), + restorePromptCancellation: vi.fn(), + waitForPromptCancellation: vi.fn(), setChatState: vi.fn(), setSessionMetadata: vi.fn(), setSessionLoadError: vi.fn(), @@ -113,6 +115,24 @@ function snapshotWithActivePrompt(activePromptAttemptId: string | null): AcpChat }; } +function pendingToolPermissionMessage(): Message & { id: string } { + return { + id: 'permission-message-1', + role: 'assistant', + created: 124, + content: [ + { + type: 'toolConfirmationRequest', + id: 'tool-call-1', + toolName: 'developer__shell', + arguments: {}, + prompt: null, + }, + ], + metadata: { userVisible: true, agentVisible: true }, + }; +} + describe('acpChatSessionController.loadSession', () => { beforeEach(() => { vi.clearAllMocks(); @@ -219,7 +239,9 @@ describe('acpChatSessionController.updateMessage', () => { beforeEach(() => { vi.clearAllMocks(); vi.mocked(acpTruncateSessionConversation).mockResolvedValue(undefined as never); + vi.mocked(acpPromptSession).mockResolvedValue({ stopReason: 'end_turn' } as never); vi.mocked(acpChatSessionStore.getSnapshot).mockReturnValue(snapshotWithActivePrompt(null)); + vi.mocked(acpChatSessionActions.waitForPromptCancellation).mockResolvedValue(undefined); }); it('rejects edits before truncating while cancellation is pending', async () => { @@ -249,7 +271,7 @@ describe('acpChatSessionController.updateMessage', () => { expect(acpPromptSession).not.toHaveBeenCalled(); }); - it('rejects edits before truncating while a prompt is active', async () => { + it('ignores edits before truncating while a prompt is active', async () => { vi.mocked(acpChatSessionStore.getSnapshot).mockReturnValue( snapshotWithActivePrompt('attempt-1') ); @@ -264,7 +286,7 @@ describe('acpChatSessionController.updateMessage', () => { getCurrentSnapshot: () => currentSnapshot, onFinish: vi.fn(), }) - ).rejects.toThrow('Cannot update message while prompt is active'); + ).resolves.toBeUndefined(); expect(acpChatSessionActions.setChatState).not.toHaveBeenCalledWith( SESSION_ID, @@ -274,4 +296,66 @@ describe('acpChatSessionController.updateMessage', () => { expect(acpChatSessionActions.setMessages).not.toHaveBeenCalled(); expect(acpPromptSession).not.toHaveBeenCalled(); }); + + it('waits for pending tool permission cancellation before truncating and rerunning', async () => { + const existingMessage = userMessage(); + const permissionMessage = pendingToolPermissionMessage(); + const activeSnapshot: AcpChatSessionSnapshot = { + ...snapshotWithActivePrompt('attempt-1'), + chatState: ChatState.WaitingForUserInput, + messages: [existingMessage, permissionMessage], + }; + let storedSnapshot = activeSnapshot; + vi.mocked(acpChatSessionStore.getSnapshot).mockImplementation(() => storedSnapshot); + vi.mocked(acpChatSessionActions.startPromptCancellation).mockReturnValue({ + ...activeSnapshot, + activePromptAttemptId: null, + pendingCancelPromptAttemptId: 'attempt-1', + }); + vi.mocked(acpCancelPrompt).mockResolvedValue(undefined); + + let resolvePromptCancellation: () => void; + const promptCancellationSettled = new Promise((resolve) => { + resolvePromptCancellation = resolve; + }); + vi.mocked(acpChatSessionActions.waitForPromptCancellation).mockReturnValue( + promptCancellationSettled + ); + + const updatePromise = acpChatSessionController.updateMessage( + SESSION_ID, + existingMessage.id, + 'Updated', + 'edit', + { + getCurrentSnapshot: () => activeSnapshot, + onFinish: vi.fn(), + } + ); + + await Promise.resolve(); + await Promise.resolve(); + + expect(acpCancelPrompt).toHaveBeenCalledWith(SESSION_ID); + expect(acpChatSessionActions.waitForPromptCancellation).toHaveBeenCalledWith( + SESSION_ID, + 'attempt-1' + ); + expect(acpTruncateSessionConversation).not.toHaveBeenCalled(); + expect(acpPromptSession).not.toHaveBeenCalled(); + + storedSnapshot = { + ...snapshotWithActivePrompt(null), + messages: [existingMessage, permissionMessage], + }; + resolvePromptCancellation!(); + await updatePromise; + + expect(acpTruncateSessionConversation).toHaveBeenCalledWith(SESSION_ID, existingMessage.created); + expect(acpPromptSession).toHaveBeenCalled(); + expect(acpChatSessionActions.clearPromptCancellation).not.toHaveBeenCalledWith( + SESSION_ID, + 'attempt-1' + ); + }); }); diff --git a/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts b/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts index 23834d0b9..68502de09 100644 --- a/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts +++ b/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts @@ -8,7 +8,9 @@ import { afterEach, describe, expect, it } from 'vitest'; import type { Message, Session } from '../../api'; import { ChatState } from '../../types/chatState'; import { + acpElicitationUserInputRequestId, acpChatSessionActions, + acpPermissionUserInputRequestId, acpChatSessionStore, useAcpChatSessionSnapshot, } from '../chatSessionStore'; @@ -294,6 +296,58 @@ describe('acpChatSessionStore', () => { expect(clearedSnapshot?.pendingCancelPromptAttemptId).toBeNull(); }); + it('restores pending user input tracking when prompt cancellation is restored', () => { + const currentSessionId = sessionId('session-1'); + + acpChatSessionActions.startPromptAttempt(currentSessionId, 'attempt-1'); + acpChatSessionActions.applyPermissionRequest(permissionRequest(currentSessionId, 'tool-1')); + acpChatSessionActions.applyElicitationRequest(elicitationRequest(currentSessionId)); + acpChatSessionActions.startPromptCancellation(currentSessionId, 'attempt-1'); + + const restoredSnapshot = acpChatSessionActions.restorePromptCancellation( + currentSessionId, + 'attempt-1' + ); + + expect(restoredSnapshot?.chatState).toBe(ChatState.WaitingForUserInput); + + const afterPermission = acpChatSessionActions.resolveUserInputRequest( + currentSessionId, + acpPermissionUserInputRequestId('tool-1') + ); + + expect(afterPermission?.chatState).toBe(ChatState.WaitingForUserInput); + + const afterElicitation = acpChatSessionActions.resolveUserInputRequest( + currentSessionId, + acpElicitationUserInputRequestId('acp_elicitation_1') + ); + + expect(afterElicitation?.chatState).toBe(ChatState.Streaming); + }); + + it('waits for prompt cancellation to clear', async () => { + const currentSessionId = sessionId('session-1'); + + acpChatSessionActions.startPromptAttempt(currentSessionId, 'attempt-1'); + acpChatSessionActions.startPromptCancellation(currentSessionId, 'attempt-1'); + + let didResolve = false; + const waitPromise = acpChatSessionActions + .waitForPromptCancellation(currentSessionId, 'attempt-1') + .then(() => { + didResolve = true; + }); + + await Promise.resolve(); + expect(didResolve).toBe(false); + + acpChatSessionActions.clearPromptCancellation(currentSessionId, 'attempt-1'); + + await waitPromise; + expect(didResolve).toBe(true); + }); + it('removes pending local steer messages when cancellation starts', () => { const currentSessionId = sessionId('session-1'); const localSteerMessage = { @@ -465,6 +519,41 @@ describe('acpChatSessionStore', () => { }); }); + it('resumes streaming only after the final pending user input request resolves', () => { + const currentSessionId = sessionId('session-1'); + + acpChatSessionActions.startPromptAttempt(currentSessionId, 'attempt-1'); + acpChatSessionActions.applyPermissionRequest(permissionRequest(currentSessionId, 'tool-1')); + acpChatSessionActions.applyElicitationRequest(elicitationRequest(currentSessionId)); + + const afterElicitation = acpChatSessionActions.resolveUserInputRequest( + currentSessionId, + acpElicitationUserInputRequestId('acp_elicitation_1') + ); + + expect(afterElicitation?.chatState).toBe(ChatState.WaitingForUserInput); + + const afterPermission = acpChatSessionActions.resolveUserInputRequest( + currentSessionId, + acpPermissionUserInputRequestId('tool-1') + ); + + expect(afterPermission?.chatState).toBe(ChatState.Streaming); + }); + + it('does not resume streaming after user input resolves without an active prompt', () => { + const currentSessionId = sessionId('session-1'); + + acpChatSessionActions.applyPermissionRequest(permissionRequest(currentSessionId, 'tool-1')); + + const snapshot = acpChatSessionActions.resolveUserInputRequest( + currentSessionId, + acpPermissionUserInputRequestId('tool-1') + ); + + expect(snapshot?.chatState).toBe(ChatState.WaitingForUserInput); + }); + it('applies elicitation requests as waiting action-required messages', () => { const currentSessionId = sessionId('session-1'); diff --git a/ui/desktop/src/acp/__tests__/elicitationRequests.test.ts b/ui/desktop/src/acp/__tests__/elicitationRequests.test.ts index c9b691b91..c84c3eb0f 100644 --- a/ui/desktop/src/acp/__tests__/elicitationRequests.test.ts +++ b/ui/desktop/src/acp/__tests__/elicitationRequests.test.ts @@ -13,8 +13,10 @@ vi.mock('../../acpChatFeatureFlag', () => ({ })); vi.mock('../chatSessionStore', () => ({ + acpElicitationUserInputRequestId: (elicitationId: string) => `elicitation:${elicitationId}`, acpChatSessionActions: { applyElicitationRequest: vi.fn(), + resolveUserInputRequest: vi.fn(), setElicitationStatus: vi.fn(), }, })); @@ -73,9 +75,8 @@ describe('ACP elicitation requests', () => { await expectStillPending(response); - const appliedRequest = vi.mocked( - acpChatSessionActions.applyElicitationRequest - ).mock.calls[0][0]; + const appliedRequest = vi.mocked(acpChatSessionActions.applyElicitationRequest).mock + .calls[0][0]; expect(appliedRequest.id).toMatch(/^acp_elicitation_/); expect(appliedRequest.sessionId).toBe('session-1'); @@ -91,6 +92,10 @@ describe('ACP elicitation requests', () => { appliedRequest.id, 'submitted' ); + expect(acpChatSessionActions.resolveUserInputRequest).toHaveBeenCalledWith( + 'session-1', + `elicitation:${appliedRequest.id}` + ); await expect(response).resolves.toEqual({ action: 'accept', @@ -142,9 +147,8 @@ describe('ACP elicitation requests', () => { vi.useFakeTimers(); try { const response = requestAcpElicitation(formRequest('session-1')); - const appliedRequest = vi.mocked( - acpChatSessionActions.applyElicitationRequest - ).mock.calls[0][0]; + const appliedRequest = vi.mocked(acpChatSessionActions.applyElicitationRequest).mock + .calls[0][0]; await expectStillPending(response); @@ -155,6 +159,10 @@ describe('ACP elicitation requests', () => { appliedRequest.id, 'cancelled' ); + expect(acpChatSessionActions.resolveUserInputRequest).toHaveBeenCalledWith( + 'session-1', + `elicitation:${appliedRequest.id}` + ); await expect(response).resolves.toEqual({ action: 'cancel' }); expect(resolveAcpElicitationRequest('session-1', appliedRequest.id, {})).toBe(false); } finally { diff --git a/ui/desktop/src/acp/__tests__/permissionRequests.test.ts b/ui/desktop/src/acp/__tests__/permissionRequests.test.ts index d2aeb5d5d..9274011fb 100644 --- a/ui/desktop/src/acp/__tests__/permissionRequests.test.ts +++ b/ui/desktop/src/acp/__tests__/permissionRequests.test.ts @@ -5,14 +5,17 @@ import { requestAcpPermission, resolveAcpPermissionRequest, } from '../permissionRequests'; +import { acpChatSessionActions } from '../chatSessionStore'; vi.mock('../../acpChatFeatureFlag', () => ({ USE_ACP_CHAT: true, })); vi.mock('../chatSessionStore', () => ({ + acpPermissionUserInputRequestId: (toolCallId: string) => `permission:${toolCallId}`, acpChatSessionActions: { applyPermissionRequest: vi.fn(), + resolveUserInputRequest: vi.fn(), }, })); @@ -57,6 +60,7 @@ async function expectStillPending(promise: Promise): describe('ACP permission requests', () => { beforeEach(() => { + vi.clearAllMocks(); for (const sessionId of TEST_SESSION_IDS) { cancelAcpPermissionRequestsForSession(sessionId); } @@ -74,6 +78,10 @@ describe('ACP permission requests', () => { await expectStillPending(response); expect(resolveAcpPermissionRequest('session-1', 'tool-1', 'allow_once')).toBe(true); + expect(acpChatSessionActions.resolveUserInputRequest).toHaveBeenCalledWith( + 'session-1', + 'permission:tool-1' + ); await expect(response).resolves.toEqual({ outcome: { outcome: 'selected', diff --git a/ui/desktop/src/acp/chatSessionController.ts b/ui/desktop/src/acp/chatSessionController.ts index cb0bf9db0..f7f2363fd 100644 --- a/ui/desktop/src/acp/chatSessionController.ts +++ b/ui/desktop/src/acp/chatSessionController.ts @@ -5,7 +5,7 @@ import { AppEvents } from '../constants/events'; import { ChatState } from '../types/chatState'; import { errorMessage } from '../utils/conversionUtils'; import { showExtensionLoadResults } from '../utils/extensionErrorUtils'; -import { createUserMessage } from '../types/message'; +import { createUserMessage, getPendingToolConfirmationIds } from '../types/message'; import { acpChatSessionActions, acpChatSessionStore, @@ -83,11 +83,21 @@ function assertNoPendingPromptCancellation(sessionId: string): void { } } -function assertNoActivePromptAttempt(sessionId: string): void { - const snapshot = acpChatSessionStore.getSnapshot(sessionId); - if (snapshot?.activePromptAttemptId) { - throw new Error('Cannot update message while prompt is active'); - } +async function forkSessionWithEditedMessage( + sessionId: string, + message: Message, + editedMessage: string +): Promise { + const targetSessionId = await acpForkSession(sessionId, message.created); + + const event = new CustomEvent(AppEvents.SESSION_FORKED, { + detail: { + newSessionId: targetSessionId, + shouldStartAgent: true, + editedMessage, + }, + }); + window.dispatchEvent(event); } async function createSession( @@ -214,37 +224,69 @@ async function updateMessage( options: AcpSubmitMessageOptions ): Promise { assertNoPendingPromptCancellation(sessionId); - assertNoActivePromptAttempt(sessionId); const resolvedEditType = editType ?? 'fork'; const currentSnapshot = options.getCurrentSnapshot(); + const storedSnapshot = acpChatSessionStore.getSnapshot(sessionId); + const activePromptAttemptId = storedSnapshot?.activePromptAttemptId; + const currentMessages = currentSnapshot?.messages ?? []; + const message = currentMessages.find((m) => m.id === messageId); + + if (!message) { + throw new Error(`Message with id ${messageId} not found in current messages`); + } + + if (resolvedEditType === 'fork') { + await forkSessionWithEditedMessage(sessionId, message, newContent); + return; + } + + const editSnapshot = currentSnapshot ?? storedSnapshot; + const isPendingToolPermission = + editSnapshot?.chatState === ChatState.WaitingForUserInput && + getPendingToolConfirmationIds(editSnapshot?.messages ?? []).size > 0; + const isIdle = editSnapshot?.chatState === ChatState.Idle; + const pendingToolPermissionPromptAttemptId = isPendingToolPermission + ? activePromptAttemptId + : undefined; + const canEditInPlace = isIdle || pendingToolPermissionPromptAttemptId != null; + + if (!canEditInPlace) { + return; + } + + if (pendingToolPermissionPromptAttemptId != null) { + const cancellation = acpChatSessionActions.startPromptCancellation( + sessionId, + pendingToolPermissionPromptAttemptId + ); + if (!cancellation) { + throw new Error('Cannot update message while prompt is active'); + } + + const promptCancellationSettled = acpChatSessionActions.waitForPromptCancellation( + sessionId, + pendingToolPermissionPromptAttemptId + ); + + try { + await acpCancelPrompt(sessionId); + } catch { + acpChatSessionActions.restorePromptCancellation( + sessionId, + pendingToolPermissionPromptAttemptId + ); + throw new Error('Cannot update message because the active prompt could not be cancelled'); + } + + cancelAcpPermissionRequestsForSession(sessionId); + cancelAcpElicitationRequestsForSession(sessionId); + await promptCancellationSettled; + } acpChatSessionActions.setChatState(sessionId, ChatState.Thinking); try { - const currentMessages = currentSnapshot?.messages ?? []; - const message = currentMessages.find((m) => m.id === messageId); - - if (!message) { - throw new Error(`Message with id ${messageId} not found in current messages`); - } - - if (resolvedEditType === 'fork') { - const targetSessionId = await acpForkSession(sessionId, message.created); - - acpChatSessionActions.setChatState(sessionId, ChatState.Idle); - const event = new CustomEvent(AppEvents.SESSION_FORKED, { - detail: { - newSessionId: targetSessionId, - shouldStartAgent: true, - editedMessage: newContent, - }, - }); - window.dispatchEvent(event); - window.electron.logInfo(`Dispatched session-forked event for session ${targetSessionId}`); - return; - } - await acpTruncateSessionConversation(sessionId, message.created); const truncatedMessages = currentMessages.filter((m) => m.created < message.created); diff --git a/ui/desktop/src/acp/chatSessionStore.ts b/ui/desktop/src/acp/chatSessionStore.ts index 08047dc45..c7901543a 100644 --- a/ui/desktop/src/acp/chatSessionStore.ts +++ b/ui/desktop/src/acp/chatSessionStore.ts @@ -29,6 +29,12 @@ type SnapshotListener = (snapshot: AcpChatSessionSnapshot) => void; interface StoreEntry extends AcpChatSessionSnapshot { adapter: AcpSessionNotificationAdapter; + promptCancellationRestoreState: { + activeRunId: string | null; + chatState: ChatState; + pendingUserInputRequestIds: Set; + } | null; + pendingUserInputRequestIds: Set; pendingLocalSteerMessageIds: Set; } @@ -72,6 +78,10 @@ export interface AcpChatSessionActions { setMessages(sessionId: string, messages: Message[]): AcpChatSessionSnapshot; addPendingLocalSteerMessage(sessionId: string, message: Message): AcpChatSessionSnapshot; setChatState(sessionId: string, chatState: ChatState): AcpChatSessionSnapshot; + resolveUserInputRequest( + sessionId: string, + userInputRequestId: string + ): AcpChatSessionSnapshot | undefined; startPromptAttempt(sessionId: string, promptAttemptId: string): AcpChatSessionSnapshot; startPromptCancellation( @@ -82,6 +92,11 @@ export interface AcpChatSessionActions { sessionId: string, promptAttemptId: string ): AcpChatSessionSnapshot | undefined; + restorePromptCancellation( + sessionId: string, + promptAttemptId: string + ): AcpChatSessionSnapshot | undefined; + waitForPromptCancellation(sessionId: string, promptAttemptId: string): Promise; finishPromptAttemptIfCurrent(sessionId: string, promptAttemptId: string, error?: string): boolean; clearActivePromptAttempt(sessionId: string): AcpChatSessionSnapshot | undefined; isCurrentPromptAttempt(sessionId: string, promptAttemptId: string): boolean; @@ -144,6 +159,8 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { activePromptAttemptId: null, activeRunId: null, pendingCancelPromptAttemptId: null, + promptCancellationRestoreState: null, + pendingUserInputRequestIds: new Set(), pendingLocalSteerMessageIds: new Set(), adapter: createAcpSessionNotificationAdapter(), }; @@ -223,6 +240,29 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { return notify(sessionId, entry); }; + const resolveUserInputRequest: AcpChatSessionActions['resolveUserInputRequest'] = ( + sessionId, + userInputRequestId + ) => { + const entry = sessionsById.get(sessionId); + if (!entry) { + return undefined; + } + + entry.pendingUserInputRequestIds.delete(userInputRequestId); + + if ( + entry.activePromptAttemptId && + entry.chatState === ChatState.WaitingForUserInput && + entry.pendingUserInputRequestIds.size === 0 + ) { + entry.chatState = ChatState.Streaming; + return notify(sessionId, entry); + } + + return snapshotFromEntry(entry); + }; + const setSessionLoadError: AcpChatSessionActions['setSessionLoadError'] = ( sessionId, sessionLoadError @@ -241,6 +281,8 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { entry.activePromptAttemptId = promptAttemptId; entry.activeRunId = null; entry.pendingCancelPromptAttemptId = null; + entry.promptCancellationRestoreState = null; + entry.pendingUserInputRequestIds.clear(); entry.chatState = ChatState.Streaming; entry.sessionLoadError = undefined; entry.notifications = []; @@ -256,9 +298,15 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { return undefined; } + entry.promptCancellationRestoreState = { + activeRunId: entry.activeRunId, + chatState: entry.chatState, + pendingUserInputRequestIds: new Set(entry.pendingUserInputRequestIds), + }; entry.activePromptAttemptId = null; entry.activeRunId = null; entry.pendingCancelPromptAttemptId = promptAttemptId; + entry.pendingUserInputRequestIds.clear(); discardPendingLocalSteerMessages(entry); entry.chatState = ChatState.Idle; return notify(sessionId, entry); @@ -274,9 +322,52 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { } entry.pendingCancelPromptAttemptId = null; + entry.promptCancellationRestoreState = null; return notify(sessionId, entry); }; + const restorePromptCancellation: AcpChatSessionActions['restorePromptCancellation'] = ( + sessionId, + promptAttemptId + ) => { + const entry = sessionsById.get(sessionId); + if ( + !entry || + entry.pendingCancelPromptAttemptId !== promptAttemptId || + !entry.promptCancellationRestoreState + ) { + return undefined; + } + + const restoreState = entry.promptCancellationRestoreState; + entry.activePromptAttemptId = promptAttemptId; + entry.activeRunId = restoreState.activeRunId; + entry.pendingCancelPromptAttemptId = null; + entry.promptCancellationRestoreState = null; + entry.pendingUserInputRequestIds = new Set(restoreState.pendingUserInputRequestIds); + entry.chatState = restoreState.chatState; + return notify(sessionId, entry); + }; + + const waitForPromptCancellation: AcpChatSessionActions['waitForPromptCancellation'] = ( + sessionId, + promptAttemptId + ) => { + const entry = sessionsById.get(sessionId); + if (!entry || entry.pendingCancelPromptAttemptId !== promptAttemptId) { + return Promise.resolve(); + } + + return new Promise((resolve) => { + const unsubscribe = subscribe(sessionId, (snapshot) => { + if (snapshot.pendingCancelPromptAttemptId !== promptAttemptId) { + unsubscribe(); + resolve(); + } + }); + }); + }; + const finishPromptAttemptIfCurrent: AcpChatSessionActions['finishPromptAttemptIfCurrent'] = ( sessionId, promptAttemptId, @@ -290,6 +381,8 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { entry.activePromptAttemptId = null; entry.activeRunId = null; entry.pendingCancelPromptAttemptId = null; + entry.promptCancellationRestoreState = null; + entry.pendingUserInputRequestIds.clear(); discardPendingLocalSteerMessages(entry); entry.chatState = ChatState.Idle; entry.sessionLoadError = error; @@ -307,6 +400,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { entry.activePromptAttemptId = null; entry.activeRunId = null; + entry.pendingUserInputRequestIds.clear(); discardPendingLocalSteerMessages(entry); entry.chatState = ChatState.Idle; return notify(sessionId, entry); @@ -338,6 +432,9 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { const entry = getOrCreateEntry(request.sessionId); const changes = entry.adapter.applyPermissionRequest(request); applyChatStateChanges(entry, changes); + entry.pendingUserInputRequestIds.add( + acpPermissionUserInputRequestId(request.toolCall.toolCallId) + ); entry.chatState = ChatState.WaitingForUserInput; return notify(request.sessionId, entry); }; @@ -346,6 +443,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { const entry = getOrCreateEntry(request.sessionId); const changes = entry.adapter.applyElicitationRequest(request); applyChatStateChanges(entry, changes); + entry.pendingUserInputRequestIds.add(acpElicitationUserInputRequestId(request.id)); entry.chatState = ChatState.WaitingForUserInput; return notify(request.sessionId, entry); }; @@ -381,9 +479,12 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { setMessages, addPendingLocalSteerMessage, setChatState, + resolveUserInputRequest, startPromptAttempt, startPromptCancellation, clearPromptCancellation, + restorePromptCancellation, + waitForPromptCancellation, finishPromptAttemptIfCurrent, clearActivePromptAttempt, isCurrentPromptAttempt, @@ -456,9 +557,12 @@ function actionsFromStore(store: AcpChatSessionStoreInternal): AcpChatSessionAct setMessages: store.setMessages, addPendingLocalSteerMessage: store.addPendingLocalSteerMessage, setChatState: store.setChatState, + resolveUserInputRequest: store.resolveUserInputRequest, startPromptAttempt: store.startPromptAttempt, startPromptCancellation: store.startPromptCancellation, clearPromptCancellation: store.clearPromptCancellation, + restorePromptCancellation: store.restorePromptCancellation, + waitForPromptCancellation: store.waitForPromptCancellation, finishPromptAttemptIfCurrent: store.finishPromptAttemptIfCurrent, clearActivePromptAttempt: store.clearActivePromptAttempt, isCurrentPromptAttempt: store.isCurrentPromptAttempt, @@ -499,10 +603,20 @@ function resetReplayState(entry: StoreEntry): void { entry.notifications = []; entry.activeRunId = null; entry.pendingCancelPromptAttemptId = null; + entry.promptCancellationRestoreState = null; + entry.pendingUserInputRequestIds.clear(); entry.pendingLocalSteerMessageIds.clear(); entry.adapter = createAcpSessionNotificationAdapter(); } +export function acpPermissionUserInputRequestId(toolCallId: string): string { + return `permission:${toolCallId}`; +} + +export function acpElicitationUserInputRequestId(elicitationId: string): string { + return `elicitation:${elicitationId}`; +} + function retainPendingLocalSteerMessageIds(entry: StoreEntry): void { if (entry.pendingLocalSteerMessageIds.size === 0) { return; diff --git a/ui/desktop/src/acp/elicitationRequests.ts b/ui/desktop/src/acp/elicitationRequests.ts index 19caf7f63..9ec139d1a 100644 --- a/ui/desktop/src/acp/elicitationRequests.ts +++ b/ui/desktop/src/acp/elicitationRequests.ts @@ -6,7 +6,7 @@ import type { } from '@agentclientprotocol/sdk'; import { v7 as uuidv7 } from 'uuid'; import { USE_ACP_CHAT } from '../acpChatFeatureFlag'; -import { acpChatSessionActions } from './chatSessionStore'; +import { acpChatSessionActions, acpElicitationUserInputRequestId } from './chatSessionStore'; type SessionScopedFormElicitationRequest = CreateElicitationRequest & { mode: 'form'; @@ -56,6 +56,10 @@ export async function requestAcpElicitation( elicitationRequest.id, 'cancelled' ); + acpChatSessionActions.resolveUserInputRequest( + elicitationRequest.sessionId, + acpElicitationUserInputRequestId(elicitationRequest.id) + ); pending.resolve(cancelledElicitationResponse()); }, ACP_ELICITATION_TIMEOUT_SECONDS * 1000); @@ -78,6 +82,10 @@ export function resolveAcpElicitationRequest( pendingRequests.delete(key); clearTimeout(pending.timeoutId); acpChatSessionActions.setElicitationStatus(sessionId, elicitationId, 'submitted'); + acpChatSessionActions.resolveUserInputRequest( + sessionId, + acpElicitationUserInputRequestId(elicitationId) + ); pending.resolve(acceptedElicitationResponse(userData)); return true; } diff --git a/ui/desktop/src/acp/permissionRequests.ts b/ui/desktop/src/acp/permissionRequests.ts index 847ad141f..c6f2167eb 100644 --- a/ui/desktop/src/acp/permissionRequests.ts +++ b/ui/desktop/src/acp/permissionRequests.ts @@ -1,7 +1,7 @@ import type { RequestPermissionRequest, RequestPermissionResponse } from '@agentclientprotocol/sdk'; import type { Permission } from '../api'; import { USE_ACP_CHAT } from '../acpChatFeatureFlag'; -import { acpChatSessionActions } from './chatSessionStore'; +import { acpChatSessionActions, acpPermissionUserInputRequestId } from './chatSessionStore'; interface PendingPermissionRequest { request: RequestPermissionRequest; @@ -41,6 +41,10 @@ export function resolveAcpPermissionRequest( } pendingRequests.delete(key); + acpChatSessionActions.resolveUserInputRequest( + sessionId, + acpPermissionUserInputRequestId(toolCallId) + ); pending.resolve(permissionResponseForAction(pending.request, action)); return true; } diff --git a/ui/desktop/src/components/ToolApprovalButtons.test.tsx b/ui/desktop/src/components/ToolApprovalButtons.test.tsx index c7a9890c3..6800c6fb5 100644 --- a/ui/desktop/src/components/ToolApprovalButtons.test.tsx +++ b/ui/desktop/src/components/ToolApprovalButtons.test.tsx @@ -6,6 +6,10 @@ import { resolveAcpPermissionRequest } from '../acp/permissionRequests'; import { IntlTestWrapper } from '../i18n/test-utils'; import ToolApprovalButtons from './ToolApprovalButtons'; +const acpChatFeatureFlagMock = vi.hoisted(() => ({ + useAcpChat: true, +})); + vi.mock('../api', () => ({ confirmToolAction: vi.fn(), })); @@ -15,7 +19,9 @@ vi.mock('../acp/permissionRequests', () => ({ })); vi.mock('../acpChatFeatureFlag', () => ({ - USE_ACP_CHAT: true, + get USE_ACP_CHAT() { + return acpChatFeatureFlagMock.useAcpChat; + }, })); const renderWithIntl = (ui: React.ReactElement, options?: RenderOptions) => @@ -27,6 +33,7 @@ const resolveAcpPermissionRequestMock = vi.mocked(resolveAcpPermissionRequest); describe('ToolApprovalButtons', () => { beforeEach(() => { vi.clearAllMocks(); + acpChatFeatureFlagMock.useAcpChat = true; }); it('marks the approval accepted when the ACP request resolves', async () => { @@ -53,11 +60,8 @@ describe('ToolApprovalButtons', () => { expect(screen.getByText('developer__shell - Allowed once')).toBeInTheDocument(); }); - it('falls back to the REST confirmation when no ACP request is pending', async () => { + it('shows a stale request error when ACP has no pending request', async () => { resolveAcpPermissionRequestMock.mockReturnValueOnce(false); - confirmToolActionMock.mockResolvedValueOnce({ error: undefined } as Awaited< - ReturnType - >); renderWithIntl( { 'tool-call-rerun', 'allow_once' ); + expect(confirmToolActionMock).not.toHaveBeenCalled(); + expect(screen.getByText('This approval request is no longer active.')).toBeInTheDocument(); + expect(screen.queryByText('developer__shell - Allowed once')).not.toBeInTheDocument(); + }); + + it('uses the REST confirmation path when ACP chat is disabled', async () => { + acpChatFeatureFlagMock.useAcpChat = false; + confirmToolActionMock.mockResolvedValueOnce({ error: undefined } as Awaited< + ReturnType + >); + + renderWithIntl( + + ); + + await userEvent.click(screen.getByRole('button', { name: 'Allow Once' })); + + expect(resolveAcpPermissionRequestMock).not.toHaveBeenCalled(); expect(confirmToolActionMock).toHaveBeenCalledWith({ body: { sessionId: 'session-1', - id: 'tool-call-rerun', + id: 'tool-call-rest', action: 'allow_once', principalType: 'Tool', }, diff --git a/ui/desktop/src/components/ToolApprovalButtons.tsx b/ui/desktop/src/components/ToolApprovalButtons.tsx index e6472c0b6..8922611ba 100644 --- a/ui/desktop/src/components/ToolApprovalButtons.tsx +++ b/ui/desktop/src/components/ToolApprovalButtons.tsx @@ -38,6 +38,10 @@ const i18n = defineMessages({ id: 'toolApprovalButtons.cancelled', defaultMessage: 'Cancelled', }, + staleApprovalRequest: { + id: 'toolApprovalButtons.staleApprovalRequest', + defaultMessage: 'This approval request is no longer active.', + }, }); const globalApprovalState = new Map< @@ -63,6 +67,13 @@ export default function ToolApprovalButtons({ data }: { data: ToolApprovalData } const storedState = globalApprovalState.get(id); const [decision, setDecision] = useState(storedState?.decision ?? null); const [isClicked, setIsClicked] = useState(storedState?.isClicked ?? initialIsClicked ?? false); + const [approvalError, setApprovalError] = useState(null); + + const setResolvedDecision = (action: Permission) => { + setDecision(action); + setIsClicked(true); + setApprovalError(null); + }; useEffect(() => { const currentState = globalApprovalState.get(id); @@ -70,6 +81,7 @@ export default function ToolApprovalButtons({ data }: { data: ToolApprovalData } setDecision(currentState.decision); setIsClicked(currentState.isClicked); } + setApprovalError(null); }, [id]); useEffect(() => { @@ -78,16 +90,16 @@ export default function ToolApprovalButtons({ data }: { data: ToolApprovalData } const handleAction = async (action: Permission) => { try { - // Edit-in-place reruns go through the legacy REST path even when ACP chat is - // enabled, so fall back to confirmToolAction when no ACP request is pending. - if (USE_ACP_CHAT && resolveAcpPermissionRequest(sessionId, id, action)) { - setDecision(action); - setIsClicked(true); + if (USE_ACP_CHAT) { + if (resolveAcpPermissionRequest(sessionId, id, action)) { + setResolvedDecision(action); + } else { + setApprovalError(intl.formatMessage(i18n.staleApprovalRequest)); + } return; } - setDecision(action); - setIsClicked(true); + setResolvedDecision(action); const response = await confirmToolAction({ body: { @@ -121,26 +133,33 @@ export default function ToolApprovalButtons({ data }: { data: ToolApprovalData } } return ( -
- - {!prompt && ( + <> +
+ {!prompt && ( + + )} + +
+ {approvalError && ( +

+ {approvalError} +

)} - -
+ ); } diff --git a/ui/desktop/src/i18n/messages/en.json b/ui/desktop/src/i18n/messages/en.json index 8933bae98..70b901ad5 100644 --- a/ui/desktop/src/i18n/messages/en.json +++ b/ui/desktop/src/i18n/messages/en.json @@ -4388,6 +4388,9 @@ "toolApprovalButtons.deny": { "defaultMessage": "Deny" }, + "toolApprovalButtons.staleApprovalRequest": { + "defaultMessage": "This approval request is no longer active." + }, "toolCallStatusIndicator.toolStatus": { "defaultMessage": "Tool status: {status}" }, diff --git a/ui/desktop/src/i18n/messages/es.json b/ui/desktop/src/i18n/messages/es.json index dda634bf1..ec10a1310 100644 --- a/ui/desktop/src/i18n/messages/es.json +++ b/ui/desktop/src/i18n/messages/es.json @@ -4388,6 +4388,9 @@ "toolApprovalButtons.deny": { "defaultMessage": "Denegar" }, + "toolApprovalButtons.staleApprovalRequest": { + "defaultMessage": "This approval request is no longer active." + }, "toolCallStatusIndicator.toolStatus": { "defaultMessage": "Estado de la herramienta: {status}" }, diff --git a/ui/desktop/src/i18n/messages/hi.json b/ui/desktop/src/i18n/messages/hi.json index d2bdbde46..312dc26b7 100644 --- a/ui/desktop/src/i18n/messages/hi.json +++ b/ui/desktop/src/i18n/messages/hi.json @@ -4388,6 +4388,9 @@ "toolApprovalButtons.deny": { "defaultMessage": "इन्कार" }, + "toolApprovalButtons.staleApprovalRequest": { + "defaultMessage": "This approval request is no longer active." + }, "toolCallStatusIndicator.toolStatus": { "defaultMessage": "उपकरण स्थिति: {status}" }, diff --git a/ui/desktop/src/i18n/messages/ja.json b/ui/desktop/src/i18n/messages/ja.json index 3feeb5b0e..08f8d39c2 100644 --- a/ui/desktop/src/i18n/messages/ja.json +++ b/ui/desktop/src/i18n/messages/ja.json @@ -4388,6 +4388,9 @@ "toolApprovalButtons.deny": { "defaultMessage": "拒否" }, + "toolApprovalButtons.staleApprovalRequest": { + "defaultMessage": "This approval request is no longer active." + }, "toolCallStatusIndicator.toolStatus": { "defaultMessage": "ツールのステータス: {status}" }, diff --git a/ui/desktop/src/i18n/messages/ko.json b/ui/desktop/src/i18n/messages/ko.json index 6de27b4df..4e440494a 100644 --- a/ui/desktop/src/i18n/messages/ko.json +++ b/ui/desktop/src/i18n/messages/ko.json @@ -4388,6 +4388,9 @@ "toolApprovalButtons.deny": { "defaultMessage": "거부" }, + "toolApprovalButtons.staleApprovalRequest": { + "defaultMessage": "This approval request is no longer active." + }, "toolCallStatusIndicator.toolStatus": { "defaultMessage": "도구 상태: {status}" }, diff --git a/ui/desktop/src/i18n/messages/ru.json b/ui/desktop/src/i18n/messages/ru.json index 3390fbf15..61ceb8cdc 100644 --- a/ui/desktop/src/i18n/messages/ru.json +++ b/ui/desktop/src/i18n/messages/ru.json @@ -4388,6 +4388,9 @@ "toolApprovalButtons.deny": { "defaultMessage": "Отклонить" }, + "toolApprovalButtons.staleApprovalRequest": { + "defaultMessage": "This approval request is no longer active." + }, "toolCallStatusIndicator.toolStatus": { "defaultMessage": "Статус инструмента: {status}" }, diff --git a/ui/desktop/src/i18n/messages/tr.json b/ui/desktop/src/i18n/messages/tr.json index 46481e610..6d983dfbf 100644 --- a/ui/desktop/src/i18n/messages/tr.json +++ b/ui/desktop/src/i18n/messages/tr.json @@ -4388,6 +4388,9 @@ "toolApprovalButtons.deny": { "defaultMessage": "Reddet" }, + "toolApprovalButtons.staleApprovalRequest": { + "defaultMessage": "This approval request is no longer active." + }, "toolCallStatusIndicator.toolStatus": { "defaultMessage": "Takım durumu: {status}" }, diff --git a/ui/desktop/src/i18n/messages/zh-CN.json b/ui/desktop/src/i18n/messages/zh-CN.json index 7c88961e8..f01a9853a 100644 --- a/ui/desktop/src/i18n/messages/zh-CN.json +++ b/ui/desktop/src/i18n/messages/zh-CN.json @@ -4388,6 +4388,9 @@ "toolApprovalButtons.deny": { "defaultMessage": "拒绝" }, + "toolApprovalButtons.staleApprovalRequest": { + "defaultMessage": "This approval request is no longer active." + }, "toolCallStatusIndicator.toolStatus": { "defaultMessage": "工具状态：{status}" },