diff --git a/pyproject.toml b/pyproject.toml index 471157d..e6fd398 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,18 @@ build-backend = "setuptools.build_meta" [tool.setuptools.packages.find] where = ['.'] -exclude = ['build', 'workspace', 'temp', 'temp.*', 'models', 'models.*', 'tests', 'tests.*'] +exclude = [ + 'build', + 'workspace', + 'temp', + 'temp.*', + 'models', + 'models.*', + 'tests', + 'tests.*', + 'uniparser_agent', + 'uniparser_agent.*', +] [tool.setuptools.dynamic] version = { attr = "uniparser_tools.__version__" } diff --git a/uniparser_agent/output_dir.py b/uniparser_agent/output_dir.py new file mode 100644 index 0000000..15e5df6 --- /dev/null +++ b/uniparser_agent/output_dir.py @@ -0,0 +1,57 @@ +"""Safe output directory resolution and replacement.""" + +from __future__ import annotations + +import shutil +from contextlib import contextmanager +from pathlib import Path +from typing import Iterator + + +def default_parse_output_dir(source_stem: str) -> Path: + """Return a contained parse output directory for ``source_stem``.""" + if ( + not source_stem + or source_stem in {".", ".."} + or "/" in source_stem + or "\\" in source_stem + or Path(source_stem).name != source_stem + ): + raise ValueError(f"Unsafe source name for output directory: {source_stem!r}") + + base = (Path.home() / "Uni-Parser-Skill").resolve() + candidate = (base / source_stem).resolve() + if candidate == base or not candidate.is_relative_to(base): + raise ValueError(f"Output directory escapes the managed root: {candidate}") + return candidate + + +def resolve_output_dir(output_dir: str | Path | None, *, default: Path) -> Path: + """Resolve an explicit output path or use the supplied safe default.""" + return Path(output_dir).expanduser().resolve() if output_dir else default.resolve() + + +def _validate_replacement_target(target: Path) -> None: + protected = { + Path(target.anchor).resolve(), + Path.home().resolve(), + Path.cwd().resolve(), + } + if target in protected: + raise ValueError(f"Refusing to replace protected directory: {target}") + if target.exists() and not target.is_dir(): + raise NotADirectoryError(f"Output path exists and is not a directory: {target}") + + +@contextmanager +def replace_output_dir(target: str | Path, *, overwrite: bool) -> Iterator[Path]: + """Create ``target``, deleting its previous contents when overwriting.""" + out = Path(target).expanduser().resolve() + _validate_replacement_target(out) + + if out.exists() and not overwrite: + raise FileExistsError(f"Output directory exists: {out}. Use --overwrite to replace.") + if out.exists(): + shutil.rmtree(out) + out.mkdir(parents=True, exist_ok=False) + yield out diff --git a/uniparser_agent/parse/api_client.py b/uniparser_agent/parse/api_client.py index 0abc6f0..02ba2fa 100644 --- a/uniparser_agent/parse/api_client.py +++ b/uniparser_agent/parse/api_client.py @@ -184,7 +184,3 @@ def resolve_input(raw: str) -> tuple[str, str, Path | None]: if path.suffix.lower() in IMAGE_SUFFIXES: return "image", stem, path return "file", stem, path - - -def default_output_dir(source_stem: str) -> Path: - return (Path.home() / "Uni-Parser-Skill" / source_stem).expanduser().resolve() diff --git a/uniparser_agent/parse/service.py b/uniparser_agent/parse/service.py index 2071d07..fbaaec5 100644 --- a/uniparser_agent/parse/service.py +++ b/uniparser_agent/parse/service.py @@ -4,11 +4,11 @@ from pathlib import Path from typing import Any +from uniparser_agent.output_dir import default_parse_output_dir, replace_output_dir, resolve_output_dir from uniparser_agent.parse.api_client import UniParserApiClient, resolve_input from uniparser_agent.parse.config import get_api_key, get_base_url from uniparser_agent.parse.storage import ( complete_parse_job, - resolve_output_dir, save_stage_error, write_trigger_meta, ) @@ -26,45 +26,47 @@ def parse_document( ) -> dict[str, Any]: """Parse with scientific-paper defaults via HTTP API (no OpenCV).""" kind, source_stem, path = resolve_input(input_path) - out = resolve_output_dir(source_stem, output_dir, overwrite=overwrite) + default_out = default_parse_output_dir(source_stem) if output_dir is None else Path(output_dir) + out = resolve_output_dir(output_dir, default=default_out) client = make_client() - if kind == "file": - trigger = client.trigger_file(str(path)) - input_type = "file" - elif kind == "image": - trigger = client.trigger_snip(str(path)) - input_type = "image" - else: - trigger = client.trigger_url(input_path) - input_type = "url" + with replace_output_dir(out, overwrite=overwrite): + if kind == "file": + trigger = client.trigger_file(str(path)) + input_type = "file" + elif kind == "image": + trigger = client.trigger_snip(str(path)) + input_type = "image" + else: + trigger = client.trigger_url(input_path) + input_type = "url" - if trigger.get("status") != "success": - save_stage_error(out, "trigger_error.json", trigger) - raise RuntimeError(trigger.get("message") or trigger.get("description") or "trigger failed") + if trigger.get("status") != "success": + save_stage_error(out, "trigger_error.json", trigger) + raise RuntimeError(trigger.get("message") or trigger.get("description") or "trigger failed") - token = trigger.get("token") - if not token: - raise RuntimeError("trigger response missing token") + token = trigger.get("token") + if not token: + raise RuntimeError("trigger response missing token") - meta_path = write_trigger_meta( - out, - token=token, - input_type=input_type, - input_value=input_path, - ) + meta_path = write_trigger_meta( + out, + token=token, + input_type=input_type, + input_value=input_path, + ) - summary = complete_parse_job(client, token, out_dir=out, source_stem=source_stem) - return { - "output_dir": summary["output_dir"], - "pages_tree_path": summary["pages_tree_path"], - "markdown_path": summary["markdown_path"], - "token": summary.get("token", ""), - "input_type": input_type, - "source_stem": source_stem, - "source": input_path, - "trigger_meta_path": str(meta_path), - } + summary = complete_parse_job(client, token, out_dir=out, source_stem=source_stem) + return { + "output_dir": summary["output_dir"], + "pages_tree_path": summary["pages_tree_path"], + "markdown_path": summary["markdown_path"], + "token": summary.get("token", ""), + "input_type": input_type, + "source_stem": source_stem, + "source": input_path, + "trigger_meta_path": str(meta_path), + } def load_pages_tree(pages_tree_path: str | Path) -> dict[str, Any]: diff --git a/uniparser_agent/parse/storage.py b/uniparser_agent/parse/storage.py index 33694b3..ce29303 100644 --- a/uniparser_agent/parse/storage.py +++ b/uniparser_agent/parse/storage.py @@ -3,7 +3,6 @@ from __future__ import annotations import json -import shutil import time from datetime import datetime, timezone from pathlib import Path @@ -16,24 +15,6 @@ POLL_TIMEOUT_SEC = 1800 -def resolve_output_dir( - source_stem: str, - output_dir: str | None, - *, - overwrite: bool, -) -> Path: - out = Path(output_dir).expanduser().resolve() if output_dir else default_output_dir(source_stem) - if out.exists() and not overwrite: - raise FileExistsError(f"Output directory exists: {out}. Use --overwrite to replace.") - if out.exists() and overwrite: - shutil.rmtree(out) - return out - - -def default_output_dir(source_stem: str) -> Path: - return (Path.home() / "Uni-Parser-Skill" / source_stem).expanduser().resolve() - - def write_trigger_meta( out_dir: Path, *, diff --git a/uniparser_agent/pdf2vqa/pipeline.py b/uniparser_agent/pdf2vqa/pipeline.py index 98b1fe4..52603f4 100644 --- a/uniparser_agent/pdf2vqa/pipeline.py +++ b/uniparser_agent/pdf2vqa/pipeline.py @@ -3,12 +3,13 @@ from __future__ import annotations import json -import shutil import time from pathlib import Path from typing import Any from uniparser_agent.llm import LLMConfig +from uniparser_agent.output_dir import replace_output_dir, resolve_output_dir +from uniparser_agent.parse.api_client import resolve_input from uniparser_agent.parse.service import load_pages_tree, parse_document from uniparser_agent.pdf2vqa.layout_adapter import adapt_pages_tree_file from uniparser_agent.pdf2vqa.llm_client import VQALLMClient @@ -19,17 +20,8 @@ from uniparser_agent.pdf2vqa.vqa_merger import jsonl_to_md, merge_vqa_pairs, write_merged_jsonl -def _resolve_output_dir(output_dir: str | Path | None, overwrite: bool) -> Path: - if output_dir: - out = Path(output_dir).expanduser().resolve() - else: - out = (Path.cwd() / "vqa_out").resolve() - if out.exists(): - if not overwrite: - raise FileExistsError(f"Output directory already exists: {out}. Pass overwrite=True or --overwrite.") - shutil.rmtree(out) - out.mkdir(parents=True, exist_ok=True) - return out +def _resolve_output_dir(output_dir: str | Path | None) -> Path: + return resolve_output_dir(output_dir, default=Path.cwd() / "vqa_out") def _require_local_pdf(path: str | Path, *, label: str) -> Path: @@ -66,28 +58,65 @@ def run_vqa_pipeline( raise ValueError("Provide input_path (pdf/url/image) or pages_tree_path.") started = time.time() - out = _resolve_output_dir(output_dir, overwrite=overwrite) - parse_dir = out / "parse" - parse_meta: dict[str, Any] = {} - merged_pdf_path: Path | None = None - + pages_tree_bytes: bytes | None = None + question_pdf: Path | None = None + answer_path: Path | None = None if pages_tree_path: src_tree = Path(pages_tree_path).expanduser().resolve() if not src_tree.is_file(): raise FileNotFoundError(f"pages_tree not found: {src_tree}") + load_pages_tree(src_tree) + pages_tree_bytes = src_tree.read_bytes() + elif answer_pdf: + assert input_path is not None + question_pdf = _require_local_pdf(input_path, label="question PDF") + answer_path = _require_local_pdf(answer_pdf, label="answer PDF") + else: + assert input_path is not None + resolve_input(input_path) + + target = _resolve_output_dir(output_dir) + with replace_output_dir(target, overwrite=overwrite) as out: + return _run_vqa_pipeline_in_dir( + out=out, + started=started, + input_path=input_path, + pages_tree_bytes=pages_tree_bytes, + question_pdf=question_pdf, + answer_path=answer_path, + strict_title_match=strict_title_match, + llm_config=llm_config, + llm_client=llm_client, + ) + + +def _run_vqa_pipeline_in_dir( + *, + out: Path, + started: float, + input_path: str | None, + pages_tree_bytes: bytes | None, + question_pdf: Path | None, + answer_path: Path | None, + strict_title_match: bool, + llm_config: LLMConfig | None, + llm_client: VQALLMClient | None, +) -> dict[str, Any]: + parse_dir = out / "parse" + parse_meta: dict[str, Any] = {} + merged_pdf_path: Path | None = None + + if pages_tree_bytes is not None: parse_dir.mkdir(parents=True, exist_ok=True) dest_tree = parse_dir / "pages_tree.json" - shutil.copy2(src_tree, dest_tree) + dest_tree.write_bytes(pages_tree_bytes) tree_path = dest_tree parse_meta = {"mode": "pages_tree", "pages_tree_path": str(tree_path)} else: assert input_path is not None parse_source = input_path - question_pdf: Path | None = None - answer_path: Path | None = None - if answer_pdf: - question_pdf = _require_local_pdf(input_path, label="question PDF") - answer_path = _require_local_pdf(answer_pdf, label="answer PDF") + if answer_path is not None: + assert question_pdf is not None merge_dir = out / "merge" merge_dir.mkdir(parents=True, exist_ok=True) merged_pdf_path = merge_pdfs( @@ -98,8 +127,8 @@ def run_vqa_pipeline( parse_result = parse_document(parse_source, output_dir=str(parse_dir), overwrite=True) tree_path = Path(parse_result["pages_tree_path"]) - if answer_pdf: - assert question_pdf is not None and answer_path is not None and merged_pdf_path is not None + if answer_path is not None: + assert question_pdf is not None and merged_pdf_path is not None parse_meta = { "mode": "dual_pdf", "question_pdf": str(question_pdf), diff --git a/uniparser_agent/pdf2vqa/prompts.py b/uniparser_agent/pdf2vqa/prompts.py index 57a3fea..5e69c61 100644 --- a/uniparser_agent/pdf2vqa/prompts.py +++ b/uniparser_agent/pdf2vqa/prompts.py @@ -1,5 +1,3 @@ -"""VQA extraction system prompt (ported from DataFlow QAExtractPrompt).""" - from __future__ import annotations diff --git a/uniparser_agent/tests/test_output_safety.py b/uniparser_agent/tests/test_output_safety.py new file mode 100644 index 0000000..3c05a96 --- /dev/null +++ b/uniparser_agent/tests/test_output_safety.py @@ -0,0 +1,167 @@ +"""Regression tests for safe output replacement.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from uniparser_agent.output_dir import default_parse_output_dir, replace_output_dir +from uniparser_agent.parse import service as parse_service +from uniparser_agent.pdf2vqa.pipeline import run_vqa_pipeline + + +class _TriggerFailureClient: + def trigger_url(self, _url: str) -> dict[str, str]: + return {"status": "error", "message": "trigger failed"} + + +class _LLMFailureClient: + def chat(self, *, system_prompt: str, user_content: str) -> str: + raise RuntimeError("LLM failed") + + +def _write_pages_tree(path: Path) -> None: + path.write_text( + json.dumps( + { + "pages_tree": [ + [ + { + "type": "paragraph", + "page": 1, + "block": 1, + "text": "Question 1", + } + ] + ] + } + ), + encoding="utf-8", + ) + + +@pytest.mark.parametrize("source_stem", ["", ".", "..", "../escape", r"..\escape"]) +def test_default_parse_output_rejects_unsafe_source_stem(source_stem: str) -> None: + with pytest.raises(ValueError, match="Unsafe source name"): + default_parse_output_dir(source_stem) + + +def test_default_parse_output_is_contained_in_managed_root( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("HOME", str(tmp_path)) + + output = default_parse_output_dir("exam") + + assert output == (tmp_path / "Uni-Parser-Skill" / "exam").resolve() + + +def test_parse_url_traversal_is_rejected_before_client_creation(monkeypatch: pytest.MonkeyPatch) -> None: + def unexpected_client() -> None: + raise AssertionError("client must not be created") + + monkeypatch.setattr(parse_service, "make_client", unexpected_client) + with pytest.raises(ValueError, match="Unsafe source name"): + parse_service.parse_document("https://example.com/..") + + +def test_output_replacement_does_not_restore_previous_directory_on_failure(tmp_path: Path) -> None: + output = tmp_path / "output" + output.mkdir() + (output / "previous.txt").write_text("previous", encoding="utf-8") + + with pytest.raises(RuntimeError, match="failed"): + with replace_output_dir(output, overwrite=True) as work_dir: + (work_dir / "partial.txt").write_text("partial", encoding="utf-8") + raise RuntimeError("failed") + + assert not (output / "previous.txt").exists() + assert (output / "partial.txt").read_text(encoding="utf-8") == "partial" + assert not list(tmp_path.glob(".output.backup-*")) + + +def test_output_replacement_commits_successful_run(tmp_path: Path) -> None: + output = tmp_path / "output" + output.mkdir() + (output / "previous.txt").write_text("previous", encoding="utf-8") + + with replace_output_dir(output, overwrite=True) as work_dir: + (work_dir / "current.txt").write_text("current", encoding="utf-8") + + assert not (output / "previous.txt").exists() + assert (output / "current.txt").read_text(encoding="utf-8") == "current" + assert not list(tmp_path.glob(".output.backup-*")) + + +def test_failed_new_output_is_left_for_diagnostics(tmp_path: Path) -> None: + output = tmp_path / "output" + + with pytest.raises(RuntimeError, match="failed"): + with replace_output_dir(output, overwrite=False) as work_dir: + (work_dir / "partial.txt").write_text("partial", encoding="utf-8") + raise RuntimeError("failed") + + assert (output / "partial.txt").read_text(encoding="utf-8") == "partial" + + +def test_parse_failure_keeps_partial_output_without_restoring_previous( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + output = tmp_path / "parse" + output.mkdir() + (output / "previous.txt").write_text("previous", encoding="utf-8") + monkeypatch.setattr(parse_service, "make_client", _TriggerFailureClient) + + with pytest.raises(RuntimeError, match="trigger failed"): + parse_service.parse_document( + "https://example.com/exam.pdf", + output_dir=str(output), + overwrite=True, + ) + + assert not (output / "previous.txt").exists() + assert (output / "trigger_error.json").is_file() + assert not list(tmp_path.glob(".parse.backup-*")) + + +def test_vqa_validates_pages_tree_before_replacing_output(tmp_path: Path) -> None: + output = tmp_path / "vqa" + output.mkdir() + (output / "previous.txt").write_text("previous", encoding="utf-8") + invalid_tree = tmp_path / "invalid.json" + invalid_tree.write_text("{}", encoding="utf-8") + + with pytest.raises(ValueError, match="missing pages_tree key"): + run_vqa_pipeline( + pages_tree_path=str(invalid_tree), + output_dir=str(output), + overwrite=True, + llm_client=_LLMFailureClient(), # type: ignore[arg-type] + ) + + assert (output / "previous.txt").read_text(encoding="utf-8") == "previous" + + +def test_vqa_failure_keeps_partial_output_without_restoring_previous(tmp_path: Path) -> None: + output = tmp_path / "vqa" + output.mkdir() + (output / "previous.txt").write_text("previous", encoding="utf-8") + pages_tree = tmp_path / "pages_tree.json" + _write_pages_tree(pages_tree) + + with pytest.raises(RuntimeError, match="LLM failed"): + run_vqa_pipeline( + pages_tree_path=str(pages_tree), + output_dir=str(output), + overwrite=True, + llm_client=_LLMFailureClient(), # type: ignore[arg-type] + ) + + assert not (output / "previous.txt").exists() + assert (output / "parse" / "pages_tree.json").is_file() + assert (output / "llm_content_list.json").is_file() + assert not list(tmp_path.glob(".vqa.backup-*"))