wave-agent-sdk 0.17.9 → 0.17.11

package/src/managers/aiManager.ts

@@ -4,7 +4,11 @@ import { convertMessagesForAPI } from "../utils/convertMessagesForAPI.js";
 import { microcompactMessages } from "../utils/microcompact.js";
 import { parseTaskNotificationXml } from "../utils/notificationXml.js";
 import { calculateComprehensiveTotalTokens } from "../utils/tokenCalculation.js";
-import * as fs from "node:fs/promises";
+import {
+  getTaskReminderTurnCounts,
+  maybeInjectTaskReminder,
+  TASK_REMINDER_CONFIG,
+} from "../utils/taskReminder.js";
 import { existsSync } from "node:fs";
 import type {
   GatewayConfig,
@@ -18,6 +22,7 @@ import type { ToolContext, ToolResult } from "../tools/types.js";
 import type { MessageManager } from "./messageManager.js";
 import type { BackgroundTaskManager } from "./backgroundTaskManager.js";
 import { ChatCompletionMessageFunctionToolCall } from "openai/resources.js";
+import type { ChatCompletionMessageParam } from "openai/resources.js";
 import type { HookManager } from "./hookManager.js";
 import type { ExtendedHookExecutionContext } from "../types/hooks.js";
 import type { PermissionManager } from "./permissionManager.js";
@@ -75,6 +80,11 @@ export class AIManager {
   private originalWorkdir: string;
   private consecutiveCompactionFailures: number = 0;
   private readonly maxTurns?: number;
+  /** Tracks file mtime/hash at read time for staleness detection on Edit/Write */
+  private readFileState = new Map<
+    string,
+    { mtime: number; hash: string; offset?: number; limit?: number }
+  >();
   /** Override tool_choice for this AI manager (e.g. for structured output) */
   public toolChoiceOverride?:
     | "auto"
@@ -238,32 +248,29 @@ export class AIManager {
   }
 
   /**
-   * Build plan mode system-reminder messages to inject into the API message stream.
-   * These are transient messages not stored in the message history.
-   * This preserves prompt caching by keeping the system prompt constant.
+   * Add plan mode reminder as a persistent meta message to the message manager.
+   * Called before getMessages() so the stored message is included in the API call.
    */
-  private buildPlanModeMessages(
+  private maybeAddPlanModeMessage(
     currentMode: PermissionMode | undefined,
-  ): import("openai/resources.js").ChatCompletionMessageParam[] {
-    const messages: import("openai/resources.js").ChatCompletionMessageParam[] =
-      [];
-    if (!this.permissionManager) return messages;
+  ): void {
+    if (!this.permissionManager) return;
 
     // Handle exit notification (one-time after leaving plan mode)
     if (this.permissionManager.getNeedsPlanModeExitAttachment()) {
       const planFilePath = this.permissionManager.getPlanFilePath();
       const planExists = planFilePath ? existsSync(planFilePath) : false;
-      messages.push({
-        role: "user",
+      this.messageManager.addUserMessage({
         content: buildExitedPlanModeReminder(planFilePath, planExists),
+        isMeta: true,
       });
       this.permissionManager.setNeedsPlanModeExitAttachment(false);
     }
 
-    if (currentMode !== "plan") return messages;
+    if (currentMode !== "plan") return;
 
     const planFilePath = this.permissionManager.getPlanFilePath();
-    if (!planFilePath) return messages;
+    if (!planFilePath) return;
 
     const planExists = existsSync(planFilePath);
 
@@ -272,25 +279,46 @@ export class AIManager {
     if (planMgr?.isPlanEntryReminderPending()) {
       if (this.permissionManager.hasExitedPlanModeInSession() && planExists) {
         // Re-entry: use small reminder
-        messages.push({
-          role: "user",
+        this.messageManager.addUserMessage({
           content: buildPlanModeReEntryReminder(planFilePath),
+          isMeta: true,
         });
       } else {
         // First entry: use full reminder
-        messages.push({
-          role: "user",
+        this.messageManager.addUserMessage({
           content: buildPlanModeReminder(
             planFilePath,
             planExists,
             !!this.subagentType,
           ),
+          isMeta: true,
         });
       }
       planMgr.consumePlanEntryReminder();
     }
+  }
+
+  private async maybeInjectTaskReminderIntoMessages(
+    messages: ChatCompletionMessageParam[],
+    toolNames: Set<string>,
+  ): Promise<void> {
+    // Guard: no task tools available
+    if (!toolNames.has("TaskUpdate")) return;
+
+    const internalMessages = this.messageManager.getMessages();
+    const turnCounts = getTaskReminderTurnCounts(internalMessages);
+
+    if (
+      turnCounts.turnsSinceLastTaskManagement <
+        TASK_REMINDER_CONFIG.TURNS_SINCE_WRITE ||
+      turnCounts.turnsSinceLastReminder <
+        TASK_REMINDER_CONFIG.TURNS_BETWEEN_REMINDERS
+    ) {
+      return;
+    }
 
-    return messages;
+    const tasks = await this.taskManager.listTasks();
+    maybeInjectTaskReminder(messages, turnCounts, tasks);
   }
 
   public setIsLoading(isLoading: boolean): void {
@@ -470,6 +498,25 @@ export class AIManager {
         compactUsage,
       );
 
+      // Re-add plan mode reminder as persistent meta message after compaction
+      const postCompactMode = this.permissionManager?.getCurrentEffectiveMode(
+        this.getModelConfig().permissionMode,
+      );
+      if (postCompactMode === "plan") {
+        const planFilePath = this.permissionManager?.getPlanFilePath();
+        if (planFilePath) {
+          const planExists = existsSync(planFilePath);
+          this.messageManager.addUserMessage({
+            content: buildPlanModeReminder(
+              planFilePath,
+              planExists,
+              !!this.subagentType,
+            ),
+            isMeta: true,
+          });
+        }
+      }
+
       // 8. Track usage
       if (compactUsage && this.callbacks?.onUsageAdded) {
         this.callbacks.onUsageAdded(compactUsage);
@@ -569,26 +616,6 @@ export class AIManager {
       if (usedTokens >= POST_COMPACT_TOKEN_BUDGET) break;
     }
 
-    // 2. Plan mode context
-    const currentMode = this.permissionManager?.getCurrentEffectiveMode(
-      this.getModelConfig().permissionMode,
-    );
-    if (currentMode === "plan") {
-      const planFilePath = this.permissionManager?.getPlanFilePath();
-      if (planFilePath) {
-        let planExists = false;
-        try {
-          await fs.access(planFilePath);
-          planExists = true;
-        } catch {
-          // Plan file doesn't exist yet
-        }
-        contextParts.push(
-          `\n\n${buildPlanModeReminder(planFilePath, planExists, !!this.subagentType)}`,
-        );
-      }
-    }
-
     // 4. Invoked skills context (with token budget, matching Claude Code)
     const POST_COMPACT_SKILLS_TOKEN_BUDGET = 25_000;
     const POST_COMPACT_MAX_TOKENS_PER_SKILL = 5_000;
@@ -784,6 +811,14 @@ export class AIManager {
       toolAbortController = this.toolAbortController!;
     }
 
+    // Get current permission mode
+    const currentMode = this.permissionManager?.getCurrentEffectiveMode(
+      this.getModelConfig().permissionMode,
+    );
+
+    // Add plan mode reminder as persistent meta message before getting messages
+    this.maybeAddPlanModeMessage(currentMode);
+
     // Get recent message history with microcompact applied
     const rawMessages = this.messageManager.getMessages();
     const microcompactedMessages = microcompactMessages(rawMessages, {
@@ -801,23 +836,12 @@ export class AIManager {
 
       logger?.debug("modelConfig in sendAIMessage", this.getModelConfig());
 
-      // Get current permission mode and plan file path
-      const currentMode = this.permissionManager?.getCurrentEffectiveMode(
-        this.getModelConfig().permissionMode,
-      );
       const toolsConfig = this.getFilteredToolsConfig();
       const toolNames = new Set(toolsConfig.map((t) => t.function.name));
       const filteredToolPlugins = this.toolManager
         .getTools()
         .filter((t) => toolNames.has(t.name));
 
-      // Inject plan mode system-reminder messages (not system prompt)
-      // This preserves prompt caching by keeping the system prompt constant
-      const planModeMessages = this.buildPlanModeMessages(currentMode);
-      if (planModeMessages.length > 0) {
-        recentMessages.push(...planModeMessages);
-      }
-
       let autoMemoryOptions: { directory: string; content: string } | undefined;
 
       if (this.getAutoMemoryEnabled()) {
@@ -857,6 +881,12 @@ export class AIManager {
         toolChoice: this.toolChoiceOverride, // Pass tool_choice override
       };
 
+      // Inject task reminder if needed (not persisted, regenerated each turn)
+      await this.maybeInjectTaskReminderIntoMessages(
+        callAgentOptions.messages,
+        toolNames,
+      );
+
       // Add streaming callbacks only if streaming is enabled
       if (this.stream) {
         callAgentOptions.onContentUpdate = (content: string) => {
@@ -1001,218 +1031,52 @@ export class AIManager {
       }
 
       if (toolCalls.length > 0) {
-        // Execute all tools in parallel using Promise.all
-        const toolExecutionPromises = toolCalls.map(
-          async (functionToolCall) => {
-            const toolId = functionToolCall.id || "";
-
-            // Check if already interrupted, skip tool execution if so
-            if (
-              abortController.signal.aborted ||
-              toolAbortController.signal.aborted
-            ) {
-              return;
-            }
-
-            const toolName = functionToolCall.function?.name || "";
-            // Safely parse tool parameters, handle tools without parameters
-            let toolArgs: Record<string, unknown> = {};
-            let jsonRecovered = false;
-            const argsString = functionToolCall.function?.arguments?.trim();
-
-            if (!argsString || argsString === "") {
-              // Tool without parameters, use empty object
-              toolArgs = {};
-            } else {
-              let recoveredArgs = argsString;
-              try {
-                toolArgs = JSON.parse(argsString);
-              } catch {
-                // Attempt to recover truncated JSON (e.g., missing closing braces)
-                recoveredArgs = recoverTruncatedJson(argsString);
-                try {
-                  toolArgs = JSON.parse(recoveredArgs);
-                  jsonRecovered = true;
-                  logger.warn(
-                    `Recovered truncated JSON for tool "${toolName}"`,
-                  );
-                } catch (parseError) {
-                  let errorMessage = `Failed to parse tool arguments`;
-                  if (result.finish_reason === "length") {
-                    errorMessage +=
-                      " (output truncated, please reduce your output)";
-                  }
-                  logger?.error(errorMessage, parseError);
-                  this.messageManager.updateToolBlock({
-                    id: toolId,
-                    parameters: argsString,
-                    result: errorMessage,
-                    success: false,
-                    error: errorMessage,
-                    stage: "end",
-                    name: toolName,
-                    compactParams: "",
-                    timestamp: Date.now(),
-                  });
-                  return;
-                }
-              }
-            }
+        // Partition tool calls into batches: consecutive concurrency-safe
+        // tools run in parallel, non-safe tools (Edit, Write, MCP) run one
+        // at a time to prevent read-modify-write races on the same file.
+        const batches: {
+          calls: ChatCompletionMessageFunctionToolCall[];
+          safe: boolean;
+        }[] = [];
+        for (const call of toolCalls) {
+          const toolName = call.function?.name || "";
+          const safe = this.toolManager.isConcurrencySafe(toolName);
+          const lastBatch = batches[batches.length - 1];
+          if (lastBatch && lastBatch.safe && safe) {
+            lastBatch.calls.push(call);
+          } else {
+            batches.push({ calls: [call], safe });
+          }
+        }
 
-            const compactParams = this.generateCompactParams(
-              toolName,
-              toolArgs,
+        // Execute batches sequentially; within a safe batch, tools run in parallel
+        for (const batch of batches) {
+          if (
+            abortController.signal.aborted ||
+            toolAbortController.signal.aborted
+          ) {
+            break;
+          }
+          if (batch.calls.length === 1) {
+            await this.executeToolCall(
+              batch.calls[0],
+              abortController,
+              toolAbortController,
+              result.finish_reason,
             );
-
-            // Emit start stage for non-streaming tool calls
-            if (!this.stream) {
-              this.messageManager.updateToolBlock({
-                id: toolId,
-                stage: "start",
-                name: toolName,
-                compactParams,
-                parameters: argsString,
-              });
-            }
-
-            // Emit running stage (tool execution about to start)
-            this.messageManager.updateToolBlock({
-              id: toolId,
-              stage: "running",
-              name: toolName,
-              compactParams,
-              parameters: argsString,
-              parametersChunk: "",
-            });
-
-            try {
-              // Execute PreToolUse hooks before tool execution
-              const shouldExecuteTool = await this.executePreToolUseHooks(
-                toolName,
-                toolArgs,
-                toolId,
-              );
-
-              // If PreToolUse hooks blocked execution, skip tool execution
-              if (!shouldExecuteTool) {
-                logger?.info(
-                  `Tool ${toolName} execution blocked by PreToolUse hooks`,
-                );
-                return; // Skip this tool and return from this map function
-              }
-
-              // Create tool execution context
-              const context: ToolContext = {
-                abortSignal: toolAbortController.signal,
-                backgroundTaskManager: this.backgroundTaskManager,
-                workdir: this.getWorkdir(),
-                originalWorkdir: this.originalWorkdir,
-                env: this.container.get<Record<string, string>>("MergedEnv"),
-                messageId: this.messageManager.getMessages().slice(-1)[0]?.id,
-                sessionId: this.messageManager.getSessionId(),
-                toolCallId: toolId,
-                taskManager: this.taskManager,
-                onShortResultUpdate: (shortResult: string) => {
-                  this.messageManager.updateToolBlock({
-                    id: toolId,
-                    shortResult,
-                    stage: "running", // Keep it in running stage while updating shortResult
-                  });
-                },
-                onResultUpdate: (result: string) => {
-                  this.messageManager.updateToolBlock({
-                    id: toolId,
-                    result,
-                    stage: "running", // Keep it in running stage while updating result
-                  });
-                },
-                onCwdChange: async (newCwd: string) => {
-                  const oldCwd = this.getWorkdir();
-                  this.container.register("Workdir", newCwd);
-                  this._onCwdChange?.(newCwd);
-                  if (this.hookManager) {
-                    const sessionId = this.messageManager.getSessionId();
-                    const transcriptPath =
-                      this.messageManager.getTranscriptPath();
-                    const env = Object.fromEntries(
-                      Object.entries(process.env).filter(
-                        (e) => e[1] !== undefined,
-                      ),
-                    ) as Record<string, string>;
-                    await this.hookManager.executeCwdChangedHooks(
-                      oldCwd,
-                      newCwd,
-                      sessionId,
-                      transcriptPath,
-                      env,
-                    );
-                  }
-                },
-              };
-
-              // Execute tool
-              const toolResult = await this.toolManager.execute(
-                functionToolCall.function?.name || "",
-                toolArgs,
-                context,
-              );
-
-              // Build result content, adding truncation warning if JSON was recovered
-              let toolResultContent =
-                toolResult.content ||
-                (toolResult.error ? `Error: ${toolResult.error}` : "");
-              if (jsonRecovered) {
-                toolResultContent +=
-                  "\n\nTool arguments were truncated (likely exceeded max output tokens). Please reduce your output or split into multiple tool calls.";
-              }
-
-              // Update message state - tool execution completed
-              this.messageManager.updateToolBlock({
-                id: toolId,
-                parameters: argsString,
-                result: toolResultContent,
-                success: toolResult.success,
-                error: toolResult.error,
-                stage: "end",
-                name: toolName,
-                compactParams,
-                shortResult: toolResult.shortResult,
-                isManuallyBackgrounded: toolResult.isManuallyBackgrounded,
-                startLineNumber: toolResult.startLineNumber,
-                timestamp: Date.now(),
-              });
-
-              // Execute PostToolUse hooks after successful tool completion
-              await this.executePostToolUseHooks(
-                toolId,
-                toolName,
-                toolArgs,
-                toolResult,
-              );
-            } catch (toolError) {
-              const errorMessage =
-                toolError instanceof Error
-                  ? toolError.message
-                  : String(toolError);
-
-              this.messageManager.updateToolBlock({
-                id: toolId,
-                parameters: JSON.stringify(toolArgs, null, 2),
-                result: `Tool execution failed: ${errorMessage}`,
-                success: false,
-                error: errorMessage,
-                stage: "end",
-                name: toolName,
-                compactParams,
-                isManuallyBackgrounded: false,
-                timestamp: Date.now(),
-              });
-            }
-          },
-        );
-
-        // Wait for all tools to complete execution in parallel
-        await Promise.all(toolExecutionPromises);
+          } else {
+            await Promise.all(
+              batch.calls.map((call) =>
+                this.executeToolCall(
+                  call,
+                  abortController,
+                  toolAbortController,
+                  result.finish_reason,
+                ),
+              ),
+            );
+          }
+        }
       }
 
       // Handle token statistics and message compaction
@@ -1584,6 +1448,206 @@ export class AIManager {
     }
   }
 
+  /**
+   * Execute a single tool call: arg parsing, block emission, PreToolUse hooks,
+   * tool execution, PostToolUse hooks, and error handling.
+   */
+  private async executeToolCall(
+    functionToolCall: ChatCompletionMessageFunctionToolCall,
+    abortController: AbortController,
+    toolAbortController: AbortController,
+    finishReason?: string | null,
+  ): Promise<void> {
+    const toolId = functionToolCall.id || "";
+
+    // Check if already interrupted, skip tool execution if so
+    if (abortController.signal.aborted || toolAbortController.signal.aborted) {
+      return;
+    }
+
+    const toolName = functionToolCall.function?.name || "";
+    // Safely parse tool parameters, handle tools without parameters
+    let toolArgs: Record<string, unknown> = {};
+    let jsonRecovered = false;
+    const argsString = functionToolCall.function?.arguments?.trim();
+
+    if (!argsString || argsString === "") {
+      // Tool without parameters, use empty object
+      toolArgs = {};
+    } else {
+      let recoveredArgs = argsString;
+      try {
+        toolArgs = JSON.parse(argsString);
+      } catch {
+        // Attempt to recover truncated JSON (e.g., missing closing braces)
+        recoveredArgs = recoverTruncatedJson(argsString);
+        try {
+          toolArgs = JSON.parse(recoveredArgs);
+          jsonRecovered = true;
+          logger.warn(`Recovered truncated JSON for tool "${toolName}"`);
+        } catch (parseError) {
+          let errorMessage = `Failed to parse tool arguments`;
+          if (finishReason === "length") {
+            errorMessage += " (output truncated, please reduce your output)";
+          }
+          logger?.error(errorMessage, parseError);
+          this.messageManager.updateToolBlock({
+            id: toolId,
+            parameters: argsString,
+            result: errorMessage,
+            success: false,
+            error: errorMessage,
+            stage: "end",
+            name: toolName,
+            compactParams: "",
+            timestamp: Date.now(),
+          });
+          return;
+        }
+      }
+    }
+
+    const compactParams = this.generateCompactParams(toolName, toolArgs);
+
+    // Emit start stage for non-streaming tool calls
+    if (!this.stream) {
+      this.messageManager.updateToolBlock({
+        id: toolId,
+        stage: "start",
+        name: toolName,
+        compactParams,
+        parameters: argsString,
+      });
+    }
+
+    // Emit running stage (tool execution about to start)
+    this.messageManager.updateToolBlock({
+      id: toolId,
+      stage: "running",
+      name: toolName,
+      compactParams,
+      parameters: argsString,
+      parametersChunk: "",
+    });
+
+    try {
+      // Execute PreToolUse hooks before tool execution
+      const shouldExecuteTool = await this.executePreToolUseHooks(
+        toolName,
+        toolArgs,
+        toolId,
+      );
+
+      // If PreToolUse hooks blocked execution, skip tool execution
+      if (!shouldExecuteTool) {
+        logger?.info(`Tool ${toolName} execution blocked by PreToolUse hooks`);
+        return;
+      }
+
+      // Create tool execution context
+      const context: ToolContext = {
+        abortSignal: toolAbortController.signal,
+        backgroundTaskManager: this.backgroundTaskManager,
+        workdir: this.getWorkdir(),
+        originalWorkdir: this.originalWorkdir,
+        messageId: this.messageManager.getMessages().slice(-1)[0]?.id,
+        sessionId: this.messageManager.getSessionId(),
+        toolCallId: toolId,
+        taskManager: this.taskManager,
+        readFileState: this.readFileState,
+        onShortResultUpdate: (shortResult: string) => {
+          this.messageManager.updateToolBlock({
+            id: toolId,
+            shortResult,
+            stage: "running",
+          });
+        },
+        onResultUpdate: (result: string) => {
+          this.messageManager.updateToolBlock({
+            id: toolId,
+            result,
+            stage: "running",
+          });
+        },
+        onCwdChange: async (newCwd: string) => {
+          const oldCwd = this.getWorkdir();
+          this.container.register("Workdir", newCwd);
+          this._onCwdChange?.(newCwd);
+          if (this.hookManager) {
+            const sessionId = this.messageManager.getSessionId();
+            const transcriptPath = this.messageManager.getTranscriptPath();
+            const env = Object.fromEntries(
+              Object.entries(process.env).filter((e) => e[1] !== undefined),
+            ) as Record<string, string>;
+            await this.hookManager.executeCwdChangedHooks(
+              oldCwd,
+              newCwd,
+              sessionId,
+              transcriptPath,
+              env,
+            );
+          }
+        },
+      };
+
+      // Execute tool
+      const toolResult = await this.toolManager.execute(
+        functionToolCall.function?.name || "",
+        toolArgs,
+        context,
+      );
+
+      // Build result content, adding truncation warning if JSON was recovered
+      let toolResultContent =
+        toolResult.content ||
+        (toolResult.error ? `Error: ${toolResult.error}` : "");
+      if (jsonRecovered) {
+        toolResultContent +=
+          "\n\nTool arguments were truncated (likely exceeded max output tokens). Please reduce your output or split into multiple tool calls.";
+      }
+
+      // Update message state - tool execution completed
+      this.messageManager.updateToolBlock({
+        id: toolId,
+        parameters: argsString,
+        result: toolResultContent,
+        success: toolResult.success,
+        error: toolResult.error,
+        stage: "end",
+        name: toolName,
+        compactParams,
+        shortResult: toolResult.shortResult,
+        isManuallyBackgrounded: toolResult.isManuallyBackgrounded,
+        startLineNumber: toolResult.startLineNumber,
+        timestamp: Date.now(),
+      });
+
+      // Execute PostToolUse hooks after successful tool completion
+      await this.executePostToolUseHooks(
+        toolId,
+        toolName,
+        toolArgs,
+        toolResult,
+      );
+    } catch (toolError) {
+      const errorMessage =
+        toolError instanceof Error ? toolError.message : String(toolError);
+
+      this.messageManager.updateToolBlock({
+        id: toolId,
+        parameters: JSON.stringify(toolArgs, null, 2),
+        result: `Tool execution failed: ${errorMessage}`,
+        success: false,
+        error: errorMessage,
+        stage: "end",
+        name: toolName,
+        compactParams,
+        isManuallyBackgrounded: false,
+        timestamp: Date.now(),
+      });
+    }
+  }
+
   /**
    * Execute PreToolUse hooks before tool execution
    * Returns true if hooks allow tool execution, false if blocked
@@ -1682,9 +1746,9 @@ export class AIManager {
         toolResponse,
         subagentType: this.subagentType, // Include subagent type in hook context
         planFilePath: this.permissionManager?.getPlanFilePath(),
-        env:
-          this.container.get<Record<string, string>>("MergedEnv") ||
-          (process.env as Record<string, string>),
+        env: Object.fromEntries(
+          Object.entries(process.env).filter((e) => e[1] !== undefined),
+        ) as Record<string, string>, // Include environment variables
       };
 
       const results = await this.hookManager.executeHooks(