diff --git a/CLAUDE.md b/CLAUDE.md index e180b0db2..4cebbec69 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -248,6 +248,70 @@ Optional web search augmentation via the Staan API, allowing the LLM to combine - `openrag/services/orchestrators/query_service.py` — `_prepare_for_web_only()`, web search logic in `_prepare_for_chat_completion()` - `openrag/api/routers/user/chat.py` — `__prepare_sources()` merges document and web sources +### Indexing Status Callbacks + +Indexing is asynchronous, so a client can either poll the task-status URL or +hand OpenRag a URL to notify once the task settles. + +**Request fields** (multipart form, on `POST` and `PUT /indexer/partition/{partition}/file/{file_id}`): +- `callback_url` — POSTed once the task reaches a terminal state +- `callback_token` — sent as `Authorization: Bearer ` on that POST + +**Body:** `{"partition", "file_id", "status": "success"|"error", "metadata"}`. `metadata` is +the upload metadata echoed back **minus** `UPLOAD_METADATA_SERVER_KEYS` (`core/utils/consts.py`: +`source`, `filename`, `original_filename`, `file_size`, `file_id`, `content_sha256`) — an exclusion, +not a fixed field list, so any caller-supplied field (cozy-stack's revision marker is `doc_rev`) +travels through unrecomputed and under whatever name the caller gave it; a key the caller never sent +is simply absent, not echoed back as `null`. The exclusion exists because the target is a +caller-supplied URL and `_build_metadata` merges those server-computed keys into the same dict — +`source` (the server's on-disk path) is the load-bearing one. `test_build_metadata_only_adds_keys_in_upload_metadata_server_keys` +(`tests/unit/services/orchestrators/test_indexing_service.py`) fails the build if `_build_metadata` +ever injects a key the constant doesn't cover. + +**Guarantees:** one attempt, no retries, a 5 s deadline on the request, never raises — a failed +callback is logged and cannot change the indexing outcome. No `callback_url` → strict no-op. A +user-cancelled task sends nothing (only a real failure notifies `"error"`), which is why the sender +keys off the return value of `set_failed_if_not_cancelled` and the pool's pre-flight handler catches +`Exception`, not `BaseException`. + +**Not guaranteed:** delivery. The send is awaited on the worker's slot, so a blackholing target costs +up to 5 s of indexing throughput per file, and an actor lost before the task settles notifies +nothing at all. Clients keep a timeout and fall back to polling the task-status URL. + +`callback_url` is checked against the SSRF guard (`is_safe_url` — scheme, loopback/private/link-local, +decimal/hex/octal/short-form IPv4 literals, all normalized through `socket.inet_aton` so a legacy +numeric spelling can't overflow `ipaddress.ip_address(int(...))` into a false-negative IPv6 address) +but not resolved: neither a DNS lookup nor an https requirement is enforced beyond that literal +check, deliberately — the URL is caller-supplied, so picking a safe target is the caller's call, not +OpenRag's to police. Checked twice — in the router (immediate `400`) and again in the sender (a +direct caller bypasses the router) — so accepting a hostname the sender would refuse never happens. +`INDEXING_CALLBACK_ALLOW_PRIVATE_URLS=true` / `indexing_callback.allow_private_urls` lifts the +*address* half of the guard for dev stacks whose target is a local instance; the scheme check always +applies. Keep it off in production: with it on, any user allowed to upload can make the server POST +to an internal address. + +**Rolling deploys:** the worker actors are named, detached and `get_if_exists`, so changing +`process_file`'s remote contract requires bumping `_INDEXER_ACTOR_PROTOCOL_VERSION` in +`indexer_pool.py` (v3 → v4 for `callback_url`/`callback_token`; v5 added worker-ref-registration wait +and TSM `set_state` fencing; v7 folds in a second, independent v6 lineage — STT-preset-aware registry +hydration plus the `_active_indexation_config` contextvar — that landed on `develop` under the same +version string while this branch's own v6 was in flight). Without the bump, new replicas attach to the +previous release's actors and every submit raises `TypeError`. Old generations are retired with +`services/workers/retire_indexer_generation.py`. + +**Key files:** +- `openrag/services/workers/indexing_callback.py` — `send_indexing_callback()` (was `webhook.py`; the + target is a normal authenticated route now, not an unauthenticated webhook trigger) +- `openrag/core/utils/url_safety.py` — `is_safe_url(url, *, allow_private_hosts=False)`, shared with + the MCP `index_url` tool (which keeps the strict default). The web-search content fetcher + (`services/websearch/content_fetcher.py`) does **not** import this — it has its own older, + un-synced `_is_safe_url`, so hardening this module does not automatically harden that one. +- `openrag/core/config/indexation.py` — `IndexingCallbackConfig` + +Both fields travel the same chain as the rest of an indexing job: `api/routers/admin/indexing.py` → +`services/orchestrators/indexing_service.py` → `core/indexing/dispatcher.py` (port) → +`services/workers/dispatcher.py` → `indexer_pool.py` → `indexer_actor.py` → `indexing_callback.py`. + ### File Quota System Per-user file quota enforcement tracked via the `file_count` and `file_quota` columns on `users`, and `created_by` on `files`. diff --git a/conf/config.yaml b/conf/config.yaml index 3e4e84ca3..be8af97d4 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -279,6 +279,13 @@ loader: concurrency_limit: 20 enable_thinking: null +# --- Indexing callback --- +# Env: INDEXING_CALLBACK_ALLOW_PRIVATE_URLS +indexing_callback: + # Dev only. Keep false in production: it disables the SSRF guard on a + # caller-supplied URL. + allow_private_urls: false + # --- Ray --- ray: indexer: diff --git a/docs/content/docs/documentation/API.mdx b/docs/content/docs/documentation/API.mdx index 95fea4f5e..e3ae47c52 100644 --- a/docs/content/docs/documentation/API.mdx +++ b/docs/content/docs/documentation/API.mdx @@ -85,11 +85,50 @@ Upload a new file to a specific partition for indexing. **Request Body (form-data):** - `file` (binary): File to upload - `metadata` (JSON string): File metadata (e.g., `{"owner": "user1"}`) +- `workspace_ids` (JSON array, optional): Workspaces to add the file to +- `callback_url` (string, optional): URL notified once indexing reaches a terminal state +- `callback_token` (string, optional): Bearer token for that notification **Responses:** - `201 Created`: Returns task status URL +- `400 Bad Request`: `callback_url` is malformed or not a public `http(s)` URL - `409 Conflict`: File already exists in partition +##### Indexing status callback + +Indexing is asynchronous. Pass a `callback_url` and OpenRAG POSTs the outcome +once the task settles, so a client can wait instead of polling: + +```json +{ + "partition": "alice.example.org", + "file_id": "file-123", + "status": "success", + "metadata": {"doc_rev": "", "datetime": "...", "doctype": "..."} +} +``` + +`status` is `"success"` or `"error"`; a user-cancelled task sends nothing. +`metadata` is whatever you sent at upload, echoed back unchanged — including +a revision-tracking field under any name you like, if that's how your receiver +orders these callbacks (the example above uses `doc_rev`). A field you didn't +send is simply absent, not sent back as `null`. Only server-computed keys are +held back: the on-disk path, content hash, file size, filenames, and +`file_id` (already a top-level field). + +For an authenticated target, pass `callback_token` — sent as `Authorization: +Bearer `, never in the URL or the payload. Picking a safe scheme for it +is on the caller: `callback_url` is otherwise unchecked beyond the public +`http(s)` requirement below. + +Best-effort: one attempt, no retries, a 5 s deadline, and any failure is +logged without affecting the indexing result. Delivery is not guaranteed, so +keep a client-side timeout and fall back to polling the task-status URL. The +`callback_url` is checked against the SSRF guard and must be a public +`http(s)` address — see +[`INDEXING_CALLBACK_ALLOW_PRIVATE_URLS`](/openrag/documentation/env_vars) to +target a local instance in development. + ##### Temporal Filtering OpenRAG supports temporal filtering to retrieve documents from specific time periods. The client can include the temporal field to allow temporal-aware search in search endpoints. diff --git a/docs/content/docs/documentation/env_vars.md b/docs/content/docs/documentation/env_vars.md index c5702d407..7dad101a8 100644 --- a/docs/content/docs/documentation/env_vars.md +++ b/docs/content/docs/documentation/env_vars.md @@ -25,6 +25,7 @@ Openrag loads all files into a pivot markdown file format before proceeding to c | `CONTENT_DEDUPLICATION_ENABLED` | `bool` | `true` | Rejects a file when identical content already exists in the same partition. Set it to `false` when a test intentionally indexes duplicates. | | `PDFLOADER` | `str` | `PyMuPDFLoader` | PDF parsing engine. `PyMuPDFLoader` (default) is a lightweight, fast, CPU-friendly backend for searchable PDFs. Switch to `MarkerLoader` for OCR / scanned documents, complex layouts and embedded images (heavier; GPU-friendly). Other options: `DoclingLoader`, `DotsOCRLoader`.| | `PARSE_TIMEOUT` | `int` | `3600` | Outer wall-clock bound (in seconds) for a single file's parse stage, whichever loader runs it. Marker and Docling self-limit via their own timeouts, but `PyMuPDFLoader` has none — this bound stops a wedged parse from stalling indexing: the file fails and is reported instead. | +| `INDEXING_CALLBACK_ALLOW_PRIVATE_URLS` | `bool` | `false` | Development-only escape hatch for the upload `callback_url`. By default a callback targeting `localhost`, a private or link-local address is refused (SSRF guard on a caller-supplied URL). Set it to `true` only in a dev stack whose callback target is a local instance (e.g. `http://cozy.localhost:8080`). Non-`http(s)` schemes stay rejected either way. | :::caution `PyMuPDFLoader` (the default) is a lightweight PDF loader that cannot process non-searchable (image-based) PDFs and does not extract or handle embedded images. Set `PDFLOADER=MarkerLoader` when you need those. diff --git a/openrag/api/routers/admin/indexing.py b/openrag/api/routers/admin/indexing.py index dc356a4ea..eee0d44d0 100644 --- a/openrag/api/routers/admin/indexing.py +++ b/openrag/api/routers/admin/indexing.py @@ -14,6 +14,7 @@ import json from pathlib import Path from typing import Any +from urllib.parse import urlparse from api.dependencies.auth import ( check_user_file_quota, @@ -36,6 +37,7 @@ from core.utils.filename import sanitize_filename from core.utils.log_tail import app_log_file from core.utils.logging import get_logger +from core.utils.url_safety import is_safe_url from di.providers import get_auth_service, get_config, get_indexing_service, get_partition_service from fastapi import ( APIRouter, @@ -52,6 +54,30 @@ logger = get_logger() +def _validate_callback_url(callback_url: str | None, config) -> None: + """Reject a callback_url the server must not POST to. + + The sender re-checks; this is only so the caller hears about it now rather + than losing the callback silently. + """ + if not callback_url: + return + allow_private = bool(getattr(getattr(config, "indexing_callback", None), "allow_private_urls", False)) + try: + # is_safe_url never touches the port; a non-numeric one raises here. + urlparse(callback_url).port + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="callback_url is not a valid URL", + ) from exc + if not is_safe_url(callback_url, allow_private_hosts=allow_private): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="callback_url must be a public http(s) URL", + ) + + def build_url(request: Request, route_name: str, *, preferred_url_scheme: str | None = None, **path_params) -> str: """Build a URL using the preferred scheme if configured.""" url = request.url_for(route_name, **path_params) @@ -126,11 +152,18 @@ async def add_file( file: UploadFile = Depends(validate_file_format), metadata: dict = Depends(validate_metadata), workspace_ids: str | None = Form(None, description="JSON array of workspace IDs to add the file to"), + callback_url: str | None = Form(None, description="Optional URL notified when async indexing finishes"), + callback_token: str | None = Form( + None, + description="Optional bearer token sent as `Authorization` on the callback_url request", + ), user=Depends(require_partition_editor), _quota_check=Depends(check_user_file_quota), config=Depends(get_config), service=Depends(get_indexing_service), ): + _validate_callback_url(callback_url, config) + if await service.file_exists(file_id, partition): raise HTTPException( status_code=status.HTTP_409_CONFLICT, @@ -195,6 +228,8 @@ async def add_file( user=user, workspace_ids=parsed_workspace_ids, content_sha256=content_sha256, + callback_url=callback_url, + callback_token=callback_token, ) except BaseException as exc: # A submission whose outcome is unknown may have left a worker running @@ -284,10 +319,17 @@ async def put_file( file_id: str = Depends(validate_file_id), file: UploadFile = Depends(validate_file_format), metadata: dict = Depends(validate_metadata), + callback_url: str | None = Form(None, description="Optional URL notified when async indexing finishes"), + callback_token: str | None = Form( + None, + description="Optional bearer token sent as `Authorization` on the callback_url request", + ), user=Depends(require_partition_editor), config=Depends(get_config), service=Depends(get_indexing_service), ): + _validate_callback_url(callback_url, config) + if not await service.file_exists(file_id, partition): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -334,6 +376,8 @@ async def put_file( user=user, replace=True, content_sha256=content_sha256, + callback_url=callback_url, + callback_token=callback_token, ) except BaseException as exc: # A submission whose outcome is unknown may have left a worker running diff --git a/openrag/core/config/indexation.py b/openrag/core/config/indexation.py index f62266c62..9370db9c2 100644 --- a/openrag/core/config/indexation.py +++ b/openrag/core/config/indexation.py @@ -201,3 +201,13 @@ class LoaderConfig(ConfigMixin): # Max depth of nested .eml-in-.eml attachments the EmlLoader will descend # into. Bounds recursion when .eml files are nested inside one another. eml_max_recursion_depth: int = 5 + + +class IndexingCallbackConfig(ConfigMixin): + """Server-side policy for the caller-supplied per-upload ``callback_url``. + + With ``allow_private_urls`` on, anyone allowed to upload can make the + server POST to an internal address. Dev stacks only. + """ + + allow_private_urls: bool = False diff --git a/openrag/core/config/loader.py b/openrag/core/config/loader.py index 0e0dcfe1c..dd9850f35 100644 --- a/openrag/core/config/loader.py +++ b/openrag/core/config/loader.py @@ -100,6 +100,8 @@ ("SAVE_MARKDOWN", "loader.save_markdown", bool), ("SAVE_UPLOADED_FILES", "loader.save_uploaded_files", bool), ("CONTENT_DEDUPLICATION_ENABLED", "loader.content_deduplication_enabled", bool), + # Indexing callback + ("INDEXING_CALLBACK_ALLOW_PRIVATE_URLS", "indexing_callback.allow_private_urls", bool), ("PDFLOADER", "loader.file_loaders.pdf", str), ("AUDIOLOADER", "loader.file_loaders.wav", str), ("MARKER_MAX_TASKS_PER_CHILD", "loader.marker_max_tasks_per_child", int), diff --git a/openrag/core/config/root.py b/openrag/core/config/root.py index 73adf8e72..4a38033e9 100644 --- a/openrag/core/config/root.py +++ b/openrag/core/config/root.py @@ -14,7 +14,7 @@ SemaphoreConfig, VLMConfig, ) -from .indexation import LoaderConfig +from .indexation import IndexingCallbackConfig, LoaderConfig from .infrastructure import ( PathsConfig, PromptsConfig, @@ -60,6 +60,7 @@ class Settings(ConfigMixin): paths: PathsConfig = Field(default_factory=PathsConfig) prompts: PromptsConfig = Field(default_factory=PromptsConfig) loader: LoaderConfig = Field(default_factory=LoaderConfig) + indexing_callback: IndexingCallbackConfig = Field(default_factory=IndexingCallbackConfig) ray: RayConfig = Field(default_factory=RayConfig) chunker: ChunkerConfig = Field(default_factory=ChunkerConfig) retriever: RetrieverConfig = Field(default_factory=SingleRetrieverConfig) diff --git a/openrag/core/indexing/dispatcher.py b/openrag/core/indexing/dispatcher.py index 3250aaa61..30c02d39a 100644 --- a/openrag/core/indexing/dispatcher.py +++ b/openrag/core/indexing/dispatcher.py @@ -37,6 +37,8 @@ async def dispatch_indexing( replace: bool, indexation_config: dict | None = None, embedder_name: str | None = None, + callback_url: str | None = None, + callback_token: str | None = None, require_existing_partition: bool = False, allow_legacy_require_existing_partition_retry: bool = False, ) -> str: diff --git a/openrag/core/utils/consts.py b/openrag/core/utils/consts.py index 3294f4329..bd1a11823 100644 --- a/openrag/core/utils/consts.py +++ b/openrag/core/utils/consts.py @@ -51,3 +51,13 @@ def strip_protected_metadata(metadata: dict | None) -> tuple[dict, list[str]]: for key in removed: del md[key] return md, removed + + +# The server-computed keys ``IndexingService._build_metadata`` merges into +# upload metadata. The indexing-status callback excludes exactly these before +# echoing metadata to a caller-supplied URL. Keep in sync with +# ``_build_metadata`` — enforced by +# ``test_build_metadata_only_adds_keys_in_upload_metadata_server_keys``. +UPLOAD_METADATA_SERVER_KEYS: frozenset[str] = frozenset( + {"source", "filename", "original_filename", "file_size", "file_id", "content_sha256"} +) diff --git a/openrag/core/utils/url_safety.py b/openrag/core/utils/url_safety.py index c50d0608f..9d033479f 100644 --- a/openrag/core/utils/url_safety.py +++ b/openrag/core/utils/url_safety.py @@ -2,18 +2,39 @@ Pure-stdlib host/address checks shared by any server-side fetcher (web-search content fetch, MCP ``index_url``). Blocks loopback / private / link-local / -reserved / non-global addresses and the decimal-integer IPv4 encoding that -resolvers accept. Regular hostnames pass the literal check; callers that follow -redirects MUST re-validate every hop, since a public hostname can redirect to a -private target. +reserved / non-global addresses and the alternate IPv4 encodings that +resolvers accept (decimal-integer, hex, octal, short form). Regular hostnames +pass the literal check; callers that follow redirects MUST re-validate every +hop, since a public hostname can redirect to a private target. """ from __future__ import annotations import ipaddress +import socket from urllib.parse import urlparse +def _bare_numeric_host_value(host: str) -> int | None: + """Parse *host* as a single C-style integer literal (decimal, ``0x`` hex, + or ``0``-prefixed octal) — the form ``inet_aton`` accepts when a host has + no dots. Returns ``None`` if *host* contains a dot or isn't purely such a + literal. + """ + if "." in host: + return None + try: + if host.lower().startswith("0x"): + return int(host, 16) + if len(host) > 1 and host[0] == "0" and host.isdigit(): + return int(host, 8) + if host.isdigit(): + return int(host, 10) + except ValueError: + return None + return None + + def is_blocked_address(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: """True for any IP a server-side fetcher must not contact. @@ -32,13 +53,17 @@ def is_blocked_address(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> b ) -def is_safe_url(url: str) -> bool: +def is_safe_url(url: str, *, allow_private_hosts: bool = False) -> bool: """Return True only if *url* is safe for a server-side fetch. Blocks non-HTTP(S) schemes, ``localhost``, IPv4/IPv6 literals in - private/loopback/link-local/reserved ranges, and decimal-integer-encoded - IPv4 (e.g. ``2130706433`` == ``127.0.0.1``). Regular hostnames pass — the - caller must re-check each redirect hop. + private/loopback/link-local/reserved ranges, and the alternate IPv4 + spellings a resolver expands (``2130706433``, ``0x7f000001``, + ``0177.0.0.1``, ``127.1`` — all ``127.0.0.1``). Regular hostnames pass — + the caller must re-check each redirect hop. + + ``allow_private_hosts`` keeps the scheme check but skips the address checks. + Never enable it for a fetch target an end user controls. """ try: parsed = urlparse(url) @@ -52,6 +77,9 @@ def is_safe_url(url: str) -> bool: if not host: return False + if allow_private_hosts: + return True + if host.lower() == "localhost": return False @@ -61,12 +89,26 @@ def is_safe_url(url: str) -> bool: except ValueError: pass - # Decimal-integer form (e.g. 2130706433 → 127.0.0.1). ip_address(int) - # interprets the value as a packed IPv4 address, matching glibc's resolver. + # A bare (undotted) numeric host that doesn't fit in 32 bits: block it + # ourselves rather than trust inet_aton's overflow handling, which is not + # portable — glibc rejects such a host outright (matching the real + # resolver, confirmed via getaddrinfo), but at least one other libc has + # been observed to silently wrap it mod 2**32 into a blocked address + # instead of raising. Relying on ip_address(int(host)) here would have + # the opposite problem: it never raises for a value this large, it just + # silently builds an unrelated, often-public-looking IPv6 address. + numeric_value = _bare_numeric_host_value(host) + if numeric_value is not None and numeric_value > 0xFFFFFFFF: + return False + + # inet_aton covers decimal/hex/octal/short-form IPv4 the same way glibc's + # resolver does, for every in-range value. try: - return not is_blocked_address(ipaddress.ip_address(int(host))) - except (ValueError, TypeError): + packed = socket.inet_aton(host) + except OSError: pass + else: + return not is_blocked_address(ipaddress.IPv4Address(packed)) # Regular hostname — passes the literal check; redirect hops re-validated. return True diff --git a/openrag/services/orchestrators/indexing_service.py b/openrag/services/orchestrators/indexing_service.py index 227f7d21f..3c9a7ecdb 100644 --- a/openrag/services/orchestrators/indexing_service.py +++ b/openrag/services/orchestrators/indexing_service.py @@ -114,7 +114,10 @@ def _build_metadata( original_filename: str | None, content_sha256: str | None, ) -> dict: - """Assemble the indexing metadata exactly as the legacy router did.""" + """Assemble the indexing metadata exactly as the legacy router did. + + The keys added below must match ``UPLOAD_METADATA_SERVER_KEYS`` exactly. + """ metadata = dict(metadata or {}) metadata.update( { @@ -243,11 +246,14 @@ async def add_file( workspace_ids: list[str] | None = None, replace: bool = False, content_sha256: str | None = None, + callback_url: str | None = None, + callback_token: str | None = None, ) -> str: """Assemble metadata and queue an (re)indexing job; return its task id. Workspace association happens inside the worker's ``add_file`` after a successful index — the router only pre-validates the ids. + *callback_url*/*callback_token* are forwarded to the worker as-is. """ if self._deduplication_enabled() and content_sha256 is None: content_sha256 = await asyncio.to_thread(_sha256_file, file_path) @@ -277,6 +283,8 @@ async def add_file( replace=replace, indexation_config=indexation_config, embedder_name=embedder_name, + callback_url=callback_url, + callback_token=callback_token, require_existing_partition=require_existing_partition, allow_legacy_require_existing_partition_retry=legacy_actor_preserves_partition_guard, ) diff --git a/openrag/services/websearch/content_fetcher.py b/openrag/services/websearch/content_fetcher.py index d13411bc8..c38b94871 100644 --- a/openrag/services/websearch/content_fetcher.py +++ b/openrag/services/websearch/content_fetcher.py @@ -1,11 +1,12 @@ import asyncio import ipaddress import socket -from urllib.parse import urljoin, urlparse +from urllib.parse import urljoin import httpx import lxml.html from core.utils.logging import get_logger +from core.utils.url_safety import is_blocked_address, is_safe_url from html_to_markdown import convert from services.websearch.base import WebResult @@ -22,70 +23,6 @@ _MAX_REDIRECTS = 10 -def _is_blocked_address(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: - """True for any IP that a server-side fetcher must not contact. - - Checks all private/reserved/non-global flags explicitly so the guard is - correct across Python minor releases (``is_global`` semantics changed - between 3.10 and 3.11 for CGNAT and some multicast ranges). - """ - return ( - addr.is_loopback - or addr.is_private - or addr.is_link_local - or addr.is_reserved - or addr.is_unspecified - or addr.is_multicast - or not addr.is_global - ) - - -def _is_safe_url(url: str) -> bool: - """Return True only if *url* is safe for a server-side fetch. - - Blocks: - - Non-HTTP(S) schemes - - ``localhost`` hostname - - IPv4/IPv6 literals in private, loopback, link-local, or reserved ranges - - Decimal-integer-encoded IPv4 addresses (e.g. ``2130706433`` == ``127.0.0.1``) - - Regular hostnames pass through; per-hop redirect validation (in - :meth:`ContentFetcher._fetch_single`) re-checks every redirect target, - covering the case where a public hostname redirects to a private address. - """ - try: - parsed = urlparse(url) - except Exception: - return False - - if parsed.scheme not in ("http", "https"): - return False - - host = parsed.hostname - if not host: - return False - - if host.lower() == "localhost": - return False - - # Dotted-decimal or IPv6 literal (e.g. "127.0.0.1", "::1", "10.0.0.1") - try: - return not _is_blocked_address(ipaddress.ip_address(host)) - except ValueError: - pass - - # Decimal-integer form (e.g. 2130706433 → 127.0.0.1). - # ipaddress.ip_address(int) interprets the value as a packed IPv4 address, - # matching how glibc's resolver (and therefore httpx) handles such hostnames. - try: - return not _is_blocked_address(ipaddress.ip_address(int(host))) - except (ValueError, TypeError): - pass - - # Regular hostname — passes initial check; every redirect hop is re-validated. - return True - - class ContentFetcher: """Fetch and extract text content from web search result URLs.""" @@ -118,7 +55,7 @@ async def _guard_request(request: httpx.Request) -> None: """httpx request hook: block requests whose host *resolves* to a non-global IP. - ``_is_safe_url`` only inspects the literal host, so a public-looking + ``is_safe_url`` only inspects the literal host, so a public-looking hostname that resolves to an internal address (DNS-rebinding-style SSRF) would otherwise slip through. This hook resolves the host and rejects any non-global resolved IP. It runs for the initial request @@ -137,7 +74,7 @@ async def _guard_request(request: httpx.Request) -> None: addr = ipaddress.ip_address(ip) except ValueError as e: raise httpx.RequestError(f"Unparseable address for {host}", request=request) from e - if _is_blocked_address(addr): + if is_blocked_address(addr): logger.warning("Blocked SSRF attempt to non-global address", host=host, ip=ip) raise httpx.RequestError(f"Blocked non-global address {ip} for {host}", request=request) @@ -145,13 +82,13 @@ async def _fetch_single(self, client: httpx.AsyncClient, url: str) -> str | None """Fetch a single URL and extract text. Returns None on any failure. Redirects are followed manually (``follow_redirects=False``) so every - hop is validated by :func:`_is_safe_url` before the request is sent. + hop is validated by :func:`is_safe_url` before the request is sent. This prevents a legitimate initial URL from redirecting to a private address after the initial check passes. """ current_url = url for _ in range(_MAX_REDIRECTS + 1): - if not _is_safe_url(current_url): + if not is_safe_url(current_url): logger.warning("Blocked unsafe URL in web search results", url=current_url) return None diff --git a/openrag/services/workers/dispatcher.py b/openrag/services/workers/dispatcher.py index b07d7f7c5..8384e2c24 100644 --- a/openrag/services/workers/dispatcher.py +++ b/openrag/services/workers/dispatcher.py @@ -239,6 +239,8 @@ async def dispatch_indexing( replace: bool, indexation_config: dict | None = None, embedder_name: str | None = None, + callback_url: str | None = None, + callback_token: str | None = None, require_existing_partition: bool = False, allow_legacy_require_existing_partition_retry: bool = False, ) -> str: @@ -321,6 +323,8 @@ async def dispatch_indexing( "replace": replace, "indexation_config": indexation_config, "embedder_name": embedder_name, + "callback_url": callback_url, + "callback_token": callback_token, } if require_existing_partition: submit_kwargs[_REQUIRE_EXISTING_PARTITION_KWARG] = True diff --git a/openrag/services/workers/indexer_actor.py b/openrag/services/workers/indexer_actor.py index 9e1e2c6e0..d2e5948bd 100644 --- a/openrag/services/workers/indexer_actor.py +++ b/openrag/services/workers/indexer_actor.py @@ -8,6 +8,7 @@ from core.models.document import Document from core.utils.logging import get_logger +from services.workers.indexing_callback import send_indexing_callback from services.workers.pipeline_builder import ( REPLACE_OLD_CHUNK_COLLECTION_ROW_KEY, REPLACE_OLD_CHUNK_IDS_ROW_KEY, @@ -20,6 +21,13 @@ logger = get_logger() +class _TaskCancelledBeforeStart(Exception): + """Internal signal: ``set_state(SERIALIZING)`` reported the task as already + fenced (cancelled) — not a failure, so it must skip failure-marking and the + error callback entirely (distinct from the TSM *raising*, which is a real + outage and does need both).""" + + class IndexerWorker: """Pure-Python core of the thin indexer actor. @@ -66,6 +74,8 @@ async def process_file( replace: bool = False, indexation_config: dict[str, Any] | None = None, embedder_name: str | None = None, + callback_url: str | None = None, + callback_token: str | None = None, require_existing_partition: bool = False, resolved_prompts: dict[str, str] | None = None, ) -> dict[str, Any]: @@ -74,16 +84,24 @@ async def process_file( Returns a plain dict ``{"stored_count": int, "stage": "stored"}`` on success. On failure, state is set to FAILED and the exception is re-raised so the Ray task is marked as errored. + + When *callback_url* is provided, a best-effort ``POST`` notification is + sent to it once the task reaches a terminal state; never affects the + indexing outcome. """ - accepted = await retry_idempotent_ray_actor_method( - lambda: self._tsm.set_state.remote(task_id, "SERIALIZING"), - task_description=f"set_state({task_id}, SERIALIZING)", - ) - if accepted is False: - raise RuntimeError(f"Task {task_id} was cancelled before indexing started") + file_id = metadata.get("file_id", "") + log = logger.bind(file_id=file_id, partition=partition, task_id=task_id) row: dict[str, Any] | None = None catalog_written = False try: + # Inside the try so a TaskStateManager outage here still reaches + # the except block below and sends the error callback. + accepted = await retry_idempotent_ray_actor_method( + lambda: self._tsm.set_state.remote(task_id, "SERIALIZING"), + task_description=f"set_state({task_id}, SERIALIZING)", + ) + if accepted is False: + raise _TaskCancelledBeforeStart document = await _load_document(path, metadata, partition) # One indexation timestamp for this file, shared by the Milvus chunks # (via the store stage) and the Postgres catalog row, so they agree. @@ -138,7 +156,10 @@ async def process_file( lambda: self._tsm.set_state.remote(task_id, "COMPLETED"), task_description=f"set_state({task_id}, COMPLETED)", ) - return {"stored_count": row.get("stored_count", 0), "stage": row.get("stage", "")} + except _TaskCancelledBeforeStart: + # The TSM already told us this task is fenced/cancelled — no need to + # ask it again, and a cancellation must not fire an error callback. + raise RuntimeError(f"Task {task_id} was cancelled before indexing started") from None except Exception: should_cleanup_vectors = row is not None and ( row.get("stored_count", 0) > 0 or row.get("stage") == "store_failed" @@ -152,11 +173,28 @@ async def process_file( task_id=task_id, ) tb = traceback.format_exc() - await retry_idempotent_ray_actor_method( - lambda: self._tsm.set_failed_if_not_cancelled.remote(task_id, tb), - task_description=f"set_failed_if_not_cancelled({task_id})", - ) + try: + was_failed = await retry_idempotent_ray_actor_method( + lambda: self._tsm.set_failed_if_not_cancelled.remote(task_id, tb), + task_description=f"set_failed_if_not_cancelled({task_id})", + ) + except Exception: + # TSM still unreachable: treat as failed so the callback still fires. + was_failed = True + if was_failed: + await send_indexing_callback( + callback_url, partition, file_id, "error", metadata, callback_token=callback_token + ) raise + else: + # else, not end-of-try: the file is already COMPLETED here, so any + # error past this point must not be caught above and reported as + # a failed run. + log.info("File indexed successfully.") + await send_indexing_callback( + callback_url, partition, file_id, "success", metadata, callback_token=callback_token + ) + return {"stored_count": row.get("stored_count", 0), "stage": row.get("stage", "")} # The raw upload is purged (when configured) by the enclosing actor, not # here: cleanup must also cover failures that happen *before* this method # runs (catalog/registry init, the SERIALIZING state update). See diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py index 212880185..614925e5f 100644 --- a/openrag/services/workers/indexer_pool.py +++ b/openrag/services/workers/indexer_pool.py @@ -3,6 +3,7 @@ import asyncio import threading import time +import traceback from contextvars import ContextVar from types import SimpleNamespace from typing import Any @@ -12,6 +13,7 @@ from core.models.catalog import CONTENT_CLAIM_TOKEN_METADATA_KEY from core.utils.exceptions import NotFoundError from services.workers.indexer_actor import IndexerWorker, _display_filename, delete_uploaded_file +from services.workers.indexing_callback import send_indexing_callback from services.workers.ray_utils import retry_idempotent_ray_actor_method from openrag.core.config.root import Settings @@ -29,7 +31,14 @@ # and workers on the same protocol generation whenever their remote contract or # cross-process indexing semantics change, so new replicas cannot attach to a # partially compatible actor fleet left by the previous release. -_INDEXER_ACTOR_PROTOCOL_VERSION = "v6" +# v4: process_file gained callback_url/callback_token; a v3 worker rejects them. +# v5 (develop): worker-ref-registration wait + TSM set_state fencing semantics. +# v6 (this branch): merge of v4 + v5. +# v6 (develop, independently): STT-preset-aware registry hydration + the +# _active_indexation_config contextvar — a different contract that happened +# to reuse the same version string on its own branch. +# v7: merge of both v6 lineages — neither alone is compatible with this one. +_INDEXER_ACTOR_PROTOCOL_VERSION = "v7" _INDEXER_POOL_DISPATCHER_ACTOR_NAME = f"IndexerPoolDispatcher-{_INDEXER_ACTOR_PROTOCOL_VERSION}" @@ -187,6 +196,9 @@ def __init__(self, namespace: str = "openrag") -> None: self._last_miss_reload_key: tuple[tuple[str, tuple[str, ...]], ...] | None = None self._registry_lock = asyncio.Lock() self._registry_reload_task: asyncio.Task[None] | None = None + # Held here too: a pre-flight failure below runs before the worker's own + # except block, so it must report its own terminal state and callback. + self._tsm = task_state_manager self._worker = IndexerWorker( pipeline=pipeline, task_state_manager=task_state_manager, @@ -403,30 +415,54 @@ async def process_file( replace: bool = False, indexation_config: dict[str, Any] | None = None, embedder_name: str | None = None, + callback_url: str | None = None, + callback_token: str | None = None, require_existing_partition: bool = False, ) -> dict[str, Any]: content_claim_token = metadata.get(CONTENT_CLAIM_TOKEN_METADATA_KEY) worker_metadata = {key: value for key, value in metadata.items() if key != CONTENT_CLAIM_TOKEN_METADATA_KEY} try: - await self._await_worker_ref_registration(task_id) - await self._ensure_catalog() - from services.workers.parsers.parser_dispatcher import routes_to_openai_audio_loader - - await self._ensure_registry_fresh( - _required_model_endpoint_names( - indexation_config, - embedder_name, - include_selected_stt=routes_to_openai_audio_loader( - self._cfg, - _display_filename(path, metadata), - ), + try: + await self._await_worker_ref_registration(task_id) + await self._ensure_catalog() + from services.workers.parsers.parser_dispatcher import routes_to_openai_audio_loader + + await self._ensure_registry_fresh( + _required_model_endpoint_names( + indexation_config, + embedder_name, + include_selected_stt=routes_to_openai_audio_loader( + self._cfg, + _display_filename(path, metadata), + ), + ) ) - ) - # Resolve the enrichment-stage prompts once for this file (partition - # override → global default → disk seed). Done here, at the job - # boundary, so per-chunk work reuses one resolved string instead of - # hitting the DB per chunk. - resolved_prompts = await self._resolve_ingest_prompts(partition, indexation_config or {}) + # Resolve the enrichment-stage prompts once for this file (partition + # override → global default → disk seed). Done here, at the job + # boundary, so per-chunk work reuses one resolved string instead of + # hitting the DB per chunk. + resolved_prompts = await self._resolve_ingest_prompts(partition, indexation_config or {}) + except Exception: + # Not BaseException: a cancellation here must not notify or be + # reported as failed (same rule as set_failed_if_not_cancelled). + tb = traceback.format_exc() + try: + was_failed = await retry_idempotent_ray_actor_method( + lambda: self._tsm.set_failed_if_not_cancelled.remote(task_id, tb), + task_description=f"set_failed_if_not_cancelled({task_id})", + ) + except Exception: + was_failed = True + if was_failed: + await send_indexing_callback( + callback_url, + partition, + worker_metadata.get("file_id", ""), + "error", + worker_metadata, + callback_token=callback_token, + ) + raise token = self._active_indexation_config.set(indexation_config) try: result = await self._worker.process_file( @@ -439,6 +475,8 @@ async def process_file( replace=replace, indexation_config=indexation_config, embedder_name=embedder_name, + callback_url=callback_url, + callback_token=callback_token, require_existing_partition=require_existing_partition, resolved_prompts=resolved_prompts, ) diff --git a/openrag/services/workers/indexing_callback.py b/openrag/services/workers/indexing_callback.py new file mode 100644 index 000000000..2cbc3e04f --- /dev/null +++ b/openrag/services/workers/indexing_callback.py @@ -0,0 +1,154 @@ +"""Best-effort status callbacks for async indexing outcomes. + +Not a webhook: the cozy-stack target is a permission-checked route that 401s +without a bearer token, hence ``callback_token``. + +Never raises, no retries, one notification per file, 5 s timeout, strict no-op +without a ``callback_url``. + +``metadata`` is echoed back verbatim except for ``UPLOAD_METADATA_SERVER_KEYS`` +(the server-computed fields, on-disk path chief among them, that must not +reach a caller-supplied URL). Any other caller-supplied key, named however the +caller likes, travels through unchanged. +""" + +from __future__ import annotations + +import asyncio +from typing import Any +from urllib.parse import ParseResult, urlparse + +import httpx +from core.config import load_config +from core.utils.consts import UPLOAD_METADATA_SERVER_KEYS +from core.utils.logging import get_logger +from core.utils.url_safety import is_safe_url + +logger = get_logger() + +_CALLBACK_TIMEOUT = 5.0 +_TIMEOUT = httpx.Timeout(_CALLBACK_TIMEOUT) + + +def _allow_private_callback_urls() -> bool: + try: + return bool(load_config().indexing_callback.allow_private_urls) + except Exception: + return False + + +def _url_credentials(callback_url: str, parsed: ParseResult) -> tuple[str, ...]: + """Credential substrings of *callback_url* to scrub from any logged message.""" + secrets: list[str] = [] + if "@" in parsed.netloc: + secrets.append(parsed.netloc.rsplit("@", 1)[0]) + for part in (parsed.password, parsed.username): + if part: + secrets.append(part) + try: + userinfo = httpx.URL(callback_url).userinfo.decode("ascii", errors="ignore") + except Exception: + userinfo = "" + if userinfo: + secrets.append(userinfo) + secrets.extend(part for part in userinfo.split(":", 1) if part) + # Longest first, so redacting "user" can't leave "REDACTED:pass" behind. + return tuple(sorted(set(secrets), key=len, reverse=True)) + + +async def send_indexing_callback( + callback_url: str | None, + partition: str, + file_id: str, + status: str, + metadata: dict[str, Any] | None = None, + callback_token: str | None = None, +) -> None: + """POST the indexing outcome to *callback_url*. + + Body: ``{"partition", "file_id", "status": "success"|"error", + "metadata": }``. + + *callback_token* is sent as ``Authorization: Bearer ``. No-op when + *callback_url* is ``None``. Any network or HTTP error is logged and + swallowed so the caller's outcome is unaffected. + """ + if not callback_url: + return + + try: + parsed = urlparse(callback_url) + netloc = parsed.hostname or "" + if parsed.port: + netloc = f"{netloc}:{parsed.port}" + # Scheme and host only: path/query/userinfo can carry secrets. + safe_url = f"{parsed.scheme}://{netloc}" + url_secrets = _url_credentials(callback_url, parsed) + except Exception as exc: + logger.warning( + "Invalid indexing callback_url; skipping callback", + partition=partition, + file_id=file_id, + status=status, + error=str(exc), + ) + return + + # Re-checked here, not just in the router: a direct caller bypasses that. + if not is_safe_url(callback_url, allow_private_hosts=_allow_private_callback_urls()): + logger.warning( + "Indexing callback_url is not a safe server-side target; skipping callback", + callback_url=safe_url, + partition=partition, + file_id=file_id, + status=status, + ) + return + + metadata = metadata or {} + echoed_metadata = {key: value for key, value in metadata.items() if key not in UPLOAD_METADATA_SERVER_KEYS} + + payload = { + "partition": partition, + "file_id": file_id, + "status": status, + "metadata": echoed_metadata, + } + + # Header only — a query-string token would land in the target's access logs. + headers = {"Authorization": f"Bearer {callback_token}"} if callback_token else {} + + try: + async with asyncio.timeout(_CALLBACK_TIMEOUT), httpx.AsyncClient(timeout=_TIMEOUT) as client: + response = await client.post(callback_url, json=payload, headers=headers) + response.raise_for_status() + except Exception as exc: + error_message = str(exc) + for secret in url_secrets: + error_message = error_message.replace(secret, "REDACTED") + if callback_token: + error_message = error_message.replace(callback_token, "REDACTED") + if len(parsed.path) > 1: + error_message = error_message.replace(parsed.path, "/REDACTED") + try: + encoded_path = httpx.URL(callback_url).raw_path.split(b"?", 1)[0].decode("ascii", errors="ignore") + if encoded_path and encoded_path != parsed.path: + error_message = error_message.replace(encoded_path, "/REDACTED") + except Exception: + pass + if parsed.query: + error_message = error_message.replace(parsed.query, "REDACTED") + try: + encoded_query = httpx.URL(callback_url).query.decode("ascii", errors="ignore") + if encoded_query: + error_message = error_message.replace(encoded_query, "REDACTED") + except Exception: + pass + logger.warning( + "Failed to send indexing callback", + callback_url=safe_url, + partition=partition, + file_id=file_id, + status=status, + error=error_message, + ) diff --git a/openrag/services/workers/task_state.py b/openrag/services/workers/task_state.py index 1443dee9c..cad558c08 100644 --- a/openrag/services/workers/task_state.py +++ b/openrag/services/workers/task_state.py @@ -408,14 +408,20 @@ async def set_error(self, task_id: str, tb_str: str) -> None: @ray.method(concurrency_group="set") async def set_failed_if_not_cancelled(self, task_id: str, tb_str: str) -> bool: - """Atomically set state to FAILED and record the traceback, unless already CANCELLED.""" + """Atomically set state to FAILED and record the traceback, unless already CANCELLED. + + Returns whether the caller should treat this as a failure to report (e.g. send an + error callback) — true even when ``task_id`` is unknown to this TSM, since that is + not a cancellation and the caller still needs to hear about the failure. + """ with self.lock: info = self.tasks.get(task_id) - if info is None or info.state == "CANCELLED": + if info is not None and info.state == "CANCELLED": return False - info.state = "FAILED" - info.error = tb_str - _save_recoverable_task(task_id, info) + if info is not None: + info.state = "FAILED" + info.error = tb_str + _save_recoverable_task(task_id, info) return True @ray.method(concurrency_group="set") diff --git a/pyproject.toml b/pyproject.toml index 64277ae76..9fd87dce4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ dependencies = [ "tenacity>=8.2.0", "aiobreaker>=1.2.0", "mcp>=1.11.0", + "httpx>=0.27.0", ] [dependency-groups] diff --git a/tests/unit/api/routers/admin/test_indexing_callback_validation.py b/tests/unit/api/routers/admin/test_indexing_callback_validation.py new file mode 100644 index 000000000..3b1886800 --- /dev/null +++ b/tests/unit/api/routers/admin/test_indexing_callback_validation.py @@ -0,0 +1,39 @@ +"""The upload routes must reject an unusable callback_url before queueing.""" + +from types import SimpleNamespace + +import pytest +from api.routers.admin.indexing import _validate_callback_url +from fastapi import HTTPException + + +def _config(*, allow_private: bool = False) -> SimpleNamespace: + return SimpleNamespace(indexing_callback=SimpleNamespace(allow_private_urls=allow_private)) + + +def test_public_https_callback_is_accepted() -> None: + _validate_callback_url("https://cozy.example.com/ai/index/status", _config()) + + +def test_private_callback_url_is_rejected_by_default() -> None: + with pytest.raises(HTTPException) as exc: + _validate_callback_url("http://127.0.0.1:8080/cb", _config()) + + assert exc.value.status_code == 400 + + +def test_private_callback_url_is_accepted_under_the_dev_opt_in() -> None: + _validate_callback_url("http://cozy.localhost:8080/cb", _config(allow_private=True)) + + +def test_no_callback_url_is_a_noop() -> None: + _validate_callback_url(None, _config()) + + +@pytest.mark.parametrize("callback_url", ["http://[::1/cb", "https://cozy.example.com:abc/cb"]) +def test_malformed_callback_url_is_a_bad_request_not_a_crash(callback_url: str) -> None: + """urlparse's own ``.port`` raises ValueError on these; unguarded that is a 500.""" + with pytest.raises(HTTPException) as exc: + _validate_callback_url(callback_url, _config()) + + assert exc.value.status_code == 400 diff --git a/tests/unit/core/utils/test_url_safety.py b/tests/unit/core/utils/test_url_safety.py index 6440bc70d..d0e607e8f 100644 --- a/tests/unit/core/utils/test_url_safety.py +++ b/tests/unit/core/utils/test_url_safety.py @@ -16,6 +16,10 @@ "http://169.254.169.254/latest/meta-data", # cloud metadata "http://[::1]/x", # IPv6 loopback "http://2130706433/x", # decimal-encoded 127.0.0.1 + "http://0x7f000001/x", # hex-encoded 127.0.0.1 + "http://0177.0.0.1/x", # octal-encoded 127.0.0.1 + "http://127.1/x", # short form of 127.0.0.1 + "http://0xa9fea9fe/x", # hex-encoded 169.254.169.254 "ftp://example.com/x", # non-http scheme "file:///etc/passwd", "not a url", @@ -35,3 +39,64 @@ def test_blocks_unsafe(url): ) def test_allows_public(url): assert is_safe_url(url) is True + + +@pytest.mark.parametrize( + "url", + [ + "http://127.0.0.1:8080/ai/index/status", + "http://localhost:8080/ai/index/status", + "http://cozy.localhost:8080/ai/index/status", + "http://192.168.1.10:8080/x", + "http://[::1]:8080/x", + ], +) +def test_allow_private_hosts_opt_in_permits_private_targets(url): + assert is_safe_url(url, allow_private_hosts=True) is True + + +@pytest.mark.parametrize( + "url", + [ + "ftp://example.com/x", + "file:///etc/passwd", + "not a url", + "https://", + ], +) +def test_allow_private_hosts_still_rejects_bad_schemes_and_hosts(url): + assert is_safe_url(url, allow_private_hosts=True) is False + + +def test_allow_private_hosts_defaults_to_off(): + """Existing callers (web-search fetch, MCP index_url) keep the strict guard.""" + assert is_safe_url("http://169.254.169.254/latest/meta-data") is False + + +def test_blocks_decimal_host_overflowing_into_a_public_ipv6_address(): + """ipaddress.ip_address(int(host)) used to accept any integer up to + 2**128-1, silently building an IPv6Address once the value exceeded + 2**32-1 instead of raising. A 38-digit decimal host crafted to land in a + public IPv6 range therefore passed as "safe". + + inet_aton's own overflow handling is not portable enough to rely on here: + glibc rejects a host this large outright (confirmed against + socket.getaddrinfo — it fails to resolve, so there is no real bypass on + Linux), but at least one other libc (BSD/macOS) silently wraps it mod + 2**32 into 127.0.0.1 instead of raising, which the guard must catch + regardless of the platform it runs on. Hence the dedicated + _bare_numeric_host_value overflow check ahead of inet_aton.""" + host = "42535295865117307932921825931101732865" + assert int(host) > 0xFFFFFFFF + assert is_safe_url(f"http://{host}/x") is False + + +def test_blocks_hex_host_overflowing_past_32_bits(): + """Same overflow class as the decimal case above, hex-spelled.""" + assert is_safe_url("http://0xfffffffff/x") is False + + +def test_decimal_ipv4_form_still_blocked_without_the_removed_branch(): + """The dedicated ip_address(int(host)) branch is gone; inet_aton alone + must still catch the ordinary in-range decimal spelling.""" + assert is_safe_url("http://2130706433/x") is False diff --git a/tests/unit/services/orchestrators/test_indexing_service.py b/tests/unit/services/orchestrators/test_indexing_service.py index 27da90814..94d19c88f 100644 --- a/tests/unit/services/orchestrators/test_indexing_service.py +++ b/tests/unit/services/orchestrators/test_indexing_service.py @@ -56,6 +56,8 @@ async def dispatch_indexing( replace, indexation_config=None, embedder_name=None, + callback_url=None, + callback_token=None, require_existing_partition=False, allow_legacy_require_existing_partition_retry=False, ): @@ -69,6 +71,8 @@ async def dispatch_indexing( "replace": replace, "indexation_config": indexation_config, "embedder_name": embedder_name, + "callback_url": callback_url, + "callback_token": callback_token, "require_existing_partition": require_existing_partition, "allow_legacy_require_existing_partition_retry": allow_legacy_require_existing_partition_retry, } @@ -774,6 +778,38 @@ async def test_copy_file_drops_protected_keys(): assert md["partition"] == "p2" +def test_build_metadata_only_adds_keys_in_upload_metadata_server_keys(tmp_path): + """The indexing-status callback echoes this dict's caller-supplied fields to + a caller-chosen URL, excluding exactly UPLOAD_METADATA_SERVER_KEYS. If this + method ever starts injecting a new server-computed key without adding it to + that constant too, the new key leaks into every callback going forward — + this test is the safety net for that drift.""" + from core.utils.consts import UPLOAD_METADATA_SERVER_KEYS + + file_path = tmp_path / "doc.pdf" + file_path.write_bytes(b"content") + + svc = _service() + caller_metadata = {"doc_rev": "3-abc", "app_tag": "x"} + full = svc._build_metadata( + metadata=caller_metadata, + file_path=str(file_path), + file_id="file-1", + sanitized_filename="doc.pdf", + original_filename="Original.pdf", + content_sha256="9f86d081", + ) + + injected_keys = set(full) - set(caller_metadata) + assert injected_keys <= UPLOAD_METADATA_SERVER_KEYS, ( + f"_build_metadata injected {injected_keys - UPLOAD_METADATA_SERVER_KEYS}, " + "not covered by UPLOAD_METADATA_SERVER_KEYS — the indexing-status " + "callback would now leak it to a caller-supplied URL" + ) + assert full["doc_rev"] == "3-abc" + assert full["app_tag"] == "x" + + class _ExplodingPresetService: """PresetService whose revision probe fails the way a DB blip does. diff --git a/tests/unit/services/websearch/test_content_fetcher.py b/tests/unit/services/websearch/test_content_fetcher.py index c5f2c730e..2cf64598d 100644 --- a/tests/unit/services/websearch/test_content_fetcher.py +++ b/tests/unit/services/websearch/test_content_fetcher.py @@ -5,7 +5,7 @@ import httpx import pytest from services.websearch.base import WebResult -from services.websearch.content_fetcher import ContentFetcher, _is_safe_url +from services.websearch.content_fetcher import ContentFetcher @pytest.fixture @@ -17,61 +17,11 @@ def _make_result(url="https://example.com", snippet="short snippet"): return WebResult(title="Test", url=url, snippet=snippet) -# --------------------------------------------------------------------------- -# _is_safe_url — unit tests (no network, no client) -# --------------------------------------------------------------------------- - - -class TestIsSafeUrl: - @pytest.mark.parametrize( - "url", - [ - "http://localhost/secret", - "http://127.0.0.1/admin", - "http://127.0.0.42/x", - "http://[::1]/admin", - "http://10.0.0.1/internal", - "http://192.168.1.1/router", - "http://169.254.169.254/metadata", # AWS/cloud metadata service - "http://0.0.0.0/x", - "http://100.64.0.1/cgnat", # RFC 6598 shared address space - "http://198.18.0.1/bench", # Benchmarking range - ], - ) - def test_blocks_private_and_reserved_addresses(self, url): - assert _is_safe_url(url) is False - - def test_blocks_decimal_encoded_loopback(self): - """2130706433 is 127.0.0.1 in decimal integer form.""" - assert _is_safe_url("http://2130706433/secret") is False - - def test_blocks_decimal_encoded_private(self): - """167772161 is 10.0.0.1 in decimal integer form.""" - assert _is_safe_url("http://167772161/secret") is False - - @pytest.mark.parametrize( - "url", - [ - "file:///etc/passwd", - "ftp://example.com/x", - "data:text/html,

