From bb719f217a77b19d89d26f96168cf463ab73b6ba Mon Sep 17 00:00:00 2001 From: "Tal.Yuan" Date: Tue, 4 Aug 2026 17:54:55 +0800 Subject: [PATCH 1/3] refactor(routes): move document domain into routes/document/ subpackage (#5885) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Slice 2m of the route-domain reorganization (#4082/#4071, per specs/architecture-runtime-inventory.md §6.3). Moves document_routes.py (1810 lines) and document_helpers.py (243 lines) into routes/document/, leaving backward-compat sys.modules shims at the old paths. Pure file reorganization, no behavior change. Both shims use sys.modules replacement so the `import ... as droutes` + `droutes.SessionLocal = ...` / `monkeypatch.setattr(droutes, ...)` pattern in multiple tests, and the `sys.modules.pop("routes.document_helpers")` + re-import pattern in test_security_regressions.py, all reach the canonical modules. The canonical document_routes.py imports helpers from the canonical path (routes.document.document_helpers), not the legacy shim. Three source-introspection test sites repointed to the new canonical path: - test_imap_mailbox_quoting.py - test_model_helper_owner_scope.py - test_vision_owner_scope.py (shared with other domains; document entry repointed) Adds tests/test_document_routes_shim.py to pin the sys.modules shim contract for both modules. Verified: compileall clean; full suite 4789 passed, 3 skipped. --- app.py | 2 +- routes/document/__init__.py | 6 + routes/document/document_helpers.py | 243 ++++ routes/document/document_routes.py | 1810 +++++++++++++++++++++++ routes/document_helpers.py | 249 +--- routes/document_routes.py | 1819 +----------------------- tests/test_document_routes_shim.py | 29 + tests/test_imap_mailbox_quoting.py | 2 +- tests/test_model_helper_owner_scope.py | 2 +- tests/test_vision_owner_scope.py | 2 +- 10 files changed, 2115 insertions(+), 2049 deletions(-) create mode 100644 routes/document/__init__.py create mode 100644 routes/document/document_helpers.py create mode 100644 routes/document/document_routes.py create mode 100644 tests/test_document_routes_shim.py diff --git a/app.py b/app.py index c85d425fb..8363ba4e9 100644 --- a/app.py +++ b/app.py @@ -739,7 +739,7 @@ app.include_router(setup_stt_routes(stt_service)) logger.info("STT service initialized (provider managed via settings)") # Documents (artifacts/canvas) -from routes.document_routes import setup_document_routes +from routes.document.document_routes import setup_document_routes document_router = setup_document_routes(session_manager, upload_handler) app.include_router(document_router) diff --git a/routes/document/__init__.py b/routes/document/__init__.py new file mode 100644 index 000000000..7f79ce1bb --- /dev/null +++ b/routes/document/__init__.py @@ -0,0 +1,6 @@ +"""Document route domain package (slice 2m, #4082/#4071). + +Contains document_routes.py and document_helpers.py, migrated from the flat +routes/ directory. Backward-compat shims at routes/document_routes.py and +routes/document_helpers.py re-export from here. +""" diff --git a/routes/document/document_helpers.py b/routes/document/document_helpers.py new file mode 100644 index 000000000..a0c2d08eb --- /dev/null +++ b/routes/document/document_helpers.py @@ -0,0 +1,243 @@ +"""document_helpers.py — Pydantic models, doc serializers, owner gating, file-locator helpers shared with document_routes.py.""" + +"""Document routes — CRUD for living documents with version history.""" + +import logging +import os +import re +from typing import Any, Dict, Optional + +from fastapi import HTTPException, Request +from pydantic import BaseModel + +from core.database import Document, DocumentVersion +from core.database import Session as DbSession +from src.auth_helpers import _auth_disabled +from src.upload_handler import UploadHandler + +logger = logging.getLogger(__name__) + + +# ---- Request schemas ---- + +class DocumentCreate(BaseModel): + session_id: Optional[str] = None + title: str = "Untitled" + language: Optional[str] = None + content: str = "" + +class DocumentUpdate(BaseModel): + content: str + summary: Optional[str] = None + force_version: bool = False + +class DocumentPatch(BaseModel): + title: Optional[str] = None + language: Optional[str] = None + session_id: Optional[str] = None # link/unlink document to a session + + +# ---- Helpers ---- + +def _doc_to_dict(doc: Document) -> Dict[str, Any]: + return { + "id": doc.id, + "session_id": doc.session_id, + "title": doc.title, + "language": doc.language, + "current_content": doc.current_content, + "version_count": doc.version_count, + "is_active": doc.is_active, + "archived": bool(getattr(doc, "archived", False)), + "created_at": (doc.created_at.isoformat() + "Z") if doc.created_at else None, + "updated_at": (doc.updated_at.isoformat() + "Z") if doc.updated_at else None, + # Source-email provenance (set when doc was created from an email + # attachment) — drives the "Send signed reply" menu item. + "source_email_uid": getattr(doc, "source_email_uid", None), + "source_email_folder": getattr(doc, "source_email_folder", None), + "source_email_account_id": getattr(doc, "source_email_account_id", None), + "source_email_message_id": getattr(doc, "source_email_message_id", None), + } + +def _version_to_dict(v: DocumentVersion) -> Dict[str, Any]: + return { + "id": v.id, + "document_id": v.document_id, + "version_number": v.version_number, + "content": v.content, + "summary": v.summary, + "source": v.source, + "created_at": v.created_at.isoformat() if v.created_at else None, + } + + +def _verify_doc_owner(db, doc: Document, user: str): + """Verify `user` owns this document. Raise 404 if not. + + Documents now carry their own `owner` column, so a doc whose session + was deleted (session_id → NULL) can still prove ownership and stay + openable / cloneable. We trust that column first and only fall back to + the session join for any not-yet-backfilled legacy row. + """ + if user is None: + if _auth_disabled(): + return # Single-user / no-auth mode: allow access + raise HTTPException(403, "Authentication required") + if doc.owner is not None: + if doc.owner != user: + raise HTTPException(404, "Document not found") + return + # Legacy fallback: derive ownership from the linked session. + if not doc.session_id: + raise HTTPException(404, "Document not found") + session = db.query(DbSession).filter(DbSession.id == doc.session_id).first() + if not session or session.owner != user: + raise HTTPException(404, "Document not found") + + +def _owner_session_filter(q, user): + """Restrict a documents query to those owned by `user`. + + Documents now carry their own `owner` column (backfilled at boot from + the linked session, or assigned to the admin user for legacy/orphaned + docs). We filter on that directly rather than on a session join, so a + document whose session was deleted (session_id → NULL) still shows up + for its owner instead of silently vanishing from the Library + search. + + The owner backfill runs in init_db before the app serves requests, so + by the time this filter is live there are no NULL-owner rows to leak; + we therefore match the owner strictly for authenticated callers.""" + if not user: + if user == "" or _auth_disabled(): + return q + return q.filter(False) + return q.filter(Document.owner == user) + + + +def _slug(name: str) -> str: + """Filesystem-friendly version of a document title. + + Whitespace becomes underscores; other unsafe punctuation is dropped. + Preserves letters, digits, dot, hyphen, underscore. Idempotent. + """ + import re as _re + s = (name or "").strip() + # Drop the trailing extension if the title happens to include one + s = _re.sub(r'\.pdf$', '', s, flags=_re.IGNORECASE) + s = _re.sub(r'\s+', '_', s) + s = _re.sub(r'[^A-Za-z0-9._-]', '', s) + s = _re.sub(r'_+', '_', s).strip('_') + return s or "form" + + +# DPI scale for the interactive PDF view. ~150 DPI (2x of 72 PDF user-units). +_PDF_RENDER_SCALE = 2.0 + + +def _upload_path_inside(upload_dir: str, path: str) -> bool: + base = os.path.realpath(upload_dir) + p = os.path.realpath(path) + try: + return os.path.commonpath([base, p]) == base + except Exception: + return False + + +def _resolve_user_upload_path( + upload_handler: Any, + upload_id: str, + owner: Optional[str], + auth_manager=None, +) -> Optional[str]: + """Resolve an upload id to a filesystem path the caller may read.""" + if upload_handler is None: + return None + resolved = upload_handler.resolve_upload( + upload_id, + owner=owner, + auth_manager=auth_manager, + ) + if not isinstance(resolved, dict) or not resolved: + return None + path = resolved.get("path") + upload_dir = getattr(upload_handler, "upload_dir", None) + if path and upload_dir and not _upload_path_inside(upload_dir, path): + logger.warning("Upload path outside upload directory: %s", path) + return None + return path + + +def _locate_upload( + upload_dir: str, + file_id: str, + owner: Optional[str] = None, + auth_manager=None, + upload_handler: Any = None, +): + """Find an upload by its filename ID via UploadHandler.resolve_upload.""" + if upload_handler is None: + from src.upload_handler import UploadHandler + + base_dir = os.path.dirname(os.path.abspath(upload_dir)) + upload_handler = UploadHandler(base_dir, upload_dir) + return _resolve_user_upload_path(upload_handler, file_id, owner, auth_manager) + + +def _assert_pdf_marker_upload_owned( + request: Request, + content: str, + user: Optional[str], + upload_handler: Any, +) -> None: + """Reject document content whose pdf_source marker points at another user's upload.""" + if upload_handler is None: + return + from src.pdf_form_doc import find_source_upload_id + + upload_id = find_source_upload_id(content or "") + if not upload_id: + return + auth_manager = getattr(getattr(request.app, "state", None), "auth_manager", None) + if not _resolve_user_upload_path(upload_handler, upload_id, user, auth_manager): + raise HTTPException( + 400, + "Document PDF marker references an upload you do not own", + ) + + +def _derive_title(content: str) -> str: + """Derive a title from document content.""" + import re + if not isinstance(content, str): + return "Untitled" + text = content.strip() + if not text: + return "Untitled" + + # Markdown header + md = re.match(r'^#{1,3}\s+(.+)', text, re.MULTILINE) + if md: + title = md.group(1).strip() + if len(title) > 50: + title = title[:48] + "…" + return title + + # HTML heading + html = re.search(r']*>([^<]+)', text, re.IGNORECASE) + if html: + title = html.group(1).strip() + if len(title) > 50: + title = title[:48] + "…" + return title + + # First non-empty line (if short enough) + for line in text.split('\n'): + line = line.strip() + if line and 2 <= len(line) <= 60: + title = re.sub(r'[:#*`]+$', '', line).strip() + if title and len(title) > 50: + title = title[:48] + "…" + return title or "Untitled" + + return "Untitled" diff --git a/routes/document/document_routes.py b/routes/document/document_routes.py new file mode 100644 index 000000000..dae8b09fa --- /dev/null +++ b/routes/document/document_routes.py @@ -0,0 +1,1810 @@ +"""Document routes — CRUD for living documents with version history.""" + +import uuid +import logging +from datetime import datetime, timezone +from typing import Dict, Any, List, Optional + +from fastapi import APIRouter, HTTPException, Query, Request, UploadFile, File, Form + +from sqlalchemy import case, func, or_ +from core.database import SessionLocal, Document, DocumentVersion +from core.database import Session as DbSession +from src.auth_helpers import get_current_user, _auth_disabled +from src.constants import MAIL_ATTACHMENTS_DIR +from src.upload_handler import reserve_upload_references + +logger = logging.getLogger(__name__) + + +def _get_session_or_404(db, session_id: str, user: Optional[str]): + session = db.query(DbSession).filter(DbSession.id == session_id).first() + if not session: + raise HTTPException(404, "Session not found") + if user and session.owner != user: + raise HTTPException(404, "Session not found") + return session + + +def _aggregate_language_facets(lang_rows): + """Sum document counts per display language for the library facet. + + NULL-language and explicit "text" rows share the "text" bucket (the + language filter treats them as one), so they must be ADDED. The old dict + comprehension keyed both to "text", silently overwriting one group and + undercounting the facet versus what the filter actually returns. + """ + out = {} + for lang, cnt in lang_rows: + key = lang or "text" + out[key] = out.get(key, 0) + cnt + return out + + +def _library_language_for_document(doc: Document) -> str: + """Return the display language used by the document library. + + PDF documents are stored as markdown wrappers so the editor can preserve + extracted text, form fields, and annotations. The library should still + identify them as PDFs instead of exposing that internal wrapper format. + """ + from src.pdf_form_doc import find_source_upload_id + + if find_source_upload_id(doc.current_content or ""): + return "pdf" + return doc.language or "text" + + +def _email_source_key(content: str) -> tuple[str, str]: + """Return the source email identity embedded in an email draft document.""" + import re + + text = content or "" + uid_m = re.search(r"(?im)^X-Source-UID:\s*(.+?)\s*$", text) + folder_m = re.search(r"(?im)^X-Source-Folder:\s*(.+?)\s*$", text) + uid = (uid_m.group(1).strip() if uid_m else "") + folder = (folder_m.group(1).strip() if folder_m else "INBOX") + return uid, folder + + +from routes.document_helpers import ( + DocumentCreate, DocumentUpdate, DocumentPatch, + _doc_to_dict, _version_to_dict, + _verify_doc_owner, _owner_session_filter, + _slug, _resolve_user_upload_path, _assert_pdf_marker_upload_owned, _derive_title, + _PDF_RENDER_SCALE, +) + + +def setup_document_routes(session_manager, upload_handler=None) -> APIRouter: + router = APIRouter(tags=["documents"]) + + def _reserve_document_uploads(user: Optional[str], content: str) -> None: + missing_id = reserve_upload_references(upload_handler, user, content) + if missing_id: + raise HTTPException( + 409, + f"Referenced upload is no longer available: {missing_id}", + ) + + def _locate_current_user_upload(request: Request, upload_id: str, user: Optional[str]): + if upload_handler is None: + return None + auth_manager = getattr(getattr(request.app, "state", None), "auth_manager", None) + return _resolve_user_upload_path(upload_handler, upload_id, user, auth_manager) + + def _load_pdf_viewer_fitz(): + from src.pdf_runtime import load_pymupdf_for_pdf_viewer + + try: + return load_pymupdf_for_pdf_viewer() + except RuntimeError as exc: + raise HTTPException(503, str(exc)) from exc + + # ---- POST /api/document ---- + @router.post("/api/document") + async def create_document(request: Request, req: DocumentCreate) -> Dict[str, Any]: + from src.auth_helpers import require_privilege + user = require_privilege(request, "can_use_documents") + db = SessionLocal() + try: + # session_id is optional: a doc can be a session-less "library" doc + # (e.g. files imported from the library) — session_id is nullable and + # the doc is owner-stamped, so it lives in the library on its own. + session = None + if req.session_id: + # Match the lenient ownership model the rest of the app uses + # (see _owner_filter): only block when an AUTHENTICATED user is + # writing into a DIFFERENT user's session. In single-user / + # unconfigured / localhost-bypass mode, falsey users preserve + # the existing lenient path. + session = _get_session_or_404(db, req.session_id, user) + + # If no language was supplied (e.g. cloning a doc whose language + # was never set), detect it from the content rather than storing + # NULL — which made the editor fall back to plain text. Defaults + # to markdown for prose. + language = req.language + if not language: + from src.agent_tools.document_tools import _looks_like_email_document, _sniff_doc_language, _coerce_email_document_content + language = _sniff_doc_language(req.content) + else: + from src.agent_tools.document_tools import _looks_like_email_document, _coerce_email_document_content + if _looks_like_email_document(req.content, req.title): + language = "email" + + _reserve_document_uploads(user, req.content) + _assert_pdf_marker_upload_owned(request, req.content, user, upload_handler) + + # Reply drafts are keyed to the source email. If a UI/tool path tries + # to create a second draft for the same email in the same chat, + # update the existing draft instead so quoted thread history stays + # attached to the visible document. + if language == "email" and req.session_id: + source_uid, source_folder = _email_source_key(req.content) + if source_uid: + candidates = ( + db.query(Document) + .filter(Document.session_id == req.session_id) + .filter(Document.is_active == True) + .filter(Document.language == "email") + .order_by(Document.updated_at.desc()) + .limit(25) + .all() + ) + for existing in candidates: + old_uid, old_folder = _email_source_key(existing.current_content or "") + if old_uid != source_uid or old_folder != source_folder: + continue + merged = _coerce_email_document_content(existing.current_content or "", req.content) + if existing.current_content != merged: + new_ver = (existing.version_count or 1) + 1 + existing.current_content = merged + existing.title = req.title or existing.title + existing.version_count = new_ver + db.add(DocumentVersion( + id=str(uuid.uuid4()), + document_id=existing.id, + version_number=new_ver, + content=merged, + summary="Updated existing email draft", + source="user", + )) + db.commit() + db.refresh(existing) + return _doc_to_dict(existing) + + doc_id = str(uuid.uuid4()) + ver_id = str(uuid.uuid4()) + + doc = Document( + id=doc_id, + session_id=req.session_id, + title=req.title, + language=language, + current_content=req.content, + version_count=1, + is_active=True, + # Stamp ownership directly so the doc survives its session + # being deleted. Fall back to the session's owner when the + # request is unauthenticated (single-user / localhost bypass). + owner=user or (session.owner if session else None), + ) + ver = DocumentVersion( + id=ver_id, + document_id=doc_id, + version_number=1, + content=req.content, + summary="Initial version", + source="user", + ) + db.add(doc) + db.add(ver) + db.commit() + db.refresh(doc) + try: + from src.event_bus import fire_event + fire_event("document_created", doc.owner) + except Exception: + logger.debug("document_created event dispatch failed", exc_info=True) + return _doc_to_dict(doc) + except HTTPException: + raise + except Exception as e: + db.rollback() + logger.error(f"Failed to create document: {e}") + raise HTTPException(500, f"Failed to create document: {e}") + finally: + db.close() + + # ---- POST /api/documents/import-pdf ---- + @router.post("/api/documents/import-pdf") + async def import_pdf( + request: Request, + file: UploadFile = File(...), + session_id: Optional[str] = Form(None), + ) -> Dict[str, Any]: + """Upload a PDF and create the matching Document. + + Detects AcroForm fields — if any, creates a form-backed markdown doc + (clickable inputs in the PDF view). Otherwise creates a plain PDF doc + with a `pdf_source` marker so the viewer renders the pages without + overlays. + """ + from src.pdf_forms import has_form_fields, extract_fields + from src.pdf_form_doc import ( + save_field_sidecar, + create_form_markdown_document, + create_plain_pdf_document, + ) + from src.document_processor import _process_pdf, strip_pdf_content_marker + import os + + from src.auth_helpers import require_privilege + user = require_privilege(request, "can_use_documents") + + # session_id is optional — a library import isn't tied to a chat. When + # given, validate it; otherwise the PDF becomes a session-less library + # doc (the doc creators below already handle a missing session). + if session_id: + db = SessionLocal() + try: + _get_session_or_404(db, session_id, user) + finally: + db.close() + + if upload_handler is None: + raise HTTPException(500, "Upload handler not configured") + + client_ip = request.client.host if request.client else "unknown" + try: + meta = upload_handler.save_upload(file, client_ip, owner=user) + except HTTPException: + raise + except Exception as e: + logger.error(f"PDF import save_upload failed: {e}") + raise HTTPException(500, f"Upload failed: {e}") + + upload_id = meta["id"] + pdf_path = _locate_current_user_upload(request, upload_id, user) + if not pdf_path: + raise HTTPException(500, "Saved PDF could not be located") + + title = os.path.splitext(meta.get("original_name") or meta.get("name") or upload_id)[0] + try: + body_text = strip_pdf_content_marker(_process_pdf(pdf_path, owner=user)) + except Exception: + body_text = None + + is_form = False + try: + is_form = has_form_fields(pdf_path) + except Exception as e: + logger.warning(f"has_form_fields failed for {pdf_path}: {e}") + + if is_form: + fields = extract_fields(pdf_path) + save_field_sidecar(pdf_path, fields) + doc_id = create_form_markdown_document( + session_id=session_id, + fields=fields, + upload_id=upload_id, + title=title, + intro_text=body_text, + ) + else: + doc_id = create_plain_pdf_document( + session_id=session_id, + upload_id=upload_id, + title=title, + body_text=body_text, + ) + + if not doc_id: + raise HTTPException(500, "Failed to create document for PDF") + + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(500, "Created document not found") + # The PDF doc creators stamp owner from the session only; a + # session-less library import leaves owner NULL, which the Library's + # owner filter then hides. Stamp the requesting user so it shows. + if not doc.owner and user: + doc.owner = user + db.commit() + db.refresh(doc) + return _doc_to_dict(doc) + finally: + db.close() + + # ---- GET /api/documents/library ---- + @router.get("/api/documents/library") + async def documents_library( + request: Request, + search: Optional[str] = Query(None), + language: Optional[str] = Query(None), + sort: str = Query("recent"), + offset: int = Query(0, ge=0), + limit: int = Query(20, ge=1, le=50), + archived: bool = Query(False), + ) -> Dict[str, Any]: + user = get_current_user(request) + db = SessionLocal() + try: + from sqlalchemy import or_ + pdf_marker_cond = or_( + Document.current_content.like('%\s*\n+#[^\n]*\n+)', re.MULTILINE) + head_match = head_re.match(content) + head = head_match.group(1) if head_match else (content.splitlines()[0] + "\n\n# " + (doc.title or "PDF") + "\n\n") + doc.current_content = head + body_text.strip() + "\n" + doc.version_count = (doc.version_count or 1) + 1 + db.add(DocumentVersion( + id=str(__import__("uuid").uuid4()), + document_id=doc_id, + version_number=doc.version_count, + content=doc.current_content, + summary="PDF text re-extracted (OCR)", + source="ocr", + )) + db.commit() + return {"ok": True, "id": doc_id, "extracted": True, "chars": len(body_text)} + finally: + db.close() + + # ---- POST /api/documents/export-zip — bundle selected docs into a .zip ---- + @router.post("/api/documents/export-zip") + async def documents_export_zip(request: Request): + """Zip the selected documents (each as a text file with the right + extension) — mirrors the gallery's bulk download-zip so multi-export + is one file instead of a blocked flood of individual downloads.""" + user = get_current_user(request) + try: + data = await request.json() + except Exception as e: + logger.warning("Failed to parse export request body, defaulting to empty", exc_info=e) + data = {} + ids = data.get("ids") or [] + if not ids: + raise HTTPException(400, "No documents specified") + _ext = { + "javascript": ".js", "python": ".py", "html": ".html", "css": ".css", + "markdown": ".md", "json": ".json", "yaml": ".yml", "bash": ".sh", + "sql": ".sql", "rust": ".rs", "go": ".go", "java": ".java", "c": ".c", + "cpp": ".cpp", "typescript": ".ts", "ruby": ".rb", "php": ".php", + "text": ".txt", "xml": ".xml", "toml": ".toml", "ini": ".ini", + } + db = SessionLocal() + try: + import io + import re + import zipfile + from fastapi import Response + docs = db.query(Document).filter(Document.id.in_(ids)).all() + buf = io.BytesIO() + used = set() + wrote = 0 + with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf: + for doc in docs: + try: + _verify_doc_owner(db, doc, user) + except HTTPException: + continue # skip docs the user doesn't own + ext = _ext.get(doc.language or "text", ".txt") + base = (doc.title or "document").strip() or "document" + base = re.sub(r"[^\w\-. ]+", "", base)[:60].strip() or doc.id + name = base if "." in base else base + ext + i = 1 + while name in used: + name = f"{base}-{i}" + ("" if "." in base else ext) + i += 1 + used.add(name) + zf.writestr(name, doc.current_content or "") + wrote += 1 + if not wrote: + raise HTTPException(404, "No documents found") + return Response( + content=buf.getvalue(), + media_type="application/zip", + headers={"Content-Disposition": 'attachment; filename="documents.zip"'}, + ) + finally: + db.close() + + # ---- PUT /api/document/{doc_id} — user manual edit ---- + # Coalesce window: if the last user version was saved within this many + # seconds, update it in-place (user is still actively editing). + # Once the gap exceeds this, the next save creates a new version. + VERSION_COALESCE_SECONDS = 60 + + @router.put("/api/document/{doc_id}") + async def update_document(request: Request, doc_id: str, req: DocumentUpdate) -> Dict[str, Any]: + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + + incoming_content = req.content + from src.agent_tools.document_tools import _coerce_email_document_content, _looks_like_email_document + is_email_doc = ( + (doc.language or "").lower() == "email" + or _looks_like_email_document(doc.current_content or "", doc.title or "") + or _looks_like_email_document(req.content or "", doc.title or "") + ) + if is_email_doc: + incoming_content = _coerce_email_document_content(doc.current_content or "", req.content) + doc.language = "email" + + # Skip if content is identical unless the caller explicitly wants + # a checkpoint version from the current editor state. + if doc.current_content == incoming_content and not req.force_version: + return _doc_to_dict(doc) + + _reserve_document_uploads(user, incoming_content) + _assert_pdf_marker_upload_owned(request, incoming_content, user, upload_handler) + + # Check if we can coalesce with the latest version + latest_ver = db.query(DocumentVersion).filter( + DocumentVersion.document_id == doc_id, + ).order_by(DocumentVersion.version_number.desc()).first() + + now = datetime.now(timezone.utc) + coalesced = False + if latest_ver and latest_ver.source == "user" and not req.force_version: + ver_time = latest_ver.created_at + if ver_time.tzinfo is None: + ver_time = ver_time.replace(tzinfo=timezone.utc) + age = (now - ver_time).total_seconds() + if age < VERSION_COALESCE_SECONDS: + # Update the existing version in-place + latest_ver.content = incoming_content + latest_ver.created_at = now + if req.summary: + latest_ver.summary = req.summary + coalesced = True + + if not coalesced: + new_ver = doc.version_count + 1 + ver = DocumentVersion( + id=str(uuid.uuid4()), + document_id=doc_id, + version_number=new_ver, + content=incoming_content, + summary=req.summary or "Manual edit", + source="user", + ) + doc.version_count = new_ver + db.add(ver) + + doc.current_content = incoming_content + db.commit() + db.refresh(doc) + return _doc_to_dict(doc) + except HTTPException: + raise + except Exception as e: + db.rollback() + raise HTTPException(500, f"Failed to update document: {e}") + finally: + db.close() + + # ---- PATCH /api/document/{doc_id} — metadata only ---- + @router.patch("/api/document/{doc_id}") + async def patch_document(request: Request, doc_id: str, req: DocumentPatch) -> Dict[str, Any]: + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + if req.title is not None: + doc.title = req.title + if req.language is not None: + doc.language = req.language + if req.session_id is not None: + # Empty string = unlink from session + if req.session_id: + _get_session_or_404(db, req.session_id, user) + doc.session_id = req.session_id if req.session_id else None + if not req.session_id: + # Tab closed / doc detached from its session — drop the + # in-memory active-doc pointer so the last-resort injection + # path doesn't re-surface this doc in a later chat (#1160). + try: + from src.agent_tools.document_tools import clear_active_document + clear_active_document(doc_id) + except Exception as e: + logger.warning("Failed to clear active document %r on detach", doc_id, exc_info=e) + db.commit() + db.refresh(doc) + return _doc_to_dict(doc) + except HTTPException: + raise + except Exception as e: + db.rollback() + raise HTTPException(500, str(e)) + finally: + db.close() + + # ---- DELETE /api/document/{doc_id} — soft delete ---- + @router.delete("/api/document/{doc_id}") + async def delete_document(request: Request, doc_id: str) -> Dict[str, str]: + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + doc.is_active = False + # Closed/deleted — drop the in-memory active-doc pointer so it isn't + # re-injected into a later, unrelated chat (#1160). + try: + from src.agent_tools.document_tools import clear_active_document + clear_active_document(doc_id) + except Exception: + pass + db.commit() + return {"status": "deleted", "id": doc_id} + except HTTPException: + raise + except Exception as e: + db.rollback() + raise HTTPException(500, str(e)) + finally: + db.close() + + # ---- GET /api/document/{doc_id}/versions ---- + @router.get("/api/document/{doc_id}/versions") + async def list_versions(request: Request, doc_id: str) -> List[Dict[str, Any]]: + user = get_current_user(request) + db = SessionLocal() + try: + # Verify ownership before listing versions + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + versions = db.query(DocumentVersion).filter( + DocumentVersion.document_id == doc_id + ).order_by(DocumentVersion.version_number.desc()).all() + return [{ + "id": v.id, + "version_number": v.version_number, + "content": v.content, + "summary": v.summary, + "source": v.source, + "created_at": v.created_at.isoformat() if v.created_at else None, + } for v in versions] + finally: + db.close() + + # ---- GET /api/document/{doc_id}/version/{num} ---- + @router.get("/api/document/{doc_id}/version/{num}") + async def get_version(request: Request, doc_id: str, num: int) -> Dict[str, Any]: + user = get_current_user(request) + db = SessionLocal() + try: + # Verify ownership + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + ver = db.query(DocumentVersion).filter( + DocumentVersion.document_id == doc_id, + DocumentVersion.version_number == num, + ).first() + if not ver: + raise HTTPException(404, "Version not found") + return _version_to_dict(ver) + finally: + db.close() + + # ---- POST /api/document/{doc_id}/restore/{num} ---- + @router.post("/api/document/{doc_id}/restore/{num}") + async def restore_version(request: Request, doc_id: str, num: int) -> Dict[str, Any]: + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + + old_ver = db.query(DocumentVersion).filter( + DocumentVersion.document_id == doc_id, + DocumentVersion.version_number == num, + ).first() + if not old_ver: + raise HTTPException(404, "Version not found") + + new_ver_num = doc.version_count + 1 + ver = DocumentVersion( + id=str(uuid.uuid4()), + document_id=doc_id, + version_number=new_ver_num, + content=old_ver.content, + summary=f"Restored from v{num}", + source="user", + ) + doc.current_content = old_ver.content + doc.version_count = new_ver_num + db.add(ver) + db.commit() + db.refresh(doc) + return _doc_to_dict(doc) + except HTTPException: + raise + except Exception as e: + db.rollback() + raise HTTPException(500, str(e)) + finally: + db.close() + + # ---- POST /api/documents/tidy — clean up broken/empty documents ---- + @router.post("/api/documents/tidy") + async def tidy_documents(request: Request) -> Dict[str, Any]: + """Fix empty titles and remove broken/empty documents (user's docs only).""" + user = get_current_user(request) + db = SessionLocal() + try: + q = ( + db.query(Document) + .outerjoin(DbSession, Document.session_id == DbSession.id) + .filter(Document.is_active == True) + .filter((Document.archived == False) | (Document.archived.is_(None))) + ) + q = _owner_session_filter(q, user) + docs = q.all() + fixed_titles = 0 + deleted = 0 + + # Same junk-detection logic as the scheduled tidy_documents + # action (src/document_actions.py). Keep these two in sync. + import re as _re + from src.document_actions import _JUNK_TITLES + + to_delete = [] + now = datetime.now(timezone.utc) + for doc in docs: + created = doc.created_at + if created and created.tzinfo is None: + created = created.replace(tzinfo=timezone.utc) + + # Skip freshly created documents to avoid deleting them while the user is actively editing + if created and (now - created).total_seconds() < 900: # 15 minutes + continue + + content = (doc.current_content or "").strip() + title_raw = (doc.title or "").strip() + title = title_raw.lower() + is_fresh_empty = ( + not content + and created is not None + and (now - created).total_seconds() < 1800 + ) + if is_fresh_empty: + continue + + # Strip markdown noise to get a "real" character count + stripped = _re.sub(r"^#{1,6}\s+", "", content, flags=_re.MULTILINE) + stripped = _re.sub(r"[*_`>\-=]+", "", stripped) + stripped = _re.sub(r"\s+", " ", stripped).strip() + real_len = len(stripped) + + # Detect email-scaffold stubs: "To: \nSubject: \n---\n" style + # bodies with nothing typed in. Stub = every meaningful line + # is a header label (To:/From:/Subject:/...) with no real + # value (blank, "empty", "(empty)", "-", "none", "n/a"). + _is_email_stub = False + _HEADER_RE = _re.compile(r"^(to|from|cc|bcc|subject|reply-to):\s*(.*)$", _re.I) + _PLACEHOLDER_VALS = {"", "empty", "(empty)", "-", "—", "none", "n/a", "na", "tbd"} + if title in ("new email", "new mail", "new message") or doc.language == "email": + body_lines = [ln.strip() for ln in content.split("\n") + if ln.strip() and ln.strip() != "---"] + def _is_filler(ln): + m = _HEADER_RE.match(ln) + if not m: + return False + val = (m.group(2) or "").strip().lower() + return val in _PLACEHOLDER_VALS + has_real_body = any(not _is_filler(ln) for ln in body_lines) + if body_lines and not has_real_body: + _is_email_stub = True + + # Hard-delete obviously empty / junk documents + if not content or content in ("", "# Untitled"): + to_delete.append(doc); deleted += 1; continue + if _is_email_stub: + to_delete.append(doc); deleted += 1; continue + if title in _JUNK_TITLES: + to_delete.append(doc); deleted += 1; continue + + # Fix empty or placeholder titles on survivors + if not title_raw or title_raw == "Untitled": + new_title = _derive_title(content) + if new_title and new_title != "Untitled": + doc.title = new_title + fixed_titles += 1 + + for doc in to_delete: + db.delete(doc) + + # Also clean up inactive empty docs from previous soft-deletes + inactive_q = ( + db.query(Document) + .outerjoin(DbSession, Document.session_id == DbSession.id) + .filter(Document.is_active == False) + .filter((Document.current_content == None) | (Document.current_content == "")) + ) + inactive_q = _owner_session_filter(inactive_q, user) + inactive_docs = inactive_q.all() + for doc in inactive_docs: + db.delete(doc) + deleted += len(inactive_docs) + + db.commit() + return { + "fixed_titles": fixed_titles, + "deleted": deleted, + "message": f"Fixed {fixed_titles} title{'s' if fixed_titles != 1 else ''}, removed {deleted} empty document{'s' if deleted != 1 else ''}", + } + except Exception as e: + db.rollback() + logger.error(f"Document tidy failed: {e}") + raise HTTPException(500, f"Tidy failed: {e}") + finally: + db.close() + + # ---- POST /api/documents/ai-tidy — AI-powered cleanup of junk/test documents ---- + @router.post("/api/documents/ai-tidy") + async def ai_tidy_documents(request: Request) -> Dict[str, Any]: + """Use AI to judge if documents are junk/test/accidental, then delete them. + Caches verdicts so previously-reviewed docs are skipped.""" + from src.task_endpoint import resolve_task_endpoint + from src.endpoint_resolver import resolve_endpoint + from src.llm_core import llm_call_async + + user = get_current_user(request) + url, model, headers = resolve_task_endpoint(owner=user or None) + if not url or not model: + # Fall back to default endpoint + url, model, headers = resolve_endpoint("default", owner=user or None) + if not url or not model: + raise HTTPException(500, "No endpoint configured for AI tidy") + + db = SessionLocal() + try: + q = ( + db.query(Document) + .outerjoin(DbSession, Document.session_id == DbSession.id) + .filter(Document.is_active == True) + .filter((Document.archived == False) | (Document.archived.is_(None))) + ) + q = _owner_session_filter(q, user) + docs = q.all() + + # Only review docs that haven't been reviewed yet + to_review = [d for d in docs if not d.tidy_verdict] + if not to_review: + return {"deleted": 0, "reviewed": 0, "message": "All documents already reviewed"} + + # Build a batch prompt — review up to 30 at a time + batch = to_review[:30] + doc_list = [] + for i, doc in enumerate(batch): + preview = (doc.current_content or "")[:300].strip() + doc_list.append(f"[{i}] title=\"{doc.title}\" lang={doc.language or 'text'} content_preview=\"{preview}\"") + + prompt = ( + "You are a document library cleaner. For each document below, decide if it is JUNK " + "(test, accidental, placeholder, empty-ish, tool-test, throwaway) or KEEP (real content worth saving).\n\n" + "Respond with ONLY a JSON array of verdicts, one per document, like: [\"junk\",\"keep\",\"junk\",...]\n" + "No explanation, no markdown, just the JSON array.\n\n" + + "\n".join(doc_list) + ) + + response = await llm_call_async( + url, model, + [{"role": "system", "content": "You classify documents as junk or keep. Respond only with a JSON array."}, + {"role": "user", "content": prompt}], + temperature=0.1, + max_tokens=200, + headers=headers, + timeout=30, + ) + + # Parse verdicts + import re + match = re.search(r'\[.*?\]', response, re.DOTALL) + if not match: + raise HTTPException(500, "AI returned invalid response") + + import json as _json + verdicts = _json.loads(match.group()) + + deleted = 0 + reviewed = 0 + for i, doc in enumerate(batch): + if i >= len(verdicts): + break + verdict = str(verdicts[i] or "").lower().strip() + if verdict == "junk": + doc.tidy_verdict = "junk" + db.delete(doc) + deleted += 1 + else: + doc.tidy_verdict = "keep" + reviewed += 1 + + db.commit() + return { + "deleted": deleted, + "reviewed": reviewed, + "remaining": len(to_review) - len(batch), + "message": f"Reviewed {reviewed}, removed {deleted} junk document{'s' if deleted != 1 else ''}", + } + except HTTPException: + raise + except Exception as e: + db.rollback() + logger.error(f"AI tidy failed: {e}") + raise HTTPException(500, f"AI tidy failed: {e}") + finally: + db.close() + + # ---- POST /api/document/{doc_id}/export-pdf/preview ---- + @router.post("/api/document/{doc_id}/export-pdf/preview") + async def export_pdf_preview(doc_id: str, request: Request) -> Dict[str, Any]: + """Return the field-value mapping that would be written to the PDF. + + Frontend shows this in a confirmation modal so the user can spot/fix + any wrong values before triggering the actual download. + """ + from src.pdf_form_doc import find_source_upload_id, parse_markdown_to_values, load_field_sidecar + + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + + upload_id = find_source_upload_id(doc.current_content or "") + if not upload_id: + raise HTTPException(400, "Document is not linked to a source PDF") + + pdf_path = _locate_current_user_upload(request, upload_id, user) + if not pdf_path: + raise HTTPException(404, f"Source PDF {upload_id} not found in uploads") + + fields = load_field_sidecar(pdf_path) + if not fields: + raise HTTPException(404, "Field schema sidecar missing for source PDF") + + values = parse_markdown_to_values(doc.current_content or "") + field_meta = {f["name"]: f for f in fields} + + preview = [] + for name, current in values.items(): + meta = field_meta.get(name) + if not meta: + continue + preview.append({ + "name": name, + "label": meta.get("label") or name, + "type": meta.get("type"), + "options": meta.get("options") or [], + "page": meta.get("page"), + "value": current, + }) + + unknown = [ + name for name in values + if name not in field_meta + ] + return { + "doc_id": doc_id, + "upload_id": upload_id, + "fields": preview, + "unknown_fields": unknown, + "total": len(fields), + "filled": sum(1 for p in preview if p["value"] not in ("", False, None)), + } + finally: + db.close() + + # ---- GET /api/document/{doc_id}/render-pages ---- + @router.get("/api/document/{doc_id}/render-pages") + async def render_pages(doc_id: str, request: Request) -> Dict[str, Any]: + """Return per-page metadata for the interactive PDF view. + + Each page entry has its rendered-image dimensions (matching what + /page/{n}.png returns at the same DPI) plus the list of form fields + on that page with their rects translated to image-pixel coordinates. + Frontend overlays HTML form controls at those positions. + """ + from src.pdf_form_doc import find_source_upload_id, parse_markdown_to_values, load_field_sidecar + + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + upload_id = find_source_upload_id(doc.current_content or "") + if not upload_id: + raise HTTPException(400, "Document is not linked to a source PDF") + pdf_path = _locate_current_user_upload(request, upload_id, user) + if not pdf_path: + raise HTTPException(404, f"Source PDF {upload_id} not found") + + fitz = _load_pdf_viewer_fitz() + schema = load_field_sidecar(pdf_path) or [] + values = parse_markdown_to_values(doc.current_content or "") + + # Group fields by page + by_page: Dict[int, list] = {} + for f in schema: + by_page.setdefault(f["page"], []).append(f) + + scale = _PDF_RENDER_SCALE + pdf_doc = fitz.open(pdf_path) + try: + pages_out = [] + for page_index in range(pdf_doc.page_count): + page = pdf_doc[page_index] + page_no = page_index + 1 + pw, ph = page.rect.width, page.rect.height + img_w = int(pw * scale) + img_h = int(ph * scale) + fields_out = [] + for f in by_page.get(page_no, []): + x0, y0, x1, y1 = f["rect"] + fields_out.append({ + "name": f["name"], + "type": f["type"], + "label": f.get("label") or "", + "options": f.get("options") or [], + "value": values.get(f["name"], f.get("value", "")), + "rect_px": [ + int(x0 * scale), int(y0 * scale), + int(x1 * scale), int(y1 * scale), + ], + }) + pages_out.append({ + "page": page_no, + "width": img_w, + "height": img_h, + "fields": fields_out, + }) + return {"doc_id": doc_id, "scale": scale, "pages": pages_out} + finally: + pdf_doc.close() + finally: + db.close() + + # ---- GET /api/document/{doc_id}/page/{n}.png ---- + @router.get("/api/document/{doc_id}/page/{page_no}.png") + async def render_page_png(doc_id: str, page_no: int, request: Request): + """Render one page of the source PDF as a PNG (no values stamped — the + frontend overlays HTML form inputs on top).""" + from fastapi.responses import Response + from src.pdf_form_doc import find_source_upload_id + + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + upload_id = find_source_upload_id(doc.current_content or "") + if not upload_id: + raise HTTPException(400, "Document is not linked to a source PDF") + pdf_path = _locate_current_user_upload(request, upload_id, user) + if not pdf_path: + raise HTTPException(404, "Source PDF not found") + finally: + db.close() + + fitz = _load_pdf_viewer_fitz() + pdf_doc = fitz.open(pdf_path) + try: + if page_no < 1 or page_no > pdf_doc.page_count: + raise HTTPException(404, "Page out of range") + page = pdf_doc[page_no - 1] + mat = fitz.Matrix(_PDF_RENDER_SCALE, _PDF_RENDER_SCALE) + pix = page.get_pixmap(matrix=mat, alpha=False) + png_bytes = pix.tobytes("png") + return Response( + content=png_bytes, + media_type="image/png", + headers={"Cache-Control": "public, max-age=3600"}, + ) + finally: + pdf_doc.close() + + # ---- POST /api/document/{doc_id}/ai-fill-annotations ---- + @router.post("/api/document/{doc_id}/ai-fill-annotations") + async def ai_fill_annotations(doc_id: str, request: Request) -> Dict[str, Any]: + """Ask a vision-capable LLM to locate fillable areas on a flat PDF and + propose annotation values for each, given a free-form user instruction. + + Returns a list of annotations: [{page, x, y, w, h, value}] where x/y/w/h + are page-percentages (0–100) — same coordinate system as the freeform + annotations the frontend already renders. + """ + import base64 + import json + import fitz + from src.pdf_form_doc import find_source_upload_id + from src.document_processor import _resolve_vl_model, _load_vl_settings + from src.llm_core import llm_call_async + + body = await request.json() if request.headers.get("content-type", "").startswith("application/json") else {} + instruction = (body or {}).get("instruction", "").strip() + if not instruction: + raise HTTPException(400, "instruction is required") + + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + upload_id = find_source_upload_id(doc.current_content or "") + if not upload_id: + raise HTTPException(400, "Document is not linked to a source PDF") + pdf_path = _locate_current_user_upload(request, upload_id, user) + if not pdf_path: + raise HTTPException(404, "Source PDF not found") + finally: + db.close() + + # Resolve VL model (admin-configured or auto-detected vision-capable) + settings = _load_vl_settings() + vl_model = settings.get("vision_model", "") + try: + url, model_id, headers = _resolve_vl_model(vl_model, owner=user) + except Exception as e: + raise HTTPException(503, f"No vision model available: {e}") + + system_prompt = ( + "You analyze rendered PDF page images and propose values to fill in. " + "For each blank line, box, underscore, or labeled space on the page that " + "should be filled given the user's instruction, output one annotation. " + "Coordinates are percentages (0-100) of the page width/height with the " + "origin at top-left. Width/height should match the visible blank box. " + "Return ONLY a JSON array, no prose, no markdown fences. Each entry: " + '{"x": number, "y": number, "w": number, "h": number, "value": string}. ' + "If a region should not be filled, omit it. If nothing should be filled, " + "return []." + ) + + all_annotations = [] + pdf_doc = fitz.open(pdf_path) + try: + for page_index in range(pdf_doc.page_count): + page = pdf_doc[page_index] + mat = fitz.Matrix(_PDF_RENDER_SCALE, _PDF_RENDER_SCALE) + pix = page.get_pixmap(matrix=mat, alpha=False) + png_bytes = pix.tobytes("png") + b64 = base64.b64encode(png_bytes).decode("ascii") + + messages = [ + {"role": "system", "content": system_prompt}, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + f"User instruction:\n{instruction}\n\n" + f"This is page {page_index + 1} of {pdf_doc.page_count}. " + "Return JSON array of annotations to add to this page." + ), + }, + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{b64}"}, + }, + ], + }, + ] + try: + raw = await llm_call_async( + url, model_id, messages, + temperature=0.1, max_tokens=2000, headers=headers, + ) + except Exception as e: + logger.error(f"VL call failed on page {page_index + 1}: {e}") + continue + + raw = (raw or "").strip() + if raw.startswith("```"): + raw = raw.split("\n", 1)[-1].rsplit("```", 1)[0].strip() + try: + parsed = json.loads(raw) + except Exception: + logger.warning(f"AI fill: page {page_index + 1} returned non-JSON: {raw[:200]}") + continue + if not isinstance(parsed, list): + continue + for item in parsed: + if not isinstance(item, dict): + continue + try: + x = float(item.get("x", 0)) + y = float(item.get("y", 0)) + w = float(item.get("w", 0)) + h = float(item.get("h", 0)) + value = str(item.get("value", "") or "") + except Exception: + continue + # Clamp + reject zero-size entries + if w <= 0.5 or h <= 0.3: + continue + x = max(0.0, min(99.0, x)) + y = max(0.0, min(99.0, y)) + w = max(0.5, min(100.0 - x, w)) + h = max(0.3, min(100.0 - y, h)) + if not value.strip(): + continue + all_annotations.append({ + "page": page_index + 1, + "x": round(x, 2), + "y": round(y, 2), + "w": round(w, 2), + "h": round(h, 2), + "value": value, + }) + finally: + pdf_doc.close() + + return {"annotations": all_annotations} + + # ---- GET /api/document/{doc_id}/render-pdf ---- + @router.get("/api/document/{doc_id}/render-pdf") + async def render_pdf(doc_id: str, request: Request): + """Inline PDF preview filled with the current markdown values. + + Same plumbing as the export route, but no signature stamping and + served inline (Content-Disposition: inline) so the browser can + embed it in an iframe. Cache-busted by the caller via query string. + """ + import base64 + import os + import tempfile + from fastapi.responses import FileResponse + from starlette.background import BackgroundTask + from src.pdf_form_doc import find_source_upload_id, parse_markdown_to_values, parse_markdown_annotations + from src.pdf_forms import fill_fields, stamp_annotations + from core.database import Signature + + # Track temp files for this request so they get unlinked AFTER + # the response is fully sent (BackgroundTask runs post-send). + _to_unlink: list[str] = [] + def _cleanup_temps(): + for _p in _to_unlink: + try: + os.unlink(_p) + except FileNotFoundError: + pass + except Exception as _e: + logger.warning(f"Could not unlink temp PDF {_p}: {_e}") + + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + upload_id = find_source_upload_id(doc.current_content or "") + if not upload_id: + raise HTTPException(400, "Document is not linked to a source PDF") + pdf_path = _locate_current_user_upload(request, upload_id, user) + if not pdf_path: + raise HTTPException(404, f"Source PDF {upload_id} not found") + + # Fail fast with a clear 503 if the optional PyMuPDF dependency + # is missing — fill_fields/stamp_annotations will otherwise + # raise RuntimeError deep inside and bubble out as a 500. + # Mirrors the convention in _load_pdf_viewer_fitz above. + _load_pdf_viewer_fitz() + + values = parse_markdown_to_values(doc.current_content or "") + out_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name + _to_unlink.append(out_path) + try: + fill_fields(pdf_path, out_path, values) + except Exception as e: + logger.error(f"render_pdf fill_fields failed for {doc_id}: {e}") + _cleanup_temps() + raise HTTPException(500, f"PDF render failed: {e}") + + annotations = parse_markdown_annotations(doc.current_content or "") + if annotations: + ann_sig_ids = [ + a["value"][len("signature:"):].strip() + for a in annotations + if a.get("kind") == "signature" + and isinstance(a.get("value"), str) + and a["value"].startswith("signature:") + ] + ann_signature_pngs: dict[str, bytes] = {} + if ann_sig_ids: + # SECURITY: filter by owner so a caller can't reference + # someone else's signature ID from doc markdown and have + # it stamped/exported. + _sig_q = db.query(Signature).filter(Signature.id.in_(ann_sig_ids)) + if user: + _sig_q = _sig_q.filter(Signature.owner == user) + sig_rows = _sig_q.all() + for s in sig_rows: + try: + ann_signature_pngs[s.id] = base64.b64decode(s.data_png) + except Exception as e: + logger.warning(f"Bad annotation signature data for {s.id}: {e}") + annotated_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name + _to_unlink.append(annotated_path) + try: + stamp_annotations(out_path, annotated_path, annotations, ann_signature_pngs) + out_path = annotated_path + except Exception as e: + logger.error(f"stamp_annotations (render) failed for {doc_id}: {e}") + + return FileResponse( + out_path, + media_type="application/pdf", + headers={"Content-Disposition": "inline"}, + background=BackgroundTask(_cleanup_temps), + ) + finally: + db.close() + + # ---- GET /api/document/{doc_id}/export-pdf ---- + @router.get("/api/document/{doc_id}/export-pdf") + async def export_pdf(doc_id: str, request: Request): + """Stream the filled PDF for download. + + Reads field values and signature selections from the markdown — there + is no separate confirmation step. Signature fields contain their + chosen signature ID encoded as `signature:` in the value. + """ + import base64 + import os + import tempfile + from fastapi.responses import FileResponse + from starlette.background import BackgroundTask + from src.pdf_form_doc import find_source_upload_id, parse_markdown_to_values, load_field_sidecar, parse_markdown_annotations + from src.pdf_forms import fill_fields, stamp_signatures, stamp_annotations + from core.database import Signature + + _to_unlink: list[str] = [] + def _cleanup_temps(): + for _p in _to_unlink: + try: + os.unlink(_p) + except FileNotFoundError: + pass + except Exception as _e: + logger.warning(f"Could not unlink temp PDF {_p}: {_e}") + + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + + upload_id = find_source_upload_id(doc.current_content or "") + if not upload_id: + raise HTTPException(400, "Document is not linked to a source PDF") + + pdf_path = _locate_current_user_upload(request, upload_id, user) + if not pdf_path: + raise HTTPException(404, f"Source PDF {upload_id} not found in uploads") + + schema = load_field_sidecar(pdf_path) or [] + sig_field_names = {f["name"] for f in schema if f.get("type") == "signature"} + + all_values = parse_markdown_to_values(doc.current_content or "") + # Split: signature fields go to stamps, everything else to fill_fields + text_values: dict = {} + sig_ids: dict[str, str] = {} + for name, raw in all_values.items(): + if name in sig_field_names and isinstance(raw, str) and raw.startswith("signature:"): + sig_ids[name] = raw[len("signature:"):].strip() + elif name not in sig_field_names: + text_values[name] = raw + + stamps: dict = {} + if sig_ids: + # SECURITY: filter by owner — same reason as render_pdf. + _sig_q2 = db.query(Signature).filter(Signature.id.in_(list(sig_ids.values()))) + if user: + _sig_q2 = _sig_q2.filter(Signature.owner == user) + rows = _sig_q2.all() + by_id = {s.id: s for s in rows} + for field_name, sid in sig_ids.items(): + s = by_id.get(sid) + if not s: + continue + try: + stamps[field_name] = base64.b64decode(s.data_png) + except Exception as e: + logger.warning(f"Bad signature data for {sid}: {e}") + + filled_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name + _to_unlink.append(filled_path) + try: + fill_fields(pdf_path, filled_path, text_values) + except Exception as e: + logger.error(f"fill_fields failed for doc {doc_id}: {e}") + _cleanup_temps() + raise HTTPException(500, f"PDF fill failed: {e}") + + out_path = filled_path + if stamps: + stamped_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name + _to_unlink.append(stamped_path) + try: + stamp_signatures(filled_path, stamped_path, stamps) + out_path = stamped_path + except Exception as e: + logger.error(f"stamp_signatures failed for doc {doc_id}: {e}") + + # Burn freeform annotations (Text/Check/Sign drops) on top. + annotations = parse_markdown_annotations(doc.current_content or "") + if annotations: + # Resolve any signature annotations to their PNG bytes. + ann_sig_ids = [ + a["value"][len("signature:"):].strip() + for a in annotations + if a.get("kind") == "signature" + and isinstance(a.get("value"), str) + and a["value"].startswith("signature:") + ] + ann_signature_pngs: dict[str, bytes] = {} + if ann_sig_ids: + # SECURITY: filter by owner so a caller can't reference + # someone else's signature ID from doc markdown and have + # it stamped/exported. + _sig_q = db.query(Signature).filter(Signature.id.in_(ann_sig_ids)) + if user: + _sig_q = _sig_q.filter(Signature.owner == user) + sig_rows = _sig_q.all() + for s in sig_rows: + try: + ann_signature_pngs[s.id] = base64.b64decode(s.data_png) + except Exception as e: + logger.warning(f"Bad annotation signature data for {s.id}: {e}") + annotated_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name + _to_unlink.append(annotated_path) + try: + stamp_annotations(out_path, annotated_path, annotations, ann_signature_pngs) + out_path = annotated_path + except Exception as e: + logger.error(f"stamp_annotations failed for doc {doc_id}: {e}") + + download_name = _slug(doc.title or "form") + "_annotated.pdf" + return FileResponse( + out_path, + media_type="application/pdf", + filename=download_name, + background=BackgroundTask(_cleanup_temps), + ) + finally: + db.close() + + # ---- POST /api/document/{doc_id}/prepare-signed-reply ---- + @router.post("/api/document/{doc_id}/prepare-signed-reply") + async def prepare_signed_reply(doc_id: str, request: Request): + """Bake the current PDF state (form fields + signature stamps + + annotations) into a flattened PDF, drop it in COMPOSE_UPLOADS_DIR + and return the reply context (To/Subject/threading headers) so the + frontend can open a reply draft with this attachment pre-loaded. + + Requires the document to have source_email_* metadata (set when the + doc was created via /api/email/attachment-as-doc). Otherwise 400. + """ + import base64 + import tempfile + import shutil + import uuid as _uuid + import email as _email_mod + from src.pdf_form_doc import ( + find_source_upload_id, parse_markdown_to_values, + load_field_sidecar, parse_markdown_annotations, + ) + from src.pdf_forms import fill_fields, stamp_signatures, stamp_annotations + from core.database import Signature + # COMPOSE_UPLOADS_DIR lives in email_routes — re-derive here so we + # don't import from a routes file (cycle-prone). Same env override + # as email_routes (ODYSSEUS_MAIL_ATTACHMENTS_DIR). + from pathlib import Path as _Path + _COMPOSE_DIR = _Path(MAIL_ATTACHMENTS_DIR) / "_compose" + _COMPOSE_DIR.mkdir(parents=True, exist_ok=True) + + user = get_current_user(request) + db = SessionLocal() + try: + doc = db.query(Document).filter(Document.id == doc_id).first() + if not doc: + raise HTTPException(404, "Document not found") + _verify_doc_owner(db, doc, user) + + if not (doc.source_email_uid and doc.source_email_folder): + raise HTTPException(400, "Document has no source email — cannot reply") + + # 1) Build the flattened PDF (same pipeline as export_pdf) + upload_id = find_source_upload_id(doc.current_content or "") + if not upload_id: + raise HTTPException(400, "Document is not linked to a source PDF") + pdf_path = _locate_current_user_upload(request, upload_id, user) + if not pdf_path: + raise HTTPException(404, f"Source PDF {upload_id} not found") + + schema = load_field_sidecar(pdf_path) or [] + sig_field_names = {f["name"] for f in schema if f.get("type") == "signature"} + all_values = parse_markdown_to_values(doc.current_content or "") + text_values: dict = {} + sig_ids: dict[str, str] = {} + for name, raw in all_values.items(): + if name in sig_field_names and isinstance(raw, str) and raw.startswith("signature:"): + sig_ids[name] = raw[len("signature:"):].strip() + elif name not in sig_field_names: + text_values[name] = raw + + stamps: dict = {} + if sig_ids: + # SECURITY: filter by owner — same reason as render_pdf. + _sig_q2 = db.query(Signature).filter(Signature.id.in_(list(sig_ids.values()))) + if user: + _sig_q2 = _sig_q2.filter(Signature.owner == user) + rows = _sig_q2.all() + by_id = {s.id: s for s in rows} + for fname, sid in sig_ids.items(): + s = by_id.get(sid) + if not s: + continue + try: + stamps[fname] = base64.b64decode(s.data_png) + except Exception: + pass + + import os + _to_unlink: list[str] = [] + filled_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name + _to_unlink.append(filled_path) + fill_fields(pdf_path, filled_path, text_values) + out_path = filled_path + if stamps: + stamped_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name + _to_unlink.append(stamped_path) + try: + stamp_signatures(filled_path, stamped_path, stamps) + out_path = stamped_path + except Exception as e: + logger.warning(f"stamp_signatures failed for {doc_id}: {e}") + + annotations = parse_markdown_annotations(doc.current_content or "") + if annotations: + ann_sig_ids = [ + a["value"][len("signature:"):].strip() + for a in annotations + if a.get("kind") == "signature" + and isinstance(a.get("value"), str) + and a["value"].startswith("signature:") + ] + ann_signature_pngs: dict[str, bytes] = {} + if ann_sig_ids: + # SECURITY: filter by owner so a caller can't reference + # someone else's signature ID from doc markdown and have + # it stamped/exported. + _sig_q = db.query(Signature).filter(Signature.id.in_(ann_sig_ids)) + if user: + _sig_q = _sig_q.filter(Signature.owner == user) + sig_rows = _sig_q.all() + for s in sig_rows: + try: + ann_signature_pngs[s.id] = base64.b64decode(s.data_png) + except Exception: + pass + annotated_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name + _to_unlink.append(annotated_path) + try: + stamp_annotations(out_path, annotated_path, annotations, ann_signature_pngs) + out_path = annotated_path + except Exception as e: + logger.warning(f"stamp_annotations failed for {doc_id}: {e}") + + # 2) Move/copy into COMPOSE_UPLOADS_DIR with the token format + # `_` that /api/email/send expects. + filename = _slug(doc.title or "signed") + "_signed.pdf" + token = f"{_uuid.uuid4().hex}_{filename}" + dest = _COMPOSE_DIR / token + shutil.copyfile(out_path, str(dest)) + # Unlink the intermediate temp PDFs now that they've been + # copied into COMPOSE_UPLOADS_DIR. + for _p in _to_unlink: + try: + os.unlink(_p) + except FileNotFoundError: + pass + except Exception as _e: + logger.warning(f"Could not unlink temp PDF {_p}: {_e}") + + # 3) Fetch the source email's headers so we can build a clean reply + # context (To/Subject/In-Reply-To/References). + try: + from routes.email_routes import _imap, _decode_header + from routes.email_helpers import _q + except Exception: + _imap = None + _decode_header = lambda x: x or "" + _q = lambda x: x or "" + + to_addr = "" + from_name = "" + subject = "" + in_reply_to = doc.source_email_message_id or "" + references = in_reply_to + if _imap: + try: + with _imap(doc.source_email_account_id or None) as conn: + conn.select(_q(doc.source_email_folder), readonly=True) + status, data = conn.fetch(doc.source_email_uid.encode(), "(RFC822.HEADER)") + if status == "OK" and data and data[0]: + raw_hdr = data[0][1] + m = _email_mod.message_from_bytes(raw_hdr) + sender = _decode_header(m.get("From", "")) + from_name, to_addr = _email_mod.utils.parseaddr(sender) + if not to_addr: + to_addr = sender + subject = _decode_header(m.get("Subject", "") or "") + if subject and not subject.lower().startswith("re:"): + subject = "Re: " + subject + msg_refs = (m.get("References") or "").strip() + msg_in_reply = (m.get("Message-ID") or "").strip() or in_reply_to + in_reply_to = msg_in_reply + references = (msg_refs + " " + msg_in_reply).strip() if msg_refs else msg_in_reply + except Exception as e: + logger.warning(f"prepare-signed-reply header fetch failed: {e}") + + return { + "ok": True, + "attachment": { + "token": token, + "filename": filename, + "size": dest.stat().st_size, + }, + "reply": { + "to": to_addr, + "to_name": from_name, + "subject": subject, + "in_reply_to": in_reply_to, + "references": references, + "account_id": doc.source_email_account_id or None, + "source_uid": doc.source_email_uid, + "source_folder": doc.source_email_folder, + "source_message_id": doc.source_email_message_id, + }, + } + finally: + db.close() + + return router diff --git a/routes/document_helpers.py b/routes/document_helpers.py index a0c2d08eb..c1f68ca51 100644 --- a/routes/document_helpers.py +++ b/routes/document_helpers.py @@ -1,243 +1,14 @@ -"""document_helpers.py — Pydantic models, doc serializers, owner gating, file-locator helpers shared with document_routes.py.""" +"""Backward-compat shim — canonical location is routes/document/document_helpers.py. -"""Document routes — CRUD for living documents with version history.""" +This module is replaced in ``sys.modules`` by the canonical module object so +that ``import routes.document_helpers``, ``from routes.document_helpers import +X``, and the ``sys.modules.pop("routes.document_helpers")`` + re-import +pattern used by test_security_regressions.py all operate on the *same* object. +Keeps existing import paths working after slice 2m (#4082/#4071). +""" -import logging -import os -import re -from typing import Any, Dict, Optional +import sys as _sys -from fastapi import HTTPException, Request -from pydantic import BaseModel +from routes.document import document_helpers as _canonical # noqa: F401 -from core.database import Document, DocumentVersion -from core.database import Session as DbSession -from src.auth_helpers import _auth_disabled -from src.upload_handler import UploadHandler - -logger = logging.getLogger(__name__) - - -# ---- Request schemas ---- - -class DocumentCreate(BaseModel): - session_id: Optional[str] = None - title: str = "Untitled" - language: Optional[str] = None - content: str = "" - -class DocumentUpdate(BaseModel): - content: str - summary: Optional[str] = None - force_version: bool = False - -class DocumentPatch(BaseModel): - title: Optional[str] = None - language: Optional[str] = None - session_id: Optional[str] = None # link/unlink document to a session - - -# ---- Helpers ---- - -def _doc_to_dict(doc: Document) -> Dict[str, Any]: - return { - "id": doc.id, - "session_id": doc.session_id, - "title": doc.title, - "language": doc.language, - "current_content": doc.current_content, - "version_count": doc.version_count, - "is_active": doc.is_active, - "archived": bool(getattr(doc, "archived", False)), - "created_at": (doc.created_at.isoformat() + "Z") if doc.created_at else None, - "updated_at": (doc.updated_at.isoformat() + "Z") if doc.updated_at else None, - # Source-email provenance (set when doc was created from an email - # attachment) — drives the "Send signed reply" menu item. - "source_email_uid": getattr(doc, "source_email_uid", None), - "source_email_folder": getattr(doc, "source_email_folder", None), - "source_email_account_id": getattr(doc, "source_email_account_id", None), - "source_email_message_id": getattr(doc, "source_email_message_id", None), - } - -def _version_to_dict(v: DocumentVersion) -> Dict[str, Any]: - return { - "id": v.id, - "document_id": v.document_id, - "version_number": v.version_number, - "content": v.content, - "summary": v.summary, - "source": v.source, - "created_at": v.created_at.isoformat() if v.created_at else None, - } - - -def _verify_doc_owner(db, doc: Document, user: str): - """Verify `user` owns this document. Raise 404 if not. - - Documents now carry their own `owner` column, so a doc whose session - was deleted (session_id → NULL) can still prove ownership and stay - openable / cloneable. We trust that column first and only fall back to - the session join for any not-yet-backfilled legacy row. - """ - if user is None: - if _auth_disabled(): - return # Single-user / no-auth mode: allow access - raise HTTPException(403, "Authentication required") - if doc.owner is not None: - if doc.owner != user: - raise HTTPException(404, "Document not found") - return - # Legacy fallback: derive ownership from the linked session. - if not doc.session_id: - raise HTTPException(404, "Document not found") - session = db.query(DbSession).filter(DbSession.id == doc.session_id).first() - if not session or session.owner != user: - raise HTTPException(404, "Document not found") - - -def _owner_session_filter(q, user): - """Restrict a documents query to those owned by `user`. - - Documents now carry their own `owner` column (backfilled at boot from - the linked session, or assigned to the admin user for legacy/orphaned - docs). We filter on that directly rather than on a session join, so a - document whose session was deleted (session_id → NULL) still shows up - for its owner instead of silently vanishing from the Library + search. - - The owner backfill runs in init_db before the app serves requests, so - by the time this filter is live there are no NULL-owner rows to leak; - we therefore match the owner strictly for authenticated callers.""" - if not user: - if user == "" or _auth_disabled(): - return q - return q.filter(False) - return q.filter(Document.owner == user) - - - -def _slug(name: str) -> str: - """Filesystem-friendly version of a document title. - - Whitespace becomes underscores; other unsafe punctuation is dropped. - Preserves letters, digits, dot, hyphen, underscore. Idempotent. - """ - import re as _re - s = (name or "").strip() - # Drop the trailing extension if the title happens to include one - s = _re.sub(r'\.pdf$', '', s, flags=_re.IGNORECASE) - s = _re.sub(r'\s+', '_', s) - s = _re.sub(r'[^A-Za-z0-9._-]', '', s) - s = _re.sub(r'_+', '_', s).strip('_') - return s or "form" - - -# DPI scale for the interactive PDF view. ~150 DPI (2x of 72 PDF user-units). -_PDF_RENDER_SCALE = 2.0 - - -def _upload_path_inside(upload_dir: str, path: str) -> bool: - base = os.path.realpath(upload_dir) - p = os.path.realpath(path) - try: - return os.path.commonpath([base, p]) == base - except Exception: - return False - - -def _resolve_user_upload_path( - upload_handler: Any, - upload_id: str, - owner: Optional[str], - auth_manager=None, -) -> Optional[str]: - """Resolve an upload id to a filesystem path the caller may read.""" - if upload_handler is None: - return None - resolved = upload_handler.resolve_upload( - upload_id, - owner=owner, - auth_manager=auth_manager, - ) - if not isinstance(resolved, dict) or not resolved: - return None - path = resolved.get("path") - upload_dir = getattr(upload_handler, "upload_dir", None) - if path and upload_dir and not _upload_path_inside(upload_dir, path): - logger.warning("Upload path outside upload directory: %s", path) - return None - return path - - -def _locate_upload( - upload_dir: str, - file_id: str, - owner: Optional[str] = None, - auth_manager=None, - upload_handler: Any = None, -): - """Find an upload by its filename ID via UploadHandler.resolve_upload.""" - if upload_handler is None: - from src.upload_handler import UploadHandler - - base_dir = os.path.dirname(os.path.abspath(upload_dir)) - upload_handler = UploadHandler(base_dir, upload_dir) - return _resolve_user_upload_path(upload_handler, file_id, owner, auth_manager) - - -def _assert_pdf_marker_upload_owned( - request: Request, - content: str, - user: Optional[str], - upload_handler: Any, -) -> None: - """Reject document content whose pdf_source marker points at another user's upload.""" - if upload_handler is None: - return - from src.pdf_form_doc import find_source_upload_id - - upload_id = find_source_upload_id(content or "") - if not upload_id: - return - auth_manager = getattr(getattr(request.app, "state", None), "auth_manager", None) - if not _resolve_user_upload_path(upload_handler, upload_id, user, auth_manager): - raise HTTPException( - 400, - "Document PDF marker references an upload you do not own", - ) - - -def _derive_title(content: str) -> str: - """Derive a title from document content.""" - import re - if not isinstance(content, str): - return "Untitled" - text = content.strip() - if not text: - return "Untitled" - - # Markdown header - md = re.match(r'^#{1,3}\s+(.+)', text, re.MULTILINE) - if md: - title = md.group(1).strip() - if len(title) > 50: - title = title[:48] + "…" - return title - - # HTML heading - html = re.search(r']*>([^<]+)', text, re.IGNORECASE) - if html: - title = html.group(1).strip() - if len(title) > 50: - title = title[:48] + "…" - return title - - # First non-empty line (if short enough) - for line in text.split('\n'): - line = line.strip() - if line and 2 <= len(line) <= 60: - title = re.sub(r'[:#*`]+$', '', line).strip() - if title and len(title) > 50: - title = title[:48] + "…" - return title or "Untitled" - - return "Untitled" +_sys.modules[__name__] = _canonical diff --git a/routes/document_routes.py b/routes/document_routes.py index dae8b09fa..dd13e3c60 100644 --- a/routes/document_routes.py +++ b/routes/document_routes.py @@ -1,1810 +1,17 @@ -"""Document routes — CRUD for living documents with version history.""" +"""Backward-compat shim — canonical location is routes/document/document_routes.py. -import uuid -import logging -from datetime import datetime, timezone -from typing import Dict, Any, List, Optional +This module is replaced in ``sys.modules`` by the canonical module object so +that ``import routes.document_routes``, ``from routes.document_routes import +X``, ``importlib.import_module("routes.document_routes")``, and the +``import ... as droutes`` + ``droutes.SessionLocal = ...`` / +``monkeypatch.setattr(droutes, ...)`` pattern used by multiple tests all +operate on the *same* object the application actually uses. Keeps existing +import paths working after slice 2m (#4082/#4071). Source-introspection tests +read the canonical file by path. +""" -from fastapi import APIRouter, HTTPException, Query, Request, UploadFile, File, Form +import sys as _sys -from sqlalchemy import case, func, or_ -from core.database import SessionLocal, Document, DocumentVersion -from core.database import Session as DbSession -from src.auth_helpers import get_current_user, _auth_disabled -from src.constants import MAIL_ATTACHMENTS_DIR -from src.upload_handler import reserve_upload_references +from routes.document import document_routes as _canonical # noqa: F401 -logger = logging.getLogger(__name__) - - -def _get_session_or_404(db, session_id: str, user: Optional[str]): - session = db.query(DbSession).filter(DbSession.id == session_id).first() - if not session: - raise HTTPException(404, "Session not found") - if user and session.owner != user: - raise HTTPException(404, "Session not found") - return session - - -def _aggregate_language_facets(lang_rows): - """Sum document counts per display language for the library facet. - - NULL-language and explicit "text" rows share the "text" bucket (the - language filter treats them as one), so they must be ADDED. The old dict - comprehension keyed both to "text", silently overwriting one group and - undercounting the facet versus what the filter actually returns. - """ - out = {} - for lang, cnt in lang_rows: - key = lang or "text" - out[key] = out.get(key, 0) + cnt - return out - - -def _library_language_for_document(doc: Document) -> str: - """Return the display language used by the document library. - - PDF documents are stored as markdown wrappers so the editor can preserve - extracted text, form fields, and annotations. The library should still - identify them as PDFs instead of exposing that internal wrapper format. - """ - from src.pdf_form_doc import find_source_upload_id - - if find_source_upload_id(doc.current_content or ""): - return "pdf" - return doc.language or "text" - - -def _email_source_key(content: str) -> tuple[str, str]: - """Return the source email identity embedded in an email draft document.""" - import re - - text = content or "" - uid_m = re.search(r"(?im)^X-Source-UID:\s*(.+?)\s*$", text) - folder_m = re.search(r"(?im)^X-Source-Folder:\s*(.+?)\s*$", text) - uid = (uid_m.group(1).strip() if uid_m else "") - folder = (folder_m.group(1).strip() if folder_m else "INBOX") - return uid, folder - - -from routes.document_helpers import ( - DocumentCreate, DocumentUpdate, DocumentPatch, - _doc_to_dict, _version_to_dict, - _verify_doc_owner, _owner_session_filter, - _slug, _resolve_user_upload_path, _assert_pdf_marker_upload_owned, _derive_title, - _PDF_RENDER_SCALE, -) - - -def setup_document_routes(session_manager, upload_handler=None) -> APIRouter: - router = APIRouter(tags=["documents"]) - - def _reserve_document_uploads(user: Optional[str], content: str) -> None: - missing_id = reserve_upload_references(upload_handler, user, content) - if missing_id: - raise HTTPException( - 409, - f"Referenced upload is no longer available: {missing_id}", - ) - - def _locate_current_user_upload(request: Request, upload_id: str, user: Optional[str]): - if upload_handler is None: - return None - auth_manager = getattr(getattr(request.app, "state", None), "auth_manager", None) - return _resolve_user_upload_path(upload_handler, upload_id, user, auth_manager) - - def _load_pdf_viewer_fitz(): - from src.pdf_runtime import load_pymupdf_for_pdf_viewer - - try: - return load_pymupdf_for_pdf_viewer() - except RuntimeError as exc: - raise HTTPException(503, str(exc)) from exc - - # ---- POST /api/document ---- - @router.post("/api/document") - async def create_document(request: Request, req: DocumentCreate) -> Dict[str, Any]: - from src.auth_helpers import require_privilege - user = require_privilege(request, "can_use_documents") - db = SessionLocal() - try: - # session_id is optional: a doc can be a session-less "library" doc - # (e.g. files imported from the library) — session_id is nullable and - # the doc is owner-stamped, so it lives in the library on its own. - session = None - if req.session_id: - # Match the lenient ownership model the rest of the app uses - # (see _owner_filter): only block when an AUTHENTICATED user is - # writing into a DIFFERENT user's session. In single-user / - # unconfigured / localhost-bypass mode, falsey users preserve - # the existing lenient path. - session = _get_session_or_404(db, req.session_id, user) - - # If no language was supplied (e.g. cloning a doc whose language - # was never set), detect it from the content rather than storing - # NULL — which made the editor fall back to plain text. Defaults - # to markdown for prose. - language = req.language - if not language: - from src.agent_tools.document_tools import _looks_like_email_document, _sniff_doc_language, _coerce_email_document_content - language = _sniff_doc_language(req.content) - else: - from src.agent_tools.document_tools import _looks_like_email_document, _coerce_email_document_content - if _looks_like_email_document(req.content, req.title): - language = "email" - - _reserve_document_uploads(user, req.content) - _assert_pdf_marker_upload_owned(request, req.content, user, upload_handler) - - # Reply drafts are keyed to the source email. If a UI/tool path tries - # to create a second draft for the same email in the same chat, - # update the existing draft instead so quoted thread history stays - # attached to the visible document. - if language == "email" and req.session_id: - source_uid, source_folder = _email_source_key(req.content) - if source_uid: - candidates = ( - db.query(Document) - .filter(Document.session_id == req.session_id) - .filter(Document.is_active == True) - .filter(Document.language == "email") - .order_by(Document.updated_at.desc()) - .limit(25) - .all() - ) - for existing in candidates: - old_uid, old_folder = _email_source_key(existing.current_content or "") - if old_uid != source_uid or old_folder != source_folder: - continue - merged = _coerce_email_document_content(existing.current_content or "", req.content) - if existing.current_content != merged: - new_ver = (existing.version_count or 1) + 1 - existing.current_content = merged - existing.title = req.title or existing.title - existing.version_count = new_ver - db.add(DocumentVersion( - id=str(uuid.uuid4()), - document_id=existing.id, - version_number=new_ver, - content=merged, - summary="Updated existing email draft", - source="user", - )) - db.commit() - db.refresh(existing) - return _doc_to_dict(existing) - - doc_id = str(uuid.uuid4()) - ver_id = str(uuid.uuid4()) - - doc = Document( - id=doc_id, - session_id=req.session_id, - title=req.title, - language=language, - current_content=req.content, - version_count=1, - is_active=True, - # Stamp ownership directly so the doc survives its session - # being deleted. Fall back to the session's owner when the - # request is unauthenticated (single-user / localhost bypass). - owner=user or (session.owner if session else None), - ) - ver = DocumentVersion( - id=ver_id, - document_id=doc_id, - version_number=1, - content=req.content, - summary="Initial version", - source="user", - ) - db.add(doc) - db.add(ver) - db.commit() - db.refresh(doc) - try: - from src.event_bus import fire_event - fire_event("document_created", doc.owner) - except Exception: - logger.debug("document_created event dispatch failed", exc_info=True) - return _doc_to_dict(doc) - except HTTPException: - raise - except Exception as e: - db.rollback() - logger.error(f"Failed to create document: {e}") - raise HTTPException(500, f"Failed to create document: {e}") - finally: - db.close() - - # ---- POST /api/documents/import-pdf ---- - @router.post("/api/documents/import-pdf") - async def import_pdf( - request: Request, - file: UploadFile = File(...), - session_id: Optional[str] = Form(None), - ) -> Dict[str, Any]: - """Upload a PDF and create the matching Document. - - Detects AcroForm fields — if any, creates a form-backed markdown doc - (clickable inputs in the PDF view). Otherwise creates a plain PDF doc - with a `pdf_source` marker so the viewer renders the pages without - overlays. - """ - from src.pdf_forms import has_form_fields, extract_fields - from src.pdf_form_doc import ( - save_field_sidecar, - create_form_markdown_document, - create_plain_pdf_document, - ) - from src.document_processor import _process_pdf, strip_pdf_content_marker - import os - - from src.auth_helpers import require_privilege - user = require_privilege(request, "can_use_documents") - - # session_id is optional — a library import isn't tied to a chat. When - # given, validate it; otherwise the PDF becomes a session-less library - # doc (the doc creators below already handle a missing session). - if session_id: - db = SessionLocal() - try: - _get_session_or_404(db, session_id, user) - finally: - db.close() - - if upload_handler is None: - raise HTTPException(500, "Upload handler not configured") - - client_ip = request.client.host if request.client else "unknown" - try: - meta = upload_handler.save_upload(file, client_ip, owner=user) - except HTTPException: - raise - except Exception as e: - logger.error(f"PDF import save_upload failed: {e}") - raise HTTPException(500, f"Upload failed: {e}") - - upload_id = meta["id"] - pdf_path = _locate_current_user_upload(request, upload_id, user) - if not pdf_path: - raise HTTPException(500, "Saved PDF could not be located") - - title = os.path.splitext(meta.get("original_name") or meta.get("name") or upload_id)[0] - try: - body_text = strip_pdf_content_marker(_process_pdf(pdf_path, owner=user)) - except Exception: - body_text = None - - is_form = False - try: - is_form = has_form_fields(pdf_path) - except Exception as e: - logger.warning(f"has_form_fields failed for {pdf_path}: {e}") - - if is_form: - fields = extract_fields(pdf_path) - save_field_sidecar(pdf_path, fields) - doc_id = create_form_markdown_document( - session_id=session_id, - fields=fields, - upload_id=upload_id, - title=title, - intro_text=body_text, - ) - else: - doc_id = create_plain_pdf_document( - session_id=session_id, - upload_id=upload_id, - title=title, - body_text=body_text, - ) - - if not doc_id: - raise HTTPException(500, "Failed to create document for PDF") - - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(500, "Created document not found") - # The PDF doc creators stamp owner from the session only; a - # session-less library import leaves owner NULL, which the Library's - # owner filter then hides. Stamp the requesting user so it shows. - if not doc.owner and user: - doc.owner = user - db.commit() - db.refresh(doc) - return _doc_to_dict(doc) - finally: - db.close() - - # ---- GET /api/documents/library ---- - @router.get("/api/documents/library") - async def documents_library( - request: Request, - search: Optional[str] = Query(None), - language: Optional[str] = Query(None), - sort: str = Query("recent"), - offset: int = Query(0, ge=0), - limit: int = Query(20, ge=1, le=50), - archived: bool = Query(False), - ) -> Dict[str, Any]: - user = get_current_user(request) - db = SessionLocal() - try: - from sqlalchemy import or_ - pdf_marker_cond = or_( - Document.current_content.like('%\s*\n+#[^\n]*\n+)', re.MULTILINE) - head_match = head_re.match(content) - head = head_match.group(1) if head_match else (content.splitlines()[0] + "\n\n# " + (doc.title or "PDF") + "\n\n") - doc.current_content = head + body_text.strip() + "\n" - doc.version_count = (doc.version_count or 1) + 1 - db.add(DocumentVersion( - id=str(__import__("uuid").uuid4()), - document_id=doc_id, - version_number=doc.version_count, - content=doc.current_content, - summary="PDF text re-extracted (OCR)", - source="ocr", - )) - db.commit() - return {"ok": True, "id": doc_id, "extracted": True, "chars": len(body_text)} - finally: - db.close() - - # ---- POST /api/documents/export-zip — bundle selected docs into a .zip ---- - @router.post("/api/documents/export-zip") - async def documents_export_zip(request: Request): - """Zip the selected documents (each as a text file with the right - extension) — mirrors the gallery's bulk download-zip so multi-export - is one file instead of a blocked flood of individual downloads.""" - user = get_current_user(request) - try: - data = await request.json() - except Exception as e: - logger.warning("Failed to parse export request body, defaulting to empty", exc_info=e) - data = {} - ids = data.get("ids") or [] - if not ids: - raise HTTPException(400, "No documents specified") - _ext = { - "javascript": ".js", "python": ".py", "html": ".html", "css": ".css", - "markdown": ".md", "json": ".json", "yaml": ".yml", "bash": ".sh", - "sql": ".sql", "rust": ".rs", "go": ".go", "java": ".java", "c": ".c", - "cpp": ".cpp", "typescript": ".ts", "ruby": ".rb", "php": ".php", - "text": ".txt", "xml": ".xml", "toml": ".toml", "ini": ".ini", - } - db = SessionLocal() - try: - import io - import re - import zipfile - from fastapi import Response - docs = db.query(Document).filter(Document.id.in_(ids)).all() - buf = io.BytesIO() - used = set() - wrote = 0 - with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf: - for doc in docs: - try: - _verify_doc_owner(db, doc, user) - except HTTPException: - continue # skip docs the user doesn't own - ext = _ext.get(doc.language or "text", ".txt") - base = (doc.title or "document").strip() or "document" - base = re.sub(r"[^\w\-. ]+", "", base)[:60].strip() or doc.id - name = base if "." in base else base + ext - i = 1 - while name in used: - name = f"{base}-{i}" + ("" if "." in base else ext) - i += 1 - used.add(name) - zf.writestr(name, doc.current_content or "") - wrote += 1 - if not wrote: - raise HTTPException(404, "No documents found") - return Response( - content=buf.getvalue(), - media_type="application/zip", - headers={"Content-Disposition": 'attachment; filename="documents.zip"'}, - ) - finally: - db.close() - - # ---- PUT /api/document/{doc_id} — user manual edit ---- - # Coalesce window: if the last user version was saved within this many - # seconds, update it in-place (user is still actively editing). - # Once the gap exceeds this, the next save creates a new version. - VERSION_COALESCE_SECONDS = 60 - - @router.put("/api/document/{doc_id}") - async def update_document(request: Request, doc_id: str, req: DocumentUpdate) -> Dict[str, Any]: - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - - incoming_content = req.content - from src.agent_tools.document_tools import _coerce_email_document_content, _looks_like_email_document - is_email_doc = ( - (doc.language or "").lower() == "email" - or _looks_like_email_document(doc.current_content or "", doc.title or "") - or _looks_like_email_document(req.content or "", doc.title or "") - ) - if is_email_doc: - incoming_content = _coerce_email_document_content(doc.current_content or "", req.content) - doc.language = "email" - - # Skip if content is identical unless the caller explicitly wants - # a checkpoint version from the current editor state. - if doc.current_content == incoming_content and not req.force_version: - return _doc_to_dict(doc) - - _reserve_document_uploads(user, incoming_content) - _assert_pdf_marker_upload_owned(request, incoming_content, user, upload_handler) - - # Check if we can coalesce with the latest version - latest_ver = db.query(DocumentVersion).filter( - DocumentVersion.document_id == doc_id, - ).order_by(DocumentVersion.version_number.desc()).first() - - now = datetime.now(timezone.utc) - coalesced = False - if latest_ver and latest_ver.source == "user" and not req.force_version: - ver_time = latest_ver.created_at - if ver_time.tzinfo is None: - ver_time = ver_time.replace(tzinfo=timezone.utc) - age = (now - ver_time).total_seconds() - if age < VERSION_COALESCE_SECONDS: - # Update the existing version in-place - latest_ver.content = incoming_content - latest_ver.created_at = now - if req.summary: - latest_ver.summary = req.summary - coalesced = True - - if not coalesced: - new_ver = doc.version_count + 1 - ver = DocumentVersion( - id=str(uuid.uuid4()), - document_id=doc_id, - version_number=new_ver, - content=incoming_content, - summary=req.summary or "Manual edit", - source="user", - ) - doc.version_count = new_ver - db.add(ver) - - doc.current_content = incoming_content - db.commit() - db.refresh(doc) - return _doc_to_dict(doc) - except HTTPException: - raise - except Exception as e: - db.rollback() - raise HTTPException(500, f"Failed to update document: {e}") - finally: - db.close() - - # ---- PATCH /api/document/{doc_id} — metadata only ---- - @router.patch("/api/document/{doc_id}") - async def patch_document(request: Request, doc_id: str, req: DocumentPatch) -> Dict[str, Any]: - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - if req.title is not None: - doc.title = req.title - if req.language is not None: - doc.language = req.language - if req.session_id is not None: - # Empty string = unlink from session - if req.session_id: - _get_session_or_404(db, req.session_id, user) - doc.session_id = req.session_id if req.session_id else None - if not req.session_id: - # Tab closed / doc detached from its session — drop the - # in-memory active-doc pointer so the last-resort injection - # path doesn't re-surface this doc in a later chat (#1160). - try: - from src.agent_tools.document_tools import clear_active_document - clear_active_document(doc_id) - except Exception as e: - logger.warning("Failed to clear active document %r on detach", doc_id, exc_info=e) - db.commit() - db.refresh(doc) - return _doc_to_dict(doc) - except HTTPException: - raise - except Exception as e: - db.rollback() - raise HTTPException(500, str(e)) - finally: - db.close() - - # ---- DELETE /api/document/{doc_id} — soft delete ---- - @router.delete("/api/document/{doc_id}") - async def delete_document(request: Request, doc_id: str) -> Dict[str, str]: - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - doc.is_active = False - # Closed/deleted — drop the in-memory active-doc pointer so it isn't - # re-injected into a later, unrelated chat (#1160). - try: - from src.agent_tools.document_tools import clear_active_document - clear_active_document(doc_id) - except Exception: - pass - db.commit() - return {"status": "deleted", "id": doc_id} - except HTTPException: - raise - except Exception as e: - db.rollback() - raise HTTPException(500, str(e)) - finally: - db.close() - - # ---- GET /api/document/{doc_id}/versions ---- - @router.get("/api/document/{doc_id}/versions") - async def list_versions(request: Request, doc_id: str) -> List[Dict[str, Any]]: - user = get_current_user(request) - db = SessionLocal() - try: - # Verify ownership before listing versions - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - versions = db.query(DocumentVersion).filter( - DocumentVersion.document_id == doc_id - ).order_by(DocumentVersion.version_number.desc()).all() - return [{ - "id": v.id, - "version_number": v.version_number, - "content": v.content, - "summary": v.summary, - "source": v.source, - "created_at": v.created_at.isoformat() if v.created_at else None, - } for v in versions] - finally: - db.close() - - # ---- GET /api/document/{doc_id}/version/{num} ---- - @router.get("/api/document/{doc_id}/version/{num}") - async def get_version(request: Request, doc_id: str, num: int) -> Dict[str, Any]: - user = get_current_user(request) - db = SessionLocal() - try: - # Verify ownership - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - ver = db.query(DocumentVersion).filter( - DocumentVersion.document_id == doc_id, - DocumentVersion.version_number == num, - ).first() - if not ver: - raise HTTPException(404, "Version not found") - return _version_to_dict(ver) - finally: - db.close() - - # ---- POST /api/document/{doc_id}/restore/{num} ---- - @router.post("/api/document/{doc_id}/restore/{num}") - async def restore_version(request: Request, doc_id: str, num: int) -> Dict[str, Any]: - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - - old_ver = db.query(DocumentVersion).filter( - DocumentVersion.document_id == doc_id, - DocumentVersion.version_number == num, - ).first() - if not old_ver: - raise HTTPException(404, "Version not found") - - new_ver_num = doc.version_count + 1 - ver = DocumentVersion( - id=str(uuid.uuid4()), - document_id=doc_id, - version_number=new_ver_num, - content=old_ver.content, - summary=f"Restored from v{num}", - source="user", - ) - doc.current_content = old_ver.content - doc.version_count = new_ver_num - db.add(ver) - db.commit() - db.refresh(doc) - return _doc_to_dict(doc) - except HTTPException: - raise - except Exception as e: - db.rollback() - raise HTTPException(500, str(e)) - finally: - db.close() - - # ---- POST /api/documents/tidy — clean up broken/empty documents ---- - @router.post("/api/documents/tidy") - async def tidy_documents(request: Request) -> Dict[str, Any]: - """Fix empty titles and remove broken/empty documents (user's docs only).""" - user = get_current_user(request) - db = SessionLocal() - try: - q = ( - db.query(Document) - .outerjoin(DbSession, Document.session_id == DbSession.id) - .filter(Document.is_active == True) - .filter((Document.archived == False) | (Document.archived.is_(None))) - ) - q = _owner_session_filter(q, user) - docs = q.all() - fixed_titles = 0 - deleted = 0 - - # Same junk-detection logic as the scheduled tidy_documents - # action (src/document_actions.py). Keep these two in sync. - import re as _re - from src.document_actions import _JUNK_TITLES - - to_delete = [] - now = datetime.now(timezone.utc) - for doc in docs: - created = doc.created_at - if created and created.tzinfo is None: - created = created.replace(tzinfo=timezone.utc) - - # Skip freshly created documents to avoid deleting them while the user is actively editing - if created and (now - created).total_seconds() < 900: # 15 minutes - continue - - content = (doc.current_content or "").strip() - title_raw = (doc.title or "").strip() - title = title_raw.lower() - is_fresh_empty = ( - not content - and created is not None - and (now - created).total_seconds() < 1800 - ) - if is_fresh_empty: - continue - - # Strip markdown noise to get a "real" character count - stripped = _re.sub(r"^#{1,6}\s+", "", content, flags=_re.MULTILINE) - stripped = _re.sub(r"[*_`>\-=]+", "", stripped) - stripped = _re.sub(r"\s+", " ", stripped).strip() - real_len = len(stripped) - - # Detect email-scaffold stubs: "To: \nSubject: \n---\n" style - # bodies with nothing typed in. Stub = every meaningful line - # is a header label (To:/From:/Subject:/...) with no real - # value (blank, "empty", "(empty)", "-", "none", "n/a"). - _is_email_stub = False - _HEADER_RE = _re.compile(r"^(to|from|cc|bcc|subject|reply-to):\s*(.*)$", _re.I) - _PLACEHOLDER_VALS = {"", "empty", "(empty)", "-", "—", "none", "n/a", "na", "tbd"} - if title in ("new email", "new mail", "new message") or doc.language == "email": - body_lines = [ln.strip() for ln in content.split("\n") - if ln.strip() and ln.strip() != "---"] - def _is_filler(ln): - m = _HEADER_RE.match(ln) - if not m: - return False - val = (m.group(2) or "").strip().lower() - return val in _PLACEHOLDER_VALS - has_real_body = any(not _is_filler(ln) for ln in body_lines) - if body_lines and not has_real_body: - _is_email_stub = True - - # Hard-delete obviously empty / junk documents - if not content or content in ("", "# Untitled"): - to_delete.append(doc); deleted += 1; continue - if _is_email_stub: - to_delete.append(doc); deleted += 1; continue - if title in _JUNK_TITLES: - to_delete.append(doc); deleted += 1; continue - - # Fix empty or placeholder titles on survivors - if not title_raw or title_raw == "Untitled": - new_title = _derive_title(content) - if new_title and new_title != "Untitled": - doc.title = new_title - fixed_titles += 1 - - for doc in to_delete: - db.delete(doc) - - # Also clean up inactive empty docs from previous soft-deletes - inactive_q = ( - db.query(Document) - .outerjoin(DbSession, Document.session_id == DbSession.id) - .filter(Document.is_active == False) - .filter((Document.current_content == None) | (Document.current_content == "")) - ) - inactive_q = _owner_session_filter(inactive_q, user) - inactive_docs = inactive_q.all() - for doc in inactive_docs: - db.delete(doc) - deleted += len(inactive_docs) - - db.commit() - return { - "fixed_titles": fixed_titles, - "deleted": deleted, - "message": f"Fixed {fixed_titles} title{'s' if fixed_titles != 1 else ''}, removed {deleted} empty document{'s' if deleted != 1 else ''}", - } - except Exception as e: - db.rollback() - logger.error(f"Document tidy failed: {e}") - raise HTTPException(500, f"Tidy failed: {e}") - finally: - db.close() - - # ---- POST /api/documents/ai-tidy — AI-powered cleanup of junk/test documents ---- - @router.post("/api/documents/ai-tidy") - async def ai_tidy_documents(request: Request) -> Dict[str, Any]: - """Use AI to judge if documents are junk/test/accidental, then delete them. - Caches verdicts so previously-reviewed docs are skipped.""" - from src.task_endpoint import resolve_task_endpoint - from src.endpoint_resolver import resolve_endpoint - from src.llm_core import llm_call_async - - user = get_current_user(request) - url, model, headers = resolve_task_endpoint(owner=user or None) - if not url or not model: - # Fall back to default endpoint - url, model, headers = resolve_endpoint("default", owner=user or None) - if not url or not model: - raise HTTPException(500, "No endpoint configured for AI tidy") - - db = SessionLocal() - try: - q = ( - db.query(Document) - .outerjoin(DbSession, Document.session_id == DbSession.id) - .filter(Document.is_active == True) - .filter((Document.archived == False) | (Document.archived.is_(None))) - ) - q = _owner_session_filter(q, user) - docs = q.all() - - # Only review docs that haven't been reviewed yet - to_review = [d for d in docs if not d.tidy_verdict] - if not to_review: - return {"deleted": 0, "reviewed": 0, "message": "All documents already reviewed"} - - # Build a batch prompt — review up to 30 at a time - batch = to_review[:30] - doc_list = [] - for i, doc in enumerate(batch): - preview = (doc.current_content or "")[:300].strip() - doc_list.append(f"[{i}] title=\"{doc.title}\" lang={doc.language or 'text'} content_preview=\"{preview}\"") - - prompt = ( - "You are a document library cleaner. For each document below, decide if it is JUNK " - "(test, accidental, placeholder, empty-ish, tool-test, throwaway) or KEEP (real content worth saving).\n\n" - "Respond with ONLY a JSON array of verdicts, one per document, like: [\"junk\",\"keep\",\"junk\",...]\n" - "No explanation, no markdown, just the JSON array.\n\n" - + "\n".join(doc_list) - ) - - response = await llm_call_async( - url, model, - [{"role": "system", "content": "You classify documents as junk or keep. Respond only with a JSON array."}, - {"role": "user", "content": prompt}], - temperature=0.1, - max_tokens=200, - headers=headers, - timeout=30, - ) - - # Parse verdicts - import re - match = re.search(r'\[.*?\]', response, re.DOTALL) - if not match: - raise HTTPException(500, "AI returned invalid response") - - import json as _json - verdicts = _json.loads(match.group()) - - deleted = 0 - reviewed = 0 - for i, doc in enumerate(batch): - if i >= len(verdicts): - break - verdict = str(verdicts[i] or "").lower().strip() - if verdict == "junk": - doc.tidy_verdict = "junk" - db.delete(doc) - deleted += 1 - else: - doc.tidy_verdict = "keep" - reviewed += 1 - - db.commit() - return { - "deleted": deleted, - "reviewed": reviewed, - "remaining": len(to_review) - len(batch), - "message": f"Reviewed {reviewed}, removed {deleted} junk document{'s' if deleted != 1 else ''}", - } - except HTTPException: - raise - except Exception as e: - db.rollback() - logger.error(f"AI tidy failed: {e}") - raise HTTPException(500, f"AI tidy failed: {e}") - finally: - db.close() - - # ---- POST /api/document/{doc_id}/export-pdf/preview ---- - @router.post("/api/document/{doc_id}/export-pdf/preview") - async def export_pdf_preview(doc_id: str, request: Request) -> Dict[str, Any]: - """Return the field-value mapping that would be written to the PDF. - - Frontend shows this in a confirmation modal so the user can spot/fix - any wrong values before triggering the actual download. - """ - from src.pdf_form_doc import find_source_upload_id, parse_markdown_to_values, load_field_sidecar - - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - - upload_id = find_source_upload_id(doc.current_content or "") - if not upload_id: - raise HTTPException(400, "Document is not linked to a source PDF") - - pdf_path = _locate_current_user_upload(request, upload_id, user) - if not pdf_path: - raise HTTPException(404, f"Source PDF {upload_id} not found in uploads") - - fields = load_field_sidecar(pdf_path) - if not fields: - raise HTTPException(404, "Field schema sidecar missing for source PDF") - - values = parse_markdown_to_values(doc.current_content or "") - field_meta = {f["name"]: f for f in fields} - - preview = [] - for name, current in values.items(): - meta = field_meta.get(name) - if not meta: - continue - preview.append({ - "name": name, - "label": meta.get("label") or name, - "type": meta.get("type"), - "options": meta.get("options") or [], - "page": meta.get("page"), - "value": current, - }) - - unknown = [ - name for name in values - if name not in field_meta - ] - return { - "doc_id": doc_id, - "upload_id": upload_id, - "fields": preview, - "unknown_fields": unknown, - "total": len(fields), - "filled": sum(1 for p in preview if p["value"] not in ("", False, None)), - } - finally: - db.close() - - # ---- GET /api/document/{doc_id}/render-pages ---- - @router.get("/api/document/{doc_id}/render-pages") - async def render_pages(doc_id: str, request: Request) -> Dict[str, Any]: - """Return per-page metadata for the interactive PDF view. - - Each page entry has its rendered-image dimensions (matching what - /page/{n}.png returns at the same DPI) plus the list of form fields - on that page with their rects translated to image-pixel coordinates. - Frontend overlays HTML form controls at those positions. - """ - from src.pdf_form_doc import find_source_upload_id, parse_markdown_to_values, load_field_sidecar - - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - upload_id = find_source_upload_id(doc.current_content or "") - if not upload_id: - raise HTTPException(400, "Document is not linked to a source PDF") - pdf_path = _locate_current_user_upload(request, upload_id, user) - if not pdf_path: - raise HTTPException(404, f"Source PDF {upload_id} not found") - - fitz = _load_pdf_viewer_fitz() - schema = load_field_sidecar(pdf_path) or [] - values = parse_markdown_to_values(doc.current_content or "") - - # Group fields by page - by_page: Dict[int, list] = {} - for f in schema: - by_page.setdefault(f["page"], []).append(f) - - scale = _PDF_RENDER_SCALE - pdf_doc = fitz.open(pdf_path) - try: - pages_out = [] - for page_index in range(pdf_doc.page_count): - page = pdf_doc[page_index] - page_no = page_index + 1 - pw, ph = page.rect.width, page.rect.height - img_w = int(pw * scale) - img_h = int(ph * scale) - fields_out = [] - for f in by_page.get(page_no, []): - x0, y0, x1, y1 = f["rect"] - fields_out.append({ - "name": f["name"], - "type": f["type"], - "label": f.get("label") or "", - "options": f.get("options") or [], - "value": values.get(f["name"], f.get("value", "")), - "rect_px": [ - int(x0 * scale), int(y0 * scale), - int(x1 * scale), int(y1 * scale), - ], - }) - pages_out.append({ - "page": page_no, - "width": img_w, - "height": img_h, - "fields": fields_out, - }) - return {"doc_id": doc_id, "scale": scale, "pages": pages_out} - finally: - pdf_doc.close() - finally: - db.close() - - # ---- GET /api/document/{doc_id}/page/{n}.png ---- - @router.get("/api/document/{doc_id}/page/{page_no}.png") - async def render_page_png(doc_id: str, page_no: int, request: Request): - """Render one page of the source PDF as a PNG (no values stamped — the - frontend overlays HTML form inputs on top).""" - from fastapi.responses import Response - from src.pdf_form_doc import find_source_upload_id - - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - upload_id = find_source_upload_id(doc.current_content or "") - if not upload_id: - raise HTTPException(400, "Document is not linked to a source PDF") - pdf_path = _locate_current_user_upload(request, upload_id, user) - if not pdf_path: - raise HTTPException(404, "Source PDF not found") - finally: - db.close() - - fitz = _load_pdf_viewer_fitz() - pdf_doc = fitz.open(pdf_path) - try: - if page_no < 1 or page_no > pdf_doc.page_count: - raise HTTPException(404, "Page out of range") - page = pdf_doc[page_no - 1] - mat = fitz.Matrix(_PDF_RENDER_SCALE, _PDF_RENDER_SCALE) - pix = page.get_pixmap(matrix=mat, alpha=False) - png_bytes = pix.tobytes("png") - return Response( - content=png_bytes, - media_type="image/png", - headers={"Cache-Control": "public, max-age=3600"}, - ) - finally: - pdf_doc.close() - - # ---- POST /api/document/{doc_id}/ai-fill-annotations ---- - @router.post("/api/document/{doc_id}/ai-fill-annotations") - async def ai_fill_annotations(doc_id: str, request: Request) -> Dict[str, Any]: - """Ask a vision-capable LLM to locate fillable areas on a flat PDF and - propose annotation values for each, given a free-form user instruction. - - Returns a list of annotations: [{page, x, y, w, h, value}] where x/y/w/h - are page-percentages (0–100) — same coordinate system as the freeform - annotations the frontend already renders. - """ - import base64 - import json - import fitz - from src.pdf_form_doc import find_source_upload_id - from src.document_processor import _resolve_vl_model, _load_vl_settings - from src.llm_core import llm_call_async - - body = await request.json() if request.headers.get("content-type", "").startswith("application/json") else {} - instruction = (body or {}).get("instruction", "").strip() - if not instruction: - raise HTTPException(400, "instruction is required") - - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - upload_id = find_source_upload_id(doc.current_content or "") - if not upload_id: - raise HTTPException(400, "Document is not linked to a source PDF") - pdf_path = _locate_current_user_upload(request, upload_id, user) - if not pdf_path: - raise HTTPException(404, "Source PDF not found") - finally: - db.close() - - # Resolve VL model (admin-configured or auto-detected vision-capable) - settings = _load_vl_settings() - vl_model = settings.get("vision_model", "") - try: - url, model_id, headers = _resolve_vl_model(vl_model, owner=user) - except Exception as e: - raise HTTPException(503, f"No vision model available: {e}") - - system_prompt = ( - "You analyze rendered PDF page images and propose values to fill in. " - "For each blank line, box, underscore, or labeled space on the page that " - "should be filled given the user's instruction, output one annotation. " - "Coordinates are percentages (0-100) of the page width/height with the " - "origin at top-left. Width/height should match the visible blank box. " - "Return ONLY a JSON array, no prose, no markdown fences. Each entry: " - '{"x": number, "y": number, "w": number, "h": number, "value": string}. ' - "If a region should not be filled, omit it. If nothing should be filled, " - "return []." - ) - - all_annotations = [] - pdf_doc = fitz.open(pdf_path) - try: - for page_index in range(pdf_doc.page_count): - page = pdf_doc[page_index] - mat = fitz.Matrix(_PDF_RENDER_SCALE, _PDF_RENDER_SCALE) - pix = page.get_pixmap(matrix=mat, alpha=False) - png_bytes = pix.tobytes("png") - b64 = base64.b64encode(png_bytes).decode("ascii") - - messages = [ - {"role": "system", "content": system_prompt}, - { - "role": "user", - "content": [ - { - "type": "text", - "text": ( - f"User instruction:\n{instruction}\n\n" - f"This is page {page_index + 1} of {pdf_doc.page_count}. " - "Return JSON array of annotations to add to this page." - ), - }, - { - "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{b64}"}, - }, - ], - }, - ] - try: - raw = await llm_call_async( - url, model_id, messages, - temperature=0.1, max_tokens=2000, headers=headers, - ) - except Exception as e: - logger.error(f"VL call failed on page {page_index + 1}: {e}") - continue - - raw = (raw or "").strip() - if raw.startswith("```"): - raw = raw.split("\n", 1)[-1].rsplit("```", 1)[0].strip() - try: - parsed = json.loads(raw) - except Exception: - logger.warning(f"AI fill: page {page_index + 1} returned non-JSON: {raw[:200]}") - continue - if not isinstance(parsed, list): - continue - for item in parsed: - if not isinstance(item, dict): - continue - try: - x = float(item.get("x", 0)) - y = float(item.get("y", 0)) - w = float(item.get("w", 0)) - h = float(item.get("h", 0)) - value = str(item.get("value", "") or "") - except Exception: - continue - # Clamp + reject zero-size entries - if w <= 0.5 or h <= 0.3: - continue - x = max(0.0, min(99.0, x)) - y = max(0.0, min(99.0, y)) - w = max(0.5, min(100.0 - x, w)) - h = max(0.3, min(100.0 - y, h)) - if not value.strip(): - continue - all_annotations.append({ - "page": page_index + 1, - "x": round(x, 2), - "y": round(y, 2), - "w": round(w, 2), - "h": round(h, 2), - "value": value, - }) - finally: - pdf_doc.close() - - return {"annotations": all_annotations} - - # ---- GET /api/document/{doc_id}/render-pdf ---- - @router.get("/api/document/{doc_id}/render-pdf") - async def render_pdf(doc_id: str, request: Request): - """Inline PDF preview filled with the current markdown values. - - Same plumbing as the export route, but no signature stamping and - served inline (Content-Disposition: inline) so the browser can - embed it in an iframe. Cache-busted by the caller via query string. - """ - import base64 - import os - import tempfile - from fastapi.responses import FileResponse - from starlette.background import BackgroundTask - from src.pdf_form_doc import find_source_upload_id, parse_markdown_to_values, parse_markdown_annotations - from src.pdf_forms import fill_fields, stamp_annotations - from core.database import Signature - - # Track temp files for this request so they get unlinked AFTER - # the response is fully sent (BackgroundTask runs post-send). - _to_unlink: list[str] = [] - def _cleanup_temps(): - for _p in _to_unlink: - try: - os.unlink(_p) - except FileNotFoundError: - pass - except Exception as _e: - logger.warning(f"Could not unlink temp PDF {_p}: {_e}") - - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - upload_id = find_source_upload_id(doc.current_content or "") - if not upload_id: - raise HTTPException(400, "Document is not linked to a source PDF") - pdf_path = _locate_current_user_upload(request, upload_id, user) - if not pdf_path: - raise HTTPException(404, f"Source PDF {upload_id} not found") - - # Fail fast with a clear 503 if the optional PyMuPDF dependency - # is missing — fill_fields/stamp_annotations will otherwise - # raise RuntimeError deep inside and bubble out as a 500. - # Mirrors the convention in _load_pdf_viewer_fitz above. - _load_pdf_viewer_fitz() - - values = parse_markdown_to_values(doc.current_content or "") - out_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name - _to_unlink.append(out_path) - try: - fill_fields(pdf_path, out_path, values) - except Exception as e: - logger.error(f"render_pdf fill_fields failed for {doc_id}: {e}") - _cleanup_temps() - raise HTTPException(500, f"PDF render failed: {e}") - - annotations = parse_markdown_annotations(doc.current_content or "") - if annotations: - ann_sig_ids = [ - a["value"][len("signature:"):].strip() - for a in annotations - if a.get("kind") == "signature" - and isinstance(a.get("value"), str) - and a["value"].startswith("signature:") - ] - ann_signature_pngs: dict[str, bytes] = {} - if ann_sig_ids: - # SECURITY: filter by owner so a caller can't reference - # someone else's signature ID from doc markdown and have - # it stamped/exported. - _sig_q = db.query(Signature).filter(Signature.id.in_(ann_sig_ids)) - if user: - _sig_q = _sig_q.filter(Signature.owner == user) - sig_rows = _sig_q.all() - for s in sig_rows: - try: - ann_signature_pngs[s.id] = base64.b64decode(s.data_png) - except Exception as e: - logger.warning(f"Bad annotation signature data for {s.id}: {e}") - annotated_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name - _to_unlink.append(annotated_path) - try: - stamp_annotations(out_path, annotated_path, annotations, ann_signature_pngs) - out_path = annotated_path - except Exception as e: - logger.error(f"stamp_annotations (render) failed for {doc_id}: {e}") - - return FileResponse( - out_path, - media_type="application/pdf", - headers={"Content-Disposition": "inline"}, - background=BackgroundTask(_cleanup_temps), - ) - finally: - db.close() - - # ---- GET /api/document/{doc_id}/export-pdf ---- - @router.get("/api/document/{doc_id}/export-pdf") - async def export_pdf(doc_id: str, request: Request): - """Stream the filled PDF for download. - - Reads field values and signature selections from the markdown — there - is no separate confirmation step. Signature fields contain their - chosen signature ID encoded as `signature:` in the value. - """ - import base64 - import os - import tempfile - from fastapi.responses import FileResponse - from starlette.background import BackgroundTask - from src.pdf_form_doc import find_source_upload_id, parse_markdown_to_values, load_field_sidecar, parse_markdown_annotations - from src.pdf_forms import fill_fields, stamp_signatures, stamp_annotations - from core.database import Signature - - _to_unlink: list[str] = [] - def _cleanup_temps(): - for _p in _to_unlink: - try: - os.unlink(_p) - except FileNotFoundError: - pass - except Exception as _e: - logger.warning(f"Could not unlink temp PDF {_p}: {_e}") - - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - - upload_id = find_source_upload_id(doc.current_content or "") - if not upload_id: - raise HTTPException(400, "Document is not linked to a source PDF") - - pdf_path = _locate_current_user_upload(request, upload_id, user) - if not pdf_path: - raise HTTPException(404, f"Source PDF {upload_id} not found in uploads") - - schema = load_field_sidecar(pdf_path) or [] - sig_field_names = {f["name"] for f in schema if f.get("type") == "signature"} - - all_values = parse_markdown_to_values(doc.current_content or "") - # Split: signature fields go to stamps, everything else to fill_fields - text_values: dict = {} - sig_ids: dict[str, str] = {} - for name, raw in all_values.items(): - if name in sig_field_names and isinstance(raw, str) and raw.startswith("signature:"): - sig_ids[name] = raw[len("signature:"):].strip() - elif name not in sig_field_names: - text_values[name] = raw - - stamps: dict = {} - if sig_ids: - # SECURITY: filter by owner — same reason as render_pdf. - _sig_q2 = db.query(Signature).filter(Signature.id.in_(list(sig_ids.values()))) - if user: - _sig_q2 = _sig_q2.filter(Signature.owner == user) - rows = _sig_q2.all() - by_id = {s.id: s for s in rows} - for field_name, sid in sig_ids.items(): - s = by_id.get(sid) - if not s: - continue - try: - stamps[field_name] = base64.b64decode(s.data_png) - except Exception as e: - logger.warning(f"Bad signature data for {sid}: {e}") - - filled_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name - _to_unlink.append(filled_path) - try: - fill_fields(pdf_path, filled_path, text_values) - except Exception as e: - logger.error(f"fill_fields failed for doc {doc_id}: {e}") - _cleanup_temps() - raise HTTPException(500, f"PDF fill failed: {e}") - - out_path = filled_path - if stamps: - stamped_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name - _to_unlink.append(stamped_path) - try: - stamp_signatures(filled_path, stamped_path, stamps) - out_path = stamped_path - except Exception as e: - logger.error(f"stamp_signatures failed for doc {doc_id}: {e}") - - # Burn freeform annotations (Text/Check/Sign drops) on top. - annotations = parse_markdown_annotations(doc.current_content or "") - if annotations: - # Resolve any signature annotations to their PNG bytes. - ann_sig_ids = [ - a["value"][len("signature:"):].strip() - for a in annotations - if a.get("kind") == "signature" - and isinstance(a.get("value"), str) - and a["value"].startswith("signature:") - ] - ann_signature_pngs: dict[str, bytes] = {} - if ann_sig_ids: - # SECURITY: filter by owner so a caller can't reference - # someone else's signature ID from doc markdown and have - # it stamped/exported. - _sig_q = db.query(Signature).filter(Signature.id.in_(ann_sig_ids)) - if user: - _sig_q = _sig_q.filter(Signature.owner == user) - sig_rows = _sig_q.all() - for s in sig_rows: - try: - ann_signature_pngs[s.id] = base64.b64decode(s.data_png) - except Exception as e: - logger.warning(f"Bad annotation signature data for {s.id}: {e}") - annotated_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name - _to_unlink.append(annotated_path) - try: - stamp_annotations(out_path, annotated_path, annotations, ann_signature_pngs) - out_path = annotated_path - except Exception as e: - logger.error(f"stamp_annotations failed for doc {doc_id}: {e}") - - download_name = _slug(doc.title or "form") + "_annotated.pdf" - return FileResponse( - out_path, - media_type="application/pdf", - filename=download_name, - background=BackgroundTask(_cleanup_temps), - ) - finally: - db.close() - - # ---- POST /api/document/{doc_id}/prepare-signed-reply ---- - @router.post("/api/document/{doc_id}/prepare-signed-reply") - async def prepare_signed_reply(doc_id: str, request: Request): - """Bake the current PDF state (form fields + signature stamps + - annotations) into a flattened PDF, drop it in COMPOSE_UPLOADS_DIR - and return the reply context (To/Subject/threading headers) so the - frontend can open a reply draft with this attachment pre-loaded. - - Requires the document to have source_email_* metadata (set when the - doc was created via /api/email/attachment-as-doc). Otherwise 400. - """ - import base64 - import tempfile - import shutil - import uuid as _uuid - import email as _email_mod - from src.pdf_form_doc import ( - find_source_upload_id, parse_markdown_to_values, - load_field_sidecar, parse_markdown_annotations, - ) - from src.pdf_forms import fill_fields, stamp_signatures, stamp_annotations - from core.database import Signature - # COMPOSE_UPLOADS_DIR lives in email_routes — re-derive here so we - # don't import from a routes file (cycle-prone). Same env override - # as email_routes (ODYSSEUS_MAIL_ATTACHMENTS_DIR). - from pathlib import Path as _Path - _COMPOSE_DIR = _Path(MAIL_ATTACHMENTS_DIR) / "_compose" - _COMPOSE_DIR.mkdir(parents=True, exist_ok=True) - - user = get_current_user(request) - db = SessionLocal() - try: - doc = db.query(Document).filter(Document.id == doc_id).first() - if not doc: - raise HTTPException(404, "Document not found") - _verify_doc_owner(db, doc, user) - - if not (doc.source_email_uid and doc.source_email_folder): - raise HTTPException(400, "Document has no source email — cannot reply") - - # 1) Build the flattened PDF (same pipeline as export_pdf) - upload_id = find_source_upload_id(doc.current_content or "") - if not upload_id: - raise HTTPException(400, "Document is not linked to a source PDF") - pdf_path = _locate_current_user_upload(request, upload_id, user) - if not pdf_path: - raise HTTPException(404, f"Source PDF {upload_id} not found") - - schema = load_field_sidecar(pdf_path) or [] - sig_field_names = {f["name"] for f in schema if f.get("type") == "signature"} - all_values = parse_markdown_to_values(doc.current_content or "") - text_values: dict = {} - sig_ids: dict[str, str] = {} - for name, raw in all_values.items(): - if name in sig_field_names and isinstance(raw, str) and raw.startswith("signature:"): - sig_ids[name] = raw[len("signature:"):].strip() - elif name not in sig_field_names: - text_values[name] = raw - - stamps: dict = {} - if sig_ids: - # SECURITY: filter by owner — same reason as render_pdf. - _sig_q2 = db.query(Signature).filter(Signature.id.in_(list(sig_ids.values()))) - if user: - _sig_q2 = _sig_q2.filter(Signature.owner == user) - rows = _sig_q2.all() - by_id = {s.id: s for s in rows} - for fname, sid in sig_ids.items(): - s = by_id.get(sid) - if not s: - continue - try: - stamps[fname] = base64.b64decode(s.data_png) - except Exception: - pass - - import os - _to_unlink: list[str] = [] - filled_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name - _to_unlink.append(filled_path) - fill_fields(pdf_path, filled_path, text_values) - out_path = filled_path - if stamps: - stamped_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name - _to_unlink.append(stamped_path) - try: - stamp_signatures(filled_path, stamped_path, stamps) - out_path = stamped_path - except Exception as e: - logger.warning(f"stamp_signatures failed for {doc_id}: {e}") - - annotations = parse_markdown_annotations(doc.current_content or "") - if annotations: - ann_sig_ids = [ - a["value"][len("signature:"):].strip() - for a in annotations - if a.get("kind") == "signature" - and isinstance(a.get("value"), str) - and a["value"].startswith("signature:") - ] - ann_signature_pngs: dict[str, bytes] = {} - if ann_sig_ids: - # SECURITY: filter by owner so a caller can't reference - # someone else's signature ID from doc markdown and have - # it stamped/exported. - _sig_q = db.query(Signature).filter(Signature.id.in_(ann_sig_ids)) - if user: - _sig_q = _sig_q.filter(Signature.owner == user) - sig_rows = _sig_q.all() - for s in sig_rows: - try: - ann_signature_pngs[s.id] = base64.b64decode(s.data_png) - except Exception: - pass - annotated_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name - _to_unlink.append(annotated_path) - try: - stamp_annotations(out_path, annotated_path, annotations, ann_signature_pngs) - out_path = annotated_path - except Exception as e: - logger.warning(f"stamp_annotations failed for {doc_id}: {e}") - - # 2) Move/copy into COMPOSE_UPLOADS_DIR with the token format - # `_` that /api/email/send expects. - filename = _slug(doc.title or "signed") + "_signed.pdf" - token = f"{_uuid.uuid4().hex}_{filename}" - dest = _COMPOSE_DIR / token - shutil.copyfile(out_path, str(dest)) - # Unlink the intermediate temp PDFs now that they've been - # copied into COMPOSE_UPLOADS_DIR. - for _p in _to_unlink: - try: - os.unlink(_p) - except FileNotFoundError: - pass - except Exception as _e: - logger.warning(f"Could not unlink temp PDF {_p}: {_e}") - - # 3) Fetch the source email's headers so we can build a clean reply - # context (To/Subject/In-Reply-To/References). - try: - from routes.email_routes import _imap, _decode_header - from routes.email_helpers import _q - except Exception: - _imap = None - _decode_header = lambda x: x or "" - _q = lambda x: x or "" - - to_addr = "" - from_name = "" - subject = "" - in_reply_to = doc.source_email_message_id or "" - references = in_reply_to - if _imap: - try: - with _imap(doc.source_email_account_id or None) as conn: - conn.select(_q(doc.source_email_folder), readonly=True) - status, data = conn.fetch(doc.source_email_uid.encode(), "(RFC822.HEADER)") - if status == "OK" and data and data[0]: - raw_hdr = data[0][1] - m = _email_mod.message_from_bytes(raw_hdr) - sender = _decode_header(m.get("From", "")) - from_name, to_addr = _email_mod.utils.parseaddr(sender) - if not to_addr: - to_addr = sender - subject = _decode_header(m.get("Subject", "") or "") - if subject and not subject.lower().startswith("re:"): - subject = "Re: " + subject - msg_refs = (m.get("References") or "").strip() - msg_in_reply = (m.get("Message-ID") or "").strip() or in_reply_to - in_reply_to = msg_in_reply - references = (msg_refs + " " + msg_in_reply).strip() if msg_refs else msg_in_reply - except Exception as e: - logger.warning(f"prepare-signed-reply header fetch failed: {e}") - - return { - "ok": True, - "attachment": { - "token": token, - "filename": filename, - "size": dest.stat().st_size, - }, - "reply": { - "to": to_addr, - "to_name": from_name, - "subject": subject, - "in_reply_to": in_reply_to, - "references": references, - "account_id": doc.source_email_account_id or None, - "source_uid": doc.source_email_uid, - "source_folder": doc.source_email_folder, - "source_message_id": doc.source_email_message_id, - }, - } - finally: - db.close() - - return router +_sys.modules[__name__] = _canonical diff --git a/tests/test_document_routes_shim.py b/tests/test_document_routes_shim.py new file mode 100644 index 000000000..68d049a62 --- /dev/null +++ b/tests/test_document_routes_shim.py @@ -0,0 +1,29 @@ +"""Regression test for the document route shim (slice 2m, #4082/#4071). + +The backward-compat shims at ``routes/document_routes.py`` and +``routes/document_helpers.py`` use ``sys.modules`` replacement so the legacy +import paths and the canonical ``routes.document.*`` paths resolve to the +*same* module objects. This is required because multiple tests do +``import routes.document_routes as droutes`` followed by +``droutes.SessionLocal = ...`` / ``monkeypatch.setattr(droutes, ...)`` and +``sys.modules.pop("routes.document_helpers")`` + re-import — for those to +take effect at runtime, the legacy and canonical module objects must be +identical. +""" + +import importlib + +import routes.document_routes as _shim_routes # noqa: F401 +import routes.document_helpers as _shim_helpers # noqa: F401 + + +def test_legacy_and_canonical_routes_are_same_object(): + legacy = importlib.import_module("routes.document_routes") + canonical = importlib.import_module("routes.document.document_routes") + assert legacy is canonical + + +def test_legacy_and_canonical_helpers_are_same_object(): + legacy = importlib.import_module("routes.document_helpers") + canonical = importlib.import_module("routes.document.document_helpers") + assert legacy is canonical diff --git a/tests/test_imap_mailbox_quoting.py b/tests/test_imap_mailbox_quoting.py index 7c5bb1645..636270a56 100644 --- a/tests/test_imap_mailbox_quoting.py +++ b/tests/test_imap_mailbox_quoting.py @@ -87,7 +87,7 @@ def test_known_imap_mailbox_call_sites_are_quoted(): assert "conn.select(sent_name" not in pollers assert "imap.append(sent_folder" not in pollers - document_routes = Path("routes/document_routes.py").read_text() + document_routes = Path("routes/document/document_routes.py").read_text() assert "conn.select(doc.source_email_folder" not in document_routes diff --git a/tests/test_model_helper_owner_scope.py b/tests/test_model_helper_owner_scope.py index dafbad594..f48a1f7e2 100644 --- a/tests/test_model_helper_owner_scope.py +++ b/tests/test_model_helper_owner_scope.py @@ -14,7 +14,7 @@ def _function_source(path: str, name: str) -> str: def test_document_ai_tidy_resolves_with_owner_scope(): - body = _function_source("routes/document_routes.py", "ai_tidy_documents") + body = _function_source("routes/document/document_routes.py", "ai_tidy_documents") assert "resolve_task_endpoint(owner=user or None)" in body assert 'resolve_endpoint("default", owner=user or None)' in body diff --git a/tests/test_vision_owner_scope.py b/tests/test_vision_owner_scope.py index f0d3a184d..29de101a3 100644 --- a/tests/test_vision_owner_scope.py +++ b/tests/test_vision_owner_scope.py @@ -88,7 +88,7 @@ def test_request_vision_call_sites_pass_owner(): chat_source = (ROOT / "src" / "chat_handler.py").read_text() processor_source = (ROOT / "src" / "document_processor.py").read_text() upload_source = (ROOT / "routes" / "upload_routes.py").read_text() - document_source = (ROOT / "routes" / "document_routes.py").read_text() + document_source = (ROOT / "routes" / "document" / "document_routes.py").read_text() gallery_source = (ROOT / "routes" / "gallery" / "gallery_routes.py").read_text() memory_source = (ROOT / "routes" / "memory" / "memory_routes.py").read_text() From 9d686180dd20e6ef842f0c23f9ef2f0ce39cee4f Mon Sep 17 00:00:00 2001 From: Ashvin <76151462+ashvinctrl@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:47:41 +0530 Subject: [PATCH 2/3] fix(integrations): pin api_call to the SSRF-validated IP (#5727) * fix(integrations): pin api_call to the SSRF-validated IP execute_api_call runs check_outbound_url on the target, but that guard only resolves the host to answer (ok, reason) and hands back no address. The request right after it opened a plain httpx.AsyncClient, which resolves the host again at connect time. A base_url host on a low TTL can pass the guard as a public IP and then flip to 169.254.169.254 for the connect, so the call lands on cloud metadata with the integration's stored auth headers attached. Resolve once, remember the IPs the guard actually validated, and pin the client's socket to that set through a small AnyIO-backed transport. SNI and the Host header still come from the URL, so TLS and vhost routing are unchanged; connect-time fallback stays inside the approved address set over one shared deadline. This is the same pinning the webhook sender and web-fetch paths already do -- api_call was the last outbound path that skipped it. Fixes #5513 * fix(integrations): de-duplicate the pinned IP list _default_resolver calls getaddrinfo(host, None) with no socktype filter, so glibc returns one record per socktype and a single-homed host comes back three times over. _validated_ips kept every entry, so the transport pinned the same address repeatedly and the connect fallback could spend its shared deadline retrying one dead address instead of moving on to a genuinely different one. Windows getaddrinfo collapses those duplicate records, which is why the ip-literal pin test only failed on CI and not locally. --- src/integrations.py | 175 ++++++++++++- tests/test_integration_api_call_ssrf.py | 240 ++++++++++++++++++ .../test_integrations_api_call_truncation.py | 14 +- 3 files changed, 420 insertions(+), 9 deletions(-) diff --git a/src/integrations.py b/src/integrations.py index aa6c4982e..52dd4b2d1 100644 --- a/src/integrations.py +++ b/src/integrations.py @@ -1,11 +1,14 @@ +import ipaddress import json import os +import time import uuid import logging import re from typing import Dict, List, Optional, Any from urllib.parse import urljoin, urlparse, urlunparse +import httpcore import httpx from fastapi import HTTPException @@ -354,6 +357,152 @@ def _find_integration(identifier: str) -> Optional[Dict[str, Any]]: return None +# httpcore raises its own exception hierarchy; map the ones a simple request can +# surface back to their httpx equivalents so the caller's `except httpx.*` blocks +# below behave exactly as they did with the default transport. +_HTTPCORE_TO_HTTPX_EXC = { + httpcore.ConnectError: httpx.ConnectError, + httpcore.ConnectTimeout: httpx.ConnectTimeout, + httpcore.NetworkError: httpx.NetworkError, + httpcore.PoolTimeout: httpx.PoolTimeout, + httpcore.ProtocolError: httpx.ProtocolError, + httpcore.ReadError: httpx.ReadError, + httpcore.ReadTimeout: httpx.ReadTimeout, + httpcore.RemoteProtocolError: httpx.RemoteProtocolError, + httpcore.TimeoutException: httpx.TimeoutException, + httpcore.WriteError: httpx.WriteError, + httpcore.WriteTimeout: httpx.WriteTimeout, +} + + +class _PinnedAsyncBackend(httpcore.AsyncNetworkBackend): + """Network backend that connects only to the pre-validated IPs, in order. + + Every address here came out of the single SSRF resolution, so moving to the + next one after a connect failure is not re-resolution — it's ordinary + multi-address fallback restricted to the set the guard already approved. + httpcore takes TLS SNI and the ``Host`` header from the request URL rather + than the connect host, so pinning the socket destination leaves certificate + validation and vhost routing pointed at the original hostname. + """ + + def __init__(self, ips: List[ipaddress._BaseAddress]): + self._ips = [str(ip) for ip in ips] + self._real = httpcore.AnyIOBackend() + + async def connect_tcp(self, host, port, timeout=None, local_address=None, + socket_options=None): + # One shared connect budget: each attempt gets the time left until the + # original deadline, so N dead addresses can't stretch the connect + # phase to N * timeout. + deadline = None if timeout is None else time.monotonic() + timeout + last_exc: Optional[Exception] = None + for ip in self._ips: + remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) + try: + return await self._real.connect_tcp( + ip, port, remaining, local_address, socket_options + ) + except (httpcore.ConnectError, httpcore.ConnectTimeout) as exc: + last_exc = exc + if deadline is not None and time.monotonic() >= deadline: + break + raise last_exc + + async def connect_unix_socket(self, path, timeout=None, socket_options=None): + return await self._real.connect_unix_socket(path, timeout, socket_options) + + async def sleep(self, seconds: float) -> None: + return await self._real.sleep(seconds) + + +class _PinnedAsyncTransport(httpx.AsyncBaseTransport): + """httpx transport that pins the TCP connect to the pre-resolved IP(s). + + Kept local, mirroring the per-module pinned transports web fetch and + webhook delivery already carry, rather than coupling api_call to the + webhook subsystem. The request URL passes through unchanged, so SNI and the + ``Host`` header stay the original hostname; only the socket destination is + pinned, which is what closes the rebinding window. + """ + + def __init__(self, ips: List[ipaddress._BaseAddress]): + self._pinned_ips = list(ips) + self._pool = httpcore.AsyncConnectionPool( + # Reuse the CA trust the default httpx client would build (certifi + # plus SSL_CERT_FILE / SSL_CERT_DIR when trust_env is set) so + # swapping in this transport doesn't quietly change which chains + # verify. ssl.create_default_context() would use system roots. + ssl_context=httpx.create_ssl_context(), + http1=True, + http2=False, + network_backend=_PinnedAsyncBackend(ips), + ) + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + core_req = httpcore.Request( + method=request.method, + url=httpcore.URL( + scheme=request.url.raw_scheme, + host=request.url.raw_host, + port=request.url.port, + target=request.url.raw_path, + ), + headers=request.headers.raw, + content=request.stream, + extensions=request.extensions, + ) + try: + core_resp = await self._pool.handle_async_request(core_req) + content = b"".join([chunk async for chunk in core_resp.aiter_stream()]) + await core_resp.aclose() + except Exception as exc: + mapped = _HTTPCORE_TO_HTTPX_EXC.get(type(exc)) + if mapped is not None: + raise mapped(str(exc)) from exc + raise + return httpx.Response( + status_code=core_resp.status, + headers=core_resp.headers, + content=content, + extensions=core_resp.extensions, + ) + + async def aclose(self) -> None: + await self._pool.aclose() + + +def _validated_ips(raw_ips: List[str]) -> List[ipaddress._BaseAddress]: + """Return every entry that parses as an IP address, de-duplicated, order + preserved. + + check_outbound_url only reports ok when *all* of these classify as safe, so + the whole list is guard-approved and any of them is a legitimate connect + target. Skipping unparseable entries mirrors how the guard walks the same + resolver output. + + De-duplication matters because the resolver is getaddrinfo(host, None) with + no socktype filter, so glibc reports the same address once per socktype + (SOCK_STREAM/SOCK_DGRAM/SOCK_RAW) — a single-homed host comes back three + times. Without this, the connect fallback would spend the shared deadline + retrying one dead address instead of moving on to a genuinely different one. + """ + ips: List[ipaddress._BaseAddress] = [] + seen = set() + for raw in raw_ips: + if not isinstance(raw, str): + continue + try: + ip = ipaddress.ip_address(raw.split("%")[0]) # strip IPv6 zone id + except ValueError: + continue + if ip in seen: + continue + seen.add(ip) + ips.append(ip) + return ips + + async def execute_api_call( integration_id: str, method: str, @@ -409,13 +558,31 @@ async def execute_api_call( # loopback for locked-down deployments. Private stays allowed by default # because LAN integrations (Home Assistant, Miniflux, ntfy) are the # primary use case. - from src.url_safety import check_outbound_url + from src.url_safety import check_outbound_url, _default_resolver block_private = os.getenv( "INTEGRATION_API_BLOCK_PRIVATE_IPS", "false" ).lower() == "true" - ok, reason = check_outbound_url(url, block_private=block_private) + # Resolve the host exactly once and remember the IPs the guard validated so + # the request below can be pinned to them. check_outbound_url only reports + # (ok, reason); a plain httpx client re-resolves the host at connect time, + # which reopens a DNS-rebinding TOCTOU — a base_url host that answers with a + # public IP for the guard and then flips to 169.254.169.254 for the connect + # would reach cloud metadata with the integration's auth headers attached. + resolved_ips: List[str] = [] + + def _recording_resolver(host: str) -> List[str]: + ips = _default_resolver(host) + resolved_ips[:] = ips + return ips + + ok, reason = check_outbound_url( + url, block_private=block_private, resolver=_recording_resolver + ) if not ok: return {"error": f"URL rejected: {reason}", "exit_code": 1} + pinned_ips = _validated_ips(resolved_ips) + if not pinned_ips: + return {"error": "URL rejected: host did not resolve to a usable address", "exit_code": 1} method = method.upper() @@ -455,7 +622,9 @@ async def execute_api_call( auth = httpx.BasicAuth(parts[0], parts[1]) try: - async with httpx.AsyncClient(timeout=30.0) as client: + async with httpx.AsyncClient( + timeout=30.0, transport=_PinnedAsyncTransport(pinned_ips) + ) as client: response = await client.request( method, url, diff --git a/tests/test_integration_api_call_ssrf.py b/tests/test_integration_api_call_ssrf.py index 53dc671c5..f23cc40de 100644 --- a/tests/test_integration_api_call_ssrf.py +++ b/tests/test_integration_api_call_ssrf.py @@ -9,8 +9,13 @@ link-local/metadata is always rejected; RFC-1918/loopback only when INTEGRATION_API_BLOCK_PRIVATE_IPS=true (LAN integrations are the primary use case, so private stays allowed by default). """ +import asyncio +import ipaddress +import ssl from unittest.mock import AsyncMock, MagicMock, patch +import httpcore +import httpx import pytest from src import integrations @@ -97,3 +102,238 @@ async def test_private_base_url_allowed_by_default_blocked_with_knob(monkeypatch assert result["exit_code"] == 1 assert "rejected" in result["error"].lower() client.request.assert_not_called() + + +async def _call_capturing_transport(base_url, path="/items"): + """Drive execute_api_call and return (result, transport) where transport is + the object passed to httpx.AsyncClient(transport=...).""" + resp = MagicMock() + resp.status_code = 200 + resp.headers = {"content-type": "application/json"} + resp.json.return_value = {"ok": True} + resp.text = '{"ok": true}' + + client = AsyncMock() + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=None) + client.request = AsyncMock(return_value=resp) + + captured = {} + + def _fake_async_client(*args, **kwargs): + captured.update(kwargs) + return client + + with ( + patch.object(integrations, "_find_integration", + return_value=_integration(base_url)), + patch("httpx.AsyncClient", side_effect=_fake_async_client), + ): + result = await integrations.execute_api_call("test_integ", "GET", path) + return result, captured.get("transport"), client + + +@pytest.mark.asyncio +async def test_connection_is_pinned_to_the_validated_ip(monkeypatch): + """DNS-rebinding defense: the guard resolves the host once to a benign + public IP, and the request must be pinned to *that* IP so a host that + rebinds to the metadata range at connect time can't be reached with the + integration's auth headers. Static resolution passing the guard is not + enough — a plain client would re-resolve at connect.""" + monkeypatch.setattr("src.url_safety._default_resolver", + lambda host: ["93.184.216.34"]) + result, transport, client = await _call_capturing_transport( + "http://rebinding.attacker.example") + + assert result.get("exit_code") == 0 + client.request.assert_called_once() + assert isinstance(transport, integrations._PinnedAsyncTransport) + assert [str(ip) for ip in transport._pinned_ips] == ["93.184.216.34"] + + +@pytest.mark.asyncio +async def test_pin_carries_the_whole_validated_ip_set(monkeypatch): + """When a host resolves to several records the transport keeps all of them + (check_outbound_url validated every one), in resolver order, so it can fall + back past a dead first address instead of failing the whole call.""" + monkeypatch.setattr("src.url_safety._default_resolver", + lambda host: ["93.184.216.34", "198.51.100.7"]) + result, transport, _ = await _call_capturing_transport("http://multi.example") + + assert result.get("exit_code") == 0 + assert [str(ip) for ip in transport._pinned_ips] == ["93.184.216.34", "198.51.100.7"] + + +class _FakeStream: + """Stand-in for the connected socket the real backend returns.""" + + +class _RecordingBackend: + """Fake httpcore backend: connect_tcp fails for the addresses in `dead` + and succeeds for the rest, recording the order it was asked to connect.""" + + def __init__(self, dead): + self.dead = set(dead) + self.attempts = [] + + async def connect_tcp(self, host, port, timeout=None, local_address=None, + socket_options=None): + self.attempts.append((host, timeout)) + if host in self.dead: + raise httpcore.ConnectError(f"connection refused: {host}") + return _FakeStream() + + +def _pinned_backend(ips, dead): + """A _PinnedAsyncBackend whose underlying connect is the recording fake.""" + backend = integrations._PinnedAsyncBackend(ips) + backend._real = _RecordingBackend(dead) + return backend + + +@pytest.mark.asyncio +async def test_connect_falls_back_from_dead_first_to_live_second(): + """first-dead / second-live: the pinned backend must try the next validated + address when the first refuses, rather than surfacing the failure. It also + ignores the `host` httpcore passes (the original hostname) and connects to + the pinned IPs, which is what keeps TLS SNI / Host on the real hostname.""" + ips = [ipaddress.ip_address("203.0.113.10"), ipaddress.ip_address("198.51.100.7")] + backend = _pinned_backend(ips, dead={"203.0.113.10"}) + + stream = await backend.connect_tcp("original.hostname.example", 443, timeout=5.0) + + assert isinstance(stream, _FakeStream) + # Tried the dead address first, then the live one — never the hostname. + assert [host for host, _ in backend._real.attempts] == ["203.0.113.10", "198.51.100.7"] + # Fallback shared one budget: the second attempt got the time left, not a fresh 5s. + assert backend._real.attempts[1][1] <= 5.0 + + +@pytest.mark.asyncio +async def test_connect_raises_when_every_validated_address_is_dead(): + ips = [ipaddress.ip_address("203.0.113.10"), ipaddress.ip_address("198.51.100.7")] + backend = _pinned_backend(ips, dead={"203.0.113.10", "198.51.100.7"}) + + with pytest.raises(httpcore.ConnectError): + await backend.connect_tcp("original.hostname.example", 443, timeout=5.0) + assert [host for host, _ in backend._real.attempts] == ["203.0.113.10", "198.51.100.7"] + + +@pytest.mark.asyncio +async def test_pinned_transport_reuses_httpx_ca_trust(monkeypatch): + """TLS trust must come from the same builder the default httpx client uses + (certifi + SSL_CERT_FILE / SSL_CERT_DIR via trust_env), not from + ssl.create_default_context()'s system roots — otherwise chains that verified + under the old default client can silently stop verifying.""" + sentinel = ssl.create_default_context() + calls = [] + + def _fake_create(*args, **kwargs): + calls.append(kwargs) + return sentinel + + monkeypatch.setattr(httpx, "create_ssl_context", _fake_create) + transport = integrations._PinnedAsyncTransport([ipaddress.ip_address("93.184.216.34")]) + try: + assert calls, "transport did not build its context via httpx.create_ssl_context" + assert transport._pool._ssl_context is sentinel + finally: + await transport.aclose() + + +@pytest.mark.asyncio +async def test_real_socket_falls_back_from_dead_first_to_live_second(): + """End-to-end over real loopback sockets: pin [127.0.0.2 (nothing + listening), 127.0.0.1 (live)], and the request must succeed by falling back + to the second address while the Host header stays the original hostname — + i.e. only the socket destination moved, vhost/SNI routing did not.""" + captured = {} + + async def handle(reader, writer): + request = await reader.read(4096) + for line in request.split(b"\r\n"): + if line.lower().startswith(b"host:"): + captured["host"] = line.split(b":", 1)[1].strip().decode() + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nhi") + await writer.drain() + writer.close() + + server = await asyncio.start_server(handle, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + async with server: + await server.start_serving() + transport = integrations._PinnedAsyncTransport( + [ipaddress.ip_address("127.0.0.2"), ipaddress.ip_address("127.0.0.1")] + ) + try: + async with httpx.AsyncClient(transport=transport) as client: + resp = await client.get(f"http://pinned.example:{port}/health") + finally: + await transport.aclose() + + assert resp.status_code == 200 + assert resp.text == "hi" + assert captured.get("host") == f"pinned.example:{port}" + + +@pytest.mark.asyncio +async def test_ip_literal_base_url_still_pins_and_is_not_rejected(): + """A base_url that is already an IP has nothing to rebind, but it must not + trip the "did not resolve" guard either. + + check_outbound_url resolves even a literal (getaddrinfo returns the address + itself), so the captured list is populated and the pin is a no-op rather + than a rejection. Uses the real resolver on purpose — no monkeypatch — so + this would catch the fail-closed branch firing on a literal. + """ + result, transport, client = await _call_capturing_transport( + "http://93.184.216.34") + + assert result.get("exit_code") == 0 + assert isinstance(transport, integrations._PinnedAsyncTransport) + assert [str(ip) for ip in transport._pinned_ips] == ["93.184.216.34"] + + +@pytest.mark.asyncio +async def test_ipv6_base_url_pins_every_validated_address(monkeypatch): + """IPv6 goes down the same path as v4. + + Resolution is stubbed rather than using a literal so this doesn't depend on + the runner having IPv6 configured. + """ + v6 = "2606:2800:220:1:248:1893:25c8:1946" + monkeypatch.setattr("src.url_safety._default_resolver", lambda host: [v6]) + result, transport, client = await _call_capturing_transport("http://v6.example") + + assert result.get("exit_code") == 0 + assert isinstance(transport, integrations._PinnedAsyncTransport) + assert [str(ip) for ip in transport._pinned_ips] == [v6] + + +def test_validated_ips_strips_zone_id_and_drops_junk(): + """getaddrinfo can hand back a scoped v6 address like 'fe80::1%eth0'.""" + got = integrations._validated_ips( + ["93.184.216.34", "fe80::1%eth0", "not-an-ip", None, "2001:db8::5"] + ) + assert [str(ip) for ip in got] == ["93.184.216.34", "fe80::1", "2001:db8::5"] + + +def test_validated_ips_deduplicates_repeated_addresses(): + """The resolver is getaddrinfo(host, None) with no socktype filter, so glibc + returns one record per socktype and a single-homed host arrives three times + over. Duplicates must collapse (first-seen order kept) or the connect + fallback wastes its shared deadline retrying one dead address.""" + got = integrations._validated_ips( + ["93.184.216.34", "93.184.216.34", "93.184.216.34"] + ) + assert [str(ip) for ip in got] == ["93.184.216.34"] + + # Order is first-seen, and distinct addresses all survive. + got = integrations._validated_ips( + ["198.51.100.7", "93.184.216.34", "198.51.100.7", "2001:db8::5"] + ) + assert [str(ip) for ip in got] == ["198.51.100.7", "93.184.216.34", "2001:db8::5"] + + # A zone-id variant is the same address once stripped, so it collapses too. + got = integrations._validated_ips(["fe80::1%eth0", "fe80::1%eth1", "fe80::1"]) + assert [str(ip) for ip in got] == ["fe80::1"] diff --git a/tests/test_integrations_api_call_truncation.py b/tests/test_integrations_api_call_truncation.py index bf1ec7d05..a0ad61b4a 100644 --- a/tests/test_integrations_api_call_truncation.py +++ b/tests/test_integrations_api_call_truncation.py @@ -83,9 +83,10 @@ async def _call(json_data, status=200): with ( patch.object(integrations, "_find_integration", return_value=DUMMY_INTEGRATION), patch("httpx.AsyncClient", return_value=mock_client), - # api.example.com doesn't resolve; the SSRF guard would fail closed. - # These tests are about truncation, so stub the guard open. - patch("src.url_safety.check_outbound_url", return_value=(True, "ok")), + # api.example.com doesn't resolve. Point the resolver at a public + # address instead of stubbing the guard open, so the real check (and + # the connect-IP pinning that reads its result) still runs. + patch("src.url_safety._default_resolver", lambda host: ["93.184.216.34"]), ): return await integrations.execute_api_call("test_integ", "GET", "/items") @@ -101,9 +102,10 @@ async def _call_with_integration(integration, path="/items"): with ( patch.object(integrations, "_find_integration", return_value=integration), patch("httpx.AsyncClient", return_value=mock_client), - # api.example.com doesn't resolve; the SSRF guard would fail closed. - # These tests are about URL joining, so stub the guard open. - patch("src.url_safety.check_outbound_url", return_value=(True, "ok")), + # api.example.com doesn't resolve. Point the resolver at a public + # address instead of stubbing the guard open, so the real check (and + # the connect-IP pinning that reads its result) still runs. + patch("src.url_safety._default_resolver", lambda host: ["93.184.216.34"]), ): result = await integrations.execute_api_call("test_integ", "GET", path) return result, mock_client From 20e7fc0164286e1521569d9edc17a4ae4d0d2e22 Mon Sep 17 00:00:00 2001 From: adabarbulescu <94562950+adabarbulescu@users.noreply.github.com> Date: Tue, 4 Aug 2026 13:17:45 +0300 Subject: [PATCH 3/3] fix(skills): require manage_skills action (#5856) --- src/tools/system.py | 6 ++++-- tests/test_manage_skills_action_required.py | 24 +++++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) create mode 100644 tests/test_manage_skills_action_required.py diff --git a/src/tools/system.py b/src/tools/system.py index 813d57df2..c2eb9ceab 100644 --- a/src/tools/system.py +++ b/src/tools/system.py @@ -46,7 +46,9 @@ async def do_manage_skills(content: str, owner: Optional[str] = None) -> Dict: except ValueError: return {"error": "Invalid JSON arguments", "exit_code": 1} - action = (args.get("action") or "").lower() + action = (args.get("action") or "").strip().lower() + if not action: + return {"error": "action is required (list|view|view_ref|add|edit|patch|publish|delete|search)", "exit_code": 1} from services.memory.skills import SkillsManager from services.memory.skill_format import Skill, slugify from src.constants import DATA_DIR @@ -55,7 +57,7 @@ async def do_manage_skills(content: str, owner: Optional[str] = None) -> Dict: # Accept legacy `skill_id` as an alias for `name`. name = (args.get("name") or args.get("skill_id") or "").strip() - if action in ("list", "index", ""): + if action in ("list", "index"): all_skills = sm.load(owner=owner) if not all_skills: return {"results": "No skills yet. Create one with action='add'."} diff --git a/tests/test_manage_skills_action_required.py b/tests/test_manage_skills_action_required.py new file mode 100644 index 000000000..4efae8026 --- /dev/null +++ b/tests/test_manage_skills_action_required.py @@ -0,0 +1,24 @@ +import json + +import pytest + +from src.tools.system import do_manage_skills + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [ + {}, + {"action": ""}, + {"action": " "}, + {"name": "demo", "description": "x", "procedure": ["step"]}, + ], +) +async def test_manage_skills_requires_action(payload): + result = await do_manage_skills(json.dumps(payload), owner="test") + + assert result == { + "error": "action is required (list|view|view_ref|add|edit|patch|publish|delete|search)", + "exit_code": 1, + }