1 files changed, 120 insertions, 15 deletions

diff --git a/packages/core/src/test-utils/tools.ts b/packages/core/src/test-utils/tools.ts
index da642212..0e3e6b86 100644
--- a/packages/core/src/test-utils/tools.ts
+++ b/packages/core/src/test-utils/tools.ts
@@ -6,17 +6,67 @@
 
 import { vi } from 'vitest';
 import {
-  BaseTool,
+  BaseDeclarativeTool,
+  BaseToolInvocation,
   ToolCallConfirmationDetails,
+  ToolInvocation,
   ToolResult,
   Kind,
 } from '../tools/tools.js';
 import { Schema, Type } from '@google/genai';
+import {
+  ModifiableDeclarativeTool,
+  ModifyContext,
+} from '../tools/modifiable-tool.js';
+
+class MockToolInvocation extends BaseToolInvocation<
+  { [key: string]: unknown },
+  ToolResult
+> {
+  constructor(
+    private readonly tool: MockTool,
+    params: { [key: string]: unknown },
+  ) {
+    super(params);
+  }
+
+  async execute(_abortSignal: AbortSignal): Promise<ToolResult> {
+    const result = this.tool.executeFn(this.params);
+    return (
+      result ?? {
+        llmContent: `Tool ${this.tool.name} executed successfully.`,
+        returnDisplay: `Tool ${this.tool.name} executed successfully.`,
+      }
+    );
+  }
+
+  override async shouldConfirmExecute(
+    _abortSignal: AbortSignal,
+  ): Promise<ToolCallConfirmationDetails | false> {
+    if (this.tool.shouldConfirm) {
+      return {
+        type: 'exec' as const,
+        title: `Confirm ${this.tool.displayName}`,
+        command: this.tool.name,
+        rootCommand: this.tool.name,
+        onConfirm: async () => {},
+      };
+    }
+    return false;
+  }
+
+  getDescription(): string {
+    return `A mock tool invocation for ${this.tool.name}`;
+  }
+}
 
 /**
  * A highly configurable mock tool for testing purposes.
  */
-export class MockTool extends BaseTool<{ [key: string]: unknown }, ToolResult> {
+export class MockTool extends BaseDeclarativeTool<
+  { [key: string]: unknown },
+  ToolResult
+> {
   executeFn = vi.fn();
   shouldConfirm = false;
 
@@ -32,32 +82,87 @@ export class MockTool extends BaseTool<{ [key: string]: unknown }, ToolResult> {
     super(name, displayName ?? name, description, Kind.Other, params);
   }
 
-  async execute(
-    params: { [key: string]: unknown },
-    _abortSignal: AbortSignal,
-  ): Promise<ToolResult> {
-    const result = this.executeFn(params);
+  protected createInvocation(params: {
+    [key: string]: unknown;
+  }): ToolInvocation<{ [key: string]: unknown }, ToolResult> {
+    return new MockToolInvocation(this, params);
+  }
+}
+
+export class MockModifiableToolInvocation extends BaseToolInvocation<
+  Record<string, unknown>,
+  ToolResult
+> {
+  constructor(
+    private readonly tool: MockModifiableTool,
+    params: Record<string, unknown>,
+  ) {
+    super(params);
+  }
+
+  async execute(_abortSignal: AbortSignal): Promise<ToolResult> {
+    const result = this.tool.executeFn(this.params);
     return (
       result ?? {
-        llmContent: `Tool ${this.name} executed successfully.`,
-        returnDisplay: `Tool ${this.name} executed successfully.`,
+        llmContent: `Tool ${this.tool.name} executed successfully.`,
+        returnDisplay: `Tool ${this.tool.name} executed successfully.`,
       }
     );
   }
 
   override async shouldConfirmExecute(
-    _params: { [key: string]: unknown },
     _abortSignal: AbortSignal,
   ): Promise<ToolCallConfirmationDetails | false> {
-    if (this.shouldConfirm) {
+    if (this.tool.shouldConfirm) {
       return {
-        type: 'exec' as const,
-        title: `Confirm ${this.displayName}`,
-        command: this.name,
-        rootCommand: this.name,
+        type: 'edit',
+        title: 'Confirm Mock Tool',
+        fileName: 'test.txt',
+        filePath: 'test.txt',
+        fileDiff: 'diff',
+        originalContent: 'originalContent',
+        newContent: 'newContent',
         onConfirm: async () => {},
       };
     }
     return false;
   }
+
+  getDescription(): string {
+    return `A mock modifiable tool invocation for ${this.tool.name}`;
+  }
+}
+
+/**
+ * Configurable mock modifiable tool for testing.
+ */
+export class MockModifiableTool
+  extends MockTool
+  implements ModifiableDeclarativeTool<Record<string, unknown>>
+{
+  constructor(name = 'mockModifiableTool') {
+    super(name);
+    this.shouldConfirm = true;
+  }
+
+  getModifyContext(
+    _abortSignal: AbortSignal,
+  ): ModifyContext<Record<string, unknown>> {
+    return {
+      getFilePath: () => 'test.txt',
+      getCurrentContent: async () => 'old content',
+      getProposedContent: async () => 'new content',
+      createUpdatedParams: (
+        _oldContent: string,
+        modifiedProposedContent: string,
+        _originalParams: Record<string, unknown>,
+      ) => ({ newContent: modifiedProposedContent }),
+    };
+  }
+
+  protected override createInvocation(
+    params: Record<string, unknown>,
+  ): ToolInvocation<Record<string, unknown>, ToolResult> {
+    return new MockModifiableToolInvocation(this, params);
+  }
 }