Refactor: Use tool wrapper for init wait instead of per-tool calls

ammar-agent · ammar-agent · commit 90777111b69b · 2025-10-27T18:29:07.000Z
- Created wrapWithInitWait() wrapper to centralize init waiting
- Removed waitForWorkspaceInit() calls from bash, file_read, file_edit_insert, file_edit_operation
- Deleted toolHelpers.ts (no longer needed)
- Updated tools.ts to wrap runtime tools (bash, file_read, file_edit_*)
- Non-runtime tools (propose_plan, todo, web_search) unchanged
- Clearer separation: explicit list of which tools wait for init
- Net: -36 lines of code, better architecture
diff --git a/src/services/tools/bash.ts b/src/services/tools/bash.ts
@@ -16,7 +16,6 @@ import { EXIT_CODE_ABORTED, EXIT_CODE_TIMEOUT } from "@/constants/exitCodes";
 import type { BashToolResult } from "@/types/tools";
 import type { ToolConfiguration, ToolFactory } from "@/utils/tools/tools";
 import { TOOL_DEFINITIONS } from "@/utils/tools/toolDefinitions";
-import { waitForWorkspaceInit } from "./toolHelpers";
 
 /**
  * Bash execution tool factory for AI assistant
@@ -39,9 +38,6 @@ export const createBashTool: ToolFactory = (config: ToolConfiguration) => {
     description: TOOL_DEFINITIONS.bash.description + "\nRuns in " + config.cwd + " - no cd needed",
     inputSchema: TOOL_DEFINITIONS.bash.schema,
     execute: async ({ script, timeout_secs }, { abortSignal }): Promise<BashToolResult> => {
-      // Wait for workspace initialization to complete (no-op if already complete or not needed)
-      await waitForWorkspaceInit(config);
-
       // Validate script is not empty - likely indicates a malformed tool call
 
       if (!script || script.trim().length === 0) {
diff --git a/src/services/tools/file_edit_insert.ts b/src/services/tools/file_edit_insert.ts
@@ -8,7 +8,6 @@ import { executeFileEditOperation } from "./file_edit_operation";
 import { RuntimeError } from "@/runtime/Runtime";
 import { fileExists } from "@/utils/runtime/fileExists";
 import { writeFileString } from "@/utils/runtime/helpers";
-import { waitForWorkspaceInit } from "./toolHelpers";
 
 /**
  * File edit insert tool factory for AI assistant
@@ -25,9 +24,6 @@ export const createFileEditInsertTool: ToolFactory = (config: ToolConfiguration)
       content,
       create,
     }): Promise<FileEditInsertToolResult> => {
-      // Wait for workspace initialization to complete (no-op if already complete or not needed)
-      await waitForWorkspaceInit(config);
-
       try {
         // Validate no redundant path prefix (must come first to catch absolute paths)
         const redundantPrefixValidation = validateNoRedundantPrefix(
diff --git a/src/services/tools/file_edit_operation.ts b/src/services/tools/file_edit_operation.ts
@@ -9,7 +9,6 @@ import {
 } from "./fileCommon";
 import { RuntimeError } from "@/runtime/Runtime";
 import { readFileString, writeFileString } from "@/utils/runtime/helpers";
-import { waitForWorkspaceInit } from "./toolHelpers";
 
 type FileEditOperationResult<TMetadata> =
   | {
@@ -41,9 +40,6 @@ export async function executeFileEditOperation<TMetadata>({
 }: ExecuteFileEditOperationOptions<TMetadata>): Promise<
   FileEditErrorResult | (FileEditDiffSuccessBase & TMetadata)
 > {
-  // Wait for workspace initialization to complete (no-op if already complete or not needed)
-  await waitForWorkspaceInit(config);
-
   try {
     // Validate no redundant path prefix (must come first to catch absolute paths)
     const redundantPrefixValidation = validateNoRedundantPrefix(
diff --git a/src/services/tools/file_read.ts b/src/services/tools/file_read.ts
@@ -5,7 +5,6 @@ import { TOOL_DEFINITIONS } from "@/utils/tools/toolDefinitions";
 import { validatePathInCwd, validateFileSize, validateNoRedundantPrefix } from "./fileCommon";
 import { RuntimeError } from "@/runtime/Runtime";
 import { readFileString } from "@/utils/runtime/helpers";
-import { waitForWorkspaceInit } from "./toolHelpers";
 
 /**
  * File read tool factory for AI assistant
@@ -22,9 +21,6 @@ export const createFileReadTool: ToolFactory = (config: ToolConfiguration) => {
     ): Promise<FileReadToolResult> => {
       // Note: abortSignal available but not used - file reads are fast and complete quickly
 
-      // Wait for workspace initialization to complete (no-op if already complete or not needed)
-      await waitForWorkspaceInit(config);
-
       try {
         // Validate no redundant path prefix (must come first to catch absolute paths)
         const redundantPrefixValidation = validateNoRedundantPrefix(
diff --git a/src/services/tools/toolHelpers.ts b/src/services/tools/toolHelpers.ts
diff --git a/src/services/tools/wrapWithInitWait.ts b/src/services/tools/wrapWithInitWait.ts
@@ -0,0 +1,36 @@
+import type { Tool } from "ai";
+import type { ToolConfiguration } from "@/utils/tools/tools";
+
+/**
+ * Wraps a tool to wait for workspace initialization before execution.
+ * 
+ * This wrapper handles the cross-cutting concern of init state waiting,
+ * keeping individual tools simple and focused on their core functionality.
+ * 
+ * Only runtime-dependent tools (bash, file_read, file_edit_*) need this wrapper.
+ * Non-runtime tools (propose_plan, todo, web_search) execute immediately.
+ * 
+ * @param tool The tool to wrap (returned from a tool factory)
+ * @param config Tool configuration containing initStateManager
+ * @returns Wrapped tool that waits for init before executing
+ */
+export function wrapWithInitWait<TParameters, TResult>(
+  tool: Tool<TParameters, TResult>,
+  config: ToolConfiguration
+): Tool<TParameters, TResult> {
+  return {
+    ...tool,
+    execute: async (args: TParameters, options) => {
+      // Wait for workspace initialization to complete (no-op if not needed)
+      // This never throws - tools proceed regardless of init outcome
+      await config.initStateManager.waitForInit(config.workspaceId);
+
+      // Execute the actual tool with all arguments
+      if (!tool.execute) {
+        throw new Error("Tool does not have an execute function");
+      }
+      return tool.execute(args, options);
+    },
+  } as Tool<TParameters, TResult>;
+}
+
diff --git a/src/utils/tools/tools.ts b/src/utils/tools/tools.ts
@@ -6,6 +6,7 @@ import { createFileEditReplaceStringTool } from "@/services/tools/file_edit_repl
 import { createFileEditInsertTool } from "@/services/tools/file_edit_insert";
 import { createProposePlanTool } from "@/services/tools/propose_plan";
 import { createTodoWriteTool, createTodoReadTool } from "@/services/tools/todo";
