bug: harden pdf2vqa output handling and package isolation.

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__" }

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

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()

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]:

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,
     *,

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),

uniparser_agent/pdf2vqa/prompts.py

@@ -1,5 +1,3 @@
-"""VQA extraction system prompt (ported from DataFlow QAExtractPrompt)."""
-
 from __future__ import annotations