wave-agent-sdk 1.0.0 → 1.0.2

package/src/managers/aiManager.ts

@@ -62,6 +62,22 @@ import {
 import { logOTelEvent } from "../telemetry/events.js";
 import type { BackgroundTask } from "../types/processes.js";
 
+/** Result of a fork-path agent loop (compaction or auto-memory extraction). */
+interface ForkLoopResult {
+  content?: string;
+  usage?: {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+  };
+}
+
+/** Max turns for the compaction fork: the model should summarize, not act. */
+const MAX_FORK_TURNS = 3;
+
+/** Max turns for the auto-memory extraction fork. */
+const MAX_AUTO_MEMORY_FORK_TURNS = 5;
+
 // Truncate text to `max` chars and append a "… [+N chars]" marker when exceeded.
 // Used for background_tasks description/command fields (≤1000 chars per spec FR-063).
 function truncateWithMarker(text: string, max: number): string {
@@ -694,8 +710,7 @@ export class AIManager {
     } catch (compactError) {
       this.consecutiveCompactionFailures++;
       logger?.error(
-        `Failed to compact messages (${this.consecutiveCompactionFailures} consecutive):`,
-        compactError,
+        `Failed to compact messages (${this.consecutiveCompactionFailures} consecutive): ${compactError instanceof Error ? compactError.message : String(compactError)}`,
       );
       this.messageManager.addErrorBlock(
         `Failed to compact conversation history: ${compactError instanceof Error ? compactError.message : String(compactError)}. You may encounter context limit issues.`,
@@ -745,27 +760,26 @@ export class AIManager {
   }
 
   /**
-   * Fork-path compaction: run a bounded agent loop over a copy of the
-   * conversation using the same system prompt, tools, model, and generation
-   * params as the main loop, so the forked request prefix matches exactly
-   * and the prompt cache is reused. Tool calls are denied locally (the model
-   * is told to summarize, not act) and their rejections are fed back for
-   * another turn. Returns undefined content when the model never produces
-   * text; the caller treats that as a compaction failure.
+   * Fork-path loop: run a bounded agent loop over a copy of the conversation
+   * using the same system prompt, tools, model, and generation params as the
+   * main loop, so the forked request prefix matches exactly and the prompt
+   * cache is reused. A `canUseTool` gate decides whether each tool call
+   * executes locally (with a stripped context) or is denied and fed back to
+   * the model for another turn. Returns undefined content when the model never
+   * produces text; the caller treats that as a failure.
    */
-  private async runCompactFork(
+  private async runForkLoop(
     historyMessages: ChatCompletionMessageParam[],
-    compactPrompt: string,
+    prompt: string,
+    options: {
+      maxTurns: number;
+      /** Gate deciding which tool calls execute locally. When undefined, every tool call is denied. */
+      canUseTool?: (name: string, args: Record<string, unknown>) => boolean;
+      /** Message fed back to the model when a tool call is denied. */
+      deniedToolMessage?: string;
+    },
     abortSignal?: AbortSignal,
-  ): Promise<{
-    content?: string;
-    usage?: {
-      prompt_tokens: number;
-      completion_tokens: number;
-      total_tokens: number;
-    };
-  }> {
-    const MAX_FORK_TURNS = 3;
+  ): Promise<ForkLoopResult> {
     const modelConfig = this.getModelConfig();
     const gatewayConfig = this.getGatewayConfig();
     const sessionId = this.messageManager.getSessionId();
@@ -783,21 +797,22 @@ export class AIManager {
       });
     }
 
-    forkMessages.push({ role: "user", content: compactPrompt });
+    forkMessages.push({ role: "user", content: prompt });
 
     const { toolsConfig, filteredToolPlugins } = this.resolveFilteredTools();
     const systemPrompt = await this.buildMainSystemPrompt(filteredToolPlugins);
 
-    let totalUsage:
-      | {
-          prompt_tokens: number;
-          completion_tokens: number;
-          total_tokens: number;
-        }
-      | undefined;
+    // Fresh read-state map so Read/Edit state built up inside the fork never
+    // leaks into the main session's dedup and staleness tracking.
+    const forkReadFileState = new Map<
+      string,
+      { mtime: number; hash: string; offset?: number; limit?: number }
+    >();
+
+    let totalUsage: ForkLoopResult["usage"];
     let content: string | undefined;
 
-    for (let turn = 0; turn < MAX_FORK_TURNS; turn++) {
+    for (let turn = 0; turn < options.maxTurns; turn++) {
       const result = await aiService.callAgent({
         gatewayConfig,
         modelConfig,
@@ -810,7 +825,7 @@ export class AIManager {
         toolChoice: this.toolChoiceOverride,
         // Stream so a slow reasoning model emits first bytes before the
         // gateway's idle timeout fires (non-streaming waits for the full
-        // summary, which exceeds the timeout on large contexts).
+        // response, which exceeds the timeout on large contexts).
         stream: true,
       });
 
@@ -832,31 +847,329 @@ export class AIManager {
       }
 
       if (result.tool_calls && result.tool_calls.length > 0) {
-        // Deny all tool calls locally and feed the rejections back so the
-        // model gets another turn to produce the summary text.
+        const functionCalls = result.tool_calls.filter(
+          (tc) => tc.type === "function",
+        );
+        if (functionCalls.length === 0) break;
+
         forkMessages.push({
           role: "assistant",
           content: result.content ?? null,
-          tool_calls: result.tool_calls,
+          tool_calls: functionCalls,
         });
-        for (const toolCall of result.tool_calls) {
+
+        for (const toolCall of functionCalls) {
+          const name = toolCall.function?.name || "";
+          const args = this.parseForkToolArgs(toolCall.function?.arguments);
+          let toolContent: string;
+          if (options.canUseTool && options.canUseTool(name, args)) {
+            toolContent = await this.executeForkTool(
+              name,
+              args,
+              workdir,
+              sessionId,
+              abortSignal,
+              forkReadFileState,
+            );
+          } else {
+            toolContent =
+              options.deniedToolMessage ??
+              "Tool use is not allowed in this context";
+          }
           forkMessages.push({
             role: "tool",
             tool_call_id: toolCall.id,
-            content: "Tool use is not allowed during compaction",
+            content: toolContent,
           });
         }
         continue;
       }
 
       // Neither text nor tool calls: retrying the identical request is
-      // pointless, bail out and let the caller fail the compaction.
+      // pointless, bail out and let the caller fail.
       break;
     }
 
     return { content, usage: totalUsage };
   }
 
+  /**
+   * Fork-path compaction: deny all tool calls locally (the model is told to
+   * summarize, not act) and feed the rejections back for another turn.
+   */
+  private async runCompactFork(
+    historyMessages: ChatCompletionMessageParam[],
+    compactPrompt: string,
+    abortSignal?: AbortSignal,
+  ): Promise<ForkLoopResult> {
+    return this.runForkLoop(
+      historyMessages,
+      compactPrompt,
+      {
+        maxTurns: MAX_FORK_TURNS,
+        deniedToolMessage: "Tool use is not allowed during compaction",
+      },
+      abortSignal,
+    );
+  }
+
+  /**
+   * Auto-memory extraction via the perfect fork: the extraction prompt is run
+   * against the same request prefix as the main conversation (same system
+   * prompt, tools, model, and message history) so the prompt cache is reused.
+   * Gate-approved tools execute locally in a stripped context; everything else
+   * is denied. Usage is reported with operation_type "agent" so extraction
+   * token costs stay visible in session accounting.
+   */
+  public async runAutoMemoryFork(
+    messages: Message[],
+    prompt: string,
+    options: {
+      canUseTool: (name: string, args: Record<string, unknown>) => boolean;
+      deniedToolMessage?: string;
+      maxTurns?: number;
+    },
+    abortSignal?: AbortSignal,
+  ): Promise<ForkLoopResult> {
+    const modelConfig = this.getModelConfig();
+    const historyMessages = convertMessagesForAPI(messages, {
+      supportsVision: supportsVision(modelConfig.capabilities),
+    });
+    // Give the fork a real signal even when the caller has none, so tools
+    // (e.g. Bash's foreground path) always receive a well-formed context.
+    const signal = abortSignal ?? new AbortController().signal;
+    const result = await this.runForkLoop(
+      historyMessages,
+      prompt,
+      {
+        maxTurns: options.maxTurns ?? MAX_AUTO_MEMORY_FORK_TURNS,
+        canUseTool: options.canUseTool,
+        deniedToolMessage: options.deniedToolMessage,
+      },
+      signal,
+    );
+
+    if (result.usage && this.callbacks?.onUsageAdded) {
+      this.callbacks.onUsageAdded({
+        ...result.usage,
+        model: modelConfig.model,
+        operation_type: "agent",
+      });
+    }
+    return result;
+  }
+
+  /**
+   * Parse a fork tool call's JSON arguments, recovering truncated JSON the
+   * same way the main loop does. Unparseable arguments fall back to `{}` and
+   * are rejected by the gate or the tool's own parameter validation.
+   */
+  private parseForkToolArgs(
+    argsString: string | undefined,
+  ): Record<string, unknown> {
+    if (!argsString?.trim()) return {};
+    try {
+      return JSON.parse(argsString) as Record<string, unknown>;
+    } catch {
+      try {
+        return JSON.parse(recoverTruncatedJson(argsString)) as Record<
+          string,
+          unknown
+        >;
+      } catch {
+        return {};
+      }
+    }
+  }
+
+  /**
+   * Execute a single tool call inside a fork with a stripped context: no
+   * permission manager (never prompts the user), no message manager (no
+   * conditional-rule triggering), no messageId (no file-history snapshots),
+   * and no background task manager (commands run in the foreground). Only
+   * gate-approved tool names reach this path.
+   */
+  private async executeForkTool(
+    name: string,
+    args: Record<string, unknown>,
+    workdir: string,
+    sessionId: string | undefined,
+    abortSignal: AbortSignal | undefined,
+    readFileState: ToolContext["readFileState"],
+  ): Promise<string> {
+    const plugin = this.toolManager.getTools().find((t) => t.name === name);
+    if (!plugin) {
+      return `Tool '${name}' not found`;
+    }
+    const context: ToolContext = {
+      abortSignal,
+      workdir,
+      originalWorkdir: this.originalWorkdir,
+      sessionId,
+      taskManager: this.taskManager,
+      readFileState,
+      onShortResultUpdate: () => {},
+      onResultUpdate: () => {},
+      onCwdChange: () => {},
+    };
+    try {
+      const result = await plugin.execute(args, context);
+      if (result.content) return result.content;
+      if (result.error) return `Error: ${result.error}`;
+      return "";
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      logger?.error(`Fork tool execution failed for ${name}:`, error);
+      return `Tool execution failed: ${message}`;
+    }
+  }
+
+  /**
+   * Fork-path side question ("/btw"): run a single-turn fork of the
+   * conversation using the same system prompt, tools, model, and generation
+   * params as the main loop, so the forked request prefix matches exactly and
+   * the prompt cache is reused. The in-progress assistant message (if any) is
+   * stripped so the fork starts from the last completed request prefix. Tools
+   * are never executed — the wrapped question instructs the model to answer
+   * directly; an attempted tool call is surfaced as an error string.
+   */
+  async runBtwFork(
+    question: string,
+    abortSignal?: AbortSignal,
+    onContent?: (content: string) => void,
+    onReasoning?: (content: string) => void,
+  ): Promise<{ content?: string; error?: string }> {
+    const modelConfig = this.getModelConfig();
+    const gatewayConfig = this.getGatewayConfig();
+    const sessionId = this.messageManager.getSessionId();
+    const workdir = this.getWorkdir();
+
+    const rawMessages = this.messageManager.getMessages();
+
+    // Strip the in-progress assistant message (a block still in "streaming"
+    // stage) so the fork's request prefix matches the last completed
+    // main-loop request and the prompt cache is reused.
+    const lastMessage = rawMessages[rawMessages.length - 1];
+    const hasInProgressMessage =
+      lastMessage?.role === "assistant" &&
+      lastMessage.blocks.some(
+        (b) => "stage" in b && (b as { stage?: string }).stage === "streaming",
+      );
+
+    const forkMessages: ChatCompletionMessageParam[] = convertMessagesForAPI(
+      hasInProgressMessage ? rawMessages.slice(0, -1) : rawMessages,
+      { supportsVision: supportsVision(modelConfig.capabilities) },
+    );
+
+    // Mirror the main loop's memory injection so the request prefix matches.
+    const { prependContent } =
+      await this.messageManager.getMemoryForInjection();
+    if (prependContent.trim()) {
+      forkMessages.unshift({
+        role: "user",
+        content: wrapInSystemReminder(prependContent),
+      });
+    }
+
+    // Wrap the question with the side-question instructions (verbatim
+    // Claude Code sideQuestion.ts) so the model answers directly.
+    const wrappedQuestion = `<system-reminder>This is a side question from the user. You must answer this question directly in a single response.
+
+IMPORTANT CONTEXT:
+- You are a separate, lightweight agent spawned to answer this one question
+- The main agent is NOT interrupted - it continues working independently in the background
+- You share the conversation context but are a completely separate instance
+- Do NOT reference being interrupted or what you were "previously doing" - that framing is incorrect
+
+CRITICAL CONSTRAINTS:
+- You have NO tools available - you cannot read files, run commands, search, or take any actions
+- This is a one-off response - there will be no follow-up turns
+- You can ONLY provide information based on what you already know from the conversation context
+- NEVER say things like "Let me try...", "I'll now...", "Let me check...", or promise to take any action
+- If you don't know the answer, say so - do not offer to look it up or investigate
+
+Simply answer the question with the information you have.</system-reminder>
+
+${question}`;
+    forkMessages.push({ role: "user", content: wrappedQuestion });
+
+    const { toolsConfig, filteredToolPlugins } = this.resolveFilteredTools();
+    const systemPrompt = await this.buildMainSystemPrompt(filteredToolPlugins);
+
+    try {
+      // Surface partial output to the caller (e.g. the /btw overlay's
+      // streaming display) as it arrives. Reasoning chunks from thinking
+      // models stream through a separate channel when the caller supplies
+      // one (webview hosts distinguish thinking from content so the panel
+      // can drop thinking text once content starts); otherwise they fall
+      // back to the content channel (CLI overlay mixes both).
+      const streamToOverlay = (text: string) => {
+        if (text.trim()) {
+          onContent?.(text);
+        }
+      };
+      const streamReasoning = (text: string) => {
+        if (text.trim()) {
+          if (onReasoning) {
+            onReasoning(text);
+          } else {
+            onContent?.(text);
+          }
+        }
+      };
+      const result = await aiService.callAgent({
+        gatewayConfig,
+        modelConfig,
+        messages: forkMessages,
+        sessionId,
+        abortSignal,
+        workdir,
+        tools: toolsConfig,
+        systemPrompt,
+        toolChoice: this.toolChoiceOverride,
+        // Stream so a slow reasoning model emits first bytes before the
+        // gateway's idle timeout fires (same rationale as runCompactFork).
+        stream: true,
+        onContentUpdate: streamToOverlay,
+        onReasoningUpdate: streamReasoning,
+      });
+
+      if (result.content?.trim()) {
+        return { content: result.content };
+      }
+
+      // A thinking model may emit only reasoning content (e.g. the stream
+      // is truncated before the final answer); surface that instead of
+      // falling through to "No response received".
+      if (result.reasoning_content?.trim()) {
+        return { content: result.reasoning_content };
+      }
+
+      if (result.tool_calls && result.tool_calls.length > 0) {
+        const firstFunctionCall = result.tool_calls.find(
+          (call) => call.type === "function",
+        );
+        const toolName = firstFunctionCall?.function?.name ?? "a tool";
+        return {
+          error: `(The model tried to call ${toolName} instead of answering directly. Try rephrasing or ask in the main conversation.)`,
+        };
+      }
+
+      // Neither text nor tool calls.
+      return { error: "No response received" };
+    } catch (error) {
+      // aiService.callAgent converts AbortError into a plain Error, so the
+      // abort is detected via the signal itself; rethrow so the UI can
+      // silently dismiss instead of showing an error.
+      if (abortSignal?.aborted) {
+        throw error;
+      }
+      return {
+        error: `(API error: ${error instanceof Error ? error.message : String(error)})`,
+      };
+    }
+  }
+
   /**
    * Build post-compact context restoration content.
    * Restores file reads, working directory, plan mode, skills, and background tasks.
@@ -1150,6 +1463,12 @@ export class AIManager {
             systemPrompt: mainSystemPrompt, // Pass custom system prompt
             maxTokens: maxTokens, // Pass max tokens override
             toolChoice: this.toolChoiceOverride, // Pass tool_choice override
+            // Fast-model subagents send disable-thinking params only when
+            // explicitly configured (never in the agent loop).
+            disableThinkingOptions:
+              this.modelOverride === "fastModel"
+                ? this.getModelConfig().disableThinkingOptions
+                : undefined,
           };
 
           // Prepend: AGENTS.md + user memory + unconditional rules as system-reminder
@@ -1578,77 +1897,19 @@ export class AIManager {
             }
           }
 
-          // Goal evaluation — supersedes Stop hooks when active
-          const goalManager = this.container.has("GoalManager")
-            ? this.container.get<import("./goalManager.js").GoalManager>(
-                "GoalManager",
-              )
-            : undefined;
-
-          let goalContinuing = false;
-
-          if (goalManager?.isGoalActive() && !this.subagentType) {
-            // 1. Increment turn count and check circuit breakers
-            goalManager.incrementTurnCount();
-            const circuitBreaker = goalManager.checkCircuitBreakers();
-
-            if (circuitBreaker) {
-              goalManager.clearGoal();
-              logger?.info(`[Goal] ${circuitBreaker}`);
-              this.messageManager.addUserMessage({
-                content: `<system-reminder>${circuitBreaker}</system-reminder>`,
-                isMeta: true,
-              });
-              // Fall through to normal Stop hooks on the final turn
-            } else {
-              // 2. Evaluate goal
-              const evaluation = await goalManager.evaluateGoal(
-                abortController.signal,
-              );
+          // Execute Stop hooks
+          const shouldContinue = await this.executeStopHooks();
 
-              if (evaluation.isMet) {
-                goalManager.clearGoal();
-                logger?.info(`[Goal] Goal achieved: ${evaluation.reason}`);
-                this.messageManager.addUserMessage({
-                  content: `<system-reminder>Goal achieved: ${evaluation.reason}</system-reminder>`,
-                  isMeta: true,
-                });
-                // Fall through to normal Stop hooks on the final turn
-              } else {
-                const goal = goalManager.getGoal()!;
-                goal.lastReason = evaluation.reason;
-                logger?.info(`[Goal] Not yet met: ${evaluation.reason}`);
-                this.messageManager.addUserMessage({
-                  content: `<system-reminder>Goal not yet met: ${evaluation.reason}. Continue working toward: ${goal.condition}</system-reminder>`,
-                  isMeta: true,
-                });
-                // Keep loading state active to prevent UI flicker
-                this.setIsLoading(true);
-                goalContinuing = true;
-                // Restart outer loop to continue goal pursuit
-                shouldRestart = true;
-                turnOffset = 0;
-              }
-            }
-          }
-
-          // Skip Stop hooks when goal evaluator is continuing the conversation
-          if (goalContinuing) {
-            // Goal evaluator supersedes Stop hooks
-          } else {
-            const shouldContinue = await this.executeStopHooks();
-
-            // If Stop/SubagentStop hooks indicate we should continue (due to blocking errors),
-            // restart the AI conversation cycle
-            if (shouldContinue) {
-              logger?.info(
-                `${this.subagentType ? "SubagentStop" : "Stop"} hooks indicate issues need fixing, continuing conversation...`,
-              );
+          // If Stop/SubagentStop hooks indicate we should continue (due to blocking errors),
+          // restart the AI conversation cycle
+          if (shouldContinue) {
+            logger?.info(
+              `${this.subagentType ? "SubagentStop" : "Stop"} hooks indicate issues need fixing, continuing conversation...`,
+            );
 
-              // Restart the conversation to let AI fix the issues
-              shouldRestart = true;
-              turnOffset = 0;
-            }
+            // Restart the conversation to let AI fix the issues
+            shouldRestart = true;
+            turnOffset = 0;
           }
         }
 
@@ -1934,6 +2195,8 @@ export class AIManager {
             id: toolId,
             shortResult,
             stage: "running",
+            compactParams,
+            name: toolName,
           });
         },
         onResultUpdate: (result: string) => {
@@ -1941,6 +2204,8 @@ export class AIManager {
             id: toolId,
             result,
             stage: "running",
+            compactParams,
+            name: toolName,
           });
         },
         onCwdChange: async (newCwd: string) => {