From ee675a233cd6b0f92f5e3e34fe770f0250c1a1eb Mon Sep 17 00:00:00 2001
From: Xiangzhe <bakryun0718@proton.me>
Date: Thu, 13 Aug 2026 15:56:14 -0300
Subject: [PATCH] fix(ocr): fail fast on non-ok poll responses instead of
 misleading 504

pollOcrOperation now checks pollRes.ok and returns a sanitized 502
immediately (logging the upstream status via console.error) instead of
looping until the 30-attempt cap and surfacing a misleading timeout for
what was actually an auth/upstream error during polling.
---
 open-sse/handlers/ocr.ts                |  4 +++
 tests/unit/ocr-handler-dispatch.test.ts | 41 +++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/open-sse/handlers/ocr.ts b/open-sse/handlers/ocr.ts
index 8e672ab56c..3edf0b5e61 100644
--- a/open-sse/handlers/ocr.ts
+++ b/open-sse/handlers/ocr.ts
@@ -135,6 +135,10 @@ async function pollOcrOperation({ pollUrl, authHeader, fetchImpl, sleepImpl }) {
       method: "GET",
       headers: authHeader,
     });
+    if (!pollRes.ok) {
+      console.error("[OCR] poll error", pollRes.status);
+      return errorResponse(502, "OCR analysis failed");
+    }
     const json = await pollRes.json();
     if (json.status === "succeeded") {
       return json;
diff --git a/tests/unit/ocr-handler-dispatch.test.ts b/tests/unit/ocr-handler-dispatch.test.ts
index 2ad1e0cec9..2474f6b0b2 100644
--- a/tests/unit/ocr-handler-dispatch.test.ts
+++ b/tests/unit/ocr-handler-dispatch.test.ts
@@ -90,3 +90,44 @@ test("azure DI poll returns failed status maps to 502", async () => {
   const body = await res.json();
   assert.ok(!body.error.message.includes("at /"));
 });
+
+test("azure DI poll returns a non-ok response (401) and fails fast without exhausting the loop", async () => {
+  const { impl, calls } = fetchStub([
+    { status: 202, headers: { "Operation-Location": "https://poll/op/1" } },
+    { status: 401, json: { error: "unauthorized" } },
+  ]);
+  const res = await handleOcr({
+    body: {
+      model: "azure-document-intelligence/prebuilt-read",
+      document: { type: "document_url", document_url: "https://x/d.pdf" },
+    },
+    credentials: { apiKey: "azkey", baseUrl: "https://r.cognitiveservices.azure.com" },
+    fetchImpl: impl,
+    sleepImpl: noSleep,
+  });
+  assert.equal(res.status, 502);
+  // 1 initial POST + 1 poll: the loop stopped immediately, it did not run all 30 attempts.
+  assert.equal(calls.length, 2);
+  const body = await res.json();
+  assert.ok(!body.error.message.includes("at /"));
+});
+
+test("azure DI poll never resolves and times out after 30 attempts with a 504", async () => {
+  const script = [{ status: 202, headers: { "Operation-Location": "https://poll/op/1" } }];
+  for (let i = 0; i < 30; i++) {
+    script.push({ status: 200, json: { status: "running" } });
+  }
+  const { impl, calls } = fetchStub(script);
+  const res = await handleOcr({
+    body: {
+      model: "azure-document-intelligence/prebuilt-read",
+      document: { type: "document_url", document_url: "https://x/d.pdf" },
+    },
+    credentials: { apiKey: "azkey", baseUrl: "https://r.cognitiveservices.azure.com" },
+    fetchImpl: impl,
+    sleepImpl: noSleep,
+  });
+  assert.equal(res.status, 504);
+  // 1 initial POST + 30 poll attempts (the max cap), no more.
+  assert.equal(calls.length, 31);
+});