Inject web crawl user messages correctly (#2129)

src/pro/main/ipc/handlers/local_agent/local_agent_handler.ts

@@ -47,6 +47,10 @@ import {
   escapeXmlContent,
   UserMessageContentPart,
 } from "./tools/types";
+import {
+  prepareStepMessages,
+  type InjectedMessage,
+} from "./prepare_step_utils";
 import { TOOL_DEFINITIONS } from "./tool_definitions";
 import { parseAiMessagesJson } from "@/ipc/utils/ai_messages_utils";
 import { parseMcpToolKey, sanitizeMcpName } from "@/ipc/utils/mcp_tool_utils";
@@ -143,6 +147,8 @@ export async function handleLocalAgentStream(
 
   // Track pending user messages to inject after tool results
   const pendingUserMessages: UserMessageContentPart[][] = [];
+  // Store injected messages with their insertion index to re-inject at the same spot each step
+  const allInjectedMessages: InjectedMessage[] = [];
 
   try {
     // Get model client
@@ -227,28 +233,11 @@ export async function handleLocalAgentStream(
       stopWhen: stepCountIs(25), // Allow multiple tool call rounds
       abortSignal: abortController.signal,
       // Inject pending user messages (e.g., images from web_crawl) between steps
-      prepareStep: ({ messages, ...rest }) => {
-        if (pendingUserMessages.length === 0) {
-          return undefined;
-        }
-        // Build user messages from pending content
-        const newMessages = [...messages];
-        for (const content of pendingUserMessages) {
-          newMessages.push({
-            role: "user" as const,
-            content: content.map((part) => {
-              if (part.type === "text") {
-                return { type: "text" as const, text: part.text };
-              }
-              // part.type === "image-url"
-              return { type: "image" as const, image: new URL(part.url) };
-            }),
-          });
-        }
-        // Clear pending messages after injection
-        pendingUserMessages.length = 0;
-        return { messages: newMessages, ...rest };
-      },
+      // We must re-inject all accumulated messages each step because the AI SDK
+      // doesn't persist dynamically injected messages in its internal state.
+      // We track the insertion index so messages appear at the same position each step.
+      prepareStep: (options) =>
+        prepareStepMessages(options, pendingUserMessages, allInjectedMessages),
       onFinish: async (response) => {
         const totalTokens = response.usage?.totalTokens;
         const inputTokens = response.usage?.inputTokens;

src/pro/main/ipc/handlers/local_agent/prepare_step_utils.ts

+/**
+ * Utility for preparing step messages with injected user content.
+ *
+ * This module contains pure functions extracted from the prepareStep callback
+ * in local_agent_handler.ts, enabling isolated unit testing.
+ */
+
+import {
+  ImagePart,
+  ModelMessage,
+  TextPart,
+  UserModelMessage,
+} from "node_modules/ai/dist";
+import type { UserMessageContentPart } from "./tools/types";
+
+/**
+ * A message that has been processed and is ready to inject.
+ */
+export interface InjectedMessage {
+  insertAtIndex: number;
+  /** Sequence number to preserve FIFO order for same-index messages */
+  sequence: number;
+  message: UserModelMessage;
+}
+
+/**
+ * Transform a UserMessageContentPart to the format expected by the AI SDK.
+ */
+export function transformContentPart(
+  part: UserMessageContentPart,
+): TextPart | ImagePart {
+  if (part.type === "text") {
+    return { type: "text", text: part.text };
+  }
+  // part.type === "image-url"
+  return { type: "image", image: new URL(part.url) };
+}
+
+/**
+ * Process pending user messages and add them to the injected messages list.
+ * Each message is recorded with the current message count as its insertion index.
+ *
+ * @param pendingUserMessages - Queue of pending messages (will be mutated/emptied)
+ * @param allInjectedMessages - List of already injected messages (will be mutated)
+ * @param currentMessageCount - The current number of messages in the conversation
+ */
+export function processPendingMessages(
+  pendingUserMessages: UserMessageContentPart[][],
+  allInjectedMessages: InjectedMessage[],
+  currentMessageCount: number,
+): void {
+  while (pendingUserMessages.length > 0) {
+    const content = pendingUserMessages.shift()!;
+    allInjectedMessages.push({
+      insertAtIndex: currentMessageCount,
+      sequence: allInjectedMessages.length, // Track insertion order
+      message: {
+        role: "user" as const,
+        content: content.map(transformContentPart),
+      },
+    });
+  }
+}
+
+/**
+ * Build a new messages array with injected messages inserted at their recorded positions.
+ * Messages are processed in reverse order of insertion index to avoid shifting issues.
+ * For messages with the same index, we process in reverse sequence order to preserve FIFO.
+ *
+ * @param messages - The original messages array
+ * @param injectedMessages - Messages to inject with their target indices
+ * @returns New array with injected messages inserted at correct positions
+ */
+export function injectMessagesAtPositions<T>(
+  messages: T[],
+  injectedMessages: InjectedMessage[],
+): (T | InjectedMessage["message"])[] {
+  if (injectedMessages.length === 0) {
+    return messages;
+  }
+
+  // Type as union from the start to allow inserting InjectedMessage["message"]
+  const newMessages: (T | InjectedMessage["message"])[] = [...messages];
+
+  // Sort by insertion index descending, then by sequence descending.
+  // The sequence descending ensures that for same-index messages,
+  // we splice the LAST-added first, so after all splices the FIRST-added
+  // ends up in front (preserving FIFO order).
+  const sortedInjections = [...injectedMessages].sort((a, b) => {
+    if (a.insertAtIndex !== b.insertAtIndex) {
+      return b.insertAtIndex - a.insertAtIndex;
+    }
+    return b.sequence - a.sequence;
+  });
+
+  for (const injection of sortedInjections) {
+    newMessages.splice(injection.insertAtIndex, 0, injection.message);
+  }
+
+  return newMessages;
+}
+
+/**
+ * The complete prepareStep logic as a pure function.
+ *
+ * @param options - The step options containing messages and other properties
+ * @param pendingUserMessages - Queue of pending messages to process
+ * @param allInjectedMessages - Accumulated list of injected messages
+ * @returns Modified options with injected messages, or undefined if no changes needed
+ */
+export function prepareStepMessages<
+  TMessage extends ModelMessage,
+  T extends { messages: TMessage[]; [key: string]: unknown },
+>(
+  options: T,
+  pendingUserMessages: UserMessageContentPart[][],
+  allInjectedMessages: InjectedMessage[],
+): (Omit<T, "messages"> & { messages: TMessage[] }) | undefined {
+  const { messages, ...rest } = options;
+
+  // Move any new pending messages to the permanent injected list
+  processPendingMessages(
+    pendingUserMessages,
+    allInjectedMessages,
+    messages.length,
+  );
+
+  // If no messages to inject, don't modify
+  if (allInjectedMessages.length === 0) {
+    return undefined;
+  }
+
+  // Build the new messages array with injections
+  // Cast is safe because InjectedMessage["message"] is a valid ModelMessage
+  const newMessages = injectMessagesAtPositions(
+    messages,
+    allInjectedMessages,
+  ) as TMessage[];
+
+  return { messages: newMessages, ...rest };
+}

src/pro/main/ipc/handlers/local_agent/tools/web_crawl.ts

@@ -35,7 +35,7 @@ Trigger a crawl ONLY if BOTH conditions are true:
 
 const CLONE_INSTRUCTIONS = `
 
-Replicate the website from the provided HTML, markdown, and screenshot.
+Replicate the website from the provided screenshot image and markdown.
 
 **Use the screenshot as your primary visual reference** to understand the layout, colors, typography, and overall design of the website. The screenshot shows exactly how the page should look.
 
@@ -118,9 +118,6 @@ export const webCrawlTool: ToolDefinition<z.infer<typeof webCrawlSchema>> = {
       throw new Error("No content available from web crawl");
     }
 
-    if (!result.html) {
-      throw new Error("No HTML available from web crawl");
-    }
     if (!result.screenshot) {
       throw new Error("No screenshot available from web crawl");
     }
@@ -133,10 +130,6 @@ export const webCrawlTool: ToolDefinition<z.infer<typeof webCrawlSchema>> = {
         type: "text",
         text: formatSnippet("Markdown snapshot:", result.markdown, "markdown"),
       },
-      {
-        type: "text",
-        text: formatSnippet("HTML snapshot:", result.html, "html"),
-      },
     ]);
 
     return "Web crawl completed.";