From 3702cd7c569f57578b85ef81b792b235e32b2bec Mon Sep 17 00:00:00 2001 From: myhloli Date: Wed, 21 Oct 2026 14:59:49 +0800 Subject: [PATCH] feat(api): support URL inputs in v1 tasks --- docs/next/api/tasks.md | 52 ++++-- mineru/kit/router/tasks.py | 37 +++- mineru/parser/api_server.py | 80 ++++++-- mineru/parser/task_api.py | 99 ++++++++-- tests/unittest/test_task_api.py | 301 ++++++++++++++++++++++++++++++- tests/unittest/test_task_http.py | 117 +++++++++++- 6 files changed, 627 insertions(+), 59 deletions(-) diff --git a/docs/next/api/tasks.md b/docs/next/api/tasks.md index 4655d985..57e18cf3 100644 --- a/docs/next/api/tasks.md +++ b/docs/next/api/tasks.md @@ -1,9 +1,10 @@ -# V1 tasks：逐文件上传与结果下载链接 +# V1 tasks：逐文件上传、URL 输入与结果下载链接 状态: Implemented（自部署 API server 与 Router） 读者: HTTP 客户端、SaaS 接口与自部署服务开发者 -共同流程是：**提交文件清单 → 上传缺失文件 → 轮询任务 → 使用返回链接下载 ZIP**。 +共同流程是：**提交文件清单 → 上传缺失的客户端文件 → 轮询任务 → 使用返回链接下载 ZIP**。 +清单也支持 URL 输入，由执行服务在后台下载，无需客户端提前获取文件大小或上传字节。 SaaS 可以把上传、下载地址指向 OSS，本地服务可以指向自身的文件接口。客户端只使用服务返回的链接，不拼接存储路径，也不要求 SaaS 实现本地 Files API。 SaaS 的 Go 服务、OSS 配置及部署由其开发方实现；这里定义客户端共同契约。 @@ -19,7 +20,8 @@ SaaS 的 Go 服务、OSS 配置及部署由其开发方实现；这里定义客 任务首次创建就获得稳定的 `task_id`；执行提交后，它也是当前部署的公开 `job_id`。等待上传阶段使用 tasks 查询，既有 `/v1/parse/jobs` 执行资源尚未创建。 状态为 `waiting_upload`、`queued`、`running`、`completed`、`partial`、`failed`、`canceled`。 -全部文件校验通过后自动启动，不要求客户端发送完成确认。等待上传不占解析并发，也不预占 Router worker。 +全部客户端上传项校验通过后自动启动，不要求客户端发送完成确认。等待上传不占解析并发，也不预占 Router worker。 +纯 URL 任务直接进入 `queued`；URL 下载在后台执行阶段进行，沿用 Job 并发控制。 `completed/partial` 只在最终 ZIP 已保存后发布；打包时仍为 `running`。`partial` 保留成功文件产物及失败文件的错误。 ## 表单请求 @@ -43,8 +45,9 @@ SaaS 的 Go 服务、OSS 配置及部署由其开发方实现；这里定义客 |---|---|---| | `client_file_id` | 是 | 一次任务内唯一，1–64 个字母、数字、下划线或短横线。 | | `filename` | 是 | 逻辑文件名，保留扩展名；不同条目允许同名。 | -| `bytes` | 是 | 文件实际字节数，正整数。 | -| `sha256sum` | 否 | 64 位小写十六进制 SHA-256；省略时视为需要上传。 | +| `url` | 否 | 文件下载地址，支持 HTTP(S)，具体协议受部署策略约束；省略或 `null` 表示客户端上传。 | +| `bytes` | 条件 | 上传项必填；URL 项可省略或为 `null`。提供时必须为严格正整数，用于核对实际字节数。 | +| `sha256sum` | 否 | 64 位小写十六进制 SHA-256；上传项用于缓存预检与完整性校验，URL 项仅用于下载完整性校验。 | | `page_range` | 否 | 逐文件页范围，未提供时继承全局值；非 PDF 仍只支持整本解析。 | 用真实文件生成清单，再发送表单： @@ -82,8 +85,32 @@ curl --fail-with-body "$MINERU_API_URL/v1/tasks" \ 匿名本地服务可省略 Authorization。无需 hash 时从清单中省略 `sha256sum`。 `client_file_id` 用于关联文件和上传项，不能按文件名判断对应关系。 +## URL 提交与混合输入 + +```bash +curl --fail-with-body "$MINERU_API_URL/v1/tasks" \ + -H "Authorization: Bearer $MINERU_API_KEY" \ + --data-urlencode 'tier=flash' \ + --data-urlencode 'ocr_mode=txt' \ + --data-urlencode 'manifest=[{"client_file_id":"f1","filename":"document.pdf","url":"https://example.com/download?signature=example"}]' +``` + +`filename` 仍然必填，用于文件类型识别、展示和产物命名，不要求 URL 路径有扩展名。 +URL 项的 `upload_status` 为 `not_required`，不产生上传链接，也不接受 PUT；源 `file_id` 为 `null`。 +未声明的 `bytes` 在下载前为 `null`，完整下载并校验通过后返回实际大小及 SHA-256。 +声明的大小或 hash 不匹配时，该文件失败，不进入解析；空文件同样失败。 + +同一 manifest 可以混合 URL 与上传项，`uploads` 仅包含需要客户端上传的条目。 +混合任务在这些上传完成后进入执行阶段，URL 下载不需要客户端参与。URL 与 multipart 字节不能同时指定给同一项。 +URL 每次重新下载，即使声明的 hash 已在源文件缓存中，也不跳过下载；不要求 HEAD 或 `Content-Length`。 + +下载失败、超时或完整性校验失败按文件记录错误，其他文件继续执行；有成功项则返回 `partial` 和 ZIP，全部失败则为 `failed` 且无 `result_url`。 +URL 大小不一致使用 `bytes_mismatch`，hash 不一致使用 `file_hash_mismatch`。请求格式和部署策略错误仍按接单校验处理。 + ## 缓存、上传与恢复 +以下缓存规则适用于无 `url` 的客户端上传项： + - 全部 hash 命中：返回空 `uploads`，直接提交解析。 - 部分命中：仅返回未命中文件的上传项，命中文件已经是 `ready`。 - 无 hash 或未命中：上传完整文件；服务端计算实际 hash、校验大小与可选 hash，并去重保存。 @@ -96,7 +123,7 @@ curl --fail-with-body "$MINERU_API_URL/v1/tasks" \ | 字段 | 说明 | |---|---| | `task_id`、`status`、`status_url` | 稳定任务 ID、当前状态和查询地址。 | -| `files` | 按输入顺序返回清单、`upload_status: pending/ready`、逐文件解析状态及错误。 | +| `files` | 按输入顺序返回清单、`upload_status: pending/ready/not_required`、逐文件解析状态及错误；URL 项始终为 `not_required`。 | | `uploads` | 尚需上传的文件，每项包含 `client_file_id`、`upload_url`、`upload_method: PUT`、`upload_headers`、`expires_at`。 | | `result_url` | 结果未就绪时为 `null`；`completed/partial` 时为最终 ZIP 下载地址。 | | `error` | 任务级错误，例如上传过期或结果打包失败。 | @@ -134,7 +161,7 @@ curl --fail-with-body "$MINERU_API_URL/v1/tasks" \ - Router 返回自身的 Files 地址并流式代理 worker 产物，不泄露内部 worker 地址。 所有 tasks 自动生成 ZIP；单文件保留现有自包含产物，批量使用 `0001/`、`0002/` 等目录并包含 `manifest.json`。 -manifest 使用 `client_file_id` 关联输入，`directory` 指向包内目录，不写入 worker 的私有 task/file ID 或短期链接。任务 ID 以 HTTP 响应为准。 +manifest 使用 `client_file_id` 关联输入，`directory` 指向包内目录，不写入 worker 的私有 task/file ID 或短期下载链接，也不写入源 `url` 字段。已校验的 URL 输入记录实际大小与 hash。任务 ID 以 HTTP 响应为准。 ZIP 在后台只生成一次；重复查询或下载不重新打包。打包失败返回 `failed/result_packaging_failed`，不会发布虚假的下载链接。 签名链接过期但产物仍在保留期内时，重新查询任务获得有效地址，不重新提交解析。 @@ -144,6 +171,7 @@ ZIP 在后台只生成一次；重复查询或下载不重新打包。打包失 ## 完整客户端示例 `scripts/http_task_example.sh` 调用同目录的 `http_task_example.py`，只需要 Python 3.10+ 标准库，无需安装 MinerU 或云存储 SDK。 +该脚本接受本地文件路径；URL 提交使用上面的 curl 示例。 默认顺序上传，每个文件最多尝试三次，仅重试网络及临时服务错误；创建请求不自动重试。 脚本校验任务身份、状态和文件对应关系，轮询有预算，下载先写临时文件并验证 ZIP 后再保存最终文件。 @@ -181,12 +209,14 @@ UPLOAD_MODE=direct MINERU_API_URL=http://127.0.0.1:8000 \ ## 本地限制与 Router 调度 -本地不设固定上传字节数或 PDF 页数上限；`usage.limits.max_file_size_bytes` 与 `max_pages_per_file` 为 `null`。 -每任务最多 100 文件。既有 Jobs 的 URL/inline 来源预算仍按其专用配置生效。 +本地不设固定客户端上传字节数或 PDF 页数上限；`usage.limits.max_file_size_bytes` 与 `max_pages_per_file` 为 `null`。 +每任务最多 100 文件。URL 下载沿用 Jobs 的配置：默认只允许 HTTPS，HTTP 需启用 `--allow-http-source`，超时由 `--url-timeout` 控制（默认 60 秒），大小由 `--max-url-bytes` 控制（默认 200 MiB）。这些预算不要求客户端提前声明 `bytes`。 +既有 Jobs 的 inline 来源预算仍按其专用配置生效。 源文件、上传索引和任务记录仍是进程内状态，重启不保证恢复；完成后的任务及产物沿用部署的保留策略。 Router 的 tasks 与 V1 Uploads 复用按调用方隔离的源字节缓存。 -全部客户端文件齐备后才选择并预占 worker，再向该 worker 提交 tasks 和传输其缺失文件。 +全部客户端上传项齐备后才选择并预占 worker，再向该 worker 提交 tasks 和传输其缺失文件。 +URL 原样交给支持该来源的 worker 下载，Router 不下载源 URL；URL 下载策略由 worker 配置决定。 worker 负责最终 ZIP，Router 只查询任务与 File 元数据并映射公开下载地址。 同负载轮询，缓存命中不把任务固定到繁忙 worker；一个批量任务仍交给同一个 worker。 失败、取消、终态和后台回收通过同一机制释放预占，等待上传不计入执行负载。 @@ -195,5 +225,5 @@ worker 负责最终 ZIP，Router 只查询任务与 File 元数据并映射公 旧 `/v1/file_parse` 和 `/v1/tasks/{task_id}/result` 已移除；不提供同步等待 HTTP 路由或结果正文路由。 客户端改为轮询 `status_url` 并下载 `result_url`，删除旧 `MODE=sync`、`RESPONSE_FORMAT`、`wait_timeout` 配置。 -旧 tasks JSON 来源、原始字节创建请求和 POST 查询参数不再接受；URL、inline、local、file_id 等资源来源继续使用既有 `/v1/parse/jobs` 协议。 +旧 tasks JSON 来源、原始字节创建请求和 POST 查询参数不再接受；URL 可通过 tasks 表单 manifest 提交。inline、local、file_id 等资源来源继续使用既有 `/v1/parse/jobs` 协议。 独立 `/v1/uploads` 与 `/v1/parse/jobs` 的 JSON 请求协议保持不变。 diff --git a/mineru/kit/router/tasks.py b/mineru/kit/router/tasks.py index 6ba263ca..ed776419 100644 --- a/mineru/kit/router/tasks.py +++ b/mineru/kit/router/tasks.py @@ -20,7 +20,7 @@ from fastapi.responses import JSONResponse from pydantic import ValidationError from ...filetypes import is_flash_only_parse_extension -from ...parser.api_server import ApiServerError, ErrorDetail, _raise_api_error +from ...parser.api_server import ApiServerError, ErrorDetail, _JobInputMetadata, _raise_api_error from ...parser.task_api import ( TERMINAL_TASK_STATUSES, TaskExecution, @@ -142,7 +142,7 @@ def _lost_execution(route: ResourceRoute, pool: WorkerPool) -> TaskExecution: file.update(status="failed", error=error.model_dump()) route.metadata.update(payload=payload, upstream_lost=True) release_job_load(route, pool) - return TaskExecution(payload, error=error) + return TaskExecution(payload, error=error, input_metadata=dict(route.metadata.get("task_input_metadata") or {})) async def read_task_execution( @@ -214,11 +214,17 @@ async def read_task_execution( ) route.metadata["task_result_file_id"] = result_file_id error = ErrorDetail.model_validate(payload["error"]) if payload.get("error") else None - return TaskExecution(rewritten, result_file_id, error) + input_metadata = { + index: _JobInputMetadata(item["bytes"], item["sha256sum"]) + for index, item in enumerate(payload["files"]) + if item.get("url") is not None and item["bytes"] is not None and item["sha256sum"] is not None + } + route.metadata["task_input_metadata"] = input_metadata + return TaskExecution(rewritten, result_file_id, error, input_metadata) class RouterTaskBackend: - """输入全部就绪后才选择 worker，并通过同一 tasks 协议提交和上传缺失文件。""" + """客户端上传项就绪后才选择 worker；URL 原样交由 worker 下载。""" def __init__(self, request: Request, owner_scope: str, finalize: FinalizeJob) -> None: """保存轻量上下文与显式收尾回调，不跨请求保存上传表单。""" @@ -239,7 +245,7 @@ class RouterTaskBackend: selection_tier = "flash" if all(lightweight) and options.tier not in tiers else options.tier candidates = self.pool.capable_workers( tier=selection_tier, - required_sources={"file_id"}, + required_sources={"url" if item.url is not None else "file_id" for item in manifest}, required_outputs=set(options.output_formats), require_flash=any(lightweight), ) @@ -293,7 +299,7 @@ class RouterTaskBackend: tier = "flash" if lightweight and options.tier not in self.pool.available_tiers() else options.tier reservation = self.pool.reserve_job( tier=tier, - required_sources={"file_id"}, + required_sources={"url" if item.url is not None else "file_id" for item in manifest}, required_outputs=set(options.output_formats), require_flash=any(is_flash_only_parse_extension(item.filename) for item in manifest), ) @@ -308,7 +314,10 @@ class RouterTaskBackend: transferred = False try: worker_manifest = [ - {**item.model_dump(), "sha256sum": uploads[index].sha256sum} for index, item in enumerate(manifest) + item.model_dump(exclude_none=True) + if item.url is not None + else {**item.model_dump(exclude_none=True), "sha256sum": uploads[index].sha256sum} + for index, item in enumerate(manifest) ] form = [("manifest", json.dumps(worker_manifest)), ("tier", options.tier), ("ocr_mode", options.ocr_mode)] form.extend(("output_formats", value) for value in options.output_formats) @@ -342,7 +351,7 @@ class RouterTaskBackend: by_client_id = {item.client_file_id: index for index, item in enumerate(manifest)} for link in payload["uploads"]: index = by_client_id.get(link["client_file_id"]) - if index is None or link["upload_method"] != "PUT": + if index not in uploads or link["upload_method"] != "PUT": _raise_api_error( 502, error_type="api_error", code="invalid_upstream_response", message="Unexpected worker upload" ) @@ -374,8 +383,18 @@ class RouterTaskBackend: ) aliases: dict[str, str] = {} for index, item in enumerate(ready["files"]): + if item["client_file_id"] != manifest[index].client_file_id: + _raise_api_error( + 502, error_type="api_error", code="invalid_upstream_response", message="Worker input order changed" + ) + if manifest[index].url is not None: + if item["upload_status"] != "not_required": + _raise_api_error( + 502, error_type="api_error", code="invalid_upstream_response", message="Unexpected URL upload" + ) + continue upstream_file_id = item.get("file_id") - if not isinstance(upstream_file_id, str) or item["client_file_id"] != manifest[index].client_file_id: + if not isinstance(upstream_file_id, str): _raise_api_error( 502, error_type="api_error", code="invalid_upstream_response", message="Missing worker source File" ) diff --git a/mineru/parser/api_server.py b/mineru/parser/api_server.py index 71e9e857..b5779352 100644 --- a/mineru/parser/api_server.py +++ b/mineru/parser/api_server.py @@ -1361,6 +1361,23 @@ def _validate_job_source_policy( ) +@dataclass(frozen=True) +class _JobInputSpec: + """仅在服务内部传递 tasks 的逻辑文件名及可选完整性声明。""" + + filename: str + bytes: int | None = None + sha256sum: str | None = None + + +@dataclass(frozen=True) +class _JobInputMetadata: + """记录完整下载且校验通过的实际输入信息，不扩展公开 Jobs 协议。""" + + bytes: int + sha256sum: str + + @dataclass class _JobRecord: id: str @@ -1374,6 +1391,8 @@ class _JobRecord: files: list[JobFileResult] = None # type: ignore[assignment] links: JobLinks | None = None input_blobs: dict[str, str] = field(default_factory=dict) + input_specs: dict[int, _JobInputSpec] = field(default_factory=dict) + input_metadata: dict[int, _JobInputMetadata] = field(default_factory=dict) file_store: FileStore | None = None usage_counted: bool = False task_result_file_id: str | None = None @@ -1506,7 +1525,14 @@ class JobStore: def _now() -> str: return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") - def create(self, req: CreateJobRequest, file_store: FileStore, *, job_id: str | None = None) -> _JobRecord: + def create( + self, + req: CreateJobRequest, + file_store: FileStore, + *, + job_id: str | None = None, + input_specs: dict[int, _JobInputSpec] | None = None, + ) -> _JobRecord: """创建唯一执行记录，允许 tasks 使用上传准备阶段预先分配的 ID。""" if req.tier is None: raise ValueError("CreateJobRequest.tier must be resolved before job creation") @@ -1520,13 +1546,14 @@ class JobStore: created_at=now, tier=req.tier, output_formats=req.output_formats, + input_specs=dict(input_specs or {}), links=JobLinks( self=f"/v1/parse/jobs/{job_id}", cancel=f"/v1/parse/jobs/{job_id}", ), ) - for entry in req.files: - name = _source_name(entry.source, file_store) + for index, entry in enumerate(req.files): + name = _source_name(entry.source, file_store, input_spec=rec.input_specs.get(index)) rec.files.append( JobFileResult( file_id=entry.source.file_id if isinstance(entry.source, FileIdSource) else None, @@ -1676,7 +1703,10 @@ class JobStore: # ── source helpers ─────────────────────────────────────────────────── -def _source_name(source: FileSource, file_store: FileStore | None = None) -> str: +def _source_name(source: FileSource, file_store: FileStore | None = None, *, input_spec: _JobInputSpec | None = None) -> str: + """优先使用 tasks 的逻辑文件名，普通 Jobs 保持既有来源命名规则。""" + if input_spec is not None: + return input_spec.filename if isinstance(source, FileIdSource) and file_store is not None: try: return file_store.get_file(source.file_id).filename @@ -1692,17 +1722,25 @@ def _source_name(source: FileSource, file_store: FileStore | None = None) -> str return "unknown" -def _job_can_default_to_flash(req: CreateJobRequest, file_store: FileStore) -> bool: +def _job_can_default_to_flash( + req: CreateJobRequest, file_store: FileStore, input_specs: dict[int, _JobInputSpec] | None = None +) -> bool: """判断整批输入是否已知为 Flash-only，或需下载后识别类型的无扩展名 URL。""" return all( - is_flash_only_parse_extension(_source_name(entry.source, file_store)) - or (isinstance(entry.source, UrlSource) and not pathlib.Path(_source_name(entry.source)).suffix) - for entry in req.files + is_flash_only_parse_extension(name) or (isinstance(entry.source, UrlSource) and not pathlib.Path(name).suffix) + for index, entry in enumerate(req.files) + for name in [_source_name(entry.source, file_store, input_spec=(input_specs or {}).get(index))] ) -def _job_has_any_flash_only_inputs(req: CreateJobRequest, file_store: FileStore) -> bool: - return any(is_flash_only_parse_extension(_source_name(entry.source, file_store)) for entry in req.files) +def _job_has_any_flash_only_inputs( + req: CreateJobRequest, file_store: FileStore, input_specs: dict[int, _JobInputSpec] | None = None +) -> bool: + """使用与执行一致的逻辑文件名检查批次中是否包含 Flash-only 格式。""" + return any( + is_flash_only_parse_extension(_source_name(entry.source, file_store, input_spec=(input_specs or {}).get(index))) + for index, entry in enumerate(req.files) + ) def _page_range_from_result_pages(pages: list[PageInfo]) -> str: @@ -1812,6 +1850,7 @@ async def _run_job( fr = rec.files[i] try: file_started = time.monotonic() + fr.status = "running" extracted = await _extract_bytes( entry.source, file_store, @@ -1823,6 +1862,16 @@ async def _run_job( input_blobs=rec.input_blobs, ) data = extracted.data + input_spec = rec.input_specs.get(i) + if input_spec is not None and isinstance(entry.source, UrlSource): + if not data: + raise MineruError("invalid_request", "Empty files are not supported") + if input_spec.bytes is not None and len(data) != input_spec.bytes: + raise MineruError("bytes_mismatch", "Declared size does not match downloaded bytes") + sha = hashlib.sha256(data).hexdigest() + if input_spec.sha256sum is not None and sha != input_spec.sha256sum: + raise MineruError("file_hash_mismatch", "SHA-256 mismatch") + rec.input_metadata[i] = _JobInputMetadata(bytes=len(data), sha256sum=sha) suffix = pathlib.Path(fr.name).suffix stype = _suffix_type(fr.name) if not stype and not suffix and isinstance(entry.source, UrlSource) and extracted.media_type == "text/html": @@ -1967,7 +2016,9 @@ Parse documents into structured Middle JSON via self-hosted or remote capacity. Submit a form manifest to `/v1/tasks`, upload only missing files to their returned URLs, then poll the task and download its `result_url`. Local servers also accept -one multipart file. Task status contains metadata only, never inline result bytes. +one multipart file. Manifest entries can instead provide a source `url`, with +optional `bytes`; execution workers download these inputs in the background. +Task status contains metadata only, never inline result bytes. The resource workflow remains: create an upload, complete it, submit a parse job for the resulting file, then poll the job and fetch its outputs. Health and usage endpoints report server readiness and quota state. Tiers (`flash`, `basic`, @@ -2283,6 +2334,7 @@ async def submit_parse_job( *, job_id: str | None = None, finalize: Callable[[_JobRecord], Awaitable[None]] | None = None, + input_specs: dict[int, _JobInputSpec] | None = None, ) -> _JobRecord: """统一校验并提交后台解析任务，供 V1 Job 与便捷入口共同调用。""" for index, entry in enumerate(body.files): @@ -2320,8 +2372,8 @@ async def submit_parse_job( # 按请求 tier 选择启动时预先解析好的 runtime；全轻量输入可按 ADR-0024 直接归一到 flash 执行。 default_tier: Tier | None = request.app.state.default_tier - can_default_to_flash = _job_can_default_to_flash(body, file_store) - if not request.app.state.flash_enabled and _job_has_any_flash_only_inputs(body, file_store): + can_default_to_flash = _job_can_default_to_flash(body, file_store, input_specs) + if not request.app.state.flash_enabled and _job_has_any_flash_only_inputs(body, file_store, input_specs): _raise_api_error( 400, error_type="invalid_request_error", @@ -2365,7 +2417,7 @@ async def submit_parse_job( if job_store._closing: _raise_api_error(503, error_type="engine_error", code="server_shutting_down", message="Server is shutting down") - rec = job_store.create(body, file_store, job_id=job_id) + rec = job_store.create(body, file_store, job_id=job_id, input_specs=input_specs) url_timeout_val: int = request.app.state.url_timeout allow_local_source_val: bool = request.app.state.allow_local_source max_inline_bytes_val: int = request.app.state.max_inline_bytes diff --git a/mineru/parser/task_api.py b/mineru/parser/task_api.py index 798f3bd7..1c76efa4 100644 --- a/mineru/parser/task_api.py +++ b/mineru/parser/task_api.py @@ -21,7 +21,7 @@ from urllib.parse import quote from fastapi import APIRouter, Request, Response from fastapi.responses import JSONResponse -from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field, ValidationError, model_validator from starlette.datastructures import UploadFile from starlette.exceptions import HTTPException as StarletteHTTPException from starlette.requests import ClientDisconnect @@ -42,9 +42,13 @@ from .api_server import ( JobProgress, JobStore, OutputFiles, + UrlSource, + _JobInputMetadata, + _JobInputSpec, _JobRecord, _raise_api_error, _require_model_preload_ready, + _validate_source_policy, submit_parse_job, ) from .page_range import normalize_page_range_input @@ -57,15 +61,26 @@ logger = logging.getLogger(__name__) class TaskManifestItem(BaseModel): - """按客户端标识关联独立上传，不依赖文件名唯一性。""" + """按客户端标识关联上传或 URL 输入，不依赖文件名唯一性。""" model_config = ConfigDict(extra="forbid") client_file_id: str = Field(pattern=r"^[A-Za-z0-9_-]{1,64}$") filename: str = Field(min_length=1, max_length=1024) - bytes: int = Field(gt=0, strict=True) + bytes: int | None = Field(default=None, gt=0, strict=True) + url: str | None = None sha256sum: str | None = Field(default=None, pattern=r"^[a-f0-9]{64}$") page_range: str | None = None + @model_validator(mode="after") + def validate_source(self) -> TaskManifestItem: + """上传项必须声明大小；URL 只校验合法性，保留签名地址的原始字符串。""" + if self.url is None: + if self.bytes is None: + raise ValueError("bytes is required when url is not supplied") + else: + AnyHttpUrl(self.url) + return self + class UploadTaskOptions(BaseModel): """解析选项仅来自表单，文件元数据独立保存在 manifest 中。""" @@ -92,7 +107,7 @@ class TaskUploadLink(BaseModel): class TaskFileStatus(TaskManifestItem): """同时提供上传准备情况与逐文件解析状态，禁止内联解析正文。""" - upload_status: Literal["pending", "ready"] + upload_status: Literal["pending", "ready", "not_required"] file_id: str | None = None name: str status: Literal["waiting_upload", "queued", "running", "completed", "failed", "canceled"] @@ -133,11 +148,12 @@ class TaskUpload: @dataclass(frozen=True) class TaskExecution: - """在后台适配器之间传递执行快照及本部署的 ZIP File 标识。""" + """在后台适配器之间传递执行快照、实际输入元数据及本部署的 ZIP File 标识。""" job: dict[str, Any] result_file_id: str | None = None error: ErrorDetail | None = None + input_metadata: dict[int, _JobInputMetadata] = field(default_factory=dict) class TaskBackend(Protocol): @@ -235,13 +251,13 @@ async def _save_upload( if canceled is not None and canceled(): _invalid("Task was canceled during upload", status_code=409, code="task_already_terminal") size += len(chunk) - if expected is not None and size > expected.bytes: + if expected is not None and expected.bytes is not None and size > expected.bytes: _invalid("Declared size does not match uploaded bytes", code="upload_size_mismatch") await run_sync(target.write, chunk) hasher.update(chunk) if size == 0: _invalid("Empty files are not supported", param="files") - if expected is not None and size != expected.bytes: + if expected is not None and expected.bytes is not None and size != expected.bytes: _invalid("Declared size does not match uploaded bytes", code="upload_size_mismatch") sha = hasher.hexdigest() if expected is not None and expected.sha256sum is not None and sha != expected.sha256sum: @@ -255,7 +271,7 @@ async def _save_upload( async def _read_submission( request: Request, directory: Path ) -> tuple[UploadTaskOptions, list[TaskManifestItem], TaskUpload | None]: - """解析表单清单或单文件直传，旧 JSON、原始字节及 URL 业务参数明确拒绝。""" + """解析含上传或 URL 来源的表单清单；拒绝旧 JSON、原始字节及查询参数。""" if request.query_params: _invalid("POST parameters must be supplied in the form body, not the URL") media_type = request.headers.get("content-type", "").split(";", 1)[0].strip().lower() @@ -288,6 +304,8 @@ async def _read_submission( if files: if len(files) != 1 or (manifest and len(manifest) != 1): _invalid("Multipart direct upload supports one file; submit a manifest for multiple files", param="files") + if manifest and manifest[0].url is not None: + _invalid("url and multipart file bytes cannot be supplied together", param="files") file = files[0] if not isinstance(file, UploadFile) or not file.filename: _invalid("files must contain one named file", param="files") @@ -410,7 +428,7 @@ class TaskStore: async def create( self, backend: TaskBackend, options: UploadTaskOptions, manifest: list[TaskManifestItem], direct: TaskUpload | None ) -> _TaskRecord: - """校验并固定已命中源文件；只有全部就绪才提交后台执行。""" + """校验并固定已命中源文件；客户端上传项齐备后提交后台执行。""" if self._closing: _invalid("Server is shutting down", status_code=503, code="server_shutting_down") await backend.validate(options, manifest) @@ -418,6 +436,8 @@ class TaskStore: record = _TaskRecord(task_id, backend, options, manifest, self._root / task_id, int(time.time()) + TASK_UPLOAD_TTL) try: for index, item in enumerate(manifest): + if item.url is not None: + continue upload = ( direct if direct is not None @@ -437,7 +457,9 @@ class TaskStore: def _maybe_start(self, record: _TaskRecord) -> None: """事件循环内原子切换准备状态，最后多个上传并发完成也只创建一个执行句柄。""" - if record.status == "waiting_upload" and len(record.uploads) == len(record.manifest): + if record.status == "waiting_upload" and all( + item.url is not None or index in record.uploads for index, item in enumerate(record.manifest) + ): record.status = "queued" record.operation = asyncio.create_task(self._dispatch(record), name=f"mineru-task-{record.id}") @@ -523,6 +545,8 @@ class TaskStore: index = next((index for index, item in enumerate(record.manifest) if item.client_file_id == client_file_id), None) if index is None: _invalid("Unknown client_file_id", status_code=404, code="upload_not_found") + if record.manifest[index].url is not None: + _invalid("URL inputs do not accept uploads", status_code=400, code="invalid_request") lock = record.upload_locks.setdefault(index, asyncio.Lock()) async with lock: if record.status in {"failed", "canceled"}: @@ -620,16 +644,18 @@ def task_payload(record: _TaskRecord, request: Request) -> dict[str, Any]: file = item.model_dump() file.update( name=item.filename, - upload_status="ready" if upload is not None else "pending", - status="queued" if upload else "waiting_upload", + upload_status="not_required" if item.url is not None else ("ready" if upload is not None else "pending"), + status="queued" if upload is not None or item.url is not None else "waiting_upload", file_id=upload.file_id if upload else None, ) if index < len(results): file.update(results[index]) + if record.execution is not None and (metadata := record.execution.input_metadata.get(index)) is not None: + file.update(bytes=metadata.bytes, sha256sum=metadata.sha256sum) if record.status == "canceled" and file["status"] not in {"completed", "failed"}: file["status"] = "canceled" files.append(file) - if upload is None and record.status == "waiting_upload": + if item.url is None and upload is None and record.status == "waiting_upload": links.append( { "client_file_id": item.client_file_id, @@ -666,7 +692,11 @@ def _request_schema() -> dict[str, Any]: # 清单使用 JSON 文本表单字段，客户端无需发送 application/json 请求。 properties["manifest"] = { "type": "string", - "description": "JSON array of 1-100 files: client_file_id, filename, bytes, optional sha256sum and page_range.", + "description": ( + "JSON array of 1-100 files: client_file_id, filename, optional url, sha256sum and page_range. " + "bytes must be a positive integer for uploads; it is optional/null for URL inputs. " + "URL inputs have upload_status=not_required and never receive upload links." + ), "example": '[{"client_file_id":"f1","filename":"document.pdf","bytes":12345}]', } plain = {"type": "object", "required": ["manifest"], "additionalProperties": False, "properties": properties} @@ -678,7 +708,7 @@ def _request_schema() -> dict[str, Any]: "minItems": 1, "maxItems": 1, "items": {"type": "string", "format": "binary"}, - "description": "Local single-file direct upload only; a supplied manifest must also contain exactly one file.", + "description": ("Local single-file direct upload only; a supplied manifest must contain exactly one file without url."), } return { "requestBody": { @@ -756,7 +786,12 @@ def _combine_archives(store: FileStore, job: _JobRecord, path: Path, manifest: l # ZIP 可独立消费；不把 worker 的私有 task/file ID 或短期下载地址写入产物。 "files": [ { - **item.model_dump(), + **item.model_dump(exclude={"url"}), + **( + {"bytes": job.input_metadata[index].bytes, "sha256sum": job.input_metadata[index].sha256sum} + if index in job.input_metadata + else {} + ), "name": result.name, "page_range": result.page_range, "status": result.status, @@ -811,6 +846,17 @@ class ApiTaskBackend: if self.jobs._closing: _invalid("Server is shutting down", code="server_shutting_down", status_code=503) state = self.request.app.state + for index, item in enumerate(manifest): + if item.url is not None: + try: + _validate_source_policy( + UrlSource(url=item.url), + allow_local_source=False, + max_inline_bytes=state.max_inline_bytes, + allow_http_source=state.allow_http_source, + ) + except ValueError as exc: + _invalid(str(exc), param=f"manifest.{index}.url") options.tier = resolve_task_tier(options, manifest, set(state.tier_runtime_options), state.default_tier) async def find_upload(self, sha256sum: str, filename: str) -> TaskUpload | None: @@ -850,7 +896,7 @@ class ApiTaskBackend: async def submit( self, task_id: str, options: UploadTaskOptions, manifest: list[TaskManifestItem], uploads: dict[int, TaskUpload] ) -> TaskExecution: - """把已经校验的输入交给现有队列，在发布成功前完成最终 ZIP。""" + """把已上传文件和 URL 来源交给现有队列，在发布成功前完成最终 ZIP。""" # 接单时已持有源字节；用户删除文件视图也不能破坏尚在等待其他文件的任务。 for upload in uploads.values(): await run_sync( @@ -858,7 +904,10 @@ class ApiTaskBackend: ) body = CreateJobRequest( files=[ - JobFileEntry(source=FileIdSource(file_id=uploads[index].file_id), page_range=item.page_range or None) + JobFileEntry( + source=UrlSource(url=item.url) if item.url is not None else FileIdSource(file_id=uploads[index].file_id), + page_range=item.page_range or None, + ) for index, item in enumerate(manifest) ], tier=options.tier, @@ -893,7 +942,18 @@ class ApiTaskBackend: job.status = "failed" logger.exception("Task ZIP creation failed: %s", task_id) - await submit_parse_job(body, self.request, self.files, self.jobs, self.access, job_id=task_id, finalize=finalize) + await submit_parse_job( + body, + self.request, + self.files, + self.jobs, + self.access, + job_id=task_id, + finalize=finalize, + input_specs={ + index: _JobInputSpec(item.filename, item.bytes, item.sha256sum) for index, item in enumerate(manifest) + }, + ) return await self.get(task_id) async def get(self, task_id: str) -> TaskExecution: @@ -903,6 +963,7 @@ class ApiTaskBackend: self.jobs.build_response(record, access_level=self.access).model_dump(by_alias=True), record.task_result_file_id, record.task_error, + dict(record.input_metadata), ) async def cancel(self, task_id: str) -> None: diff --git a/tests/unittest/test_task_api.py b/tests/unittest/test_task_api.py index 5b5b437a..8f40de0f 100644 --- a/tests/unittest/test_task_api.py +++ b/tests/unittest/test_task_api.py @@ -9,8 +9,9 @@ import json import threading import time import zipfile -from collections.abc import Iterator +from collections.abc import AsyncIterator, Iterator from pathlib import Path +from types import SimpleNamespace from typing import Any import httpx @@ -26,6 +27,45 @@ HTML = b"

