diff --git a/openrag/api/schemas/admin/model_endpoint_schemas.py b/openrag/api/schemas/admin/model_endpoint_schemas.py index d98dbd6cf..15789957d 100644 --- a/openrag/api/schemas/admin/model_endpoint_schemas.py +++ b/openrag/api/schemas/admin/model_endpoint_schemas.py @@ -2,6 +2,7 @@ from __future__ import annotations +import re from datetime import datetime from typing import Any, Literal @@ -15,12 +16,43 @@ # ``extra`` — validated here so a typo can't persist a nonsensical value. _LLM_TOKEN_EXTRA_KEYS = (LLM_CONTEXT_SIZE_KEY, LLM_OUTPUT_TOKENS_KEY) +# Allowlist, not a denylist: `name` is a single path segment in every +# single-endpoint route (see `_normalize_name`), and enumerating unsafe values +# one at a time as they're discovered — first `/` (#768), then the RFC 3986 +# dot-segments `.`/`..` — never closes the class. Anchoring both ends on +# alphanumeric rules out `/`, `.`, `..`, and any leading/trailing separator by +# construction, while `.`/`_`/`-` stay available in the middle for realistic +# names like `gpt-4.1` or `jina_v3`. +_NAME_PATTERN = re.compile(r"[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?") +_NAME_MAX_LENGTH = 128 + def _normalize_name(value: str) -> str: - """Trim a user-facing registry name and reject blank values.""" + """Trim a user-facing registry name and reject any value unsafe as a URL path segment. + + ``name`` is embedded as a single path segment in every single-endpoint route + (``GET/PUT/DELETE /model-endpoints/{model_type}/{name}``, ``.../set-default``, + ``.../reveal-api-key``, ``.../validate``). A value outside ``_NAME_PATTERN`` + — a ``/`` (splits across path segments), the exact values ``.``/``..`` + (RFC 3986 dot-segments: browsers and HTTP clients normalize these out of + the URL before the request is even sent, resolving to the collection route + or dropping the ``model_type`` segment entirely), or anything else that + doesn't start/end alphanumeric — would leave the row visible in the list + endpoint but permanently unreachable by get/update/delete/set-default, + surfacing as a spurious "not found". Percent-encoding never helps: ASGI + servers decode ``%2F``/dot-segment escapes before Starlette's router sees + the path. + """ value = value.strip() if not value: raise ValueError("name must be non-empty") + if len(value) > _NAME_MAX_LENGTH: + raise ValueError(f"name must be at most {_NAME_MAX_LENGTH} characters") + if not _NAME_PATTERN.fullmatch(value): + raise ValueError( + "name must start and end with a letter or digit, and contain only " + "letters, digits, '.', '_', or '-' (it is used as a URL path segment)" + ) return value diff --git a/openrag/di/container.py b/openrag/di/container.py index 2510b5638..28a8e8ae7 100644 --- a/openrag/di/container.py +++ b/openrag/di/container.py @@ -437,6 +437,7 @@ def model_endpoint_service(self) -> ModelEndpointService: model_endpoint_repo=self.model_endpoint_repo, config=self._require_settings(), partition_service=self.partition_service, + preset_service=self.preset_service, client_caches={ "embedder": self._embedder_cache, "reranker": self._reranker_cache, diff --git a/openrag/services/orchestrators/model_endpoint_service.py b/openrag/services/orchestrators/model_endpoint_service.py index 058abfcc1..825e111db 100644 --- a/openrag/services/orchestrators/model_endpoint_service.py +++ b/openrag/services/orchestrators/model_endpoint_service.py @@ -114,11 +114,13 @@ def __init__( model_endpoint_repo: ModelEndpointRepository, config: Settings, partition_service: Any = None, + preset_service: Any = None, client_caches: dict[str, dict[str, Any]] | None = None, ) -> None: self._repo = model_endpoint_repo self._config = config self._partition_service = partition_service + self._preset_service = preset_service self._client_caches: dict[str, dict[str, Any]] = client_caches or {} # ------------------------------------------------------------------ @@ -396,6 +398,29 @@ async def update_model_endpoint(self, name: str, model_type: str, **fields: obje Pass ``new_name=`` to rename. After any change the in-memory config is reloaded and the stale cached client instance is evicted so the next request builds a fresh client against the updated config. + + A rename also cascades to every stored reference — ``partitions.embedder`` + / ``partitions.chat_llm`` and endpoint-name fields embedded in + ``pipeline_presets.config`` — inside the repo's own rename transaction + (see ``PgModelEndpointRepository.rename``, #770). Those writes are + invisible until the referencing services reload their in-memory caches, + which is why a rename also refreshes presets then partitions here — + the same order ``PresetService.update_preset`` uses, since partition + resolution reads the presets dict. + + Both reload calls ``await``, so a concurrent request can run between + them — and the DB rename has *already* committed by that point. Without + ``_alias_renamed_name``, a request landing in that window could resolve + a partition/preset that the cascade already repointed at ``new_name`` + against a registry that (until the final ``load_all()`` below) still + only knows ``name`` — a bare ``KeyError``. The alias makes both ``name`` + and ``new_name`` resolve immediately, built from the row this call just + wrote — not whatever the in-memory bucket held before it — so a rename + combined with a field change (e.g. a new ``endpoint``) aliases the + *updated* config, not a stale pre-update one. That also covers a reload + call above raising: the registry stays queryable under both names, + correctly, instead of the update's failure leaving it stuck on a stale + config until process restart. """ existing = await self._repo.get(name, model_type) if existing is None: @@ -423,6 +448,21 @@ async def update_model_endpoint(self, name: str, model_type: str, **fields: obje await self._repo.rename(name, model_type, new_name) effective_name = new_name renamed_from = name + self._alias_renamed_name(model_type, name, new_name, updated or existing) + # A cached *client instance* under either name would otherwise survive + # this alias — the factory checks its cache before consulting the + # config registry, so a stale pre-rename/pre-update client would keep + # serving until the eviction at the end of this method, which a + # reload call below raising would skip entirely. The config alias + # above is already fresh, so evicting now is safe: anything rebuilt + # from either name resolves through the up-to-date config, not stale + # cached state. + self._invalidate_client_cache(model_type, name) + self._invalidate_client_cache(model_type, new_name) + if self._preset_service is not None: + await self._preset_service.load_all() + if self._partition_service is not None: + await self._partition_service.load_partitions() if promote_to_default: # Clears any prior default and sets this row in one transaction, then @@ -540,6 +580,39 @@ async def validate_endpoint( # Internals # ------------------------------------------------------------------ + def _alias_renamed_name(self, model_type: str, old_name: str, new_name: str, row: ModelEndpointRow) -> None: + """Make both ``old_name`` and ``new_name`` resolve to *row* before any reload runs. + + Runs synchronously right after the rename ``await`` returns — no + further ``await`` happens before this executes, so no concurrent + request can observe the DB already renamed while the registry still + only answers to ``old_name``. + + Built from ``row`` — the just-written DB state — rather than copying + whatever the in-memory bucket currently holds under ``old_name``: a + rename can land in the same call as a field update (e.g. a new + ``endpoint`` URL), applied to the DB *before* this runs, so the stale + in-memory entry would alias both names to the pre-update config. If a + reload below then raises, that staleness would never get corrected + by the final ``load_all()`` this call never reaches — the registry + would keep serving the old endpoint under the new (DB-authoritative) + name until process restart. The next full ``load_all()`` (below, or + from any later CRUD call) rebuilds the bucket straight from DB and + drops the ``old_name`` entry on its own. + """ + bucket: dict[str, Any] | None = getattr(self._config.models, model_type, None) + if bucket is None: + return + cfg = ModelEndpointConfig( + endpoint=row.endpoint, + model_name=row.model_name, + batch_size=row.batch_size, + timeout=row.timeout, + extra=row.extra, + ) + bucket[old_name] = cfg + bucket[new_name] = cfg + def _invalidate_client_cache(self, model_type: str, name: str) -> None: """Evict ``name`` from the component-factory cache for ``model_type``.""" cache = self._client_caches.get(model_type) diff --git a/openrag/services/orchestrators/retrieval_service.py b/openrag/services/orchestrators/retrieval_service.py index 0fcbf4862..eea37049a 100644 --- a/openrag/services/orchestrators/retrieval_service.py +++ b/openrag/services/orchestrators/retrieval_service.py @@ -164,6 +164,71 @@ def _require_partition_config(self, partition: str): def _legacy_retriever_value(self, name: str, default: Any) -> Any: return getattr(self._config.retriever, name, default) + def _resolve_reranker(self, reranker_name: str | None, partition: str) -> Reranker | None: + """Effective reranker for one partition's retrieval pipeline. + + Resolution order — mirrors ``QueryService._resolve_llm``: + + 1. The partition's configured ``reranker`` preset, resolved fresh via + the model-endpoint catalog factory so a rename/promotion of that + endpoint takes effect immediately. + 2. The **catalog default** endpoint (``is_default=True``) when the + partition sets no preset, or its preset name has gone stale (the + endpoint was renamed/deleted after assignment — unlike + ``chat_llm``, this field has no create/PATCH-time validation, so a + stale name reaching here is expected, not a bug). + 3. The static reranker built at startup from ``settings.reranker``, + only when no factory is wired (unit tests) or the catalog has no + default reranker endpoint yet. + + The resolved endpoint name is always logged (at debug), including for + the default, so "which reranker ran?" is answerable from the logs. + """ + if self._reranker_factory is None: + return self._legacy_reranker + if reranker_name: + try: + reranker = self._reranker_factory(reranker_name) + except KeyError: + logger.bind(reranker=reranker_name, partition=partition).warning( + "Partition reranker preset not found in the model-endpoint catalog — " + "falling back to the default reranker" + ) + else: + logger.bind(reranker=reranker_name, partition=partition).debug( + "Reranking with the partition's reranker preset" + ) + return reranker + try: + reranker = self._reranker_factory("default") + except KeyError: + pass + else: + logger.bind(reranker=self._default_reranker_name(), partition=partition).debug( + "Reranking with the default reranker preset" + ) + return reranker + logger.bind(partition=partition).debug( + "Reranking with the static default reranker (no catalog default endpoint)" + ) + return self._legacy_reranker + + def _default_reranker_name(self) -> str: + """Real endpoint name behind the catalog reranker ``"default"`` alias, for logging. + + Same identity-lookup trick as ``QueryService._default_llm_name``: the + ``"default"`` alias config object is the *same* object as its real-named + entry, so the name is recovered by identity. Returns ``"default"`` when + it can't be resolved (e.g. the alias isn't populated yet). + """ + rerankers = self._config.models.reranker + default_cfg = rerankers.get("default") + if default_cfg is not None: + for name, cfg in rerankers.items(): + if name != "default" and cfg is default_cfg: + return name + return "default" + def _pipeline_for_partition(self, partition: str) -> tuple[RetrieverPipeline, int | None]: # Callers only ever pass a concrete partition name — the "all" sentinel is # expanded to concrete keys by _pipeline_groups_for_partitions before this @@ -182,12 +247,7 @@ def _pipeline_for_partition(self, partition: str) -> tuple[RetrieverPipeline, in if rtype in {"multiQuery", "hyde"} and self._llm_factory is not None: llm = self._llm_factory(pipeline_cfg.llm or partition_cfg.chat_llm or "default") - reranker = None - if pipeline_cfg.enable_reranker: - if self._reranker_factory is not None: - reranker = self._reranker_factory(pipeline_cfg.reranker or "default") - else: - reranker = self._legacy_reranker + reranker = self._resolve_reranker(pipeline_cfg.reranker, partition) if pipeline_cfg.enable_reranker else None retriever = self._build_retriever( rtype=rtype, diff --git a/openrag/services/persistence/model_endpoint_repo.py b/openrag/services/persistence/model_endpoint_repo.py index 933352b88..57b8164b8 100644 --- a/openrag/services/persistence/model_endpoint_repo.py +++ b/openrag/services/persistence/model_endpoint_repo.py @@ -27,6 +27,20 @@ # transaction; ModelEndpointService.update_model_endpoint routes is_default there. _ALLOWED_UPDATE_FIELDS = frozenset({"endpoint", "model_name", "batch_size", "timeout", "extra"}) +# Endpoint names are referenced by value elsewhere, and nothing updates those +# references when an endpoint is renamed (#770) — so ``rename()`` cascades to +# every known reference in the same transaction as the name change. Direct +# endpoint-name columns on ``partitions``, keyed by the model_type they hold: +_PARTITION_COLUMN_BY_TYPE = {"embedder": "embedder", "llm": "chat_llm"} +# Endpoint-name keys embedded in ``pipeline_presets.config`` (JSONB), by the +# preset_type that carries them — see core/config/retrieval_pipeline.py and +# core/config/indexation_pipeline.py for the field definitions. +_RETRIEVAL_PRESET_KEYS_BY_TYPE = {"llm": ("llm",), "reranker": ("reranker",)} +_INDEXATION_PRESET_KEYS_BY_TYPE = { + "llm": ("contextualization_llm", "metadata_extraction_llm", "topic_tagging_llm"), + "vlm": ("vlm",), +} + class PgModelEndpointRepository(ModelEndpointRepository): """asyncpg-backed implementation of :class:`ModelEndpointRepository`.""" @@ -124,12 +138,80 @@ async def update(self, name: str, model_type: str, **fields: object) -> ModelEnd return self._to_model(rec) if rec else None async def rename(self, name: str, model_type: str, new_name: str) -> None: - await self.pool.execute( - "UPDATE model_endpoints SET name = $3, updated_at = now() WHERE name = $1 AND model_type = $2", - name, - model_type, - new_name, - ) + """Rename an endpoint and cascade the new name to every stored reference. + + Left alone, a rename silently strands every partition or preset that + pointed at the old name — ``partitions.embedder`` / ``partitions.chat_llm``, + and the endpoint-name fields embedded in ``pipeline_presets.config`` + (JSONB) — since nothing else in the schema updates those when the + referenced row's name changes (#770). All writes run in this one + transaction so a partial cascade can never leave the registry and its + referents disagreeing. + + The caller (``ModelEndpointService.update_model_endpoint``) still owns + refreshing the in-memory partition/preset caches afterwards — this + method only makes the DB-side references consistent. + + Raises :class:`NotFoundError` if ``name`` vanished between the + service's existence check and this transaction (a concurrent delete) + — mirroring ``PgPipelinePresetRepository.rename``. Without the + ``RETURNING`` check, a lost race would still run the cascade below, + repointing partitions/presets at a ``new_name`` that was never + actually created. + + Also ``LOCK``s ``partitions`` ``IN SHARE MODE`` before touching + anything — the same lock :meth:`PgPresetRepository.delete` takes, and + the same table :meth:`PgPartitionRepository.update_partition` writes + to *before* its own DB-authoritative ``chat_llm`` re-check. Without + this, a partition PATCH could validate ``name`` against the in-memory + catalog, then block behind this transaction's cascade on the exact + row it's about to write, and resume writing the now-renamed-away + ``name`` straight back once this commits — permanently stranding that + partition the moment the service's temporary alias (see + ``ModelEndpointService._alias_renamed_name``) drops. Locking first, in + the same order both call sites use, makes the two block on each other + instead of interleaving: whichever transaction's ``partitions`` write + commits first is the one the other observes. + """ + async with self.pool.acquire() as conn: + async with conn.transaction(): + await conn.execute("LOCK TABLE partitions IN SHARE MODE") + + renamed = await conn.fetchrow( + "UPDATE model_endpoints SET name = $3, updated_at = now() " + "WHERE name = $1 AND model_type = $2 RETURNING name", + name, + model_type, + new_name, + ) + if renamed is None: + raise NotFoundError(f"Endpoint '{name}' of type '{model_type}' not found.") + + partition_col = _PARTITION_COLUMN_BY_TYPE.get(model_type) + if partition_col: + await conn.execute( + f"UPDATE partitions SET {partition_col} = $2 WHERE {partition_col} = $1", + name, + new_name, + ) + + for preset_type, keys in ( + ("retrieval", _RETRIEVAL_PRESET_KEYS_BY_TYPE.get(model_type, ())), + ("indexation", _INDEXATION_PRESET_KEYS_BY_TYPE.get(model_type, ())), + ): + for key in keys: + await conn.execute( + """ + UPDATE pipeline_presets + SET config = jsonb_set(config, $1::text[], to_jsonb($2::text)), updated_at = now() + WHERE preset_type = $3 AND config->>$4 = $5 + """, + [key], + new_name, + preset_type, + key, + name, + ) async def delete(self, name: str, model_type: str) -> bool: result = await self.pool.execute( diff --git a/openrag/services/persistence/partition_repo.py b/openrag/services/persistence/partition_repo.py index 1d4770c33..1931f2f42 100644 --- a/openrag/services/persistence/partition_repo.py +++ b/openrag/services/persistence/partition_repo.py @@ -35,6 +35,16 @@ "indexation_preset": "indexation", "retrieval_preset": "retrieval", } +# Partition columns that reference a model_endpoints row, mapped to the +# model_type they point at. Only `chat_llm` is assignment-validated today +# (PartitionService._validate_chat_llm_ref checks the in-memory catalog); +# `embedder` carries no such check, so it is deliberately not listed here. +# Assigning chat_llm must be guarded against a concurrent rename the same way +# a preset assignment is guarded against a concurrent preset delete — see +# update_partition and PgModelEndpointRepository.rename. +_MODEL_ENDPOINT_COLUMN_TYPES = { + "chat_llm": "llm", +} _PARTITION_UPDATE_COLUMNS = frozenset( { "description", @@ -56,6 +66,19 @@ def _partition_updates(fields: dict[str, object]) -> dict[str, object]: return {key: value for key, value in fields.items() if key in _PARTITION_UPDATE_COLUMNS} +def _endpoint_refs(updates: dict[str, object]) -> dict[str, str]: + """Model-endpoint-referencing columns in *updates*, mapped to their model_type. + + A ``None`` value clears the (nullable) column rather than pointing it at a + name, so it needs no existence check. + """ + return { + col: model_type + for col, model_type in _MODEL_ENDPOINT_COLUMN_TYPES.items() + if col in updates and updates[col] is not None + } + + class _PartitionOperationGuard: def __init__(self, repo: PgPartitionRepository, conn: asyncpg.Connection) -> None: self._repo = repo @@ -282,24 +305,35 @@ async def update_partition(self, name: str, **fields: object) -> dict | None: """Update a partition's config columns. When the update assigns a preset column (``indexation_preset`` / - ``retrieval_preset``), the write and a DB-authoritative existence check - run in one transaction that touches ``partitions`` before - ``pipeline_presets`` — the same lock order :meth:`PgPresetRepository.delete` - uses (it ``LOCK``s ``partitions`` ``IN SHARE MODE``, then ``DELETE``s the - preset). That makes an assign and a concurrent delete of the same preset + ``retrieval_preset``) or ``chat_llm``, the write and a DB-authoritative + existence check run in one transaction that touches ``partitions`` + before ``pipeline_presets`` / ``model_endpoints`` — the same lock order + :meth:`PgPresetRepository.delete` and :meth:`PgModelEndpointRepository. + rename` use (both ``LOCK`` ``partitions`` ``IN SHARE MODE`` first). That + makes an assign and a concurrent delete/rename of the same reference serialize without deadlocking, so a partition can never end up pointing - at a preset that no longer exists: - - * if this UPDATE commits first, the delete's ``COUNT`` sees the reference - and refuses with 409; - * if the delete commits first, this UPDATE blocks on its ``SHARE`` lock, - then the follow-up ``SELECT`` sees the vanished preset and the - transaction rolls the write back (raising ``PRESET_NOT_FOUND``). + at a preset or model endpoint that no longer exists under that name: + + * if this UPDATE commits first, the delete/rename's own guard against + ``partitions`` (a ``COUNT`` for presets, the ``SHARE`` lock itself for + renames) sees the reference and blocks or refuses accordingly; + * if the delete/rename commits first, this UPDATE blocks on its + ``SHARE``-conflicting write, then the follow-up ``SELECT`` sees the + vanished name and the transaction rolls the write back (raising + ``PRESET_NOT_FOUND`` / ``MODEL_ENDPOINT_NOT_FOUND``) — instead of + silently writing back a name a concurrent rename already moved on + from, which is what a validate-in-memory-then-blind-UPDATE sequence + could otherwise do. + + ``embedder`` carries no such check — it has no assignment-time + validation at all today (see ``_MODEL_ENDPOINT_COLUMN_TYPES``), so + there is nothing here for a concurrent rename to race against. """ updates = _partition_updates(fields) if updates: preset_refs = {col: _PRESET_COLUMN_TYPES[col] for col in updates if col in _PRESET_COLUMN_TYPES} - if preset_refs: + endpoint_refs = _endpoint_refs(updates) + if preset_refs or endpoint_refs: async with self.pool.acquire() as conn: return await self._update_partition_on_conn(conn, name, **fields) return await self._update_partition_on_conn(self.pool, name, **fields) @@ -323,13 +357,14 @@ async def _update_partition_on_conn( sql = f"UPDATE partitions SET {', '.join(sets)}, updated_at = now() WHERE partition = $1 RETURNING *" preset_refs = {col: _PRESET_COLUMN_TYPES[col] for col in updates if col in _PRESET_COLUMN_TYPES} - if not preset_refs: + endpoint_refs = _endpoint_refs(updates) + if not preset_refs and not endpoint_refs: row = await conn.fetchrow(sql, *params) return self._row_to_full_dict(row) if row else None transaction = getattr(conn, "transaction", None) if transaction is None: - raise TypeError("preset-reference updates require a connection transaction") + raise TypeError("preset/model-endpoint reference updates require a connection transaction") async with transaction(): row = await conn.fetchrow(sql, *params) if row is None: @@ -345,6 +380,17 @@ async def _update_partition_on_conn( f"{preset_type.capitalize()} preset '{updates[col]}' does not exist.", code="PRESET_NOT_FOUND", ) + for col, model_type in endpoint_refs.items(): + exists = await conn.fetchval( + "SELECT 1 FROM model_endpoints WHERE name = $1 AND model_type = $2", + updates[col], + model_type, + ) + if not exists: + raise ValidationError( + f"{model_type.upper()} endpoint '{updates[col]}' referenced by {col} not found.", + code="MODEL_ENDPOINT_NOT_FOUND", + ) return self._row_to_full_dict(row) # ── Legacy method names used by the Phase 7C shim ──────────────── diff --git a/tests/unit/api/schemas/admin/test_phase14_schemas.py b/tests/unit/api/schemas/admin/test_phase14_schemas.py index eb5697d8b..f7cd50f90 100644 --- a/tests/unit/api/schemas/admin/test_phase14_schemas.py +++ b/tests/unit/api/schemas/admin/test_phase14_schemas.py @@ -46,6 +46,45 @@ def test_create_model_endpoint_rejects_empty_normalized_endpoint(endpoint): CreateModelEndpointRequest(name="default", model_type="llm", endpoint=endpoint) +_UNSAFE_NAMES = [ + "owner/model", # splits across the {model_type}/{name} route segment (#768) + ".", # RFC 3986 dot-segment: normalizes to the collection route + "..", # RFC 3986 dot-segment: normalizes away the model_type segment too + "-leading-dash", + "trailing-dash-", + ".leading-dot", + "trailing-dot.", + "_leading_underscore", + "trailing_underscore_", + "has space", + "has%percent", + "a" * 129, # over _NAME_MAX_LENGTH +] + + +@pytest.mark.parametrize("bad_name", _UNSAFE_NAMES) +def test_create_model_endpoint_rejects_unsafe_name(bad_name): + """Any name outside the URL-path-segment allowlist is rejected, not just '/'.""" + with pytest.raises(ValidationError): + CreateModelEndpointRequest(name=bad_name, model_type="reranker", endpoint="http://host") + + +@pytest.mark.parametrize("bad_name", _UNSAFE_NAMES) +def test_update_model_endpoint_rejects_unsafe_name(bad_name): + """Same allowlist applies to renames via the update schema — kept in sync + with the create matrix above (a separate, optional-name validator) so a + regression in one can't go uncovered by the other.""" + with pytest.raises(ValidationError): + UpdateModelEndpointRequest(name=bad_name) + + +@pytest.mark.parametrize("good_name", ["default", "gpt-4.1", "jina_v3", "LocalReranker.prod", "a", "a" * 128]) +def test_create_model_endpoint_accepts_realistic_names(good_name): + """Interior '.', '_', '-' stay available for realistic names.""" + request = CreateModelEndpointRequest(name=good_name, model_type="reranker", endpoint="http://host") + assert request.name == good_name + + def test_update_model_endpoint_requires_at_least_one_field(): """Endpoint updates must contain at least one field.""" with pytest.raises(ValidationError): diff --git a/tests/unit/services/orchestrators/test_model_endpoint_service.py b/tests/unit/services/orchestrators/test_model_endpoint_service.py index a06590457..55ba765b9 100644 --- a/tests/unit/services/orchestrators/test_model_endpoint_service.py +++ b/tests/unit/services/orchestrators/test_model_endpoint_service.py @@ -103,13 +103,15 @@ async def delete_and_promote_default(self, name: str, model_type: str) -> tuple[ return ("ok", promoted) -def _make_service(repo=None, rows=None, settings=None): +def _make_service(repo=None, rows=None, settings=None, partition_service=None, preset_service=None): from core.config.root import Settings from services.orchestrators.model_endpoint_service import ModelEndpointService return ModelEndpointService( model_endpoint_repo=repo or _FakeEndpointRepo(rows), config=settings or Settings(), + partition_service=partition_service, + preset_service=preset_service, ) @@ -913,6 +915,159 @@ async def test_update_model_endpoint_renames_and_evicts_cache(): assert ("new-name", "embedder") in repo._store +class _FakePresetServiceForReload: + def __init__(self): + self.load_all_calls = 0 + + async def load_all(self): + self.load_all_calls += 1 + + +class _FakePartitionServiceForReload: + def __init__(self): + self.load_partitions_calls = 0 + + async def load_partitions(self): + self.load_partitions_calls += 1 + + +@pytest.mark.asyncio +async def test_update_model_endpoint_rename_reloads_presets_then_partitions(): + """A rename cascades DB-side references (#770) inside the repo's own rename + transaction (see PgModelEndpointRepository.rename), but those writes are + invisible until PresetService / PartitionService reload their in-memory + caches — pin that both get refreshed on a rename.""" + existing = _make_row(name="old-name", model_type="llm") + repo = _FakeEndpointRepo(rows=[existing]) + preset_service = _FakePresetServiceForReload() + partition_service = _FakePartitionServiceForReload() + svc = _make_service(repo, partition_service=partition_service, preset_service=preset_service) + + await svc.update_model_endpoint("old-name", "llm", new_name="new-name") + + assert preset_service.load_all_calls == 1 + assert partition_service.load_partitions_calls == 1 + + +@pytest.mark.asyncio +async def test_update_model_endpoint_without_rename_skips_preset_and_partition_reload(): + """A plain field update (no rename) touches no cross-referenced name, so + it must not pay for a presets/partitions reload it doesn't need.""" + existing = _make_row(name="jina") + repo = _FakeEndpointRepo(rows=[existing]) + preset_service = _FakePresetServiceForReload() + partition_service = _FakePartitionServiceForReload() + svc = _make_service(repo, partition_service=partition_service, preset_service=preset_service) + + await svc.update_model_endpoint("jina", "embedder", endpoint="http://new:8000/v1") + + assert preset_service.load_all_calls == 0 + assert partition_service.load_partitions_calls == 0 + + +@pytest.mark.asyncio +async def test_update_model_endpoint_rename_aliases_new_name_before_reload_awaits(): + """A request racing the rename must resolve `new_name` even before the + presets/partitions reload below completes — the DB cascade (#770) has + already repointed partitions/presets at it by the time the rename + `await` returns, so the registry can't lag behind until `load_all()`.""" + existing = _make_row(name="old-name", model_type="llm", endpoint="http://old:8000/v1") + repo = _FakeEndpointRepo(rows=[existing]) + svc = _make_service(repo) + await svc.load_all() + + seen_during_reload = {} + + class _SnoopingPresetService: + async def load_all(self): + seen_during_reload["new-name"] = svc._config.models.llm.get("new-name") + seen_during_reload["old-name"] = svc._config.models.llm.get("old-name") + + svc._preset_service = _SnoopingPresetService() + + await svc.update_model_endpoint("old-name", "llm", new_name="new-name") + + assert seen_during_reload["new-name"].endpoint == "http://old:8000/v1" + assert seen_during_reload["old-name"].endpoint == "http://old:8000/v1" + + +@pytest.mark.asyncio +async def test_update_model_endpoint_rename_keeps_both_names_resolvable_after_failed_reload(): + """If PresetService.load_all() (or PartitionService.load_partitions()) + raises mid-rename, the DB rename has already committed — the registry + must still resolve both the old and the new name afterward, instead of + being stuck answering only to the pre-rename one until process restart.""" + existing = _make_row(name="old-name", model_type="llm", endpoint="http://old:8000/v1") + repo = _FakeEndpointRepo(rows=[existing]) + svc = _make_service(repo) + await svc.load_all() + + class _FailingPresetService: + async def load_all(self): + raise RuntimeError("db blip") + + svc._preset_service = _FailingPresetService() + + with pytest.raises(RuntimeError): + await svc.update_model_endpoint("old-name", "llm", new_name="new-name") + + assert svc._config.models.llm.get("old-name").endpoint == "http://old:8000/v1" + assert svc._config.models.llm.get("new-name").endpoint == "http://old:8000/v1" + + +@pytest.mark.asyncio +async def test_update_model_endpoint_rename_with_field_change_aliases_the_new_values(): + """A rename combined with a field change (e.g. a new endpoint URL) must + alias `new_name`/`old_name` to the row this call just wrote, not to + whatever the in-memory bucket held before the update ran — otherwise a + reload failure right after would leave the registry silently serving the + stale pre-update config under the DB-authoritative new name forever.""" + existing = _make_row(name="old-name", model_type="llm", endpoint="http://old:8000/v1") + repo = _FakeEndpointRepo(rows=[existing]) + svc = _make_service(repo) + await svc.load_all() + + class _FailingPresetService: + async def load_all(self): + raise RuntimeError("db blip") + + svc._preset_service = _FailingPresetService() + + with pytest.raises(RuntimeError): + await svc.update_model_endpoint("old-name", "llm", new_name="new-name", endpoint="http://new:8000/v1") + + assert svc._config.models.llm.get("new-name").endpoint == "http://new:8000/v1" + assert svc._config.models.llm.get("old-name").endpoint == "http://new:8000/v1" + + +@pytest.mark.asyncio +async def test_update_model_endpoint_rename_evicts_stale_client_cache_before_failed_reload(): + """A cached *client instance* under old_name predates this call and can't + know about a field change baked into the same rename — the factory checks + its cache before the config registry, so it must be evicted eagerly (not + only in the post-reload cleanup a failing reload would skip), or a + request through old_name keeps getting the stale pre-update client.""" + existing = _make_row(name="old-name", model_type="llm", endpoint="http://old:8000/v1") + repo = _FakeEndpointRepo(rows=[existing]) + svc = _make_service(repo) + await svc.load_all() + stale_client = object() + cache: dict = {"old-name": stale_client} + svc._client_caches["llm"] = cache + + class _FailingPresetService: + async def load_all(self): + raise RuntimeError("db blip") + + svc._preset_service = _FailingPresetService() + + with pytest.raises(RuntimeError): + await svc.update_model_endpoint("old-name", "llm", new_name="new-name", endpoint="http://new:8000/v1") + + assert "old-name" not in cache + assert "new-name" not in cache + + @pytest.mark.asyncio async def test_update_default_endpoint_evicts_default_alias_cache(): # Updating the current default must evict the 'default' cache key too, not just diff --git a/tests/unit/services/orchestrators/test_retrieval_service.py b/tests/unit/services/orchestrators/test_retrieval_service.py index 217f737b4..67d361f49 100644 --- a/tests/unit/services/orchestrators/test_retrieval_service.py +++ b/tests/unit/services/orchestrators/test_retrieval_service.py @@ -65,6 +65,7 @@ def _config(rtype: str = "single", reranker_enabled: bool = False) -> SimpleName ), reranker=SimpleNamespace(enabled=reranker_enabled, top_k=5), partitions={}, + models=SimpleNamespace(reranker={}), ) @@ -249,6 +250,68 @@ async def test_retrieve_uses_partition_retrieval_config_and_named_reranker(): assert call["similarity_threshold"] == 0.77 +@pytest.mark.asyncio +async def test_retrieve_falls_back_to_default_reranker_when_preset_stale(): + """A partition's ``reranker`` preset can go stale (renamed/deleted after + assignment — this field has no create/PATCH-time validation, unlike + ``chat_llm``). The stale name must fall back to the catalog default + instead of raising.""" + s = FakeSearcher() + s.search_result = [_chunk("a")] + default_reranker = FakeReranker() + reranker_calls: list[str] = [] + + def factory(name: str): + reranker_calls.append(name) + if name == "default": + return default_reranker + raise KeyError(name) + + cfg = _config() + cfg.partitions = { + "tenant-a": _partition(retrieval=RetrievalPipelineConfig(enable_reranker=True, reranker="stale-ranker")) + } + + svc = RetrievalService( + searcher=s, + reranker=None, + llm=None, + config=cfg, + reranker_factory=factory, + ) + + out = await svc.retrieve(partitions=["tenant-a"], query=Query(query="hello")) + + assert [c.id for c in out] == ["a"] + assert reranker_calls == ["stale-ranker", "default"] + assert default_reranker.calls[0]["query"] == "hello" + + +@pytest.mark.asyncio +async def test_retrieve_falls_back_to_legacy_reranker_when_no_catalog_default(): + """No ``is_default`` reranker endpoint registered yet — fall back to the + static reranker built at startup instead of raising.""" + s = FakeSearcher() + s.search_result = [_chunk("a")] + legacy_reranker = FakeReranker() + + cfg = _config() + cfg.partitions = {"tenant-a": _partition(retrieval=RetrievalPipelineConfig(enable_reranker=True))} + + svc = RetrievalService( + searcher=s, + reranker=legacy_reranker, + llm=None, + config=cfg, + reranker_factory=lambda name: (_ for _ in ()).throw(KeyError(name)), + ) + + out = await svc.retrieve(partitions=["tenant-a"], query=Query(query="hello")) + + assert [c.id for c in out] == ["a"] + assert legacy_reranker.calls[0]["query"] == "hello" + + @pytest.mark.asyncio async def test_retrieve_uses_partition_searcher_factory_for_named_embedder(): default_searcher = FakeSearcher() diff --git a/tests/unit/services/persistence/test_model_endpoint_repo.py b/tests/unit/services/persistence/test_model_endpoint_repo.py index 660ca6a5a..8544e16ec 100644 --- a/tests/unit/services/persistence/test_model_endpoint_repo.py +++ b/tests/unit/services/persistence/test_model_endpoint_repo.py @@ -237,6 +237,157 @@ async def _execute(query, *params): assert await repo.delete("ghost", "embedder") is False +@pytest.mark.asyncio +async def test_rename_updates_the_model_endpoints_row(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old", "embedder", "new") + + queries = [q for q, _ in pool.conn.executed] + assert any("UPDATE model_endpoints SET name = $3" in q for q in queries) + params = next(p for q, p in pool.conn.executed if "UPDATE model_endpoints SET name" in q) + assert params == ("old", "embedder", "new") + + +@pytest.mark.asyncio +async def test_rename_locks_partitions_table_before_touching_model_endpoints(): + """rename() must LOCK partitions IN SHARE MODE before its own UPDATE — + the same order PgPartitionRepository.update_partition's chat_llm guard + touches partitions (write) then model_endpoints (check), so the two + transactions can only block on each other, never deadlock. Without this + lock, a partition PATCH could validate 'old' in-memory, block on this + transaction's cascade instead, then resume and write 'old' straight back + after this commits — see PgPartitionRepository.update_partition.""" + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old", "llm", "new") + + queries = [q for q, _ in pool.conn.executed] + lock_i = next(i for i, q in enumerate(queries) if q == "LOCK TABLE partitions IN SHARE MODE") + rename_i = next(i for i, q in enumerate(queries) if "UPDATE model_endpoints SET name" in q) + assert lock_i < rename_i + + +@pytest.mark.asyncio +async def test_rename_raises_not_found_and_skips_cascade_when_row_vanished(): + """A concurrent delete between the service's existence check and this + transaction must abort before the cascade — not repoint partitions/presets + at a `new_name` that was never actually created (mirrors + PgPipelinePresetRepository.rename's RETURNING guard).""" + from core.utils.exceptions import NotFoundError + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() # _fetchrow_result defaults to None: row is gone + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + with pytest.raises(NotFoundError): + await repo.rename("old-llm", "llm", "new-llm") + + queries = [q for q, _ in pool.conn.executed] + assert not any("UPDATE partitions SET" in q for q in queries) + assert not any("pipeline_presets" in q for q in queries) + + +@pytest.mark.asyncio +async def test_rename_embedder_cascades_to_partitions_embedder_only(): + """Renaming an embedder must update `partitions.embedder` and touch no + preset JSONB — the embedder name isn't referenced inside any preset.""" + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new-embedder"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old-embedder", "embedder", "new-embedder") + + queries = pool.conn.executed + partition_updates = [(q, p) for q, p in queries if "UPDATE partitions SET" in q] + assert len(partition_updates) == 1 + q, p = partition_updates[0] + assert "embedder = $2 WHERE embedder = $1" in q + assert p == ("old-embedder", "new-embedder") + assert not any("pipeline_presets" in q for q, _ in queries) + + +@pytest.mark.asyncio +async def test_rename_llm_cascades_to_chat_llm_and_both_preset_types(): + """Renaming an LLM endpoint must update `partitions.chat_llm`, the + retrieval preset's `llm` key, and every indexation-preset LLM field.""" + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new-llm"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old-llm", "llm", "new-llm") + + queries = pool.conn.executed + partition_updates = [(q, p) for q, p in queries if "UPDATE partitions SET" in q] + assert len(partition_updates) == 1 + q, p = partition_updates[0] + assert "chat_llm = $2 WHERE chat_llm = $1" in q + assert p == ("old-llm", "new-llm") + + preset_updates = [(q, p) for q, p in queries if "pipeline_presets" in q] + # retrieval.llm + indexation.{contextualization_llm, metadata_extraction_llm, topic_tagging_llm} + assert len(preset_updates) == 4 + keys_by_preset_type = {(p[2], p[3]) for _, p in preset_updates} + assert keys_by_preset_type == { + ("retrieval", "llm"), + ("indexation", "contextualization_llm"), + ("indexation", "metadata_extraction_llm"), + ("indexation", "topic_tagging_llm"), + } + for _, p in preset_updates: + assert p[0] == [p[3]] # jsonb_set path matches the ->> key checked in WHERE + assert p[1] == "new-llm" + assert p[4] == "old-llm" + + +@pytest.mark.asyncio +async def test_rename_reranker_cascades_to_retrieval_preset_only(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new-ranker"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old-ranker", "reranker", "new-ranker") + + queries = pool.conn.executed + assert not any("UPDATE partitions SET" in q for q, _ in queries) + preset_updates = [(q, p) for q, p in queries if "pipeline_presets" in q] + assert len(preset_updates) == 1 + q, p = preset_updates[0] + assert p == (["reranker"], "new-ranker", "retrieval", "reranker", "old-ranker") + + +@pytest.mark.asyncio +async def test_rename_vlm_cascades_to_indexation_preset_only(): + from services.persistence.model_endpoint_repo import PgModelEndpointRepository + + pool = _FakePool() + pool.conn._fetchrow_result = {"name": "new-vlm"} + repo = PgModelEndpointRepository(pool_getter=lambda: pool) + + await repo.rename("old-vlm", "vlm", "new-vlm") + + queries = pool.conn.executed + assert not any("UPDATE partitions SET" in q for q, _ in queries) + preset_updates = [(q, p) for q, p in queries if "pipeline_presets" in q] + assert len(preset_updates) == 1 + q, p = preset_updates[0] + assert p == (["vlm"], "new-vlm", "indexation", "vlm", "old-vlm") + + def _row(name, is_default): return {"name": name, "is_default": is_default} diff --git a/tests/unit/services/persistence/test_partition_repo.py b/tests/unit/services/persistence/test_partition_repo.py index 0aaf6209c..7ca31e504 100644 --- a/tests/unit/services/persistence/test_partition_repo.py +++ b/tests/unit/services/persistence/test_partition_repo.py @@ -180,10 +180,13 @@ class _UpdateFakeConn: ``preset_exists`` models whether the referenced preset row is still present when the guard's follow-up SELECT runs (False simulates a concurrent delete_preset committing while this UPDATE was blocked on its SHARE lock). + ``model_endpoint_exists`` is the same, for a ``chat_llm`` reference racing + a concurrent ``PgModelEndpointRepository.rename``. """ - def __init__(self, *, preset_exists: bool = True): + def __init__(self, *, preset_exists: bool = True, model_endpoint_exists: bool = True): self.preset_exists = preset_exists + self.model_endpoint_exists = model_endpoint_exists self.operations: list[tuple[str, tuple]] = [] self.transactions = 0 @@ -209,6 +212,8 @@ async def fetchval(self, query: str, *params): self.operations.append((query, params)) if "FROM pipeline_presets" in query: return 1 if self.preset_exists else None + if "FROM model_endpoints" in query: + return 1 if self.model_endpoint_exists else None return None # conn interface @@ -267,3 +272,76 @@ async def test_update_partition_without_preset_change_skips_the_guard(): assert result["description"] == "notes" assert conn.transactions == 0 assert not any("FROM pipeline_presets" in q for q, _ in conn.operations) + + +# ── update_partition chat_llm-assignment race guard ────────────────── + + +@pytest.mark.asyncio +async def test_update_partition_rolls_back_when_chat_llm_endpoint_renamed_concurrently(): + """A concurrent PgModelEndpointRepository.rename() moving 'old-name' away + while this UPDATE was blocked on the LOCK it also takes must not let the + partition silently keep pointing at the now-nonexistent name.""" + from services.persistence.partition_repo import PgPartitionRepository + + conn = _UpdateFakeConn(model_endpoint_exists=False) + repo = PgPartitionRepository(pool_getter=lambda: conn) + + with pytest.raises(ValidationError) as exc: + await repo.update_partition("p1", chat_llm="old-name") + + assert exc.value.code == "MODEL_ENDPOINT_NOT_FOUND" + # The write and the existence check share one transaction, and the write + # (partitions) happens before the check (model_endpoints) — the same lock + # order PgModelEndpointRepository.rename uses, keeping the two deadlock-free. + assert conn.transactions == 1 + queries = [q for q, _ in conn.operations] + update_i = next(i for i, q in enumerate(queries) if "UPDATE partitions" in q) + check_i = next(i for i, q in enumerate(queries) if "FROM model_endpoints" in q) + assert update_i < check_i + + +@pytest.mark.asyncio +async def test_update_partition_commits_when_chat_llm_endpoint_exists(): + from services.persistence.partition_repo import PgPartitionRepository + + conn = _UpdateFakeConn(model_endpoint_exists=True) + repo = PgPartitionRepository(pool_getter=lambda: conn) + + result = await repo.update_partition("p1", chat_llm="gpt-4.1") + + assert result["chat_llm"] == "gpt-4.1" + assert conn.transactions == 1 + assert any("FROM model_endpoints" in q for q, _ in conn.operations) + + +@pytest.mark.asyncio +async def test_update_partition_clearing_chat_llm_skips_the_guard(): + """chat_llm=None clears the (nullable) column — it names no endpoint to + validate, so this must take the fast, non-transactional path.""" + from services.persistence.partition_repo import PgPartitionRepository + + conn = _UpdateFakeConn(model_endpoint_exists=False) + repo = PgPartitionRepository(pool_getter=lambda: conn) + + result = await repo.update_partition("p1", chat_llm=None) + + assert result["chat_llm"] is None + assert conn.transactions == 0 + assert not any("FROM model_endpoints" in q for q, _ in conn.operations) + + +@pytest.mark.asyncio +async def test_update_partition_embedder_change_skips_the_guard(): + """embedder carries no assignment-time validation today, so assigning it + alone must not pay for a transaction or a model_endpoints lookup.""" + from services.persistence.partition_repo import PgPartitionRepository + + conn = _UpdateFakeConn(model_endpoint_exists=False) + repo = PgPartitionRepository(pool_getter=lambda: conn) + + result = await repo.update_partition("p1", embedder="some-embedder") + + assert result["embedder"] == "some-embedder" + assert conn.transactions == 0 + assert not any("FROM model_endpoints" in q for q, _ in conn.operations) diff --git a/ui/src/pages/admin/models.tsx b/ui/src/pages/admin/models.tsx index 8811682db..154bbf379 100644 --- a/ui/src/pages/admin/models.tsx +++ b/ui/src/pages/admin/models.tsx @@ -72,6 +72,14 @@ type RevealedApiKey = { const normalizeEndpointUrl = (value: string) => value.trim().replace(/\/+$/, ""); +// Mirrors the backend's `_NAME_PATTERN` allowlist (api/schemas/admin/model_endpoint_schemas.py): +// `name` is a single path segment in every single-endpoint route. An allowlist, not a +// denylist — `/` splits across path segments, and `.`/`..` are RFC 3986 dot-segments that +// browsers normalize out of the URL before the request is even sent — so anchoring both +// ends on alphanumeric rules out all of that (and any leading/trailing separator) at once. +const NAME_PATTERN = /^[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?$/; +const NAME_MAX_LENGTH = 128; + // Placeholder shown when a per-endpoint budget is left blank. The real // fallback (the MAX_LLM_CONTEXT_SIZE / MAX_OUTPUT_TOKENS env vars — see // core/config/endpoints.py:LLMContextConfig) is environment-configurable, so @@ -316,6 +324,20 @@ function EndpointDialog({ const [vendor, setVendor] = useState(""); const [extraJson, setExtraJson] = useState("{}"); + // The backend trims `name` before validating it (and before persisting it), + // so validate — and submit — the same trimmed value here, not the raw + // input; otherwise e.g. " gpt-4.1 " would be accepted by the backend but + // blocked by this form. Checked client-side too because the resulting 422 + // has no readable message — FastAPI's validation `detail` is a list, and + // ApiError only unwraps string detail. + const trimmedName = name.trim(); + const nameError = + trimmedName !== "" && trimmedName.length > NAME_MAX_LENGTH + ? `Name must be at most ${NAME_MAX_LENGTH} characters.` + : trimmedName !== "" && !NAME_PATTERN.test(trimmedName) + ? "Name must start/end with a letter or digit, and contain only letters, digits, '.', '_', or '-' — it's used in the endpoint's URL path." + : null; + const modelType = (editing ? editing.model_type : activeTab) as ModelType; // LLM token-budget fields (max context / max output) apply to LLM endpoints only. const isLlm = modelType === "llm"; @@ -572,6 +594,10 @@ function EndpointDialog({ const handleSubmit = (e: React.FormEvent) => { e.preventDefault(); + if (nameError) { + toast.error(nameError); + return; + } let extra: Record = {}; try { extra = mergeModelEndpointApiKeyExtra(JSON.parse(extraJson), apiKey, { @@ -594,13 +620,13 @@ function EndpointDialog({ timeout: numOr(timeout, 30), extra, }; - if (name !== editing.name) { - updateData.name = name; + if (trimmedName !== editing.name) { + updateData.name = trimmedName; } onUpdate(editing.model_type, editing.name, updateData); } else { onCreate({ - name, + name: trimmedName, model_type: activeTab as ModelType, endpoint, model_name: modelName || undefined, @@ -633,6 +659,7 @@ function EndpointDialog({ onChange={(e) => setName(e.target.value)} required /> + {nameError &&

{nameError}

}
@@ -800,7 +827,7 @@ function EndpointDialog({