hi

", - "javascript:alert(1)", - ], - ) - def test_blocks_non_http_schemes(self, url): - assert _is_safe_url(url) is False - - def test_allows_public_ip(self): - # 93.184.216.34 is example.com — globally routable - assert _is_safe_url("http://93.184.216.34/page") is True - - def test_allows_regular_hostname(self): - assert _is_safe_url("https://example.com/page") is True - - # --------------------------------------------------------------------------- # _fetch_single — integration tests with mock transport # --------------------------------------------------------------------------- +# SSRF literal-host checks (_is_safe_url) now delegate to +# core.utils.url_safety.is_safe_url; see tests/unit/core/utils/test_url_safety.py. class TestFetchSingleURL: diff --git a/tests/unit/services/workers/test_dispatcher.py b/tests/unit/services/workers/test_dispatcher.py index 2653ea6aa..bbd5e49b1 100644 --- a/tests/unit/services/workers/test_dispatcher.py +++ b/tests/unit/services/workers/test_dispatcher.py @@ -260,6 +260,8 @@ async def test_dispatch_indexing_relies_on_pool_worker_registration() -> None: replace=True, indexation_config={"parsing_strategy": "pymupdf"}, embedder_name="embed-fast", + callback_url="https://cozy.example.com/callback", + callback_token="jwt-token", require_existing_partition=True, ) @@ -292,6 +294,8 @@ async def test_dispatch_indexing_relies_on_pool_worker_registration() -> None: replace=True, indexation_config={"parsing_strategy": "pymupdf"}, embedder_name="embed-fast", + callback_url="https://cozy.example.com/callback", + callback_token="jwt-token", require_existing_partition=True, ) tsm.set_object_ref.remote.assert_not_called() diff --git a/tests/unit/services/workers/test_indexer_pool.py b/tests/unit/services/workers/test_indexer_pool.py index 3ec14f5c0..7be85681b 100644 --- a/tests/unit/services/workers/test_indexer_pool.py +++ b/tests/unit/services/workers/test_indexer_pool.py @@ -157,7 +157,7 @@ def fake_options(**kwargs): opts = options_calls[0] # A protocol-specific name prevents a rolling deployment from attaching to # a detached actor that still runs the previous claim implementation. - assert opts["name"] == "IndexerPoolDispatcher-v6" + assert opts["name"] == "IndexerPoolDispatcher-v7" assert opts["namespace"] == "openrag" assert opts["get_if_exists"] is True assert opts["lifetime"] == "detached" @@ -195,9 +195,9 @@ def fake_options(**kwargs): # One detached worker actor per pool_size slot, each capped at max_tasks_per_worker. assert len(pool._workers) == 3 assert {c["name"] for c in calls} == { - "IndexerWorker-v6-0", - "IndexerWorker-v6-1", - "IndexerWorker-v6-2", + "IndexerWorker-v7-0", + "IndexerWorker-v7-1", + "IndexerWorker-v7-2", } for c in calls: assert c["lifetime"] == "detached" @@ -1192,7 +1192,7 @@ async def test_pool_drain_rejects_new_work_and_reports_accepted_work(monkeypatch await pool.submit(task_id="accepted-before-drain") assert await pool.begin_drain() == { - "protocol_version": "v6", + "protocol_version": "v7", "accepting_tasks": False, "inflight_jobs": 1, "worker_names": ["test-worker-0"], @@ -1225,7 +1225,7 @@ async def test_pool_drain_rejects_new_work_and_reports_accepted_work(monkeypatch await _settle_pool_release_tasks(pool, worker.futures[0]) assert await pool.status() == { - "protocol_version": "v6", + "protocol_version": "v7", "accepting_tasks": False, "inflight_jobs": 0, "worker_names": ["test-worker-0"], @@ -1260,7 +1260,7 @@ async def test_pool_abort_drain_restores_acceptance() -> None: await pool.submit(task_id="rejected-while-draining") assert await pool.abort_drain() == { - "protocol_version": "v6", + "protocol_version": "v7", "accepting_tasks": True, "inflight_jobs": 0, "worker_names": ["test-worker-0"], @@ -1275,7 +1275,7 @@ async def test_pool_abort_drain_restores_acceptance() -> None: async def test_pool_reports_current_protocol_version() -> None: pool = _bare_pool([_FakeWorker()]) - assert await pool.protocol_version() == "v6" + assert await pool.protocol_version() == "v7" @pytest.mark.asyncio @@ -1704,6 +1704,7 @@ async def _noop(*_a, **_k): ) actor._save_uploaded_files = save_uploaded_files actor._logger = SimpleNamespace(debug=lambda *a, **k: None, warning=lambda *a, **k: None) + actor._tsm = SimpleNamespace(set_failed_if_not_cancelled=SimpleNamespace(remote=lambda *a: "ref")) actor._active_indexation_config = ContextVar("test_active_indexation_config", default=None) # These build the actor with __new__, so __init__ never runs. Captioning is # enabled by default, so ingest now resolves its prompt even for a config @@ -2101,3 +2102,132 @@ async def _boom(*_a, **_k): await actor.process_file(task_id="t", path=str(path), metadata={"file_id": "f"}, partition="p") assert path.exists() + + +@pytest.mark.asyncio +async def test_preflight_failure_sends_the_error_callback_and_sets_failed(tmp_path, monkeypatch) -> None: + """IndexerWorker.process_file owns the success-path callback and is never + reached here, but the pre-flight path must still report a terminal state — + otherwise the callback says "error" while GET task-status shows QUEUED + forever, and a client trusting either source disagrees with the other.""" + import services.workers.indexer_pool as module + + path = tmp_path / "doc.txt" + path.write_bytes(b"x") + worker = _RecordingWorker() + actor = _bare_worker_actor(save_uploaded_files=True, worker=worker) + + async def _boom(*_a, **_k): + raise RuntimeError("postgres down") + + actor._ensure_catalog = _boom + actor._await_worker_ref_registration = AsyncMock(return_value=None) + callback = AsyncMock() + monkeypatch.setattr(module, "send_indexing_callback", callback) + mark_failed = AsyncMock(return_value=True) + monkeypatch.setattr(module, "retry_idempotent_ray_actor_method", mark_failed) + + metadata = {"file_id": "f", "doc_rev": "3-abc"} + with pytest.raises(RuntimeError, match="postgres down"): + await actor.process_file( + task_id="t", + path=str(path), + metadata=metadata, + partition="p", + callback_url="https://cozy.example.com/ai/index/status", + callback_token="jwt", + ) + + assert worker.calls == 0 + mark_failed.assert_awaited_once() + assert mark_failed.await_args.kwargs["task_description"] == "set_failed_if_not_cancelled(t)" + callback.assert_awaited_once_with( + "https://cozy.example.com/ai/index/status", "p", "f", "error", metadata, callback_token="jwt" + ) + + +@pytest.mark.asyncio +async def test_preflight_failure_without_callback_url_still_sets_failed(tmp_path, monkeypatch) -> None: + import services.workers.indexer_pool as module + + path = tmp_path / "doc.txt" + path.write_bytes(b"x") + actor = _bare_worker_actor(save_uploaded_files=True, worker=_RecordingWorker()) + + async def _boom(*_a, **_k): + raise RuntimeError("postgres down") + + actor._ensure_catalog = _boom + actor._await_worker_ref_registration = AsyncMock(return_value=None) + callback = AsyncMock() + monkeypatch.setattr(module, "send_indexing_callback", callback) + mark_failed = AsyncMock(return_value=True) + monkeypatch.setattr(module, "retry_idempotent_ray_actor_method", mark_failed) + + with pytest.raises(RuntimeError, match="postgres down"): + await actor.process_file(task_id="t", path=str(path), metadata={"file_id": "f"}, partition="p") + + mark_failed.assert_awaited_once() + callback.assert_awaited_once() + assert callback.await_args[0][0] is None + + +@pytest.mark.asyncio +async def test_preflight_failure_still_notifies_when_the_tsm_is_unreachable(tmp_path, monkeypatch) -> None: + """If the TSM is down, set_failed_if_not_cancelled can't run either, but a + client waiting on the callback must not be left with neither signal.""" + import services.workers.indexer_pool as module + + path = tmp_path / "doc.txt" + path.write_bytes(b"x") + actor = _bare_worker_actor(save_uploaded_files=True, worker=_RecordingWorker()) + + async def _boom(*_a, **_k): + raise RuntimeError("postgres down") + + actor._ensure_catalog = _boom + actor._await_worker_ref_registration = AsyncMock(return_value=None) + callback = AsyncMock() + monkeypatch.setattr(module, "send_indexing_callback", callback) + monkeypatch.setattr( + module, "retry_idempotent_ray_actor_method", AsyncMock(side_effect=RuntimeError("tsm unreachable")) + ) + + with pytest.raises(RuntimeError, match="postgres down"): + await actor.process_file( + task_id="t", + path=str(path), + metadata={"file_id": "f"}, + partition="p", + callback_url="https://cozy.example.com/ai/index/status", + ) + + callback.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_cancelled_preflight_sends_no_callback(tmp_path, monkeypatch) -> None: + """A cancelled task notifies nothing — the worker's gate says the same.""" + import services.workers.indexer_pool as module + + path = tmp_path / "doc.txt" + path.write_bytes(b"x") + actor = _bare_worker_actor(save_uploaded_files=True, worker=_RecordingWorker()) + + async def _cancelled(*_a, **_k): + raise asyncio.CancelledError() + + actor._ensure_catalog = _cancelled + callback = AsyncMock() + monkeypatch.setattr(module, "send_indexing_callback", callback) + + with pytest.raises(asyncio.CancelledError): + await actor.process_file( + task_id="t", + path=str(path), + metadata={"file_id": "f"}, + partition="p", + callback_url="https://cozy.example.com/ai/index/status", + ) + + callback.assert_not_awaited() diff --git a/tests/unit/services/workers/test_indexer_worker.py b/tests/unit/services/workers/test_indexer_worker.py index 88f7f8bd7..bae2cf102 100644 --- a/tests/unit/services/workers/test_indexer_worker.py +++ b/tests/unit/services/workers/test_indexer_worker.py @@ -633,6 +633,231 @@ async def test_process_file_stores_indexation_config_snapshot_on_replace(tmp_pat assert repo.update_calls[0]["indexation_config"] == indexation_config +@pytest.mark.asyncio +async def test_process_file_success_sends_callback_with_status_and_metadata( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + processed = ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="content")]) + chunks = [Chunk(id="c1", text="content", partition="p")] + worker = IndexerWorker(pipeline=_make_pipeline(processed, chunks), task_state_manager=_fake_tsm()) + + callback_mock = AsyncMock() + monkeypatch.setattr("services.workers.indexer_actor.send_indexing_callback", callback_mock) + + metadata = {"file_id": "f1", "doc_rev": "abc123", "datetime": "2026-01-01T00:00:00Z", "doctype": "text"} + await worker.process_file( + task_id="t-cb", + path=str(path), + metadata=metadata, + partition="p", + callback_url="https://cozy.example.com/callback", + ) + + callback_mock.assert_awaited_once_with( + "https://cozy.example.com/callback", "p", "f1", "success", metadata, callback_token=None + ) + + +@pytest.mark.asyncio +async def test_a_broken_success_callback_does_not_flip_a_completed_task_to_failed( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """send_indexing_callback is documented to never raise for Exception, but if it + ever did (or the process is cancelled mid-await), the file is already COMPLETED + at that point — the except block above must not reinterpret that as a failed + indexing run, flip the state to FAILED, or fire a spurious "error" callback + right after "success" was attempted.""" + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + processed = ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="content")]) + chunks = [Chunk(id="c1", text="content", partition="p")] + tsm = _fake_tsm() + worker = IndexerWorker(pipeline=_make_pipeline(processed, chunks), task_state_manager=tsm) + + callback_mock = AsyncMock(side_effect=RuntimeError("boom")) + monkeypatch.setattr("services.workers.indexer_actor.send_indexing_callback", callback_mock) + + with pytest.raises(RuntimeError, match="boom"): + await worker.process_file( + task_id="t-cb-broken", + path=str(path), + metadata={"file_id": "f1"}, + partition="p", + callback_url="https://cozy.example.com/callback", + ) + + # set_state is also called for SERIALIZING, hence the call-list filter. + completed_calls = [ + call for call in tsm.set_state.remote.call_args_list if call.args == ("t-cb-broken", "COMPLETED") + ] + assert len(completed_calls) == 1 + tsm.set_failed_if_not_cancelled.remote.assert_not_awaited() + # Only the one (failing) "success" attempt — no follow-up "error" callback. + callback_mock.assert_awaited_once() + assert callback_mock.await_args[0][3] == "success" + + +@pytest.mark.asyncio +async def test_process_file_failure_sends_error_callback(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + path = tmp_path / "bad.txt" + path.write_bytes(b"x") + + class BrokenParser: + async def parse(self, document: Document) -> ProcessedDocument: + raise RuntimeError("parser exploded") + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + pipeline = build_indexing_pipeline( + parser=BrokenParser(), + chunker=FakeChunker([]), + embedder=FakeEmbedder(), + vector_store=FakeVectorStore(), + ) + worker = IndexerWorker(pipeline=pipeline, task_state_manager=_fake_tsm()) + + callback_mock = AsyncMock() + monkeypatch.setattr("services.workers.indexer_actor.send_indexing_callback", callback_mock) + + metadata = {"file_id": "f1", "doc_rev": "abc123"} + with pytest.raises(RuntimeError, match="parser exploded"): + await worker.process_file( + task_id="t-cb-fail", + path=str(path), + metadata=metadata, + partition="p", + callback_url="https://cozy.example.com/callback", + ) + + callback_mock.assert_awaited_once_with( + "https://cozy.example.com/callback", "p", "f1", "error", metadata, callback_token=None + ) + + +@pytest.mark.asyncio +async def test_process_file_still_reports_the_original_failure_when_the_tsm_is_unreachable( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """set_failed_if_not_cancelled itself raising (TSM down at report time — the + likely cause of the original failure too) must not replace the real error + with a TSM-unavailability one, and must not cost the client its callback.""" + path = tmp_path / "bad.txt" + path.write_bytes(b"x") + + class BrokenParser: + async def parse(self, document: Document) -> ProcessedDocument: + raise RuntimeError("parser exploded") + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + pipeline = build_indexing_pipeline( + parser=BrokenParser(), + chunker=FakeChunker([]), + embedder=FakeEmbedder(), + vector_store=FakeVectorStore(), + ) + worker = IndexerWorker(pipeline=pipeline, task_state_manager=_fake_tsm()) + + callback_mock = AsyncMock() + monkeypatch.setattr("services.workers.indexer_actor.send_indexing_callback", callback_mock) + + async def _serializing_ok_report_fails(submit, task_description: str = "") -> None: + if "SERIALIZING" in task_description: + return None + raise RuntimeError("tsm unreachable") + + monkeypatch.setattr( + "services.workers.indexer_actor.retry_idempotent_ray_actor_method", + _serializing_ok_report_fails, + ) + + metadata = {"file_id": "f1", "doc_rev": "abc123"} + with pytest.raises(RuntimeError, match="parser exploded"): + await worker.process_file( + task_id="t-cb-fail", + path=str(path), + metadata=metadata, + partition="p", + callback_url="https://cozy.example.com/callback", + ) + + callback_mock.assert_awaited_once_with( + "https://cozy.example.com/callback", "p", "f1", "error", metadata, callback_token=None + ) + + +@pytest.mark.asyncio +async def test_process_file_success_forwards_callback_token(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """The uploader's token must reach the sender, or the status is lost to a 401.""" + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + processed = ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="content")]) + chunks = [Chunk(id="c1", text="content", partition="p")] + worker = IndexerWorker(pipeline=_make_pipeline(processed, chunks), task_state_manager=_fake_tsm()) + + callback_mock = AsyncMock() + monkeypatch.setattr("services.workers.indexer_actor.send_indexing_callback", callback_mock) + + metadata = {"file_id": "f1", "doc_rev": "abc123"} + await worker.process_file( + task_id="t-cb-token", + path=str(path), + metadata=metadata, + partition="p", + callback_url="https://cozy.example.com/ai/index/status", + callback_token="jwt-token", + ) + + callback_mock.assert_awaited_once_with( + "https://cozy.example.com/ai/index/status", "p", "f1", "success", metadata, callback_token="jwt-token" + ) + # A credential, not payload: never in the metadata echoed back. + assert "callback_token" not in metadata + + +@pytest.mark.asyncio +async def test_process_file_failure_forwards_callback_token(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + path = tmp_path / "bad.txt" + path.write_bytes(b"x") + + class BrokenParser: + async def parse(self, document: Document) -> ProcessedDocument: + raise RuntimeError("parser exploded") + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + pipeline = build_indexing_pipeline( + parser=BrokenParser(), + chunker=FakeChunker([]), + embedder=FakeEmbedder(), + vector_store=FakeVectorStore(), + ) + worker = IndexerWorker(pipeline=pipeline, task_state_manager=_fake_tsm()) + + callback_mock = AsyncMock() + monkeypatch.setattr("services.workers.indexer_actor.send_indexing_callback", callback_mock) + + metadata = {"file_id": "f1"} + with pytest.raises(RuntimeError, match="parser exploded"): + await worker.process_file( + task_id="t-cb-token-fail", + path=str(path), + metadata=metadata, + partition="p", + callback_url="https://cozy.example.com/ai/index/status", + callback_token="jwt-token", + ) + + callback_mock.assert_awaited_once_with( + "https://cozy.example.com/ai/index/status", "p", "f1", "error", metadata, callback_token="jwt-token" + ) + + @pytest.mark.asyncio async def test_process_file_catalog_failure_sets_failed_state(tmp_path: Path) -> None: path = tmp_path / "doc.txt" @@ -864,3 +1089,37 @@ async def run(self, row: dict[str, Any]) -> dict[str, Any]: assert len(document_repo.add_calls) == 1 assert vector_store.deleted_filters == [] tsm.set_failed_if_not_cancelled.remote.assert_called_once() + + +@pytest.mark.asyncio +async def test_serializing_state_failure_still_sends_the_error_callback( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The pool's pre-flight handler stops short of this call — nobody else notifies.""" + path = tmp_path / "doc.txt" + path.write_bytes(b"x") + + processed = ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="x")]) + pipeline = _make_pipeline(processed, [Chunk(id="c1", text="x")]) + tsm = _fake_tsm() + # Not retryable, so the helper gives up at once instead of burning its budget. + tsm.set_state.remote = AsyncMock(side_effect=RuntimeError("task state manager is gone")) + worker = IndexerWorker(pipeline=pipeline, task_state_manager=tsm) + + callback_mock = AsyncMock() + monkeypatch.setattr("services.workers.indexer_actor.send_indexing_callback", callback_mock) + + metadata = {"file_id": "f1"} + with pytest.raises(RuntimeError, match="task state manager is gone"): + await worker.process_file( + task_id="t-serializing", + path=str(path), + metadata=metadata, + partition="p", + callback_url="https://cozy.example.com/ai/index/status", + callback_token="jwt", + ) + + callback_mock.assert_awaited_once_with( + "https://cozy.example.com/ai/index/status", "p", "f1", "error", metadata, callback_token="jwt" + ) diff --git a/tests/unit/services/workers/test_indexing_callback.py b/tests/unit/services/workers/test_indexing_callback.py new file mode 100644 index 000000000..94d23bd4e --- /dev/null +++ b/tests/unit/services/workers/test_indexing_callback.py @@ -0,0 +1,507 @@ +from __future__ import annotations + +import json +from unittest import mock + +import httpx +import pytest +from services.workers.indexing_callback import send_indexing_callback + + +def _patch_async_client(monkeypatch: pytest.MonkeyPatch, handler) -> None: + real_async_client = httpx.AsyncClient + + def factory(*args, **kwargs): + return real_async_client(transport=httpx.MockTransport(handler)) + + monkeypatch.setattr("services.workers.indexing_callback.httpx.AsyncClient", factory) + + +@pytest.fixture +def captured_body(monkeypatch: pytest.MonkeyPatch) -> dict: + """Captures the POST body in ``captured["body"]``; the transport answers 200.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response(200, json={"ok": True}) + + _patch_async_client(monkeypatch, handler) + return captured + + +@pytest.fixture +def captured_request(monkeypatch: pytest.MonkeyPatch) -> dict: + """Like ``captured_body``, but keeps the request so headers can be asserted.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["request"] = request + return httpx.Response(200, json={"ok": True}) + + _patch_async_client(monkeypatch, handler) + return captured + + +@pytest.mark.asyncio +async def test_success_callback_echoes_caller_metadata_verbatim(captured_body: dict) -> None: + """Pass-through, not a fixed schema: whatever revision key a caller uses + (cozy-stack's is doc_rev) travels through opaque and unchanged — this + module has no opinion on its name or shape.""" + doc_rev = "3-a1b2C3d4E5f6/g7+h8==" + + await send_indexing_callback( + "https://cozy.example.com/rag/callback", + "alice.mycozy.cloud", + "file-123", + "success", + {"doc_rev": doc_rev, "datetime": "2026-01-01T00:00:00Z", "doctype": "io.cozy.files"}, + ) + + body = captured_body["body"] + assert body["partition"] == "alice.mycozy.cloud" + assert body["file_id"] == "file-123" + assert body["status"] == "success" + assert body["metadata"] == { + "doc_rev": doc_rev, + "datetime": "2026-01-01T00:00:00Z", + "doctype": "io.cozy.files", + } + + +@pytest.mark.asyncio +async def test_error_callback_uses_error_status_and_echoes_metadata(captured_body: dict) -> None: + await send_indexing_callback( + "https://cozy.example.com/rag/callback", + "p", + "f1", + "error", + {"doc_rev": "deadbeef"}, + ) + + body = captured_body["body"] + assert body["status"] == "error" + assert body["metadata"] == {"doc_rev": "deadbeef"} + + +@pytest.mark.asyncio +async def test_server_computed_keys_are_excluded_from_the_echo(captured_body: dict) -> None: + """The upload metadata also holds source/content_sha256/file_size/file_id/ + filename/original_filename by the time it reaches this function — none of + those may reach a caller-supplied URL. Everything else the caller sent, + including fields this module has never heard of, passes through.""" + await send_indexing_callback( + "https://cozy.example.com/rag/callback", + "p", + "f1", + "success", + { + "doc_rev": "3-abc", + "md5sum": "d41d8cd98f00b204", + "app_specific_tag": "keep-me", + "source": "/var/data/tenant42/uploads/xyz.pdf", + "content_sha256": "9f86d081", + "file_size": "42.00 B", + "file_id": "f1", + "filename": "sanitized.pdf", + "original_filename": "Original Name.pdf", + }, + ) + + assert captured_body["body"]["metadata"] == { + "doc_rev": "3-abc", + "md5sum": "d41d8cd98f00b204", + "app_specific_tag": "keep-me", + } + + +@pytest.mark.asyncio +async def test_no_callback_url_is_a_strict_noop(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: # pragma: no cover - must never be called + raise AssertionError("no HTTP call should be made when callback_url is None") + + _patch_async_client(monkeypatch, handler) + + await send_indexing_callback(None, "p", "f1", "success", {"doc_rev": "abc"}) + + +@pytest.mark.asyncio +async def test_unsafe_callback_url_is_skipped_not_sent(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: # pragma: no cover - must never be called + raise AssertionError("no HTTP call should be made to an unsafe callback_url") + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback("http://127.0.0.1:9999/x", "p", "f1", "success", {"doc_rev": "abc"}) + mock_logger.warning.assert_called_once() + + +@pytest.mark.asyncio +async def test_callback_url_credentials_are_not_logged(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: # pragma: no cover - must never be called + raise AssertionError("no HTTP call should be made to an unsafe callback_url") + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback( + "https://leakme:hunter2@127.0.0.1/callback", "p", "f1", "success", {"doc_rev": "abc"} + ) + mock_logger.warning.assert_called_once() + logged_url = mock_logger.warning.call_args[1]["callback_url"] + assert "leakme" not in logged_url + assert "hunter2" not in logged_url + + +@pytest.mark.asyncio +async def test_callback_error_with_bad_port_and_query_does_not_crash(monkeypatch: pytest.MonkeyPatch) -> None: + """urlparse accepts a non-numeric port; httpx.URL rejects it, inside the handler.""" + + def handler(request: httpx.Request) -> httpx.Response: # pragma: no cover - never reached + raise AssertionError("should fail before an HTTP call is attempted") + + _patch_async_client(monkeypatch, handler) + + await send_indexing_callback( + "https://cozy.example.com:abc/callback?token=SECRET", "p", "f1", "success", {"doc_rev": "abc"} + ) + + +@pytest.mark.asyncio +async def test_callback_error_with_non_printable_char_does_not_crash(monkeypatch: pytest.MonkeyPatch) -> None: + """urlparse tolerates a non-printable character; httpx.URL rejects it mid-flight.""" + + def handler(request: httpx.Request) -> httpx.Response: # pragma: no cover - httpx rejects before dispatch + raise AssertionError("should fail before a response is produced") + + _patch_async_client(monkeypatch, handler) + + await send_indexing_callback( + "https://cozy.example.com/callback?token=x\x00y", "p", "f1", "success", {"doc_rev": "abc"} + ) + + +@pytest.mark.asyncio +async def test_callback_failure_is_swallowed(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500, json={"error": "boom"}) + + _patch_async_client(monkeypatch, handler) + + # Must not raise even though the remote endpoint returns an error. + await send_indexing_callback( + "https://cozy.example.com/rag/callback", + "p", + "f1", + "success", + {"doc_rev": "abc"}, + ) + + +@pytest.mark.asyncio +async def test_callback_error_redacts_query_string_from_logged_exception(monkeypatch: pytest.MonkeyPatch) -> None: + """httpx's HTTPStatusError message embeds the full request URL.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500, json={"error": "boom"}) + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback( + "https://cozy.example.com/rag/callback?token=SECRETVALUE", + "p", + "f1", + "success", + {"doc_rev": "abc"}, + ) + + mock_logger.warning.assert_called_once() + call_args = mock_logger.warning.call_args + assert call_args[1]["callback_url"] == "https://cozy.example.com" + assert "SECRETVALUE" not in call_args[1]["error"] + + +@pytest.mark.asyncio +async def test_malformed_callback_url_is_swallowed_not_raised(monkeypatch: pytest.MonkeyPatch) -> None: + """urlparse raises on some inputs, e.g. an unterminated IPv6 literal.""" + + def handler(request: httpx.Request) -> httpx.Response: # pragma: no cover - must never be called + raise AssertionError("no HTTP call should be made for a malformed callback_url") + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback("http://[::1", "p", "f1", "success", {"doc_rev": "abc"}) + mock_logger.warning.assert_called_once() + assert "callback_url" in mock_logger.warning.call_args[0][0].lower() + + +@pytest.mark.asyncio +async def test_callback_error_redacts_percent_encoded_query_from_logged_exception( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """httpx percent-encodes the query, so the message embeds the encoded form.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500, json={"error": "boom"}) + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback( + "https://cozy.example.com/rag/callback?token=SECRET VALUE!", + "p", + "f1", + "success", + {"doc_rev": "abc"}, + ) + + mock_logger.warning.assert_called_once() + error_message = mock_logger.warning.call_args[1]["error"] + assert "SECRET" not in error_message + + +@pytest.mark.asyncio +async def test_missing_metadata_is_an_empty_dict_not_null_placeholders(captured_body: dict) -> None: + """Pass-through echoes whatever the caller sent, or nothing — it no longer + synthesizes a fixed set of ``None``-valued keys for a schema that doesn't + exist any more.""" + await send_indexing_callback("https://cozy.example.com/rag/callback", "p", "f1", "success", None) + + body = captured_body["body"] + assert body["metadata"] == {} + + +@pytest.mark.asyncio +async def test_absent_or_present_caller_fields_never_log_a_warning(captured_body: dict) -> None: + """This module has no opinion on any caller field's presence or name — + whether the receiver treats a missing one as an error (e.g. cozy-stack + ordering callbacks on a revision) is entirely the receiver's call.""" + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback( + "https://cozy.example.com/rag/callback", "p", "f1", "success", {"doctype": "io.cozy.files"} + ) + await send_indexing_callback("https://cozy.example.com/rag/callback", "p", "f1", "success", {"doc_rev": "abc"}) + + mock_logger.warning.assert_not_called() + + +@pytest.mark.asyncio +async def test_callback_token_is_sent_as_bearer_header(captured_request: dict) -> None: + await send_indexing_callback( + "https://cozy.example.com/ai/index/status", + "alice.mycozy.cloud", + "file-123", + "success", + {"doc_rev": "abc"}, + callback_token="jwt-abc.def.ghi", + ) + + request = captured_request["request"] + assert request.headers["Authorization"] == "Bearer jwt-abc.def.ghi" + + +@pytest.mark.asyncio +async def test_callback_token_never_reaches_url_or_payload(captured_request: dict) -> None: + await send_indexing_callback( + "https://cozy.example.com/ai/index/status", + "p", + "f1", + "success", + {"doc_rev": "abc"}, + callback_token="s3cret-token", + ) + + request = captured_request["request"] + assert "s3cret-token" not in str(request.url) + assert "s3cret-token" not in request.content.decode() + assert json.loads(request.content).keys() == {"partition", "file_id", "status", "metadata"} + + +@pytest.mark.asyncio +async def test_no_callback_token_sends_no_authorization_header(captured_request: dict) -> None: + """Back-compat: an unauthenticated endpoint sees the request it saw before.""" + await send_indexing_callback("https://cozy.example.com/rag/callback", "p", "f1", "success", {"doc_rev": "abc"}) + + assert "Authorization" not in captured_request["request"].headers + + +@pytest.mark.asyncio +async def test_callback_token_is_redacted_from_logged_errors(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("upstream refused Bearer s3cret-token") + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback( + "https://cozy.example.com/ai/index/status", + "p", + "f1", + "success", + {"doc_rev": "abc"}, + callback_token="s3cret-token", + ) + + mock_logger.warning.assert_called_once() + assert "s3cret-token" not in mock_logger.warning.call_args[1]["error"] + assert "REDACTED" in mock_logger.warning.call_args[1]["error"] + + +@pytest.mark.asyncio +async def test_private_callback_url_is_sent_when_operator_opts_in( + monkeypatch: pytest.MonkeyPatch, captured_request: dict +) -> None: + monkeypatch.setattr("services.workers.indexing_callback._allow_private_callback_urls", lambda: True) + + await send_indexing_callback("http://localhost:8080/ai/index/status", "p", "f1", "success", {"doc_rev": "abc"}) + + assert str(captured_request["request"].url) == "http://localhost:8080/ai/index/status" + + +@pytest.mark.asyncio +async def test_private_callback_url_still_blocked_by_default(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: # pragma: no cover - must never be called + raise AssertionError("no HTTP call should be made to a private callback_url by default") + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback("http://localhost:8080/ai/index/status", "p", "f1", "success", {"doc_rev": "abc"}) + mock_logger.warning.assert_called_once() + + +@pytest.mark.asyncio +async def test_opting_in_to_private_urls_still_rejects_non_http_schemes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("services.workers.indexing_callback._allow_private_callback_urls", lambda: True) + + def handler(request: httpx.Request) -> httpx.Response: # pragma: no cover - must never be called + raise AssertionError("no HTTP call should be made for a non-http(s) scheme") + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback("file:///etc/passwd", "p", "f1", "success", {"doc_rev": "abc"}) + mock_logger.warning.assert_called_once() + + +@pytest.mark.asyncio +async def test_allow_private_callback_urls_defaults_to_false_on_config_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from services.workers import indexing_callback + + def boom() -> None: + raise RuntimeError("config unavailable") + + monkeypatch.setattr(indexing_callback, "load_config", boom) + assert indexing_callback._allow_private_callback_urls() is False + + +@pytest.mark.asyncio +async def test_url_credentials_are_redacted_from_logged_errors(monkeypatch: pytest.MonkeyPatch) -> None: + """httpx echoes the full URL in its exception messages, which we log verbatim.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500) + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback( + "https://leakme:hunter2@cozy.example.com/cb", "p", "f1", "success", {"doc_rev": "abc"} + ) + + mock_logger.warning.assert_called_once() + logged = mock_logger.warning.call_args[1] + assert "hunter2" not in logged["error"] + assert "leakme" not in logged["error"] + assert "hunter2" not in logged["callback_url"] + # The host is still there: redaction must not cost us the diagnostics. + assert "cozy.example.com" in logged["error"] + + +@pytest.mark.asyncio +async def test_url_credentials_with_reserved_chars_are_redacted(monkeypatch: pytest.MonkeyPatch) -> None: + """httpx percent-encodes, so the password appears in a spelling urlparse never produced.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500) + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback( + "https://user:p%40ss word@cozy.example.com/cb", "p", "f1", "success", {"doc_rev": "abc"} + ) + + mock_logger.warning.assert_called_once() + error = mock_logger.warning.call_args[1]["error"] + assert "p%40ss" not in error + assert "p@ss" not in error + + +@pytest.mark.asyncio +async def test_username_redaction_does_not_leave_password_behind(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500) + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback( + "https://bob:bobsecret@cozy.example.com/cb", "p", "f1", "success", {"doc_rev": "abc"} + ) + + error = mock_logger.warning.call_args[1]["error"] + assert "bobsecret" not in error + + +@pytest.mark.asyncio +async def test_a_secret_in_the_callback_path_is_not_logged(monkeypatch: pytest.MonkeyPatch) -> None: + """Webhook secrets conventionally live in the path, not the query.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500) + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback( + "https://cozy.example.com/hooks/T000/B000/SECRETPATH", "p", "f1", "success", {"doc_rev": "abc"} + ) + + mock_logger.warning.assert_called_once() + logged = mock_logger.warning.call_args[1] + assert "SECRETPATH" not in logged["callback_url"] + assert "SECRETPATH" not in logged["error"] + # The host is still there: redaction must not cost us the diagnostics. + assert "cozy.example.com" in logged["error"] + + +@pytest.mark.asyncio +async def test_a_percent_encoded_secret_in_the_callback_path_is_not_logged( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """httpx percent-encodes the path the same way it does the query — a raw, + unencoded substring match misses a secret containing a space or other + character httpx re-encodes before it ever reaches the exception message.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500) + + _patch_async_client(monkeypatch, handler) + + with mock.patch("services.workers.indexing_callback.logger") as mock_logger: + await send_indexing_callback( + "https://cozy.example.com/hooks/SECRET PATH/x", "p", "f1", "success", {"doc_rev": "abc"} + ) + + mock_logger.warning.assert_called_once() + logged = mock_logger.warning.call_args[1] + assert "SECRET" not in logged["error"] + assert "PATH" not in logged["error"] + assert "%20" not in logged["error"] + assert "cozy.example.com" in logged["error"] diff --git a/uv.lock b/uv.lock index e961ca218..cdd077119 100644 --- a/uv.lock +++ b/uv.lock @@ -2783,6 +2783,7 @@ dependencies = [ { name = "faster-whisper" }, { name = "hdbscan" }, { name = "html-to-markdown" }, + { name = "httpx" }, { name = "infinity-client" }, { name = "itsdangerous" }, { name = "langchain-community" }, @@ -2853,6 +2854,7 @@ requires-dist = [ { name = "faster-whisper", specifier = ">=1.1.0" }, { name = "hdbscan", specifier = ">=0.8.40" }, { name = "html-to-markdown", specifier = ">=2.4.0" }, + { name = "httpx", specifier = ">=0.27.0" }, { name = "infinity-client", specifier = ">=0.0.76" }, { name = "itsdangerous", specifier = ">=2.2" }, { name = "langchain-community", specifier = ">=0.3.18" },