fix: preserve streamed DOM tool calls in active turns

bridge-browser/src/content/dom_tool_turn.ts

@@ -0,0 +1,120 @@
+import type {
+  ToolRequestRegistry,
+  ToolRequestTurn,
+  UnflushedRequestBatch,
+} from "./tool_request_registry";
+
+export interface DomToolTurnLocation {
+  conversationKey: string;
+  messageIndex: number;
+}
+
+interface ActiveDomToolTurn {
+  conversationKey: string;
+  messageElement: Element;
+  messageIndex: number;
+  requests: ToolRequestTurn;
+}
+
+/**
+ * Keeps the active DOM tool turn stable while one assistant message grows across streaming scans.
+ *
+ * A turn becomes trusted only when it is first observed outside virtualized history. Once trusted,
+ * the same message can keep adding calls even if response growth leaves the viewport far behind the
+ * live bottom. Unrelated history messages never replace a pending active turn.
+ */
+export class DomToolTurnController {
+  private activeTurn: ActiveDomToolTurn | null = null;
+
+  public constructor(private readonly requestRegistry: ToolRequestRegistry) {}
+
+  /**
+   * Observe the latest rendered assistant message and report whether it belongs to the active turn.
+   */
+  public observeMessage(
+    messageElement: Element,
+    location: DomToolTurnLocation,
+    viewingVirtualizedHistory: boolean
+  ): boolean {
+    if (this.activeTurn?.conversationKey !== location.conversationKey) {
+      this.activeTurn = null;
+    }
+
+    const matchingTurn = this.getActiveTurn(messageElement, location, viewingVirtualizedHistory);
+    if (matchingTurn) {
+      matchingTurn.messageElement = messageElement;
+      return true;
+    }
+
+    if (this.activeTurn?.requests.getUnflushedBatch().hasRequests) {
+      return false;
+    }
+
+    if (viewingVirtualizedHistory) {
+      return false;
+    }
+
+    this.activeTurn = {
+      conversationKey: location.conversationKey,
+      messageElement,
+      messageIndex: location.messageIndex,
+      requests: this.requestRegistry.createTurn(),
+    };
+    return true;
+  }
+
+  /**
+   * Record the current identity for a code-block slot in the active message.
+   */
+  public recordRequest(codeBlockIndex: number, requestKey: string): void {
+    this.activeTurn?.requests.set(codeBlockIndex, requestKey);
+  }
+
+  public getUnflushedBatch(): UnflushedRequestBatch {
+    return this.activeTurn?.requests.getUnflushedBatch() ?? createEmptyBatch();
+  }
+
+  /**
+   * Release a turn after its delivered request keys have been marked flushed in the registry.
+   */
+  public finalizeRequests(requestKeys: readonly string[]): void {
+    if (!this.activeTurn?.requests.hasAny(requestKeys)) {return;}
+    if (this.activeTurn.requests.getUnflushedBatch().hasRequests) {return;}
+    this.activeTurn = null;
+  }
+
+  public reset(): void {
+    this.activeTurn = null;
+  }
+
+  private getActiveTurn(
+    messageElement: Element,
+    location: DomToolTurnLocation,
+    viewingVirtualizedHistory: boolean
+  ): ActiveDomToolTurn | null {
+    const activeTurn = this.activeTurn;
+    if (activeTurn?.conversationKey !== location.conversationKey) {
+      return null;
+    }
+
+    if (activeTurn.messageIndex !== location.messageIndex) {
+      return null;
+    }
+
+    if (!viewingVirtualizedHistory || activeTurn.messageElement === messageElement) {
+      return activeTurn;
+    }
+
+    return !activeTurn.messageElement.isConnected ? activeTurn : null;
+  }
+}
+
+function createEmptyBatch(): UnflushedRequestBatch {
+  return {
+    completedCount: 0,
+    hasRequests: false,
+    ids: [],
+    isComplete: false,
+    totalCount: 0,
+  };
+}

bridge-browser/src/content/main.ts

