From: Stefan Gasser Date: Thu, 2 Jul 2026 20:45:49 +0000 (+0200) Subject: Avoid concurrent detector span scans (#135) X-Git-Tag: v0.7.4~1 X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=3fe543a14fb18e3a95955e5f1d131f943bb4c592;p=sgasser-llm-shield.git Avoid concurrent detector span scans (#135) --- diff --git a/src/pii/detect.test.ts b/src/pii/detect.test.ts index 5b18425..3ef6326 100644 --- 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 d1ebeaa..957e680 100644 --- 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();