🤖 Eliminate race conditions in init state with promise-based waiting

Replace event-based waiting with promise-based state map in InitStateManager.
This eliminates race conditions between timeout and event firing, simplifies
cleanup, and handles workspace deletion correctly.

Changes:
- Add initPromises map to store completion promises for running inits
- startInit() creates promise, endInit() resolves it
- waitForInit() uses Promise.race() for clean timeout handling
- clearInMemoryState() rejects orphaned promises on workspace deletion
- No event listener cleanup needed (promises auto-resolve)

Extended integration tests:
- Verify tools wait for init completion before executing
- Send second message after init completes to verify state persistence
- Use failed init (exit code 1) to verify tools proceed regardless of status
- Both messages succeed, second is faster (no init wait)

Benefits:
- Multiple tools share same promise (no duplicate listeners)
- Promise.race() handles timeout cleanly (no manual cleanup)
- Workspace deletion properly cancels pending waits
- Simpler code, fewer edge cases, easier to reason about

Tests: All 7 integration tests pass (19s), 111 tool tests pass (3.5s)

Author: ammar-agent

src/debug/agentSessionCli.ts

@@ -7,8 +7,8 @@ import { parseArgs } from "util";
 import { Config } from "@/config";
 import { HistoryService } from "@/services/historyService";
 import { PartialService } from "@/services/partialService";
-import { AIService } from "@/services/aiService";
 import { InitStateManager } from "@/services/initStateManager";
