From bc5388af6b888be8285896ae9dece0e654eec901 Mon Sep 17 00:00:00 2001
From: Richard R <me@richardr.dev>
Date: Wed, 3 Jun 2026 04:35:27 -0600
Subject: [PATCH] fix(pdf): add strict validation for layout model output
 shapes and extend test coverage

Add explicit error handling for invalid or inconsistent pred_boxes and logits array lengths in runLayoutModel to prevent silent failures. Expand test suite to verify correct region filtering and error scenarios, ensuring only labeled regions are returned and malformed outputs are handled robustly.
---
 compute/core/src/pdf/runLayoutModel.ts        | 26 +++++++++++--
 .../control-plane/run-layout-model.test.ts    | 39 +++++++++++++++++++
 2 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/compute/core/src/pdf/runLayoutModel.ts b/compute/core/src/pdf/runLayoutModel.ts
index c97dd6b..5c726e0 100644
--- a/compute/core/src/pdf/runLayoutModel.ts
+++ b/compute/core/src/pdf/runLayoutModel.ts
@@ -286,11 +286,29 @@ export async function runLayoutModel(input: RunLayoutInput): Promise<LayoutRegio
     const logits = output.logits?.data as Float32Array | undefined;
     const predBoxes = output.pred_boxes?.data as Float32Array | undefined;
     if (!logits || !predBoxes) return [];
+    if (predBoxes.length === 0 && logits.length === 0) return [];
+    if (predBoxes.length === 0 || logits.length === 0) {
+      throw new Error(
+        `layout-model-invalid-output-shape: pred_boxes length=${predBoxes.length}, logits length=${logits.length}`,
+      );
+    }
 
-    const numQueries = Math.floor(predBoxes.length / 4);
-    if (numQueries <= 0) return [];
-    const classCount = Math.floor(logits.length / numQueries);
-    if (classCount <= 0) return [];
+    if (predBoxes.length % 4 !== 0) {
+      throw new Error(`layout-model-invalid-pred-box-shape: length ${predBoxes.length} is not divisible by 4`);
+    }
+    const numQueries = predBoxes.length / 4;
+    if (numQueries <= 0) {
+      throw new Error(`layout-model-invalid-pred-box-shape: expected positive query count, got ${numQueries}`);
+    }
+    if (logits.length % numQueries !== 0) {
+      throw new Error(
+        `layout-model-invalid-logit-shape: length ${logits.length} is not divisible by query count ${numQueries}`,
+      );
+    }
+    const classCount = logits.length / numQueries;
+    if (classCount <= 0) {
+      throw new Error(`layout-model-invalid-logit-shape: expected positive class count, got ${classCount}`);
+    }
     const regions: LayoutRegion[] = [];
 
     for (let queryIdx = 0; queryIdx < numQueries; queryIdx += 1) {
diff --git a/compute/core/tests/control-plane/run-layout-model.test.ts b/compute/core/tests/control-plane/run-layout-model.test.ts
index 490d0b4..e951bf3 100644
--- a/compute/core/tests/control-plane/run-layout-model.test.ts
+++ b/compute/core/tests/control-plane/run-layout-model.test.ts
@@ -133,4 +133,43 @@ describe('runLayoutModel', () => {
       expect.closeTo(40, 5),
     ]);
   });
+
+  test('drops unlabeled query winners and keeps only labeled regions', async () => {
+    mockState.runOutput = {
+      logits: {
+        data: new Float32Array([
+          0.1,
+          0.2,
+          5,
+          4,
+          0.1,
+          0.1,
+        ]),
+      },
+      pred_boxes: {
+        data: new Float32Array([
+          0.25, 0.25, 0.3, 0.3,
+          0.75, 0.75, 0.3, 0.3,
+        ]),
+      },
+    };
+
+    const { runLayoutModel } = await import('../../src/pdf/runLayoutModel');
+    const regions = await runLayoutModel({
+      pageWidth: 100,
+      pageHeight: 100,
+      textItems: [{} as never],
+      pageImage: Buffer.from([1]),
+    });
+
+    expect(regions).toHaveLength(1);
+    expect(regions[0]?.label).toBe('text');
+    expect(regions[0]?.confidence).toBeCloseTo(0.96109135, 5);
+    expect(regions[0]?.bbox).toEqual([
+      expect.closeTo(60, 5),
+      expect.closeTo(60, 5),
+      expect.closeTo(90, 5),
+      expect.closeTo(90, 5),
+    ]);
+  });
 });