Implement review suggestion

Signed-off-by: Leslie Lazzarino <leslie.lazzarino@gmail.com>

Author: leslielazzarino

src/__tests__/reportTools.test.ts

@@ -558,6 +558,46 @@ describe("reportTools", () => {
   });
 
   describe("aggregate_messages", () => {
+    test("falls back to the same header matches as search_messages, incl. single-term matches", async () => {
+      vi.useFakeTimers();
+      try {
+        const query = vi.fn(async (info: Record<string, unknown>) => {
+          if (info.fullText) return new Promise(() => {});
+          return {
+            messages: [
+              { id: 1, subject: "Lunch", author: "a@x.com", recipients: [], date: new Date("2026-01-01") },
+              { id: 2, subject: "Würth Zugang", author: "b@x.com", recipients: [], date: new Date("2026-01-02") },
+              { id: 3, subject: "Passwort", author: "c@x.com", recipients: [], date: new Date("2026-01-03") },
+            ],
+          };
+        });
+        setBrowser({ query });
+        const handlers = createReportToolHandlers({ ...BASE_SCOPE });
+
+        // Three terms, but no header matches two of them: single-term matches are counted, as search does.
+        const pending = handlers.aggregate_messages({
+          groupBy: "author",
+          query: "wuerth passwort service",
+        }) as Promise<{
+          totalMatched: number;
+          groups: Array<{ key: string; count: number }>;
+          note?: string;
+        }>;
+        await vi.advanceTimersByTimeAsync(10_000);
+        const result = await pending;
+        expect(result.totalMatched).toBe(2);
+        expect(result.groups.map((g) => g.key).sort()).toEqual(["b@x.com", "c@x.com"]);
+        expect(result.note).toMatch(/bodies were NOT searched/i);
+
+        const search = (await handlers.search_messages({ query: "wuerth passwort service" })) as {
+          hits: Array<{ id: number }>;
+        };
+        expect(search.hits.map((h) => h.id).sort()).toEqual([2, 3]);
+      } finally {
+        vi.useRealTimers();
+      }
+    });
+
     test("counts messages grouped by author", async () => {
       const query = vi.fn().mockResolvedValue({
         id: undefined,

src/__tests__/reportTrace.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import type { LlmApiRequestBody, LlmTextCompletionResponse } from "../llmConnection";
 import { LlmRoles } from "../llmConnection";
 import { DEFAULT_OPTIONS, type Options } from "../optionsParams";
@@ -66,8 +66,8 @@ describe("ReportTrace", () => {
     const trace = newTrace();
     trace.llmRequest("https://llm.example", requestBody(["ask"]));
     trace.llmResponse(completion(null), 120);
-    trace.toolCall("search_messages", { query: "foo" });
-    trace.toolResult("search_messages", '{"hits":3}');
+    const startedAt = trace.toolCall("search_messages", { query: "foo" });
+    trace.toolResult("search_messages", '{"hits":3}', startedAt);
     trace.llmRequest("https://llm.example", requestBody(["ask", "tool result"]));
     trace.llmResponse(completion("the report"), 200);
     trace.finish("success", { report: "the report" });
@@ -88,6 +88,37 @@ describe("ReportTrace", () => {
     );
   });
 
+  it("times concurrent tool calls individually and counts their overlap once in toolMs", () => {
+    vi.useFakeTimers();
+    try {
+      const trace = newTrace();
+      const a = trace.toolCall("search_messages", {}); // t=0
+      vi.advanceTimersByTime(100);
+      const b = trace.toolCall("get_thread", {}); // t=100
+      vi.advanceTimersByTime(50);
+      trace.toolResult("get_thread", "{}", b); // t=150
+      vi.advanceTimersByTime(250);
+      trace.toolResult("search_messages", "{}", a); // t=400
+      vi.advanceTimersByTime(100);
+      const c = trace.toolCall("get_messages", {}); // t=500, after an idle gap
+      vi.advanceTimersByTime(30);
+      trace.toolResult("get_messages", "{}", c); // t=530
+
+      const json = trace.toJson();
+      const durations = (json.events as Array<Record<string, unknown>>)
+        .filter((e) => e.type === "tool-result")
+        .map((e) => [e.name, e.durationMs]);
+      expect(durations).toEqual([
+        ["get_thread", 50],
+        ["search_messages", 400],
+        ["get_messages", 30],
+      ]);
+      expect((json.totals as Record<string, unknown>).toolMs).toBe(430);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
   it("marks an aborted run as cancelled and keeps the error message", () => {
     const trace = newTrace();
     trace.finish("cancelled", { error: new DOMException("Report generation cancelled", "AbortError") });

src/llmConnection.ts

@@ -320,6 +320,7 @@ export async function runAgenticLlm(
   async function runToolCall(toolCall: LlmToolCall): Promise<string> {
     const handler = toolHandlers[toolCall.function.name];
     let resultContent: string;
+    let startedAt: number | undefined;
     if (!handler) {
       console.warn(`REPORT: tool '${toolCall.function.name}' is not registered`);
       resultContent = JSON.stringify({ error: `Unknown tool: ${toolCall.function.name}` });
@@ -330,16 +331,16 @@ export async function runAgenticLlm(
           `REPORT: running tool '${toolCall.function.name}' with arg keys: ${Object.keys(parsedArgs).join(",") || "(none)"}`,
         );
         reportProgress(describeToolPhase(toolCall.function.name));
-        currentTrace()?.toolCall(toolCall.function.name, parsedArgs);
+        startedAt = currentTrace()?.toolCall(toolCall.function.name, parsedArgs);
         const result = await handler(parsedArgs);
         resultContent = JSON.stringify(result ?? null);
-        currentTrace()?.toolResult(toolCall.function.name, resultContent);
+        currentTrace()?.toolResult(toolCall.function.name, resultContent, startedAt);
         console.log(`REPORT: tool '${toolCall.function.name}' completed (resultChars=${resultContent.length})`);
       } catch (e) {
         if ((e as Error).name === "AbortError") throw e;
         console.warn(`REPORT: tool '${toolCall.function.name}' failed:`, e);
         resultContent = JSON.stringify({ error: (e as Error).message });
-        currentTrace()?.toolResult(toolCall.function.name, resultContent);
+        currentTrace()?.toolResult(toolCall.function.name, resultContent, startedAt);
       }
     }
     toolCallCount++;

src/reportTools.ts

@@ -717,8 +717,7 @@ async function collectHeaderMatches(
     if (score > 0) matches.push({ score, hit: toHit(msg) });
     return ++scanned < MAX_AGGREGATE_SCAN;
   });
-  const strong = matches.filter((m) => m.score >= minTermMatches(terms));
-  const kept = strong.length > 0 ? strong : matches;
+  const kept = keepBestMatches(matches, terms);
   kept.sort((a, b) => b.score - a.score || b.hit.date.localeCompare(a.hit.date));
   return {
     hits: kept.slice(0, cap).map((m) => m.hit),
@@ -727,6 +726,15 @@ async function collectHeaderMatches(
   };
 }
 
+/**
+ * Keep the matches scoring at least {@link minTermMatches}; if none does, fall back to every match with
+ * at least one term, so a header fallback never comes back empty while some header mentions a term.
+ */
+function keepBestMatches<T extends { score: number }>(matches: T[], terms: string[]): T[] {
+  const strong = matches.filter((m) => m.score >= minTermMatches(terms));
+  return strong.length > 0 ? strong : matches;
+}
+
 /** ISO timestamp at minute precision ("2026-09-11T09:36Z"); seconds cost tokens but add no meaning. */
 function compactDate(date: Date | string | number | undefined): string {
   return date ? `${new Date(date).toISOString().slice(0, 16)}Z` : "";
@@ -982,24 +990,30 @@ async function handleAggregateMessages(
   if (folderRef?.folderId) queryInfo.folderId = folderRef.folderId;
 
   const countMatches = async (info: QueryInfo, terms: string[] | null, timeoutMs: number) => {
-    const counts = new Map<string, number>();
-    let scanned = 0;
+    // With terms, headers are scored like search_messages' fallback and only the kept ones are counted.
+    const matches: Array<{ score: number; msg: browser.messages.MessageHeader }> = [];
     let reachedCap = false;
     const { complete, timedOut } = await scanHeaders(info, abortSignal, timeoutMs, (msg) => {
       if (filters.subjectFilter && !(msg.subject ?? "").toLowerCase().includes(filters.subjectFilter)) return true;
+      let score = 1;
       if (terms) {
         const text = headerText(msg);
-        if (terms.filter((t) => text.includes(t)).length < minTermMatches(terms)) return true;
+        score = terms.filter((t) => text.includes(t)).length;
+        if (score === 0) return true;
       }
-      scanned++;
+      matches.push({ score, msg });
+      reachedCap = matches.length >= MAX_AGGREGATE_SCAN;
+      return !reachedCap;
+    });
+    const counted = terms ? keepBestMatches(matches, terms) : matches;
+    const counts = new Map<string, number>();
+    for (const { msg } of counted) {
       for (const key of groupKeys(msg, groupBy)) {
         counts.set(key, (counts.get(key) ?? 0) + 1);
       }
-      reachedCap = scanned >= MAX_AGGREGATE_SCAN;
-      return !reachedCap;
-    });
+    }
     // A failed or stalled page flags the scan as capped so a partial aggregate is not reported as complete.
-    return { counts, scanned, capped: reachedCap || !complete, timedOut };
+    return { counts, scanned: counted.length, capped: reachedCap || !complete, timedOut };
   };
 
   // Like search_messages, a stalled full-text aggregate falls back to matching terms against headers.

src/reportTrace.ts

@@ -100,7 +100,9 @@ export class ReportTrace {
   private responseChars = 0;
   private toolResultChars = 0;
   private lastLlmRequestAt = 0;
-  private lastToolCallAt = 0;
+  /** Tool calls of one step run concurrently; toolMs counts the time at least one was running. */
+  private activeTools = 0;
+  private toolsActiveSince = 0;
   private outcome: TraceOutcome = "error";
   private errorMessage?: string;
   private report?: string;
@@ -195,16 +197,24 @@ export class ReportTrace {
     });
   }
 
-  /** Record a tool invocation requested by the model. */
-  toolCall(name: string, args: Record<string, unknown>): void {
-    this.lastToolCallAt = Date.now();
+  /** Record a tool invocation requested by the model; returns its start time for {@link toolResult}. */
+  toolCall(name: string, args: Record<string, unknown>): number {
+    const startedAt = Date.now();
+    if (this.activeTools++ === 0) this.toolsActiveSince = startedAt;
     this.push("tool-call", { name, args });
+    return startedAt;
   }
 
-  /** Record the result a tool handler returned; `result` is the JSON string handed back to the model. */
-  toolResult(name: string, result: string): void {
-    const durationMs = this.lastToolCallAt ? Date.now() - this.lastToolCallAt : 0;
-    this.toolMs += durationMs;
+  /**
+   * Record the result a tool handler returned; `result` is the JSON string handed back to the model and
+   * `startedAt` the value {@link toolCall} returned (absent when the call failed before it started).
+   */
+  toolResult(name: string, result: string, startedAt?: number): void {
+    const now = Date.now();
+    const durationMs = startedAt !== undefined ? now - startedAt : 0;
+    if (startedAt !== undefined && this.activeTools > 0 && --this.activeTools === 0) {
+      this.toolMs += now - this.toolsActiveSince;
+    }
     this.toolCalls++;
     this.toolResultChars += result.length;
     this.push("tool-result", { name, durationMs, resultChars: result.length, result: clamp(result) });