diff --git a/mineru/parser/api_server.py b/mineru/parser/api_server.py index 2cc9be05..2ccbe779 100644 --- a/mineru/parser/api_server.py +++ b/mineru/parser/api_server.py @@ -82,6 +82,8 @@ from .writer import DataWriter _DEFAULT_API_SERVER_TIER: ServerTier = "standard" _API_SERVER_LANGUAGES = PUBLIC_OCR_LANGUAGES _MAX_INLINE_BYTES_DEFAULT = 1024 * 1024 +_MAX_FILES_PER_JOB_DEFAULT = 100 +_MAX_FILE_SIZE_BYTES_DEFAULT = 209715200 _LOCAL_PARSE_OUTPUT_FORMATS: tuple[OutputFormat, ...] = ( "markdown", "middle_json", @@ -472,7 +474,7 @@ class CallbackConfig(BaseModel): class CreateJobRequest(BaseModel): model_config = _PYDANTIC_CONFIG - files: list[JobFileEntry] = Field(min_length=1) + files: list[JobFileEntry] = Field(min_length=1, max_length=_MAX_FILES_PER_JOB_DEFAULT) tier: Tier | None = None ocr_mode: Literal["auto", "txt", "ocr"] = Field( default="auto", description="OCR mode for this parse job: auto-detect, native text, or forced OCR." @@ -592,8 +594,8 @@ class UsageCurrent(BaseModel): class UsageLimits(BaseModel): model_config = _PYDANTIC_CONFIG max_pages_per_file: int = 1000 - max_file_size_bytes: int = 209715200 - max_files_per_job: int = 100 + max_file_size_bytes: int = _MAX_FILE_SIZE_BYTES_DEFAULT + max_files_per_job: int = _MAX_FILES_PER_JOB_DEFAULT max_concurrent_jobs: int = 1 max_file_retention_days: int | None = None @@ -1106,7 +1108,8 @@ class _JobRecord: class JobStore: def __init__(self, concurrency: int = 1) -> None: self._jobs: dict[str, _JobRecord] = {} - self._semaphore = asyncio.Semaphore(max(1, concurrency)) + self._concurrency = max(1, concurrency) + self._semaphore = asyncio.Semaphore(self._concurrency) self._started_at = JobStore._now() self._tasks: dict[str, asyncio.Task[None]] = {} self._closing = False @@ -1265,7 +1268,7 @@ class JobStore: has_more=(start + limit) < len(recs), ) - def build_response(self, rec: _JobRecord, access_level: AccessLevel = "registered") -> JobAsyncResponse: + def build_response(self, rec: _JobRecord, *, access_level: AccessLevel) -> JobAsyncResponse: return JobAsyncResponse( job_id=rec.id, status=rec.status, @@ -1281,7 +1284,7 @@ class JobStore: ) def usage(self, access_level: AccessLevel) -> UsageResponse: - completed = sum(1 for j in self._jobs.values() if j.status in ("completed", "partial")) + files_processed = sum(1 for j in self._jobs.values() for fr in j.files if fr.status == "completed") processed_page_count = sum( sum(_count_pages_in_range(fr.page_range) for fr in j.files if fr.status == "completed") for j in self._jobs.values() ) @@ -1290,11 +1293,11 @@ class JobStore: billing_period=UsageBillingPeriod(start=self._started_at), current=UsageCurrent( pages_processed=processed_page_count, - files_processed=completed, + files_processed=files_processed, jobs_created=len(self._jobs), ), limits=UsageLimits( - max_concurrent_jobs=self._semaphore._value, + max_concurrent_jobs=self._concurrency, ), ) @@ -1362,6 +1365,7 @@ async def _extract_bytes( allow_local_source: bool = False, max_inline_bytes: int = _MAX_INLINE_BYTES_DEFAULT, allow_http_source: bool = False, + max_url_bytes: int = _MAX_FILE_SIZE_BYTES_DEFAULT, ) -> _ExtractedSource: """按既有来源策略读取字节,并保留不会由临时文件替代的来源上下文。""" _validate_source_policy( @@ -1377,14 +1381,22 @@ async def _extract_bytes( return _ExtractedSource(file_store.read_blob(rec.sha256sum)) if isinstance(source, UrlSource): async with httpx.AsyncClient(timeout=url_timeout) as cli: - r = await cli.get(source.url) - r.raise_for_status() - return _ExtractedSource( - r.content, - source_uri=str(r.url), - transport_encoding=r.charset_encoding, - media_type=(r.headers.get("content-type") or "").partition(";")[0].strip().casefold() or None, - ) + # 流式读取并在超过 max_url_bytes 时中止,避免无上限响应被整包缓冲进内存。 + async with cli.stream("GET", source.url) as r: + r.raise_for_status() + chunks: list[bytes] = [] + received = 0 + async for chunk in r.aiter_bytes(): + received += len(chunk) + if received > max_url_bytes: + raise ValueError(f"url source exceeds max_url_bytes ({max_url_bytes})") + chunks.append(chunk) + return _ExtractedSource( + b"".join(chunks), + source_uri=str(r.url), + transport_encoding=r.charset_encoding, + media_type=(r.headers.get("content-type") or "").partition(";")[0].strip().casefold() or None, + ) if isinstance(source, InlineSource): return _ExtractedSource(_decode_inline_data(source.data)) if isinstance(source, LocalSource): @@ -1411,6 +1423,7 @@ async def _run_job( allow_local_source: bool = False, max_inline_bytes: int = _MAX_INLINE_BYTES_DEFAULT, allow_http_source: bool = False, + max_url_bytes: int = _MAX_FILE_SIZE_BYTES_DEFAULT, flash_enabled: bool = True, vlm_config: VlmConfig | None = None, ) -> None: @@ -1434,6 +1447,7 @@ async def _run_job( allow_local_source=allow_local_source, max_inline_bytes=max_inline_bytes, allow_http_source=allow_http_source, + max_url_bytes=max_url_bytes, ) data = extracted.data suffix = pathlib.Path(fr.name).suffix @@ -1932,6 +1946,7 @@ async def create_job( allow_local_source_val: bool = request.app.state.allow_local_source max_inline_bytes_val: int = request.app.state.max_inline_bytes allow_http_source_val: bool = request.app.state.allow_http_source + max_url_bytes_val: int = request.app.state.max_url_bytes flash_enabled_val: bool = request.app.state.flash_enabled image_analysis_val: bool = request.app.state.image_analysis vlm_config_val: VlmConfig = request.app.state.vlm_config @@ -1947,12 +1962,15 @@ async def create_job( allow_local_source=allow_local_source_val, max_inline_bytes=max_inline_bytes_val, allow_http_source=allow_http_source_val, + max_url_bytes=max_url_bytes_val, flash_enabled=flash_enabled_val, vlm_config=vlm_config_val, ) job_store.start_task(rec, _bg_run) - return JSONResponse(content=job_store.build_response(rec).model_dump(by_alias=True), status_code=202) + return JSONResponse( + content=job_store.build_response(rec, access_level=access_level).model_dump(by_alias=True), status_code=202 + ) @_router.get( @@ -1965,9 +1983,10 @@ async def create_job( async def get_job( job_id: str = Path(description="Job ID (job_...)"), job_store: JobStore = Depends(_get_job_store), + access_level: AccessLevel = Depends(_resolve_access_level), ) -> JobAsyncResponse: """Retrieve a job's current status and results.""" - return job_store.build_response(job_store.get(job_id)) + return job_store.build_response(job_store.get(job_id), access_level=access_level) @_router.get( @@ -2219,6 +2238,7 @@ def create_app( url_timeout: int = 60, allow_local_source: bool = False, max_inline_bytes: int = _MAX_INLINE_BYTES_DEFAULT, + max_url_bytes: int = _MAX_FILE_SIZE_BYTES_DEFAULT, allow_http_source: bool = False, api_key: str | None = None, language: str = "ch", @@ -2247,6 +2267,9 @@ def create_app( Whether ``local`` sources may read paths visible to the server process. max_inline_bytes: Maximum decoded bytes accepted for ``inline`` sources. + max_url_bytes: + Maximum bytes downloaded for ``url`` sources; larger responses are + aborted mid-download instead of being buffered into memory. allow_http_source: Whether ``url`` sources may use plain HTTP. HTTPS is always allowed. api_key: @@ -2276,6 +2299,8 @@ def create_app( _upload_dir.mkdir(parents=True, exist_ok=True) if max_inline_bytes < 0: raise ValueError("max_inline_bytes must be non-negative") + if max_url_bytes < 0: + raise ValueError("max_url_bytes must be non-negative") language = validate_public_ocr_lang(language) _model_ids, _tiers = _model_ids_and_tiers_for_server_tiers(server_tiers) @@ -2297,6 +2322,7 @@ def create_app( application.state.url_timeout = url_timeout application.state.allow_local_source = allow_local_source application.state.max_inline_bytes = max_inline_bytes + application.state.max_url_bytes = max_url_bytes application.state.allow_http_source = allow_http_source application.state.api_key = _api_key application.state.language = language @@ -2361,6 +2387,7 @@ def create_app( application.state.url_timeout = url_timeout application.state.allow_local_source = allow_local_source application.state.max_inline_bytes = max_inline_bytes + application.state.max_url_bytes = max_url_bytes application.state.allow_http_source = allow_http_source application.state.api_key = _api_key application.state.language = language @@ -2508,6 +2535,12 @@ def _build_server_log_config(log_level: str) -> dict[str, Any]: type=int, help=f"Maximum decoded bytes for inline sources (default: {_MAX_INLINE_BYTES_DEFAULT})", ) +@click.option( + "--max-url-bytes", + default=_MAX_FILE_SIZE_BYTES_DEFAULT, + type=int, + help=f"Maximum bytes downloaded for url sources (default: {_MAX_FILE_SIZE_BYTES_DEFAULT})", +) @click.option( "--allow-http-source", is_flag=True, @@ -2556,6 +2589,7 @@ def main( url_timeout: int, allow_local_source: bool, max_inline_bytes: int, + max_url_bytes: int, allow_http_source: bool, language: str, disable_image_analysis: bool, @@ -2612,6 +2646,7 @@ def main( url_timeout=url_timeout, allow_local_source=allow_local_source, max_inline_bytes=max_inline_bytes, + max_url_bytes=max_url_bytes, allow_http_source=allow_http_source, api_key=api_key, language=language, diff --git a/tests/unittest/test_flash_html.py b/tests/unittest/test_flash_html.py index 2fc49658..16d4b86d 100644 --- a/tests/unittest/test_flash_html.py +++ b/tests/unittest/test_flash_html.py @@ -30,6 +30,26 @@ from mineru.render.structured_content import render_structured_content from mineru.types import BlockType, ImageBlock, ImageBodyBlock, MiddleJson +class _StreamContext: + """把已构造的 httpx.Response 包装成 cli.stream 返回的异步上下文。""" + + def __init__(self, response: httpx.Response) -> None: + self._response = response + + async def __aenter__(self) -> httpx.Response: + return self._response + + async def __aexit__(self, *args: object) -> None: + return None + + +def _url_stream_client(response: httpx.Response) -> AsyncMock: + client = AsyncMock() + client.__aenter__.return_value = client + client.stream = lambda *_args, **_kwargs: _StreamContext(response) + return client + + def _all_raw_blocks(model_pages: list[list[dict[str, object]]]) -> list[dict[str, object]]: """按页展开 raw model-list,方便断言 HTML 映射结果。""" return [block for page in model_pages for block in page] @@ -115,9 +135,7 @@ def test_html_parse_server_url_preserves_http_declared_charset(tmp_path: Path, m headers={"Content-Type": "text/html; charset=shift_jis"}, request=httpx.Request("GET", url), ) - client = AsyncMock() - client.__aenter__.return_value = client - client.get.return_value = response + client = _url_stream_client(response) monkeypatch.setattr(api_server, "httpx", SimpleNamespace(AsyncClient=lambda **_: client)) file_store = FileStore(tmp_path / "api-files") request = CreateJobRequest.model_validate( @@ -165,9 +183,7 @@ def test_html_parse_server_url_accepts_extensionless_text_html( headers={"Content-Type": "text/html; charset=utf-8"}, request=httpx.Request("GET", url), ) - client = AsyncMock() - client.__aenter__.return_value = client - client.get.return_value = response + client = _url_stream_client(response) monkeypatch.setattr(api_server, "httpx", SimpleNamespace(AsyncClient=lambda **_: client)) file_store = FileStore(tmp_path / "api-files") request = CreateJobRequest.model_validate( @@ -231,9 +247,7 @@ def test_html_parse_server_no_flash_rejects_extensionless_text_html_after_fetch( headers={"Content-Type": "text/html; charset=utf-8"}, request=httpx.Request("GET", url), ) - client = AsyncMock() - client.__aenter__.return_value = client - client.get.return_value = response + client = _url_stream_client(response) parse_async_mock = AsyncMock() monkeypatch.setattr(api_server, "httpx", SimpleNamespace(AsyncClient=lambda **_: client)) monkeypatch.setattr(api_server, "parse_async", parse_async_mock) @@ -264,6 +278,43 @@ def test_html_parse_server_no_flash_rejects_extensionless_text_html_after_fetch( parse_async_mock.assert_not_awaited() +def test_url_source_aborts_when_exceeding_max_url_bytes(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """验证 URL 源超过 max_url_bytes 时流式下载中止并使该文件失败。""" + url = "https://example.com/big.html" + response = httpx.Response( + 200, + content=b"x" * 1024, + headers={"Content-Type": "text/html; charset=utf-8"}, + request=httpx.Request("GET", url), + ) + client = _url_stream_client(response) + monkeypatch.setattr(api_server, "httpx", SimpleNamespace(AsyncClient=lambda **_: client)) + file_store = FileStore(tmp_path / "api-files") + request = CreateJobRequest.model_validate( + { + "files": [{"source": {"type": "url", "url": url}}], + "tier": "standard", + "output_formats": ["middle_json"], + } + ) + record = api_server.JobStore().create(request, file_store) + + asyncio.run( + api_server._run_job( + record, + request, + file_store, + image_analysis=True, + max_url_bytes=64, + ) + ) + + assert record.status == "failed" + assert record.files[0].status == "failed" + assert record.files[0].error is not None + assert record.files[0].error.message == "url source exceeds max_url_bytes (64)" + + def test_html_local_base_images_styles_and_escape_are_bounded(tmp_path: Path) -> None: """验证本地 base、CSS、栅格图可读取,但父目录逃逸图片只保留说明。""" assets = tmp_path / "assets" diff --git a/tests/unittest/test_parser_api_contract.py b/tests/unittest/test_parser_api_contract.py index 488f7bd9..cf3cde9a 100644 --- a/tests/unittest/test_parser_api_contract.py +++ b/tests/unittest/test_parser_api_contract.py @@ -2059,6 +2059,87 @@ def test_run_job_preserves_input_file_id_and_keeps_source_undownloadable( assert exc_info.value.status_code == 403 +def test_create_job_request_rejects_more_than_max_files_per_job(tmp_path: Path) -> None: + """超过广告的 max_files_per_job 时拒绝创建请求。""" + source = str(tmp_path / "demo.pdf") + entries = [{"source": {"type": "local", "path": source}}] * (api_server._MAX_FILES_PER_JOB_DEFAULT + 1) + + with pytest.raises(ValidationError): + CreateJobRequest.model_validate({"files": entries, "tier": "flash"}) + + +def test_job_responses_report_resolved_access_level( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """job 创建与查询响应携带按 api_key 解析的 access_level,而不是固定 registered。""" + + async def fake_parse_async(*args: object, **kwargs: object) -> ParseResult: + return ParseResult(middle_json=_full_middle_json(PageInfo(page_idx=0))) + + monkeypatch.setattr("mineru.parser.api_server.parse_async", fake_parse_async) + source = tmp_path / "demo.html" + source.write_text("

content

", encoding="utf-8") + payload = { + "files": [{"source": {"type": "local", "path": str(source)}}], + "tier": "flash", + "output_formats": ["middle_json"], + } + + anonymous_app = create_app(upload_dir=str(tmp_path / "api-anonymous"), tier="flash", allow_local_source=True) + with TestClient(anonymous_app) as client: + created = client.post("/v1/parse/jobs", json=payload) + assert created.status_code == 202 + assert created.json()["access_level"] == "anonymous" + fetched = client.get(f"/v1/parse/jobs/{created.json()['job_id']}") + assert fetched.status_code == 200 + assert fetched.json()["access_level"] == "anonymous" + + registered_app = create_app( + upload_dir=str(tmp_path / "api-registered"), tier="flash", allow_local_source=True, api_key="secret-key" + ) + with TestClient(registered_app) as client: + created = client.post("/v1/parse/jobs", json=payload, headers={"Authorization": "Bearer secret-key"}) + assert created.status_code == 202 + assert created.json()["access_level"] == "registered" + fetched = client.get(f"/v1/parse/jobs/{created.json()['job_id']}", headers={"Authorization": "Bearer secret-key"}) + assert fetched.status_code == 200 + assert fetched.json()["access_level"] == "registered" + + +def test_usage_counts_completed_files_and_reports_configured_concurrency(tmp_path: Path) -> None: + """files_processed 按完成的文件计数,max_concurrent_jobs 报告配置上限而非剩余信号量。""" + file_store = FileStore(tmp_path / "api-files") + store = api_server.JobStore(concurrency=2) + + def _request(*names: str) -> CreateJobRequest: + return CreateJobRequest.model_validate( + { + "files": [{"source": {"type": "local", "path": str(tmp_path / name)}} for name in names], + "tier": "flash", + "output_formats": ["middle_json"], + } + ) + + rec_a = store.create(_request("a1.html", "a2.html"), file_store) + rec_a.files[0].status = "completed" + rec_a.files[1].status = "failed" + rec_a.status = "partial" + rec_b = store.create(_request("b1.html"), file_store) + rec_b.files[0].status = "completed" + rec_b.status = "completed" + + async def _report_while_busy() -> "api_server.UsageResponse": + await store._semaphore.acquire() + return store.usage("anonymous") + + usage = asyncio.run(_report_while_busy()) + + assert usage.current.jobs_created == 2 + assert usage.current.files_processed == 2 + assert usage.limits.max_concurrent_jobs == 2 + + @pytest.mark.parametrize(("request_tier", "response_tier"), [("standard", "standard"), (None, "flash")]) def test_api_server_accepts_lightweight_job_without_requested_quality_runtime( monkeypatch: pytest.MonkeyPatch, @@ -2595,7 +2676,7 @@ def test_api_server_standard_jobs_use_requested_tier(tmp_path: Path, monkeypatch max_inline_bytes=app.state.max_inline_bytes, allow_http_source=app.state.allow_http_source, ) - return job_store.build_response(rec).model_dump(by_alias=True) + return job_store.build_response(rec, access_level="anonymous").model_dump(by_alias=True) default_response = asyncio.run( run_request(