Avoid concurrent detector span scans (#135)

diff --git a/src/pii/detect.test.ts b/src/pii/detect.test.ts
index 5b18425dd2c5513c34a4d80b01bf097d577a5533..3ef632606be46bd1c5d263bc3b28285cc742ae63 100644 (file)
--- a/src/pii/detect.test.ts
+++ b/src/pii/detect.test.ts
@@ -149,6 +149,44 @@ describe("PIIDetector", () => {
       expect(analyzeRequests).toEqual([expect.objectContaining({ text: "mcp-pii here" })]);
     });
 
+    test("scans spans sequentially to avoid detector-side inference queue timeouts", async () => {
+      const analyzeRequests: string[] = [];
+      let inFlight = 0;
+      let maxInFlight = 0;
+
+      globalThis.fetch = mock(async (url: string | URL | Request, init?: RequestInit) => {
+        const urlStr = url.toString();
+
+        if (urlStr.includes("/analyze") && init?.body) {
+          const body = JSON.parse(init.body as string);
+          analyzeRequests.push(body.text);
+          inFlight++;
+          maxInFlight = Math.max(maxInFlight, inFlight);
+          await new Promise((resolve) => setTimeout(resolve, 5));
+          inFlight--;
+          return new Response(JSON.stringify([]), {
+            status: 200,
+            headers: { "Content-Type": "application/json" },
+          });
+        }
+
+        return originalFetch(url, init);
+      }) as unknown as typeof fetch;
+
+      const detector = new PIIDetector();
+      const spans: TextSpan[] = [
+        { text: "first", path: "0", messageIndex: 0, partIndex: 0, role: "user" },
+        { text: "second", path: "1", messageIndex: 1, partIndex: 0, role: "user" },
+        { text: "third", path: "2", messageIndex: 2, partIndex: 0, role: "user" },
+      ];
+
+      const result = await detector.analyzeRequest(spans, spanExtractor);
+
+      expect(result.hasPII).toBe(false);
+      expect(analyzeRequests).toEqual(["first", "second", "third"]);
+      expect(maxInFlight).toBe(1);
+    });
+
     test("honors explicit scan_roles override", async () => {
       const config = getConfig();
       const previousScanRoles = config.pii_detection.scan_roles;

diff --git a/src/pii/detect.ts b/src/pii/detect.ts
index d1ebeaa5d93b67e84ee1615fbb3e7d78ffc1eea4..957e680cbbbac51d2268170b4f0296fa43bb9b20 100644 (file)
--- a/src/pii/detect.ts
+++ b/src/pii/detect.ts
@@ -225,22 +225,23 @@ export class PIIDetector {
     const allowlist = config.masking.allowlist;
     const denylist = config.masking.denylist;
 
-    const spanEntities: PIIEntity[][] = await Promise.all(
-      spans.map(async (span) => {
-        if (!span.text) return [];
+    const spanEntities: PIIEntity[][] = [];
+    for (const span of spans) {
+      if (!span.text) {
+        spanEntities.push([]);
+        continue;
+      }
 
-        if (!span.role || !scanRoles.has(span.role)) {
-          return [];
-        }
+      if (!span.role || !scanRoles.has(span.role)) {
+        spanEntities.push([]);
+        continue;
+      }
 
-        const denylistedEntities = findDenylistedEntities(span.text, denylist, knownPlaceholders);
-        const detectedEntities = config.pii_detection.enabled
-          ? await this.detectPII(span.text)
-          : [];
-        const filteredEntities = filterAllowlistedEntities(span.text, detectedEntities, allowlist);
-        return mergeDenylistEntities(filteredEntities, denylistedEntities);
-      }),
-    );
+      const denylistedEntities = findDenylistedEntities(span.text, denylist, knownPlaceholders);
+      const detectedEntities = config.pii_detection.enabled ? await this.detectPII(span.text) : [];
+      const filteredEntities = filterAllowlistedEntities(span.text, detectedEntities, allowlist);
+      spanEntities.push(mergeDenylistEntities(filteredEntities, denylistedEntities));
+    }
 
     const allEntities = spanEntities.flat();