+import { wrapWithInitWait } from "@/services/tools/wrapWithInitWait";
 import { log } from "@/services/log";
 
 import type { Runtime } from "@/runtime/Runtime";
@@ -54,22 +55,32 @@ export async function getToolsForModel(
 ): Promise<Record<string, Tool>> {
   const [provider, modelId] = modelString.split(":");
 
-  // Base tools available for all models
-  const baseTools: Record<string, Tool> = {
-    // Use snake_case for tool names to match what seems to be the convention.
-    file_read: createFileReadTool(config),
-    file_edit_replace_string: createFileEditReplaceStringTool(config),
+  // Runtime-dependent tools need to wait for workspace initialization
+  // Wrap them to handle init waiting centrally instead of in each tool
+  const runtimeTools: Record<string, Tool> = {
+    file_read: wrapWithInitWait(createFileReadTool(config), config),
+    file_edit_replace_string: wrapWithInitWait(createFileEditReplaceStringTool(config), config),
     // DISABLED: file_edit_replace_lines - causes models (particularly GPT-5-Codex)
     // to leave repository in broken state due to issues with concurrent file modifications
     // and line number miscalculations. Use file_edit_replace_string or file_edit_insert instead.
-    // file_edit_replace_lines: createFileEditReplaceLinesTool(config),
-    file_edit_insert: createFileEditInsertTool(config),
-    bash: createBashTool(config),
+    // file_edit_replace_lines: wrapWithInitWait(createFileEditReplaceLinesTool(config), config),
+    file_edit_insert: wrapWithInitWait(createFileEditInsertTool(config), config),
+    bash: wrapWithInitWait(createBashTool(config), config),
+  };
+
+  // Non-runtime tools execute immediately (no init wait needed)
+  const nonRuntimeTools: Record<string, Tool> = {
     propose_plan: createProposePlanTool(config),
     todo_write: createTodoWriteTool(config),
     todo_read: createTodoReadTool(config),
   };
 
+  // Base tools available for all models
+  const baseTools: Record<string, Tool> = {
+    ...runtimeTools,
+    ...nonRuntimeTools,
+  };
+
   // Try to add provider-specific web search tools if available
   // Lazy-load providers to avoid loading all AI SDKs at startup
   try {
diff --git a/tests/ipcMain/helpers.ts b/tests/ipcMain/helpers.ts
@@ -565,7 +565,9 @@ export function collectInitEvents(
   return env.sentEvents
     .filter((e) => e.channel === getChatChannel(workspaceId))
     .map((e) => e.data as WorkspaceChatMessage)
-    .filter((msg) => isInitStart(msg) || isInitOutput(msg) || isInitEnd(msg)) as WorkspaceInitEvent[];
+    .filter(
+      (msg) => isInitStart(msg) || isInitOutput(msg) || isInitEnd(msg)
+    ) as WorkspaceInitEvent[];
 }
 
 /**