@@ -12,6 +12,7 @@ import { AutoInitPromptController } from "./auto_init_prompt";
 import { createApprovalState, parseStoredApprovalEntries, type ApprovalState } from "./approval_policy";
 import { CompletionNotifier } from "./completion_notifier";
 import { DomToolActivityController } from "./dom_tool_activity";
+import { DomToolTurnController } from "./dom_tool_turn";
 import { hasPromptResourceChange, loadPromptsFromStorage } from "./prompt_resources";
 import { createNetworkCaptureRuntime } from "./network_capture_runtime";
 import { ResultDeliveryController } from "./result_delivery_controller";
@@ -194,6 +195,7 @@ function applySyncedSiteConfig(siteId: string, sites: SyncedAiSite[]): void {
     DOM = matchedSite.selectors;
     currentSiteName = matchedSite.name ?? matchedSite.id;
     domToolActivity.reset();
+    domToolTurns.reset();
     networkCapture.configure(getSiteNetworkCaptureConfig(matchedSite.capture));
     completionNotifier.reset();
     autoInitPrompt.setupTrigger();
@@ -206,12 +208,14 @@ function applySyncedSiteConfig(siteId: string, sites: SyncedAiSite[]): void {
   DOM = null;
   currentSiteName = null;
   domToolActivity.reset();
+  domToolTurns.reset();
   networkCapture.reset();
   console.log(`${BRANDING.productName}: Site '${siteId}' is not configured in VS Code. Idle.`);
 }
 
 function resetCurrentSite(): void {
   domToolActivity.reset();
+  domToolTurns.reset();
   networkCapture.reset();
   DOM = null;
   currentSiteName = null;
@@ -243,6 +247,7 @@ chrome.storage.onChanged.addListener((changes, namespace) => {
 const requestRegistry = new ToolRequestRegistry();
 const toolActivityTracker = new ToolActivityTracker();
 const domToolActivity = new DomToolActivityController(toolActivityTracker);
+const domToolTurns = new DomToolTurnController(requestRegistry);
 new ToolActivityOverlay(toolActivityTracker);
 let lastProgressLogTime = 0;
 let lastProgressStatus = "";
@@ -288,6 +293,7 @@ const networkCapture = createNetworkCaptureRuntime({
 const resultDelivery = new ResultDeliveryController({
   getAutoSend: () => CONFIG.autoSend,
   hasPendingTurns: () => networkCapture.hasPendingTurns(),
+  onBatchFinalized: (requestKeys) => domToolTurns.finalizeRequests(requestKeys),
   requestRegistry,
   scheduleMainLoop,
   toolActivityTracker,
@@ -340,14 +346,9 @@ function runMainLoop() {
   if (!latestCodeBlocks) { return; }
 
   const { messageIndex, messageElement, codeElements } = latestCodeBlocks;
-  const messageLocation = {
-    conversationKey: location.href,
-    messageIndex,
-  };
-  const skipNewCapturesForVirtualizedHistory = UI.isLikelyViewingVirtualizedHistory(DOM);
-
-  // 当前轮次对象只记录本次扫描看到的 requestKey；去重、排序和已回填过滤由 registry 统一处理。
-  const currentTurn = requestRegistry.createTurn();
+  const messageLocation = { conversationKey: location.href, messageIndex };
+  const viewingVirtualizedHistory = UI.isLikelyViewingVirtualizedHistory(DOM);
+  const isActiveDomTurn = domToolTurns.observeMessage(messageElement, messageLocation, viewingVirtualizedHistory);
 
   for (const [codeBlockIndex, codeEl] of codeElements.entries()) {
     const codeElement = codeEl as HTMLElement;
@@ -375,11 +376,11 @@ function runMainLoop() {
       const isProcessing = requestRegistry.isRunning(requestIdentity.requestKey);
       const isKnown = requestRegistry.hasSeen(requestIdentity.requestKey);
 
-      if (!isKnown && skipNewCapturesForVirtualizedHistory) {
-        logVirtualizedHistorySkip(payload.name);
+      if (!isKnown && !isActiveDomTurn) {
+        if (viewingVirtualizedHistory) {logVirtualizedHistorySkip(payload.name);}
         continue;
       }
-      currentTurn.add(requestIdentity.requestKey);
+      if (isActiveDomTurn) {domToolTurns.recordRequest(codeBlockIndex, requestIdentity.requestKey);}
 
       if (!isKnown) {
         // 新发现的工具调用只进入执行路径一次，后续扫描只会根据 registry 中的执行状态刷新视觉状态。
@@ -408,10 +409,8 @@ function runMainLoop() {
         }
       }
     } catch (error) {
-      const isKnown = Boolean(codeElement.dataset.mcpRequestKey && requestRegistry.hasSeen(codeElement.dataset.mcpRequestKey));
-
-      if (!isKnown && skipNewCapturesForVirtualizedHistory) {
-        logVirtualizedHistorySkip();
+      if (!isActiveDomTurn) {
+        if (viewingVirtualizedHistory) {logVirtualizedHistorySkip();}
         continue;
       }
 
@@ -423,12 +422,12 @@ function runMainLoop() {
         codeBlockIndex,
         error
       );
-      currentTurn.add(requestIdentity.requestKey);
+      domToolTurns.recordRequest(codeBlockIndex, requestIdentity.requestKey);
     }
   }
 
   // 只处理当前轮次里还没有写回过的请求。已 flush 的 requestKey 不会再次写入输入框。
-  const unflushedBatch = currentTurn.getUnflushedBatch();
+  const unflushedBatch = domToolTurns.getUnflushedBatch();
 
   if (unflushedBatch.hasRequests) {
     // 工具完成的判定由 registry 统一计算：已不在执行中，并且已经有结果可回填。
@@ -458,6 +457,7 @@ function runMainLoop() {
         // 某些路径可能没有文本输出；它们完成后也要标记为已处理。
         if (resultBatch.hasAnyResult) {
           requestRegistry.markFlushed(resultBatch.ids);
+          domToolTurns.finalizeRequests(resultBatch.ids);
           toolActivityTracker.updateDelivery(resultBatch.ids, "delivered");
         }
       }

bridge-browser/src/content/result_delivery_controller.ts

@@ -7,6 +7,7 @@ import type { BufferedResultBatch, ToolRequestRegistry } from "./tool_request_re
 interface ResultDeliveryControllerOptions {
   getAutoSend: () => boolean;
   hasPendingTurns: () => boolean;
+  onBatchFinalized?: (requestKeys: readonly string[]) => void;
   requestRegistry: ToolRequestRegistry;
   scheduleMainLoop: (delayMs: number) => void;
   toolActivityTracker: ToolActivityTracker;
@@ -30,7 +31,7 @@ export class ResultDeliveryController {
     void UI.deliverResult(resultBatch, selectors)
       .then((delivery) => {
         batchFinalized = true;
-        this.options.requestRegistry.markFlushed(resultBatch.ids);
+        this.finalizeBatch(resultBatch.ids);
         if (!delivery.delivered) {
           this.handleDeliveryFailure(resultBatch);
           return;
@@ -41,13 +42,18 @@ export class ResultDeliveryController {
       })
       .catch((error: unknown) => {
         batchFinalized = true;
-        this.options.requestRegistry.markFlushed(resultBatch.ids);
+        this.finalizeBatch(resultBatch.ids);
         this.options.toolActivityTracker.updateDelivery(resultBatch.ids, "failed");
         Logger.log(`Result delivery failed: ${getErrorMessage(error)}`, "error");
       })
       .finally(() => this.finishDelivery(batchFinalized));
   }
 
+  private finalizeBatch(requestKeys: readonly string[]): void {
+    this.options.requestRegistry.markFlushed(requestKeys);
+    this.options.onBatchFinalized?.(requestKeys);
+  }
+
   private handleDeliveryFailure(resultBatch: BufferedResultBatch): void {
     this.options.toolActivityTracker.updateDelivery(resultBatch.ids, "failed");
     Logger.log(

bridge-browser/src/content/tool_request_registry.ts

@@ -110,9 +110,9 @@ export class ToolRequestRegistry {
   private toolCallCount = 0;
 
   /**
-   * 创建一次页面扫描轮次的临时收集器。
+   * 创建一个由调用方控制生命周期的工具调用轮次收集器。
    *
-   * ToolRequestTurn 只保存本次扫描看到的 requestKey 及其顺序；跨轮次状态仍由 registry 持有。
+   * DOM 捕获会让同一个实例跨多次流式扫描存活；网络捕获仍可直接使用 registry 的批处理方法。
    */
   public createTurn(): ToolRequestTurn {
     return new ToolRequestTurn(this);
@@ -314,39 +314,35 @@ export class ToolRequestRegistry {
 }
 
 /**
- * 一次 runMainLoop 扫描过程中的 requestKey 收集器。
+ * 一个工具调用轮次中的 requestKey 收集器。
  *
- * 它只保存本轮扫描看到的 ID，并保持页面出现顺序。生命周期很短，每次 runMainLoop 都会创建
- * 新实例；跨轮次的执行/回填状态由 ToolRequestRegistry 管理。
+ * DOM 流式输出会在多次扫描中逐步补充代码块，因此收集器按代码块位置更新 requestKey，并由
+ * 上层决定何时释放。跨轮次的执行、结果和已回填状态仍由 ToolRequestRegistry 管理。
  */
 export class ToolRequestTurn {
   /**
-   * 当前扫描轮次内按页面顺序出现的 requestKey。
+   * 代码块位置到当前 requestKey 的映射。
    *
-   * 后续回填会按这个顺序合并结果，保证多工具调用结果顺序和 AI 原始请求顺序一致。
+   * 流式 JSON 可能先产生一个临时的协议错误身份，随后补全为有效调用。按位置覆盖可以避免旧身份
+   * 永久留在持久化轮次中，同时在较早代码块暂时离开 DOM 时保留已经发现的调用。
    */
-  private readonly requestKeys: string[] = [];
-
-  /**
-   * 当前扫描轮次内的去重集合。
-   *
-   * 同一个 requestKey 可能因为重复代码块、协议错误反馈或 DOM 结构变化被看到多次；Set 用来
-   * 保证 requestKeys 中只出现一次。
-   */
-  private readonly requestKeySet = new Set<string>();
+  private readonly requestKeysByCodeBlock = new Map<number, string>();
 
   public constructor(private readonly registry: ToolRequestRegistry) {}
 
   /**
-   * 记录本轮扫描看到的一个 requestKey。
-   *
-   * null 表示调用方明确没有要纳入本轮批处理的 requestKey；这种情况直接忽略。
+   * 记录指定代码块位置当前对应的 requestKey。
    */
-  public add(requestKey: string | null): void {
-    if (!requestKey || this.requestKeySet.has(requestKey)) {return;}
+  public set(codeBlockIndex: number, requestKey: string): void {
+    this.requestKeysByCodeBlock.set(codeBlockIndex, requestKey);
+  }
 
-    this.requestKeys.push(requestKey);
-    this.requestKeySet.add(requestKey);
+  /**
+   * 判断给定的一批 requestKey 是否属于这个轮次。
+   */
+  public hasAny(requestKeys: readonly string[]): boolean {
+    const candidates = new Set(requestKeys);
+    return this.getOrderedRequestKeys().some((requestKey) => candidates.has(requestKey));
   }
 
   /**
@@ -355,7 +351,14 @@ export class ToolRequestTurn {
    * 具体的已回填过滤和完成状态计算交给 registry，这个对象只提供本轮 ID 的有序列表。
    */
   public getUnflushedBatch(): UnflushedRequestBatch {
-    return this.registry.getUnflushedBatch(this.requestKeys);
+    return this.registry.getUnflushedBatch(this.getOrderedRequestKeys());
+  }
+
+  private getOrderedRequestKeys(): string[] {
+    const orderedKeys = [...this.requestKeysByCodeBlock.entries()]
+      .sort(([leftIndex], [rightIndex]) => leftIndex - rightIndex)
+      .map(([, requestKey]) => requestKey);
+    return [...new Set(orderedKeys)];
   }
 }