feat: implemented confirmation mechanism for tool invokation

2026-06-03 17:05:42 -07:00
parent 414565ecc9
commit ad1e0c2aa8
23 changed files with 860 additions and 84 deletions
@@ -1,7 +1,23 @@
 import { act, renderHook } from '@testing-library/react';
-import { describe, expect, it, vi } from 'vitest';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
 import type { TodoItem } from '../agent/core/todo';

+const { mockedStoreState } = vi.hoisted(() => {
+  const state = {
+    activeRegionId: null as string | null,
+    selectedRegionIds: [] as string[],
+    timeSignature: { numerator: 4, denominator: 4 },
+    tracks: [] as unknown[],
+    refreshProjectState: vi.fn(),
+    toolFastForwardEnabled: false,
+    setToolFastForwardEnabled: vi.fn(),
+  };
+  state.setToolFastForwardEnabled.mockImplementation((enabled: boolean) => {
+    state.toolFastForwardEnabled = enabled;
+  });
+  return { mockedStoreState: state };
+});
+
 vi.mock('../agent/core/AgentCore', () => ({
  AgentCore: {
    instance: vi.fn()
@@ -16,13 +32,7 @@ vi.mock('../core/KGCore', () => ({

 vi.mock('../stores/projectStore', () => ({
  useProjectStore: {
-    getState: vi.fn(() => ({
-      activeRegionId: null,
-      selectedRegionIds: [],
-      timeSignature: { numerator: 4, denominator: 4 },
-      tracks: [],
-      refreshProjectState: vi.fn(),
-    })),
+    getState: vi.fn(() => mockedStoreState),
  },
 }));

@@ -53,6 +63,13 @@ const flushMicrotasks = async (): Promise<void> => {
 };

 describe('useStreamProcessor', () => {
+  beforeEach(() => {
+    mockedStoreState.activeRegionId = null;
+    mockedStoreState.toolFastForwardEnabled = false;
+    mockedStoreState.setToolFastForwardEnabled.mockClear();
+    mockedStoreState.refreshProjectState.mockClear();
+  });
+
  it('switches from Thinking to Processing after the first text token arrives', async () => {
    let releaseDone!: () => void;
    const doneGate = new Promise<void>((resolve) => {
@@ -388,4 +405,172 @@ describe('useStreamProcessor', () => {
    expect(toolResultMessage?.toolRawResult).toBe('raw music result');
    expect(toolResultMessage?.toolResultDisplayContent).toBe('raw music result');
  });
+
+  it('shows a confirmation card and replaces it with a denied result when the user denies execution', async () => {
+    vi.spyOn(AgentCore, 'instance').mockReturnValue({
+      getAgentState: () => ({
+        getTodos: () => [],
+      }),
+      processUserInput: async function* (_input: string, options?: { requestToolApproval?: (toolCall: { id: string; type: 'function'; function: { name: string; arguments: string } }) => Promise<'allow' | 'always_allow' | 'deny'> }) {
+        const toolCall = {
+          id: 'add-notes-call-confirm',
+          type: 'function' as const,
+          function: {
+            name: 'add_notes',
+            arguments: JSON.stringify({
+              notes: [{ pitch: 'C4', start: 16, length: 4 }],
+            }),
+          },
+        };
+
+        yield { type: 'tool_call', content: '', toolCall };
+        const decision = await options?.requestToolApproval?.(toolCall);
+        yield {
+          type: 'tool_result',
+          content: '',
+          toolResult: {
+            toolCallId: toolCall.id,
+            name: 'add_notes',
+            success: false,
+            result: 'Execution was denied by the user.',
+            denied: decision === 'deny',
+          },
+        };
+      },
+    } as unknown as AgentCore);
+
+    const selectedRegion = new KGMidiRegion('region-1', '1', 0, 'Verse Melody');
+    vi.spyOn(KGCore, 'instance').mockReturnValue({
+      getCurrentProject: () => ({
+        getTimeSignature: () => ({ numerator: 4, denominator: 4 }),
+        getTracks: () => [
+          {
+            getId: () => '1',
+            getName: () => 'Lead',
+            getRegions: () => [selectedRegion],
+          },
+        ],
+      }),
+      getSelectedItems: () => [selectedRegion],
+    } as unknown as KGCore);
+    mockedStoreState.activeRegionId = selectedRegion.getId();
+
+    const messages = new Map<string, ChatMessage>();
+    const { result } = renderHook(() => useStreamProcessor({
+      onMessageAdd: (message) => {
+        messages.set(message.id, message);
+      },
+      onMessageUpdate: (messageId, updater) => {
+        const current = messages.get(messageId);
+        if (!current) {
+          throw new Error(`Missing message ${messageId}`);
+        }
+        messages.set(messageId, updater(current));
+      },
+      onMessageRemove: (messageId) => {
+        messages.delete(messageId);
+      },
+      onProcessingChange: () => undefined,
+    }));
+
+    let responsePromise!: Promise<string>;
+    await act(async () => {
+      responsePromise = result.current.processStream('confirm prompt');
+      await flushMicrotasks();
+    });
+
+    const confirmationMessage = [...messages.values()].find(message => message.toolConfirmation);
+    expect(confirmationMessage?.toolConfirmation?.toolName).toBe('add_notes');
+    expect(confirmationMessage?.toolConfirmation?.message).toContain('Allow creating 1 note in region **Verse Melody**');
+
+    act(() => {
+      confirmationMessage?.onToolConfirmationDecision?.('deny');
+    });
+
+    await act(async () => {
+      await responsePromise;
+    });
+
+    expect([...messages.values()].some(message => message.toolConfirmation)).toBe(false);
+    const deniedResult = [...messages.values()].find(message => message.toolName === 'add_notes');
+    expect(deniedResult?.toolDenied).toBe(true);
+    expect(deniedResult?.toolRawResult).toBe('Execution was denied by the user.');
+  });
+
+  it('skips the confirmation card when fast-forward mode is enabled', async () => {
+    mockedStoreState.toolFastForwardEnabled = true;
+    let capturedDecision: string | undefined;
+
+    vi.spyOn(AgentCore, 'instance').mockReturnValue({
+      getAgentState: () => ({
+        getTodos: () => [],
+      }),
+      processUserInput: async function* (_input: string, options?: { requestToolApproval?: (toolCall: { id: string; type: 'function'; function: { name: string; arguments: string } }) => Promise<'allow' | 'always_allow' | 'deny'> }) {
+        const toolCall = {
+          id: 'add-notes-call-auto',
+          type: 'function' as const,
+          function: {
+            name: 'add_notes',
+            arguments: JSON.stringify({
+              notes: [{ pitch: 'C4', start: 0, length: 4 }],
+            }),
+          },
+        };
+
+        yield { type: 'tool_call', content: '', toolCall };
+        capturedDecision = await options?.requestToolApproval?.(toolCall);
+        yield {
+          type: 'tool_result',
+          content: '',
+          toolResult: {
+            toolCallId: toolCall.id,
+            name: 'add_notes',
+            success: true,
+            result: 'Successfully created 1 note: C4 (beat 0, length 4)',
+          },
+        };
+      },
+    } as unknown as AgentCore);
+
+    const selectedRegion = new KGMidiRegion('region-1', '1', 0, 'Intro');
+    vi.spyOn(KGCore, 'instance').mockReturnValue({
+      getCurrentProject: () => ({
+        getTimeSignature: () => ({ numerator: 4, denominator: 4 }),
+        getTracks: () => [
+          {
+            getId: () => '1',
+            getName: () => 'Lead',
+            getRegions: () => [selectedRegion],
+          },
+        ],
+      }),
+      getSelectedItems: () => [selectedRegion],
+    } as unknown as KGCore);
+    mockedStoreState.activeRegionId = selectedRegion.getId();
+
+    const messages = new Map<string, ChatMessage>();
+    const { result } = renderHook(() => useStreamProcessor({
+      onMessageAdd: (message) => {
+        messages.set(message.id, message);
+      },
+      onMessageUpdate: (messageId, updater) => {
+        const current = messages.get(messageId);
+        if (!current) {
+          throw new Error(`Missing message ${messageId}`);
+        }
+        messages.set(messageId, updater(current));
+      },
+      onMessageRemove: (messageId) => {
+        messages.delete(messageId);
+      },
+      onProcessingChange: () => undefined,
+    }));
+
+    await act(async () => {
+      await result.current.processStream('auto allow prompt');
+    });
+
+    expect(capturedDecision).toBe('allow');
+    expect([...messages.values()].some(message => message.toolConfirmation)).toBe(false);
+  });
 });
@@ -1,8 +1,10 @@
 import { useState, useCallback } from 'react';
 import { AgentCore } from '../agent/core/AgentCore';
-import { AVAILABLE_TOOLS } from '../agent/tools';
+import { createToolInstance } from '../agent/tools';
 import { createStreamingMessage, createMessage } from '../utils/chatMessageUtils';
 import type { ChatMessage } from '../types/projectTypes';
+import type { ToolApprovalDecision } from '../agent/llm/StreamingTypes';
+import { useProjectStore } from '../stores/projectStore';

 const TODO_TOOL_NAME = 'update_todo_list';

@@ -12,29 +14,6 @@ interface PendingToolCall {
  arguments: Record<string, unknown> | null;
 }

-const buildToolResultDisplayContent = (
-  toolName: string,
-  success: boolean,
-  rawResult: string,
-  toolArgs: Record<string, unknown> | null,
-): string => {
-  if (!success) {
-    return rawResult;
-  }
-
-  const ToolClass = AVAILABLE_TOOLS[toolName as keyof typeof AVAILABLE_TOOLS];
-  if (!ToolClass) {
-    return rawResult;
-  }
-
-  try {
-    const toolInstance = new ToolClass();
-    return toolInstance.buildToolResultDisplayContent(toolArgs, { success, result: rawResult }) ?? rawResult;
-  } catch {
-    return rawResult;
-  }
-};
-
 interface StreamProcessorOptions {
  onMessageUpdate: (messageId: string, updater: (msg: ChatMessage) => ChatMessage) => void;
  onMessageAdd: (message: ChatMessage) => void;
@@ -79,7 +58,55 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce
      console.log(input);
      console.log('------------------------------');

-      for await (const chunk of agentCore.processUserInput(input)) {
+      const requestToolApproval = async (toolCall: PendingToolCall): Promise<ToolApprovalDecision> => {
+        const { toolFastForwardEnabled, setToolFastForwardEnabled } = useProjectStore.getState();
+        if (toolFastForwardEnabled) {
+          return 'allow';
+        }
+
+        const toolInstance = createToolInstance(toolCall.name);
+        const confirmationContent = toolInstance?.buildConfirmationContent(toolCall.arguments) ?? undefined;
+        if (!confirmationContent) {
+          return 'allow';
+        }
+
+        return await new Promise<ToolApprovalDecision>((resolve) => {
+          const confirmationMessage = {
+            ...createMessage('assistant', confirmationContent),
+            toolConfirmation: {
+              toolCallId: toolCall.id,
+              toolName: toolCall.name,
+              message: confirmationContent,
+            },
+            onToolConfirmationDecision: (decision: ToolApprovalDecision) => {
+              onMessageRemove(confirmationMessage.id);
+              if (decision === 'always_allow') {
+                setToolFastForwardEnabled(true);
+              }
+              resolve(decision);
+            },
+          };
+
+          onMessageAdd(confirmationMessage);
+        });
+      };
+
+      for await (const chunk of agentCore.processUserInput(input, {
+        requestToolApproval: async (toolCall) => {
+          let parsedArguments: Record<string, unknown> | null;
+          try {
+            parsedArguments = JSON.parse(toolCall.function.arguments);
+          } catch {
+            parsedArguments = null;
+          }
+
+          return requestToolApproval({
+            id: toolCall.id,
+            name: toolCall.function.name,
+            arguments: parsedArguments,
+          });
+        },
+      })) {
        if (controller.signal.aborted) {
          return '';
        }
@@ -142,17 +169,23 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce
          console.log('-------------------------------------');

          // Show tool result in UI
-          const { toolCallId, name, success, result } = chunk.toolResult;
+          const { toolCallId, name, success, result, denied } = chunk.toolResult;
          const pendingToolCallIndex = pendingToolCalls.findIndex(toolCall => toolCall.id === toolCallId);
          const pendingToolCall = pendingToolCallIndex >= 0
            ? pendingToolCalls.splice(pendingToolCallIndex, 1)[0]
            : undefined;
-          const toolResultDisplayContent = buildToolResultDisplayContent(
-            name,
-            success,
-            result,
-            pendingToolCall?.arguments ?? null,
-          );
+          let toolResultDisplayContent = result;
+          if (success) {
+            try {
+              const toolInstance = createToolInstance(name);
+              toolResultDisplayContent = toolInstance?.buildToolResultDisplayContent(
+                pendingToolCall?.arguments ?? null,
+                { success, result },
+              ) ?? result;
+            } catch {
+              toolResultDisplayContent = result;
+            }
+          }
          const toolResultMsg = name === TODO_TOOL_NAME
            ? {
              ...createMessage('assistant', result),
@@ -166,6 +199,7 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce
              toolSuccess: success,
              toolRawResult: result,
              toolResultDisplayContent,
+              toolDenied: denied,
            };
          onMessageAdd(toolResultMsg);

@@ -176,9 +210,11 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce
          performanceInfo = undefined;

          // Create a fresh streaming placeholder for the next LLM response
-          const nextMsg = createStreamingMessage();
-          currentStreamingId = nextMsg.id;
-          onMessageAdd(nextMsg);
+          if (!denied) {
+            const nextMsg = createStreamingMessage();
+            currentStreamingId = nextMsg.id;
+            onMessageAdd(nextMsg);
+          }
        } else if (chunk.type === 'done') {
          performanceInfo = chunk.performanceInfo;
          // Finalize the streaming message
@@ -209,12 +245,17 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce
      }

      console.error('Error processing stream:', error);
-      onMessageUpdate(currentStreamingId, (msg) => ({
-        ...msg,
-        content: `Error: ${error instanceof Error ? error.message : 'Failed to process message'}`,
-        isStreaming: false,
-        tokenCount: undefined
-      }));
+      const errorContent = `Error: ${error instanceof Error ? error.message : 'Failed to process message'}`;
+      try {
+        onMessageUpdate(currentStreamingId, (msg) => ({
+          ...msg,
+          content: errorContent,
+          isStreaming: false,
+          tokenCount: undefined
+        }));
+      } catch {
+        onMessageAdd(createMessage('assistant', errorContent));
+      }
      throw error;
    } finally {
      setAbortController(null);