+import { AIService } from "@/services/aiService";
 import { AgentSession, type AgentSessionChatEvent } from "@/services/agentSession";
 import {
   isCaughtUpMessage,
@@ -209,8 +209,8 @@ async function main(): Promise<void> {
 
   const historyService = new HistoryService(config);
   const partialService = new PartialService(config, historyService);
-  const aiService = new AIService(config, historyService, partialService);
   const initStateManager = new InitStateManager(config);
+  const aiService = new AIService(config, historyService, partialService, initStateManager);
   ensureProvidersConfig(config);
 
   const session = new AgentSession({

src/debug/replay-history.ts

@@ -17,6 +17,7 @@ import { parseArgs } from "util";
 import { defaultConfig } from "@/config";
 import type { CmuxMessage } from "@/types/message";
 import { createCmuxMessage } from "@/types/message";
+import { InitStateManager } from "@/services/initStateManager";
 import { AIService } from "@/services/aiService";
 import { HistoryService } from "@/services/historyService";
 import { PartialService } from "@/services/partialService";
@@ -123,7 +124,8 @@ async function main() {
   const config = defaultConfig;
   const historyService = new HistoryService(config);
   const partialService = new PartialService(config, historyService);
-  const aiService = new AIService(config, historyService, partialService);
+  const initStateManager = new InitStateManager(config);
+  const aiService = new AIService(config, historyService, partialService, initStateManager);
 
   const modelString = values.model ?? "openai:gpt-5-codex";
   const thinkingLevel = (values.thinking ?? "high") as "low" | "medium" | "high";

src/services/aiService.test.ts

@@ -6,6 +6,7 @@ import { describe, it, expect, beforeEach } from "bun:test";
 import { AIService } from "./aiService";
 import { HistoryService } from "./historyService";
 import { PartialService } from "./partialService";
+import { InitStateManager } from "./initStateManager";
 import { Config } from "@/config";
 
 describe("AIService", () => {
@@ -15,7 +16,8 @@ describe("AIService", () => {
     const config = new Config();
     const historyService = new HistoryService(config);
     const partialService = new PartialService(config, historyService);
-    service = new AIService(config, historyService, partialService);
+    const initStateManager = new InitStateManager(config);
+    service = new AIService(config, historyService, partialService, initStateManager);
   });
 
   // Note: These tests are placeholders as Bun doesn't support Jest mocking

src/services/aiService.ts

@@ -12,6 +12,7 @@ import type { CmuxMessage, CmuxTextPart } from "@/types/message";
 import { createCmuxMessage } from "@/types/message";
 import type { Config } from "@/config";
 import { StreamManager } from "./streamManager";
+import { InitStateManager } from "./initStateManager";
 import type { SendMessageError } from "@/types/errors";
 import { getToolsForModel } from "@/utils/tools/tools";
 import { createRuntime } from "@/runtime/runtimeFactory";
@@ -108,17 +109,24 @@ export class AIService extends EventEmitter {
   private readonly historyService: HistoryService;
   private readonly partialService: PartialService;
   private readonly config: Config;
+  private readonly initStateManager: InitStateManager;
   private readonly mockModeEnabled: boolean;
   private readonly mockScenarioPlayer?: MockScenarioPlayer;
 
-  constructor(config: Config, historyService: HistoryService, partialService: PartialService) {
+  constructor(
+    config: Config,
+    historyService: HistoryService,
+    partialService: PartialService,
+    initStateManager: InitStateManager
+  ) {
     super();
     // Increase max listeners to accommodate multiple concurrent workspace listeners
     // Each workspace subscribes to stream events, and we expect >10 concurrent workspaces
     this.setMaxListeners(50);
     this.config = config;
     this.historyService = historyService;
     this.partialService = partialService;
+    this.initStateManager = initStateManager;
     this.streamManager = new StreamManager(historyService, partialService);
     void this.ensureSessionsDir();
     this.setupStreamEventForwarding();
@@ -427,6 +435,8 @@ export class AIService extends EventEmitter {
         cwd: process.cwd(),
         runtime: earlyRuntime,
         runtimeTempDir: os.tmpdir(),
+        workspaceId: "", // Empty workspace ID for early stub config
+        initStateManager: this.initStateManager,
         secrets: {},
       });
       const earlyTools = applyToolPolicy(earlyAllTools, toolPolicy);
@@ -536,6 +546,8 @@ export class AIService extends EventEmitter {
       const allTools = await getToolsForModel(modelString, {
         cwd: workspacePath,
         runtime,
+        workspaceId,
+        initStateManager: this.initStateManager,
         secrets: secretsToRecord(projectSecrets),
         runtimeTempDir,
       });

src/services/initStateManager.ts

@@ -47,15 +47,29 @@ type InitHookState = InitStatus;
  * - Permanent persistence (init logs kept forever as workspace metadata)
  *
  * Lifecycle:
- * 1. startInit() - Create in-memory state, emit init-start
+ * 1. startInit() - Create in-memory state, emit init-start, create completion promise
  * 2. appendOutput() - Accumulate lines, emit init-output
- * 3. endInit() - Finalize state, write to disk, emit init-end
+ * 3. endInit() - Finalize state, write to disk, emit init-end, resolve promise
  * 4. State remains in memory until cleared or process restart
  * 5. replayInit() - Re-emit events from in-memory or disk state (via EventStore)
+ *
+ * Waiting: Tools use waitForInit() which returns a promise that resolves when
+ * init completes. This promise is stored in initPromises map and resolved by
+ * endInit(). No event listeners needed, eliminating race conditions.
  */
 export class InitStateManager extends EventEmitter {
   private readonly store: EventStore<InitHookState, WorkspaceInitEvent & { workspaceId: string }>;
 
+  /**
+   * Promise-based completion tracking for running inits.
+   * Each running init has a promise that resolves when endInit() is called.
+   * Multiple tools can await the same promise without race conditions.
+   */
+  private readonly initPromises: Map<
+    string,
+    { promise: Promise<void>; resolve: () => void; reject: (error: Error) => void }
+  > = new Map();
+
   constructor(config: Config) {
     super();
     this.store = new EventStore(
@@ -111,7 +125,7 @@ export class InitStateManager extends EventEmitter {
 
   /**
    * Start tracking a new init hook execution.
-   * Creates in-memory state and emits init-start event.
+   * Creates in-memory state, completion promise, and emits init-start event.
    */
   startInit(workspaceId: string, hookPath: string): void {
     const startTime = Date.now();
@@ -127,6 +141,21 @@ export class InitStateManager extends EventEmitter {
 
     this.store.setState(workspaceId, state);
 
+    // Create completion promise for this init
+    // This allows multiple tools to await the same init without event listeners
+    let resolve: () => void;
+    let reject: (error: Error) => void;
+    const promise = new Promise<void>((res, rej) => {
+      resolve = res;
+      reject = rej;
+    });
+
+    this.initPromises.set(workspaceId, {
+      promise,
+      resolve: resolve!,
+      reject: reject!,
+    });
+
     log.debug(`Init hook started for workspace ${workspaceId}: ${hookPath}`);
 
     // Emit init-start event
@@ -167,7 +196,7 @@ export class InitStateManager extends EventEmitter {
 
   /**
    * Finalize init hook execution.
-   * Updates state, persists to disk, and emits init-end event.
+   * Updates state, persists to disk, emits init-end event, and resolves completion promise.
    */
   async endInit(workspaceId: string, exitCode: number): Promise<void> {
     const state = this.store.getState(workspaceId);
@@ -197,6 +226,13 @@ export class InitStateManager extends EventEmitter {
       timestamp: endTime,
     } satisfies WorkspaceInitEvent & { workspaceId: string });
 
+    // Resolve completion promise for waiting tools
+    const promiseEntry = this.initPromises.get(workspaceId);
+    if (promiseEntry) {
+      promiseEntry.resolve();
+      this.initPromises.delete(workspaceId);
+    }
+
     // Keep state in memory for replay (unlike streams which delete immediately)
   }
 
@@ -244,8 +280,68 @@ export class InitStateManager extends EventEmitter {
    * Clear in-memory state for a workspace.
    * Useful for testing or cleanup after workspace deletion.
    * Does NOT delete disk file (use deleteInitStatus for that).
+   *
+   * Also cancels any running init promises to prevent orphaned waiters.
    */
   clearInMemoryState(workspaceId: string): void {
     this.store.deleteState(workspaceId);
+
+    // Cancel any running init promise for this workspace
+    const promiseEntry = this.initPromises.get(workspaceId);
+    if (promiseEntry) {
+      promiseEntry.reject(new Error(`Workspace ${workspaceId} was deleted`));
+      this.initPromises.delete(workspaceId);
+    }
+  }
+
+  /**
+   * Wait for workspace initialization to complete.
+   * Used by tools (bash, file_*) to ensure files are ready before executing.
+   *
+   * Behavior:
+   * - No init state: Returns immediately (init not needed or backwards compat)
+   * - Init succeeded/failed: Returns immediately (let tool proceed/fail naturally)
+   * - Init running: Waits for completion promise (up to 5 minutes)
+   *
+   * Promise-based approach eliminates race conditions:
+   * - Multiple tools share the same promise (no duplicate listeners)
+   * - No event cleanup needed (promise auto-resolves once)
+   * - Timeout races handled by Promise.race()
+   *
+   * @param workspaceId Workspace ID to wait for
+   * @throws Error if init times out after 5 minutes
+   */
+  async waitForInit(workspaceId: string): Promise<void> {
+    const state = this.getInitState(workspaceId);
+
+    // No init state - proceed immediately (backwards compat or init not needed)
+    if (!state) {
+      return;
+    }
+
+    // Init already completed (success or failure) - proceed immediately
+    if (state.status !== "running") {
+      return;
+    }
+
+    // Init is running - wait for completion promise with timeout
+    const promiseEntry = this.initPromises.get(workspaceId);
+
+    if (!promiseEntry) {
+      // State says running but no promise exists (shouldn't happen, but handle gracefully)
+      log.error(`Init state is running for ${workspaceId} but no promise found, proceeding`);
+      return;
+    }
+
+    const INIT_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes
+    const timeoutPromise = new Promise<void>((_, reject) => {
+      setTimeout(() => {
+        reject(new Error("Workspace initialization timed out after 5 minutes"));
+      }, INIT_TIMEOUT_MS);
+    });
+
+    // Race between completion and timeout
+    // If timeout wins, promise rejection propagates to caller
+    await Promise.race([promiseEntry.promise, timeoutPromise]);
   }
 }

src/services/ipcMain.ts

@@ -134,9 +134,14 @@ export class IpcMain {
     this.config = config;
     this.historyService = new HistoryService(config);
     this.partialService = new PartialService(config, this.historyService);
-    this.aiService = new AIService(config, this.historyService, this.partialService);
-    this.bashService = new BashExecutionService();
     this.initStateManager = new InitStateManager(config);
+    this.aiService = new AIService(
+      config,
+      this.historyService,
+      this.partialService,
+      this.initStateManager
+    );
+    this.bashService = new BashExecutionService();
   }
 
   private getOrCreateSession(workspaceId: string): AgentSession {
@@ -898,6 +903,8 @@ export class IpcMain {
           const bashTool = createBashTool({
             cwd: workspacePath, // Bash executes in the workspace directory
             runtime,
+            workspaceId,
+            initStateManager: this.initStateManager,
             secrets: secretsToRecord(projectSecrets),
             niceness: options?.niceness,
             runtimeTempDir: tempDir.path,