Shared tasks

Preserved content.Second document

Second content.

" +@pytest.fixture +def url_downloads(monkeypatch: pytest.MonkeyPatch) -> list[httpx.Request]: + """仅替换 worker 的外部下载传输，Router 和客户端仍经过真实 ASGI 协议。""" + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + """保留完整签名 URL 和请求头，并提供下载成功及失败来源。""" + requests.append(request) + if request.url.path == "/timeout": + raise httpx.ReadTimeout("source timed out", request=request) + return httpx.Response( + 404 if request.url.path == "/missing" else 200, + content=b"" if request.url.path == "/empty" else HTML, + headers={"content-type": "text/html; charset=utf-8"}, + ) + + def download_client(**kwargs: Any) -> httpx.AsyncClient: + """使用真实 httpx 流式客户端，隔离源站网络而不替换下载业务实现。""" + return httpx.AsyncClient(transport=httpx.MockTransport(respond), **kwargs) + + monkeypatch.setattr(api_server, "httpx", SimpleNamespace(AsyncClient=download_client)) + return requests + + +def _execution_apps(app: FastAPI) -> list[FastAPI]: + """取得真正执行 URL 下载的服务，用于验证部署配置及取消后的资源释放。""" + return list(app.state.test_transport.apps.values()) if hasattr(app.state, "test_transport") else [app] + + +def _url_manifest(**metadata: Any) -> dict[str, Any]: + """构造与下载地址名称无关的逻辑文件条目，默认不声明大小或 hash。""" + return { + "client_file_id": "remote", + "filename": "logical.html", + "url": "https://source.example/opaque?signature=a%2Fb&part=1", + **metadata, + } + + class _WorkerTransport(httpx.AsyncBaseTransport): """使用真实 worker 路由验证源文件权限和任务生命周期，不以宽松 fake 替代。""" @@ -476,6 +516,265 @@ def test_input_validation_and_openapi(service: tuple[TestClient, FastAPI]) -> No assert limits["max_files_per_job"] == 100 +@pytest.mark.parametrize( + "metadata", + [ + {}, + {"bytes": None}, + {"bytes": len(HTML)}, + {"sha256sum": hashlib.sha256(HTML).hexdigest()}, + {"bytes": len(HTML), "sha256sum": hashlib.sha256(HTML).hexdigest()}, + {"url": "https://source.example/opaque?signature=" + "a" * 3000}, + ], +) +def test_url_inputs_download_in_worker_and_report_actual_metadata( + service: tuple[TestClient, FastAPI], url_downloads: list[httpx.Request], metadata: dict[str, Any] +) -> None: + """无扩展名签名 URL 使用逻辑文件名，未知大小在下载后补齐且无需上传。""" + client, app = service + item = _url_manifest(**metadata) + if hasattr(app.state, "worker_pool"): + app.state.worker_pool.workers[0].features["sources"] = ["file_id"] + response = client.post("/v1/tasks", data={"manifest": json.dumps([item])}) + assert response.status_code == 202, response.text + task = response.json() + assert task["status"] == "queued" and task["uploads"] == [] + assert task["files"][0]["bytes"] == metadata.get("bytes") + assert task["files"][0]["upload_status"] == "not_required" + result = _terminal(client, task) + assert result["status"] == "completed" + file = result["files"][0] + assert file["filename"] == file["name"] == "logical.html" + assert file["bytes"] == len(HTML) and file["sha256sum"] == hashlib.sha256(HTML).hexdigest() + assert file["upload_status"] == "not_required" and file["file_id"] is None + assert file["url"] == item["url"] + assert b"Preserved content" in _archive(client, result)["markdown.md"] + assert len(url_downloads) == 1 and str(url_downloads[0].url) == item["url"] + assert url_downloads[0].method == "GET" and "authorization" not in url_downloads[0].headers + assert client.put(f"/v1/tasks/{task['task_id']}/uploads/remote", content=HTML).status_code == 400 + if hasattr(app.state, "test_transport"): + assert not any(method == "PUT" for _, method, _ in app.state.test_transport.calls) + assert all(host == "worker-b" for host, method, path in app.state.test_transport.calls if path == "/v1/tasks") + + +def test_url_manifest_validation(service: tuple[TestClient, FastAPI], url_downloads: list[httpx.Request]) -> None: + """大小只对 URL 可选，非法 URL、非正整数字节数及混合直传在接单前拒绝。""" + client, app = service + invalid = [ + {"client_file_id": "f1", "filename": "a.html"}, + {"client_file_id": "f1", "filename": "a.html", "bytes": None}, + *[_url_manifest(url=value) for value in ("", "file:///tmp/a.html", "https://", "ftp://source/a.html")], + *[_url_manifest(bytes=value) for value in (0, -1, True, "10", 10.5)], + ] + for item in invalid: + response = client.post("/v1/tasks", data={"manifest": json.dumps([item])}) + assert response.status_code == 400, response.text + response = client.post("/v1/tasks", data={"manifest": json.dumps([_url_manifest()])}, files={"files": ("a.html", HTML)}) + assert response.status_code == 400 + assert _job_count(app) == 0 and url_downloads == [] + schema = client.get("/openapi.json").json() + assert "not_required" in schema["components"]["schemas"]["TaskFileStatus"]["properties"]["upload_status"]["enum"] + + +def test_mixed_url_upload_partial_and_no_early_reservation( + service: tuple[TestClient, FastAPI], url_downloads: list[httpx.Request] +) -> None: + """混合批次仅等待本地上传，URL 逐项失败不丢失同名成功项或目录顺序。""" + client, app = service + manifest = [ + _url_manifest(filename="same.html"), + {"client_file_id": "local", "filename": "same.html", "bytes": len(SECOND)}, + _url_manifest(client_file_id="bad", url="https://source.example/missing"), + ] + response = client.post("/v1/tasks", data={"manifest": json.dumps(manifest)}) + assert response.status_code == 202, response.text + task = response.json() + assert task["status"] == "waiting_upload" + assert [link["client_file_id"] for link in task["uploads"]] == ["local"] + assert [file["upload_status"] for file in task["files"]] == ["not_required", "pending", "not_required"] + assert _job_count(app) == 0 and url_downloads == [] + if hasattr(app.state, "worker_pool"): + assert all(worker.active_jobs == 0 for worker in app.state.worker_pool.workers) + assert client.put(task["uploads"][0]["upload_url"], content=SECOND).status_code == 200 + result = _terminal(client, task) + assert result["status"] == "partial" + assert result["progress"] == {"completed": 2, "failed": 1, "total": 3} + assert result["files"][0]["bytes"] == len(HTML) + assert result["files"][2]["error"] and result["files"][2]["bytes"] is None + members = _archive(client, result) + assert b"Preserved content" in members["0001/markdown.md"] + assert b"Second content" in members["0002/markdown.md"] + assert not any(name.startswith("0003/") for name in members) + archived_inputs = json.loads(members["manifest.json"])["files"] + assert [item["client_file_id"] for item in archived_inputs] == ["remote", "local", "bad"] + assert archived_inputs[0]["bytes"] == len(HTML) + assert archived_inputs[0]["sha256sum"] == hashlib.sha256(HTML).hexdigest() + assert all("url" not in item for item in archived_inputs) + + +@pytest.mark.parametrize( + ("metadata", "code"), + [ + ({"bytes": len(HTML) + 1}, "bytes_mismatch"), + ({"sha256sum": "0" * 64}, "file_hash_mismatch"), + ({"url": "https://source.example/empty"}, "invalid_request"), + ({"url": "https://source.example/timeout"}, "parse_failed"), + ], +) +def test_url_download_failures_are_per_file( + service: tuple[TestClient, FastAPI], url_downloads: list[httpx.Request], metadata: dict[str, Any], code: str +) -> None: + """下载及完整性错误进入文件错误；全失败不生成结果链接或实际元数据。""" + client, _ = service + response = client.post("/v1/tasks", data={"manifest": json.dumps([_url_manifest(**metadata)])}) + assert response.status_code == 202 + result = _terminal(client, response.json()) + assert result["status"] == "failed" and result["result_url"] is None + assert result["files"][0]["status"] == "failed" and result["files"][0]["error"]["code"] == code + assert result["files"][0]["bytes"] == metadata.get("bytes") + assert len(url_downloads) == 1 + + +def test_url_hash_does_not_skip_download(service: tuple[TestClient, FastAPI], url_downloads: list[httpx.Request]) -> None: + """即使相同 hash 已在上传缓存中，URL 来源仍每次重新读取并校验。""" + client, _ = service + inputs = [("cached.html", HTML)] + cached = _submit(client, inputs) + _upload(client, cached, inputs) + assert _terminal(client, cached)["status"] == "completed" + for _ in range(2): + response = client.post( + "/v1/tasks", data={"manifest": json.dumps([_url_manifest(sha256sum=hashlib.sha256(HTML).hexdigest())])} + ) + assert response.status_code == 202 + assert _terminal(client, response.json())["status"] == "completed" + assert len(url_downloads) == 2 + + +def test_url_download_obeys_worker_policy_and_limit( + service: tuple[TestClient, FastAPI], url_downloads: list[httpx.Request] +) -> None: + """HTTP 开关和 URL 下载预算继续由真正执行下载的服务控制。""" + client, app = service + item = _url_manifest(url="http://source.example/opaque") + response = client.post("/v1/tasks", data={"manifest": json.dumps([item])}) + if hasattr(app.state, "worker_pool"): + assert response.status_code == 202 + assert _terminal(client, response.json())["status"] == "failed" + else: + assert response.status_code == 400 + assert url_downloads == [] + for worker in _execution_apps(app): + worker.state.allow_http_source = True + worker.state.max_url_bytes = len(HTML) - 1 + worker.state.url_timeout = 7 + response = client.post("/v1/tasks", data={"manifest": json.dumps([item])}) + result = _terminal(client, response.json()) + assert result["status"] == "failed" + assert "max_url_bytes" in result["files"][0]["error"]["message"] + assert url_downloads[0].extensions["timeout"]["read"] == 7 + for worker in _execution_apps(app): + worker.state.max_url_bytes = len(HTML) + response = client.post("/v1/tasks", data={"manifest": json.dumps([item])}) + assert _terminal(client, response.json())["status"] == "completed" + + +def test_url_html_preserves_source_context( + service: tuple[TestClient, FastAPI], url_downloads: list[httpx.Request], monkeypatch: pytest.MonkeyPatch +) -> None: + """URL HTML 解析继续使用源站地址和编码，而不是临时路径作为资源基址。""" + client, _ = service + contexts: list[Any] = [] + original = api_server.parse_async + + async def inspect_context(*args: Any, **kwargs: Any) -> api_server.ParseResult: + """观察真实解析调用的来源上下文，再继续执行原解析器生成 ZIP。""" + contexts.append(kwargs.get("source_context")) + return await original(*args, **kwargs) + + monkeypatch.setattr(api_server, "parse_async", inspect_context) + response = client.post("/v1/tasks", data={"manifest": json.dumps([_url_manifest()])}) + assert _terminal(client, response.json())["status"] == "completed" + assert len(contexts) == 1 and len(url_downloads) == 1 + assert contexts[0].source_uri == _url_manifest()["url"] + assert contexts[0].transport_encoding == "utf-8" + assert contexts[0].local_resource_root is None + + +def test_url_logical_filename_drives_flash_fallback( + service: tuple[TestClient, FastAPI], url_downloads: list[httpx.Request] +) -> None: + """逻辑 HTML 文件保留轻量格式回退，不能被 URL 路径的 PDF 后缀误导。""" + client, _ = service + response = client.post( + "/v1/tasks", + data={ + "manifest": json.dumps([_url_manifest(url="https://source.example/document.pdf")]), + "tier": "standard", + }, + ) + assert response.status_code == 202, response.text + result = _terminal(client, response.json()) + assert result["status"] == "completed", result + assert result["files"][0]["name"] == "logical.html" + assert b"Preserved content" in _archive(client, result)["markdown.md"] + assert len(url_downloads) == 1 + + +def test_cancel_url_download_closes_stream_and_releases_worker( + service: tuple[TestClient, FastAPI], monkeypatch: pytest.MonkeyPatch +) -> None: + """下载中取消立即返回，并关闭源站流、释放 Job 与 Router 预占。""" + client, app = service + started, closed = threading.Event(), threading.Event() + + class SlowSource(httpx.AsyncByteStream): + """模拟已经收到部分字节但源站迟迟不结束的下载。""" + + async def __aiter__(self) -> AsyncIterator[bytes]: + """收到取消时退出等待，避免后台下载在任务取消后继续运行。""" + try: + yield HTML[:10] + started.set() + await asyncio.Future() + finally: + closed.set() + + async def aclose(self) -> None: + """记录 httpx 已关闭下载资源。""" + closed.set() + + def respond(request: httpx.Request) -> httpx.Response: + """返回实际异步流，使取消发生在读取源字节期间。""" + return httpx.Response(200, stream=SlowSource(), headers={"content-type": "text/html"}) + + def download_client(**kwargs: Any) -> httpx.AsyncClient: + """仅隔离 URL 来源网络，不改变任务和 Router 的生命周期实现。""" + return httpx.AsyncClient(transport=httpx.MockTransport(respond), **kwargs) + + monkeypatch.setattr(api_server, "httpx", SimpleNamespace(AsyncClient=download_client)) + response = client.post("/v1/tasks", data={"manifest": json.dumps([_url_manifest()])}) + assert response.status_code == 202 + task = response.json() + assert started.wait(5) + before = time.monotonic() + canceled = client.delete(task["status_url"]) + assert canceled.status_code == 200 and canceled.json()["status"] == "canceled" + assert time.monotonic() - before < 2 + assert closed.wait(5) + deadline = time.monotonic() + 5 + while any(worker.state.job_store._tasks for worker in _execution_apps(app)): + assert time.monotonic() < deadline, "Canceled URL job retained its execution slot" + time.sleep(0.01) + if hasattr(app.state, "worker_pool"): + while any(worker.active_jobs for worker in app.state.worker_pool.workers): + assert time.monotonic() < deadline, "Canceled URL task retained its worker reservation" + time.sleep(0.01) + final = client.get(task["status_url"]).json() + assert final["status"] == "canceled" and final["result_url"] is None + assert final["files"][0]["bytes"] is None + + def test_router_private_hash_cache_and_task_scope(tmp_path: Path) -> None: """Router 的缓存、任务与文件地址均隔离调用方，不能仅凭已知 hash 使用他人输入。""" workers = {"worker-a": api_server.create_app(upload_dir=str(tmp_path / "worker"), tier="flash")} diff --git a/tests/unittest/test_task_http.py b/tests/unittest/test_task_http.py index 70c1517e..7cc5b90c 100644 --- a/tests/unittest/test_task_http.py +++ b/tests/unittest/test_task_http.py @@ -13,7 +13,7 @@ import sys import threading import time import zipfile -from collections.abc import Iterator +from collections.abc import AsyncIterator, Iterator from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -22,7 +22,7 @@ from urllib.parse import urlsplit import httpx import pytest from fastapi import FastAPI, Request, Response -from fastapi.responses import JSONResponse, RedirectResponse +from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse from pypdf import PdfWriter from starlette.types import ASGIApp, Message, Receive, Scope, Send from test_http_api_example_script import _Server @@ -69,20 +69,25 @@ class _LiveService: url: str app: FastAPI counts: dict[str, int] + execution_apps: list[FastAPI] @pytest.fixture(params=["api", "router"]) def live_service(request: pytest.FixtureRequest, tmp_path: Path) -> Iterator[_LiveService]: """启动真实 Uvicorn 服务；Router 的两个 worker 同样使用真实 TCP。""" servers: list[_Server] = [] + execution_apps: list[FastAPI] = [] try: if request.param == "api": app = api_server.create_app(upload_dir=str(tmp_path / "api"), tier="flash", api_key="test-key") + execution_apps.append(app) else: for index in range(2): - worker = _Server( - api_server.create_app(upload_dir=str(tmp_path / f"worker-{index}"), tier="flash", api_key="test-key") + worker_app = api_server.create_app( + upload_dir=str(tmp_path / f"worker-{index}"), tier="flash", api_key="test-key" ) + execution_apps.append(worker_app) + worker = _Server(worker_app) worker.start() servers.append(worker) app = create_router( @@ -97,7 +102,7 @@ def live_service(request: pytest.FixtureRequest, tmp_path: Path) -> Iterator[_Li server = _Server(app) server.start() servers.append(server) - yield _LiveService(server.url, app, counts) + yield _LiveService(server.url, app, counts, execution_apps) finally: for server in reversed(servers): server.stop() @@ -397,6 +402,108 @@ def test_real_documents_match_jobs_outputs(live_service: _LiveService, document: assert _members(first.content) == _members(second.content) +@pytest.mark.parametrize("document", ["html", "pdf"]) +def test_url_task_over_http_downloads_chunked_source(live_service: _LiveService, document: str) -> None: + """真实 TCP URL 无需 HEAD 或长度头，保留文件页范围且不向源站发送 API Key。""" + if document == "pdf": + sample = Path(__file__).resolve().parents[2] / "examples" / "demo1.pdf" + if sample.is_file(): + data = sample.read_bytes() + else: + writer = PdfWriter() + for _ in range(3): + writer.add_blank_page(width=72, height=72) + buffer = io.BytesIO() + writer.write(buffer) + data = buffer.getvalue() + else: + data = HTML + requests: list[Request] = [] + source = FastAPI() + + @source.api_route("/opaque", methods=["GET", "HEAD"]) + async def download(request: Request) -> StreamingResponse: + """使用不带 Content-Length 的分块响应，验证只以真实收到的字节计算大小。""" + requests.append(request) + + async def chunks() -> AsyncIterator[bytes]: + """分段返回源文件，不借助上传入口或提前探测文件长度。""" + for offset in range(0, len(data), 8192): + yield data[offset : offset + 8192] + + return StreamingResponse(chunks(), media_type="application/pdf" if document == "pdf" else "text/html") + + source_server = _Server(source) + source_server.start() + try: + for worker in live_service.execution_apps: + worker.state.allow_http_source = True + item = { + "client_file_id": "url-1", + "filename": f"logical.{document}", + "url": f"{source_server.url}/opaque?signature=a%2Fb&part=1", + **({"page_range": "2"} if document == "pdf" else {}), + } + with _client(live_service) as client: + response = client.post( + "/v1/tasks", + data={ + "manifest": json.dumps([item]), + "tier": "flash", + "ocr_mode": "txt", + "page_range": "1" if document == "pdf" else "", + }, + ) + assert response.status_code == 202, response.text + task = response.json() + assert task["uploads"] == [] and task["files"][0]["upload_status"] == "not_required" + result = _wait(client, task) + assert result["status"] == "completed", result + assert result["files"][0]["bytes"] == len(data) + assert result["files"][0]["sha256sum"] == hashlib.sha256(data).hexdigest() + members = _members(client.get(result["result_url"]).content) + assert live_service.counts == {} + assert len(requests) == 1 and requests[0].method == "GET" + assert requests[0].scope["query_string"] == b"signature=a%2Fb&part=1" + assert "authorization" not in requests[0].headers + if document == "html": + # HTML 的来源地址参与页面标识，因此与保留相同来源的 Jobs URL 路径比较。 + response = client.post( + "/v1/parse/jobs", + json={ + "tier": "flash", + "ocr_mode": "txt", + "output_formats": ["markdown", "zip"], + "files": [{"source": {"type": "url", "url": item["url"]}}], + }, + ) + assert response.status_code == 202, response.text + job_id = response.json()["job_id"] + deadline = time.monotonic() + 30 + while True: + baseline_result = client.get(f"/v1/parse/jobs/{job_id}").json() + if baseline_result["status"] in TERMINAL: + break + assert time.monotonic() < deadline + time.sleep(0.02) + assert baseline_result["status"] == "completed" + output_id = baseline_result["files"][0]["output_files"]["zip"]["file_id"] + baseline_url = f"/v1/files/{output_id}/content" + else: + baseline = _submit(client, [(item["filename"], data)], page_range="2") + assert client.put(baseline["uploads"][0]["upload_url"], content=data).status_code == 200 + baseline_result = _wait(client, baseline) + assert baseline_result["status"] == "completed" + baseline_url = baseline_result["result_url"] + baseline_members = _members(client.get(baseline_url).content) + assert members == baseline_members + if document == "pdf": + middle = json.loads(members["middle_json.json"]) + assert len(middle["pages"]) == 1 and middle["pages"][0]["page_idx"] == 1 + finally: + source_server.stop() + + @dataclass class _CloudState: """模拟仅处理元数据的 SaaS 网关与独立对象存储，覆盖客户端可观察的协议。"""