diff --git a/src/polygateway/transports/monkey_ocr.py b/src/polygateway/transports/monkey_ocr.py
index 9af9edc..89faae6 100644
--- a/src/polygateway/transports/monkey_ocr.py
+++ b/src/polygateway/transports/monkey_ocr.py
@@ -28,8 +28,13 @@ from polygateway.transports._http_errors import (
response_body,
summarize_body,
)
-from polygateway.transports._monkey_document import _finite_number, _parse_page_size
+from polygateway.transports._monkey_document import (
+ _finite_number,
+ _parse_page_size,
+ parse_document_package,
+)
from polygateway.types import (
+ OcrDocumentTransportResult,
OcrLayoutElement,
OcrLayoutTransportResult,
OcrTextTransportResult,
@@ -70,6 +75,11 @@ _DEFAULT_REJECTED_MARKERS: tuple[str, ...] = (
# 服务把一切异常(含它自己抛的 400、`Model not initialized` 这类真故障)都包成同一前缀,
# 所以前缀只是必要条件,必须同时命中具体标记
_PARSING_FAILED_PREFIX = "parsing failed:"
+
+_MAX_PACKAGE_BYTES = 256 * 1024 * 1024
+"""整包(含 `images/*.jpg`)体积上限,仅文档路径。4 页含 4 图约 598KB,数百页含图的文档
+也远低于此;超限即拒,防止畸形/恶意响应把内存吃满。httpx 非流式 `get` 会先整包 `aread()`,
+这正是文档路径改用流式下载并逐块累计的原因。"""
# Adobe 阅读器容忍头部前有垃圾字节(惯例而非 ISO 32000 正文);服务端实测可解析
# 前缀 200 字节垃圾的 PDF(#34 设计 E3)
_PDF_HEADER_WINDOW = 1024
@@ -244,6 +254,24 @@ def _parse_middle_json(
raise ResultInvalidError(f"OCR 结果包内容无效: {exc}") from exc
+def _package_too_large(source: SourceConfig, size: int) -> ResultInvalidError:
+ return ResultInvalidError(
+ f"{source.name} OCR 结果包体积超过上限 {_MAX_PACKAGE_BYTES}(已读 {size})",
+ source_name=source.name,
+ operation="download_result",
+ )
+
+
+def _check_declared_length(resp: httpx.Response, source: SourceConfig) -> None:
+ """声明长度超限就不必读了;缺失/非法的头部不作判据(靠累计字节兜底)。"""
+ try:
+ declared = int(resp.headers.get("content-length", ""))
+ except ValueError:
+ return
+ if declared > _MAX_PACKAGE_BYTES:
+ raise _package_too_large(source, declared)
+
+
class MonkeyOcrTransport:
"""MonkeyOCR transport: 每源一个预配 httpx client,懒创建,aclose 统一释放。"""
@@ -311,32 +339,47 @@ class MonkeyOcrTransport:
text=content, raw={k: v for k, v in data.items() if k != "content"}
)
- async def parse_layout(
- self, *, image: bytes, source: SourceConfig, call_id: str
- ) -> OcrLayoutTransportResult:
- """两段协议: POST /parse → GET download_url → ZIP → 全元素提取。"""
- del call_id
- client = self._client_for(source)
- with _translate_http_errors(source.name, "parse"):
- resp = await client.post("/parse", files=self._files(image))
+ async def _post_parse(
+ self,
+ client: httpx.AsyncClient,
+ data: bytes,
+ source: SourceConfig,
+ rejected_markers: tuple[str, ...],
+ ) -> str:
+ """两段协议的第一段: POST /parse → 校验 `success` → 返回 download_url。
+
+ layout 与 document 两条路径共用,保证"success 为假 → 附 status_code=200"等
+ 设计 §10.2 的有意修复只有一份。
+ """
+ with _translate_http_errors(source.name, "parse", rejected_markers=rejected_markers):
+ resp = await client.post("/parse", files=self._files(data))
resp.raise_for_status()
- data = self._response_json(resp, source, "parse")
- if data.get("success") is not True:
+ body = self._response_json(resp, source, "parse")
+ if body.get("success") is not True:
# 设计 §10.2 有意修复: 200 响应确证服务活着,附 status_code
# 让治理侧按"响应即健康"记熔断成功(CHS 原样会误走本地拒绝分支)
raise RequestRejectedError(
- f"{source.name} OCR 解析失败: {data.get('message')}",
+ f"{source.name} OCR 解析失败: {body.get('message')}",
source_name=source.name,
status_code=resp.status_code,
operation="parse",
)
- download_url = data.get("download_url")
+ download_url = body.get("download_url")
if not isinstance(download_url, str) or not download_url.strip():
raise TransientError(
f"{source.name} /parse 响应缺少 download_url",
source_name=source.name,
operation="parse",
)
+ return download_url
+
+ async def parse_layout(
+ self, *, image: bytes, source: SourceConfig, call_id: str
+ ) -> OcrLayoutTransportResult:
+ """两段协议: POST /parse → GET download_url → ZIP → 全元素提取。"""
+ del call_id
+ client = self._client_for(source)
+ download_url = await self._post_parse(client, image, source, ())
with _translate_http_errors(source.name, "download_result"):
# 实测 download_url 为相对路径(/static/*.zip),相对 base_url 解析;
# httpx 对绝对 URL 会忽略 base_url,两种形态天然兼容
@@ -349,6 +392,54 @@ class MonkeyOcrTransport:
raw={"success": True, "download_url": download_url},
)
+ async def parse_document(
+ self, *, document: bytes, source: SourceConfig, call_id: str
+ ) -> OcrDocumentTransportResult:
+ """整份文档: 本地预检 → POST /parse → 流式 GET → 线程内解析(#34 设计 §5)。"""
+ del call_id
+ if _sniff_media(document) == "unknown":
+ # 服务从未被打扰: HTML 付费墙页面、乱码、空串在批处理语料里最常见。不带
+ # status_code(服务没响应),治理侧只归还探针、不计熔断成败
+ raise RequestRejectedError(
+ f"{source.name} OCR 无法识别文档格式(既非 PDF 也非已知图片),拒绝发送",
+ source_name=source.name,
+ operation="parse",
+ )
+ client = self._client_for(source)
+ download_url = await self._post_parse(client, document, source, self._rejected_markers)
+ package = await self._download_package(client, download_url, source)
+ # 纯函数、可能处理大对象: 放线程里,取消与 call_deadline_s 才能及时生效
+ elements, page_sizes, markdown = await asyncio.to_thread(parse_document_package, package)
+ return OcrDocumentTransportResult(
+ elements=elements,
+ page_sizes=page_sizes,
+ markdown=markdown,
+ raw={"success": True, "download_url": download_url},
+ )
+
+ async def _download_package(
+ self, client: httpx.AsyncClient, download_url: str, source: SourceConfig
+ ) -> bytes:
+ """流式下载结果包并限制体积。
+
+ **先处理错误状态,再检查体积**: 结果包的资源检查不得改变 HTTP 错误分类
+ (一个声明超限的 401 仍是 SourceDead,响应体摘要也要保留——issue #10)。
+ """
+ chunks: list[bytes] = []
+ total = 0
+ with _translate_http_errors(source.name, "download_result"):
+ async with client.stream("GET", download_url) as resp:
+ if resp.is_error:
+ await resp.aread() # 先落缓冲:摘要口径只认已缓冲的响应体
+ resp.raise_for_status()
+ _check_declared_length(resp, source)
+ async for chunk in resp.aiter_bytes():
+ total += len(chunk)
+ if total > _MAX_PACKAGE_BYTES:
+ raise _package_too_large(source, total)
+ chunks.append(chunk)
+ return b"".join(chunks)
+
async def check_health(self, *, source: SourceConfig) -> bool:
"""GET /health 探测: 2xx → True,失败 → False,取消穿透(R10)。"""
client = self._client_for(source)
diff --git a/tests/unit/_document_fixtures.py b/tests/unit/_document_fixtures.py
new file mode 100644
index 0000000..e3f017a
--- /dev/null
+++ b/tests/unit/_document_fixtures.py
@@ -0,0 +1,71 @@
+"""OCR 文档结果包的合成 fixtures(#34)。
+
+全部**合成**: 保留 2026-10-21 实测的结构骨架(para_blocks 块类型、table/image 的嵌套
+`blocks`、span 字段),一切识别文本替换为 `SYNTH-*` 占位——第三方论文内容许可不明、
+医疗数据不入库(零业务假设 + P5),真实 ZIP 不提交。
+"""
+
+import io
+import json
+import zipfile
+
+_BBOX = (10.0, 20.0, 110.0, 40.0)
+
+
+def _span(content="SYNTH-TEXT", type_="text"):
+ return {"type": type_, "content": content, "bbox": list(_BBOX), "score": 1.0}
+
+
+def _text_block(content="SYNTH-TEXT", *, type_="text", bbox=_BBOX):
+ return {
+ "type": type_,
+ "bbox": list(bbox),
+ "index": 0,
+ "lines": [{"bbox": list(bbox), "spans": [_span(content)]}],
+ }
+
+
+def _body(type_, spans, bbox=_BBOX):
+ return {"type": type_, "bbox": list(bbox), "lines": [{"bbox": list(bbox), "spans": spans}]}
+
+
+def _table_block(html="", extra_subs=()):
+ body = _body("table_body", [{"type": "table", "html": html, "image_path": "synth.jpg"}])
+ return {"type": "table", "bbox": list(_BBOX), "index": 1, "blocks": [*extra_subs, body]}
+
+
+def _image_block():
+ body = _body("image_body", [{"type": "image", "image_path": "synth.jpg"}])
+ return {"type": "image", "bbox": list(_BBOX), "index": 2, "blocks": [body]}
+
+
+def _page(blocks=None, *, page_idx=0, page_size=(612.0, 792.0), **extra):
+ page = {"page_idx": page_idx, "page_size": list(page_size), "discarded_blocks": []}
+ if blocks is not None:
+ page["para_blocks"] = blocks
+ page.update(extra)
+ return page
+
+
+def _package(
+ pages=None,
+ *,
+ md="# synth\n",
+ middle=None,
+ middle_name="doc_middle.json",
+ md_name=None,
+ extra=(),
+):
+ """组装结果包。`middle` 给原始 bytes 时原样写入;`md_name` 默认由 middle 名配对。"""
+ if middle is None:
+ middle = json.dumps({"pdf_info": pages if pages is not None else [_page([])]}).encode()
+ if md_name is None:
+ md_name = middle_name[: -len("_middle.json")] + ".md"
+ buf = io.BytesIO()
+ with zipfile.ZipFile(buf, "w") as zf:
+ zf.writestr(middle_name, middle)
+ if md is not None:
+ zf.writestr(md_name, md if isinstance(md, bytes) else md.encode())
+ for name, data in extra:
+ zf.writestr(name, data)
+ return buf.getvalue()
diff --git a/tests/unit/test_monkey_document.py b/tests/unit/test_monkey_document.py
index f2cc4cc..ce77c84 100644
--- a/tests/unit/test_monkey_document.py
+++ b/tests/unit/test_monkey_document.py
@@ -20,67 +20,16 @@ from polygateway.errors import ResultInvalidError
from polygateway.transports import _monkey_document
from polygateway.transports._monkey_document import parse_document_package
from polygateway.types import OcrDocumentElement
-
-_BBOX = (10.0, 20.0, 110.0, 40.0)
-
-
-def _span(content="SYNTH-TEXT", type_="text"):
- return {"type": type_, "content": content, "bbox": list(_BBOX), "score": 1.0}
-
-
-def _text_block(content="SYNTH-TEXT", *, type_="text", bbox=_BBOX):
- return {
- "type": type_,
- "bbox": list(bbox),
- "index": 0,
- "lines": [{"bbox": list(bbox), "spans": [_span(content)]}],
- }
-
-
-def _body(type_, spans, bbox=_BBOX):
- return {"type": type_, "bbox": list(bbox), "lines": [{"bbox": list(bbox), "spans": spans}]}
-
-
-def _table_block(html="", extra_subs=()):
- body = _body("table_body", [{"type": "table", "html": html, "image_path": "synth.jpg"}])
- return {"type": "table", "bbox": list(_BBOX), "index": 1, "blocks": [*extra_subs, body]}
-
-
-def _image_block():
- body = _body("image_body", [{"type": "image", "image_path": "synth.jpg"}])
- return {"type": "image", "bbox": list(_BBOX), "index": 2, "blocks": [body]}
-
-
-def _page(blocks=None, *, page_idx=0, page_size=(612.0, 792.0), **extra):
- page = {"page_idx": page_idx, "page_size": list(page_size), "discarded_blocks": []}
- if blocks is not None:
- page["para_blocks"] = blocks
- page.update(extra)
- return page
-
-
-def _package(
- pages=None,
- *,
- md="# synth\n",
- middle=None,
- middle_name="doc_middle.json",
- md_name=None,
- extra=(),
-):
- """组装结果包。`middle` 给原始 bytes 时原样写入;`md_name` 默认由 middle 名配对。"""
- if middle is None:
- middle = json.dumps({"pdf_info": pages if pages is not None else [_page([])]}).encode()
- if md_name is None:
- md_name = middle_name[: -len("_middle.json")] + ".md"
- buf = io.BytesIO()
- with zipfile.ZipFile(buf, "w") as zf:
- zf.writestr(middle_name, middle)
- if md is not None:
- zf.writestr(md_name, md if isinstance(md, bytes) else md.encode())
- for name, data in extra:
- zf.writestr(name, data)
- return buf.getvalue()
+from tests.unit._document_fixtures import (
+ _BBOX,
+ _body,
+ _image_block,
+ _package,
+ _page,
+ _span,
+ _table_block,
+ _text_block,
+)
def _parse(pages=None, **kwargs):
diff --git a/tests/unit/test_monkey_ocr.py b/tests/unit/test_monkey_ocr.py
index 31eeff8..47fca66 100644
--- a/tests/unit/test_monkey_ocr.py
+++ b/tests/unit/test_monkey_ocr.py
@@ -8,6 +8,7 @@ P5 医疗数据不入库)。
import asyncio
import io
import json
+import threading
import zipfile
import httpx
@@ -19,6 +20,7 @@ from polygateway.errors import (
SourceDeadError,
TransientError,
)
+from polygateway.transports import monkey_ocr
from polygateway.transports.monkey_ocr import (
_DEFAULT_REJECTED_MARKERS,
MonkeyOcrTransport,
@@ -28,7 +30,10 @@ from polygateway.transports.monkey_ocr import (
_translate_http_errors,
_upload_part,
)
-from polygateway.types import SourceConfig
+from polygateway.types import OcrDocumentElement, OcrDocumentTransportResult, SourceConfig
+from tests.unit._document_fixtures import _package as _doc_package
+from tests.unit._document_fixtures import _page as _doc_page
+from tests.unit._document_fixtures import _table_block, _text_block
def _source(**overrides):
@@ -535,8 +540,9 @@ class TestRejectedMarkers:
assert type(err) is exc_type
def test_translate_http_errors_passes_markers_through(self):
- with pytest.raises(RequestRejectedError), _translate_http_errors(
- "monkey_1", "parse", rejected_markers=_DEFAULT_REJECTED_MARKERS
+ with (
+ pytest.raises(RequestRejectedError),
+ _translate_http_errors("monkey_1", "parse", rejected_markers=_DEFAULT_REJECTED_MARKERS),
):
raise _http_error(500, _parsing_failed(_DEFAULT_REJECTED_MARKERS[0]))
# 默认不传: 行为不变
@@ -563,6 +569,198 @@ class TestRejectedMarkers:
await t.parse_layout(image=b"jpg", source=_source(), call_id="c1")
+class TestParseDocument:
+ """#34 设计 §5: 本地预检 → POST /parse → 流式 GET(体积上限)→ 线程内解析。"""
+
+ @staticmethod
+ def _doc_zip():
+ return _doc_package(
+ [_doc_page([_text_block("SYNTH-A"), _table_block("")])]
+ )
+
+ def _transport(self, seen=None, *, zip_resp=None, parse_resp=None):
+ def handler(request: httpx.Request) -> httpx.Response:
+ if seen is not None:
+ seen.append((request.method, request.url.path, request.read()))
+ if request.url.path == "/parse":
+ return parse_resp or httpx.Response(200, json=_parse_body())
+ if request.url.path.startswith("/static/"):
+ return zip_resp or httpx.Response(200, content=self._doc_zip())
+ return httpx.Response(500)
+
+ return _transport_for(handler)
+
+ async def test_full_chain_relative_download_url(self):
+ t = self._transport()
+ result = await t.parse_document(document=_PDF_BODY, source=_source(), call_id="c1")
+ assert isinstance(result, OcrDocumentTransportResult)
+ assert [(e.type, e.text) for e in result.elements] == [
+ ("text", "SYNTH-A"),
+ ("table", ""),
+ ]
+ assert all(isinstance(e, OcrDocumentElement) for e in result.elements)
+ assert result.page_sizes == [(612.0, 792.0)]
+ assert result.markdown == "# synth\n"
+ assert result.raw == {"success": True, "download_url": "/static/sample_parsed_1.zip"}
+
+ async def test_absolute_download_url(self):
+ url = "http://10.77.0.20:7866/static/abs.zip"
+ t = self._transport(parse_resp=httpx.Response(200, json=_parse_body(url)))
+ result = await t.parse_document(document=_PDF_BODY, source=_source(), call_id="c1")
+ assert result.raw["download_url"] == url
+
+ @pytest.mark.parametrize(
+ ("document", "filename", "content_type"),
+ [
+ pytest.param(_PDF_BODY, b"doc.pdf", b"application/pdf", id="pdf"),
+ pytest.param(_JUNK + _PDF_BODY, b"doc.pdf", b"application/pdf", id="pdf-junk-prefix"),
+ pytest.param(b"\xff\xd8\xff\xe0abc", b"image.jpg", b"image/jpeg", id="jpeg"),
+ ],
+ )
+ async def test_upload_label_matrix_for_the_document_entry(
+ self, document, filename, content_type
+ ):
+ seen = []
+ await self._transport(seen).parse_document(
+ document=document, source=_source(), call_id="c1"
+ )
+ post_body = next(body for method, path, body in seen if path == "/parse")
+ assert b'filename="' + filename + b'"' in post_body
+ assert b"Content-Type: " + content_type in post_body
+
+ @pytest.mark.parametrize(
+ "document",
+ [b"paywall", b"GARBAGE-NOT-AN-IMAGE", b""],
+ ids=["html", "garbage", "empty"],
+ )
+ async def test_unrecognizable_bytes_are_rejected_locally_without_any_request(self, document):
+ seen = []
+ with pytest.raises(RequestRejectedError) as ei:
+ await self._transport(seen).parse_document(
+ document=document, source=_source(), call_id="c1"
+ )
+ assert seen == [] # 服务从未被打扰
+ assert ei.value.status_code is None # 服务未响应: 熔断侧只归还探针,不计成败
+
+ @pytest.mark.parametrize("marker", _DEFAULT_REJECTED_MARKERS)
+ async def test_known_deterministic_post_failures_are_rejected_with_status_500(self, marker):
+ body = _parsing_failed(marker).encode()
+ t = self._transport(parse_resp=httpx.Response(500, content=body))
+ with pytest.raises(RequestRejectedError) as ei:
+ await t.parse_document(document=_PDF_BODY, source=_source(), call_id="c1")
+ assert ei.value.status_code == 500 and ei.value.body_text
+
+ @pytest.mark.parametrize(
+ "reason",
+ ["CUDA out of memory", "400: Unsupported file type: .webp", "500: Model not initialized"],
+ )
+ async def test_unknown_post_500_stays_transient(self, reason):
+ t = self._transport(
+ parse_resp=httpx.Response(500, content=_parsing_failed(reason).encode())
+ )
+ with pytest.raises(TransientError):
+ await t.parse_document(document=_PDF_BODY, source=_source(), call_id="c1")
+
+ async def test_success_false_is_rejected_with_status_200(self):
+ t = self._transport(
+ parse_resp=httpx.Response(200, json={"success": False, "message": "no"})
+ )
+ with pytest.raises(RequestRejectedError) as ei:
+ await t.parse_document(document=_PDF_BODY, source=_source(), call_id="c1")
+ assert ei.value.status_code == 200
+
+ async def test_missing_download_url_and_non_json_are_transient(self):
+ body = {"success": True, "message": "ok"}
+ for resp in (httpx.Response(200, json=body), httpx.Response(200, content=b"not json")):
+ t = self._transport(parse_resp=resp)
+ with pytest.raises(TransientError):
+ await t.parse_document(document=_PDF_BODY, source=_source(), call_id="c1")
+
+ @pytest.mark.parametrize(
+ ("status", "exc_type"),
+ [
+ (401, SourceDeadError),
+ (429, TransientError),
+ (500, TransientError),
+ (404, RequestRejectedError),
+ ],
+ )
+ async def test_download_errors_keep_their_class_and_body(self, status, exc_type):
+ """GET 阶段不适用确定性窄化,且响应体摘要随每个分支(issue #10)。"""
+ body = _parsing_failed(_DEFAULT_REJECTED_MARKERS[0])
+ t = self._transport(zip_resp=httpx.Response(status, content=body.encode()))
+ with pytest.raises(exc_type) as ei:
+ await t.parse_document(document=_PDF_BODY, source=_source(), call_id="c1")
+ assert ei.value.status_code == status and ei.value.body_text == body
+
+ async def test_error_status_with_oversized_content_length_is_still_classified_by_http(
+ self, monkeypatch
+ ):
+ """结果包资源检查不得改变错误分类: 超限的 Content-Length 出现在 401 上,仍是 SourceDead。"""
+ monkeypatch.setattr(monkey_ocr, "_MAX_PACKAGE_BYTES", 100)
+ resp = httpx.Response(401, content=b"denied", headers={"content-length": "101"})
+ with pytest.raises(SourceDeadError):
+ await self._transport(zip_resp=resp).parse_document(
+ document=_PDF_BODY, source=_source(), call_id="c1"
+ )
+
+ async def test_oversized_declared_content_length_is_rejected(self, monkeypatch):
+ monkeypatch.setattr(monkey_ocr, "_MAX_PACKAGE_BYTES", 100)
+ resp = httpx.Response(200, content=b"x" * 10, headers={"content-length": "101"})
+ with pytest.raises(ResultInvalidError):
+ await self._transport(zip_resp=resp).parse_document(
+ document=_PDF_BODY, source=_source(), call_id="c1"
+ )
+
+ async def test_oversized_streamed_body_is_rejected_by_cumulative_count(self, monkeypatch):
+ """没有 Content-Length(分块传输)时靠累计字节兜底: 多个小块合起来超限。"""
+ monkeypatch.setattr(monkey_ocr, "_MAX_PACKAGE_BYTES", 100)
+
+ async def chunks():
+ for _ in range(3):
+ yield b"x" * 60
+
+ resp = httpx.Response(200, content=chunks())
+ with pytest.raises(ResultInvalidError):
+ await self._transport(zip_resp=resp).parse_document(
+ document=_PDF_BODY, source=_source(), call_id="c1"
+ )
+
+ async def test_bad_package_is_result_invalid(self):
+ t = self._transport(zip_resp=httpx.Response(200, content=b"not a zip"))
+ with pytest.raises(ResultInvalidError):
+ await t.parse_document(document=_PDF_BODY, source=_source(), call_id="c1")
+
+ async def test_cancellation_during_parse_does_not_block_the_event_loop(self, monkeypatch):
+ """解析在线程里跑: 同步解析大对象会占住事件循环,取消与 call_deadline_s 都会失效。
+
+ 线程内用 threading.Event 阻塞,并用 call_soon_threadsafe 通知(线程里不能直接用
+ asyncio.Event)。若实现退回同步解析,桩会把循环占满 1 秒后正常返回,
+ `task.cancel()` 就取不到 CancelledError——用例转红。
+ """
+ loop = asyncio.get_running_loop()
+ entered = asyncio.Event()
+ release = threading.Event()
+
+ def blocking_parse(zip_bytes):
+ loop.call_soon_threadsafe(entered.set)
+ release.wait(timeout=1)
+ return [], [(1.0, 1.0)], ""
+
+ monkeypatch.setattr(monkey_ocr, "parse_document_package", blocking_parse)
+ task = asyncio.create_task(
+ self._transport().parse_document(document=_PDF_BODY, source=_source(), call_id="c1")
+ )
+ try:
+ await asyncio.wait_for(entered.wait(), timeout=2)
+ await asyncio.sleep(0.01) # 线程还卡着,循环必须仍能调度别的协程
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ finally:
+ release.set() # 让线程收尾,不泄漏
+
+
class TestCheckHealth:
async def test_healthy(self):
t = _transport_for(_routes(health_resp=httpx.Response(200, json={"status": "healthy"})))