fix: address policy bundle review feedback

Author: erichare

src/backend/base/langflow/alembic/versions/f7a9c2d4e6b8_add_shared_policy_bundle.py

@@ -272,6 +272,9 @@ def _sync_legacy_policy(conn: sa.Connection) -> None:
         .one_or_none()
     )
     if active is None:
+        if conn.execute(sa.select(sa.literal(1)).select_from(revision_table).limit(1)).first() is not None:
+            msg = "Shared policy bundle has immutable revision history but no active singleton"
+            raise RuntimeError(msg)
         return
     active_revision = active["revision"]
     bundle = (
@@ -280,7 +283,8 @@ def _sync_legacy_policy(conn: sa.Connection) -> None:
         .one_or_none()
     )
     if bundle is None:
-        return
+        msg = "Active policy bundle points to a missing immutable revision"
+        raise RuntimeError(msg)
 
     provider_table = _legacy_provider_table()
     conn.execute(
@@ -328,8 +332,22 @@ def _sync_legacy_policy(conn: sa.Connection) -> None:
 def downgrade() -> None:
     """Copy the active bundle to legacy stores before removing new tables."""
     conn = op.get_bind()
-    if not migration.table_exists(REVISION_TABLE, conn) or not migration.table_exists(ACTIVE_TABLE, conn):
+    revision_exists = migration.table_exists(REVISION_TABLE, conn)
+    active_exists = migration.table_exists(ACTIVE_TABLE, conn)
+    if revision_exists != active_exists:
+        existing_table = _revision_table() if revision_exists else _active_table()
+        if conn.execute(sa.select(sa.literal(1)).select_from(existing_table).limit(1)).first() is not None:
+            msg = "Shared policy bundle schema is partially initialized with durable data"
+            raise RuntimeError(msg)
+        if active_exists:
+            op.drop_table(ACTIVE_TABLE)
+        if revision_exists:
+            op.drop_table(REVISION_TABLE)
+        return
+
+    if not revision_exists:
         return
+
     _sync_legacy_policy(conn)
     op.drop_table(ACTIVE_TABLE)
     op.drop_table(REVISION_TABLE)

src/backend/base/langflow/api/v1/catalog_policy.py

@@ -8,6 +8,7 @@
 from fastapi import APIRouter, Depends, HTTPException, status
 from lfx.services.catalog_policy import BaseCatalogPolicyService, CatalogPolicySnapshot
 
+from langflow.api.v1.policy_bundle_errors import policy_bundle_revision_conflict
 from langflow.api.v1.schemas.catalog_policy import CatalogPolicyBlockedSet, CatalogPolicyRead
 from langflow.services.auth.utils import get_current_active_superuser
 from langflow.services.authorization.audit import audit_decision
@@ -48,17 +49,6 @@ def _raise_if_externally_managed(service: BaseCatalogPolicyService) -> None:
         )
 
 
-def _revision_conflict(exc: PolicyBundleRevisionConflictError) -> HTTPException:
-    return HTTPException(
-        status_code=status.HTTP_409_CONFLICT,
-        detail={
-            "message": "Policy bundle revision conflict",
-            "expected_revision": exc.expected_revision,
-            "active_revision": exc.active_revision,
-        },
-    )
-
-
 async def _audit_update(
     *,
     user_id: UUID,
@@ -115,7 +105,7 @@ async def replace_component_policy(
             actor_user_id=admin.id,
         )
     except PolicyBundleRevisionConflictError as exc:
-        raise _revision_conflict(exc) from exc
+        raise policy_bundle_revision_conflict(exc) from exc
     await _audit_update(
         user_id=admin.id,
         resource_kind="component",
@@ -149,7 +139,7 @@ async def replace_template_policy(
             actor_user_id=admin.id,
         )
     except PolicyBundleRevisionConflictError as exc:
-        raise _revision_conflict(exc) from exc
+        raise policy_bundle_revision_conflict(exc) from exc
     await _audit_update(
         user_id=admin.id,
         resource_kind="template",

src/backend/base/langflow/api/v1/model_provider_policy.py

@@ -10,6 +10,7 @@
 from pydantic import BaseModel, Field, StringConstraints, field_validator
 
 from langflow.api.utils import DbSession, DbSessionReadOnly
+from langflow.api.v1.policy_bundle_errors import policy_bundle_revision_conflict
 from langflow.services.auth.utils import get_current_active_superuser
 from langflow.services.authorization.audit import AUDIT_ALLOW, audit_decision
 from langflow.services.database.models.user.model import User
@@ -145,14 +146,7 @@ async def replace_model_provider_policy(
     except PolicyBundleApplicationNotSupportedError as exc:
         raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(exc)) from exc
     except PolicyBundleRevisionConflictError as exc:
-        raise HTTPException(
-            status_code=status.HTTP_409_CONFLICT,
-            detail={
-                "message": "Policy bundle revision conflict",
-                "expected_revision": exc.expected_revision,
-                "active_revision": exc.active_revision,
-            },
-        ) from exc
+        raise policy_bundle_revision_conflict(exc) from exc
 
     try:
         await audit_decision(

src/backend/base/langflow/api/v1/policy_bundle.py

@@ -13,6 +13,8 @@
 from pydantic import BaseModel, Field, StringConstraints, field_validator
 
 from langflow.api.utils import DbSession, DbSessionReadOnly
+from langflow.api.v1.policy_bundle_errors import policy_bundle_revision_conflict
+from langflow.api.v1.schemas.catalog_policy import CatalogPolicyKeyList, normalize_catalog_policy_keys
 from langflow.services.auth.utils import get_current_active_superuser
 from langflow.services.authorization.audit import AUDIT_ALLOW, audit_decision
 from langflow.services.database.models.policy_bundle import POLICY_BUNDLE_REASON_MAX_LENGTH
@@ -34,24 +36,13 @@
 ProviderId = Annotated[str, StringConstraints(pattern=r"^[a-z0-9][a-z0-9._-]*$", max_length=255)]
 
 
-def _normalize_keys(values: list[str]) -> list[str]:
-    normalized: set[str] = set()
-    for raw_value in values:
-        value = raw_value.strip()
-        if not value:
-            msg = "Policy bundle catalog keys must not be empty"
-            raise ValueError(msg)
-        normalized.add(value)
-    return sorted(normalized)
-
-
 class PolicyBundleWrite(BaseModel):
     """Complete replacement guarded by the caller's observed revision."""
 
     expected_revision: int = Field(ge=1)
     approved_provider_ids: Annotated[list[ProviderId], Field(max_length=1000)]
-    blocked_component_keys: list[str]
-    blocked_template_keys: list[str]
+    blocked_component_keys: CatalogPolicyKeyList
+    blocked_template_keys: CatalogPolicyKeyList
     reason: str | None = Field(default=None, max_length=POLICY_BUNDLE_REASON_MAX_LENGTH)
 
     @field_validator("approved_provider_ids", mode="before")
@@ -72,7 +63,7 @@ def deduplicate_provider_ids(cls, provider_ids: list[str]) -> list[str]:
     @field_validator("blocked_component_keys", "blocked_template_keys")
     @classmethod
     def normalize_catalog_keys(cls, values: list[str]) -> list[str]:
-        return _normalize_keys(values)
+        return normalize_catalog_policy_keys(values)
 
 
 class PolicyBundleRollbackWrite(BaseModel):
@@ -131,17 +122,6 @@ def _unavailable() -> HTTPException:
     )
 
 
-def _conflict(exc: PolicyBundleRevisionConflictError) -> HTTPException:
-    return HTTPException(
-        status_code=status.HTTP_409_CONFLICT,
-        detail={
-            "message": "Policy bundle revision conflict",
-            "expected_revision": exc.expected_revision,
-            "active_revision": exc.active_revision,
-        },
-    )
-
-
 def _raise_if_externally_managed() -> None:
     if _managed_externally():
         raise HTTPException(
@@ -204,7 +184,7 @@ async def replace_policy_bundle(
     except PolicyBundleNotInitializedError as exc:
         raise _unavailable() from exc
     except PolicyBundleRevisionConflictError as exc:
-        raise _conflict(exc) from exc
+        raise policy_bundle_revision_conflict(exc) from exc
 
     try:
         await _audit_bundle(snapshot, user_id=admin.id, action="policy_bundle:replace")
@@ -251,7 +231,7 @@ async def rollback_policy_bundle(
     except PolicyBundleNotInitializedError as exc:
         raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(exc)) from exc
     except PolicyBundleRevisionConflictError as exc:
-        raise _conflict(exc) from exc
+        raise policy_bundle_revision_conflict(exc) from exc
 
     try:
         await _audit_bundle(snapshot, user_id=admin.id, action="policy_bundle:rollback")

src/backend/base/langflow/api/v1/policy_bundle_errors.py

@@ -0,0 +1,20 @@
+"""Shared HTTP error mapping for policy-bundle-backed administration APIs."""
+
+from fastapi import HTTPException, status
+
+from langflow.services.policy_bundle import PolicyBundleRevisionConflictError
+
+
+def policy_bundle_revision_conflict(exc: PolicyBundleRevisionConflictError) -> HTTPException:
+    """Preserve the structured optimistic-concurrency response across policy APIs."""
+    return HTTPException(
+        status_code=status.HTTP_409_CONFLICT,
+        detail={
+            "message": "Policy bundle revision conflict",
+            "expected_revision": exc.expected_revision,
+            "active_revision": exc.active_revision,
+        },
+    )
+
+
+__all__ = ["policy_bundle_revision_conflict"]

src/backend/base/langflow/api/v1/schemas/catalog_policy.py

@@ -2,32 +2,59 @@
 
 from __future__ import annotations
 
-from pydantic import BaseModel, Field, field_validator
+from typing import Annotated
+
+from pydantic import BaseModel, Field, StringConstraints, field_validator
+
+CATALOG_POLICY_KEY_MAX_LENGTH = 255
+CATALOG_POLICY_KEYS_MAX_LENGTH = 1000
+
+CatalogPolicyKey = Annotated[
+    str,
+    StringConstraints(
+        strip_whitespace=True,
+        min_length=1,
+        max_length=CATALOG_POLICY_KEY_MAX_LENGTH,
+    ),
+]
+CatalogPolicyKeyList = Annotated[list[CatalogPolicyKey], Field(max_length=CATALOG_POLICY_KEYS_MAX_LENGTH)]
+
+
+def normalize_catalog_policy_keys(value: list[str]) -> list[str]:
+    """Deduplicate validated catalog keys and sort them deterministically."""
+    return sorted(set(value))
 
 
 class CatalogPolicyBlockedSet(BaseModel):
     """A normalized whole-set catalog block policy."""
 
-    blocked: list[str] = Field(
+    blocked: CatalogPolicyKeyList = Field(
         ...,
         description="Complete set of blocked catalog keys for this resource kind.",
     )
 
     @field_validator("blocked")
     @classmethod
     def normalize_blocked_keys(cls, value: list[str]) -> list[str]:
-        """Trim, reject empty values, deduplicate, and sort deterministically."""
-        normalized: set[str] = set()
-        for raw_key in value:
-            key = raw_key.strip()
-            if not key:
-                msg = "Blocked catalog keys must not be empty"
-                raise ValueError(msg)
-            normalized.add(key)
-        return sorted(normalized)
-
-
-class CatalogPolicyRead(CatalogPolicyBlockedSet):
+        """Deduplicate and sort the already-trimmed, validated keys."""
+        return normalize_catalog_policy_keys(value)
+
+
+class CatalogPolicyRead(BaseModel):
     """Current catalog block policy and its ownership source."""
 
+    # Response/history schemas intentionally remain unconstrained so a policy
+    # written by an older release can still be read and repaired.
+    blocked: list[str]
     managed_externally: bool
+
+
+__all__ = [
+    "CATALOG_POLICY_KEYS_MAX_LENGTH",
+    "CATALOG_POLICY_KEY_MAX_LENGTH",
+    "CatalogPolicyBlockedSet",
+    "CatalogPolicyKey",
+    "CatalogPolicyKeyList",
+    "CatalogPolicyRead",
+    "normalize_catalog_policy_keys",
+]

src/backend/base/langflow/services/catalog_policy/service.py

@@ -18,6 +18,7 @@
     CatalogResourceKind,
 )
 from langflow.services.policy_bundle import (
+    PolicyBundleRevisionConflictError,
     apply_policy_bundle_state,
     get_policy_bundle_state,
     replace_policy_bundle_state,
@@ -32,6 +33,9 @@
     from langflow.services.database.service import DatabaseService
 
 
+_CATALOG_POLICY_CAS_ATTEMPTS = 3
+
+
 def _normalize_keys(keys: Collection[str]) -> frozenset[str]:
     """Trim and deduplicate keys while preserving their case."""
     normalized: set[str] = set()
@@ -57,6 +61,7 @@ def __init__(
         self._policy_bundle_service = policy_bundle_service
         self._legacy_snapshot = CatalogPolicySnapshot()
         self._legacy_hydrated = False
+        self._projected_snapshot: tuple[PolicyBundleSnapshot, CatalogPolicySnapshot] | None = None
         self._write_lock = asyncio.Lock()
         # Direct service-class registration does not call set_ready(), unlike
         # factory creation. The fail-open snapshot is usable immediately.
@@ -73,10 +78,17 @@ def snapshot(self) -> CatalogPolicySnapshot:
         if self._policy_bundle_service is None:
             return self._legacy_snapshot
         bundle = self._policy_bundle_service.snapshot
-        return CatalogPolicySnapshot(
-            blocked_component_keys=bundle.blocked_component_keys,
-            blocked_template_keys=bundle.blocked_template_keys,
-        )
+        projected = self._projected_snapshot
+        if projected is None or projected[0] is not bundle:
+            projected = (
+                bundle,
+                CatalogPolicySnapshot(
+                    blocked_component_keys=bundle.blocked_component_keys,
+                    blocked_template_keys=bundle.blocked_template_keys,
+                ),
+            )
+            self._projected_snapshot = projected
+        return projected[1]
 
     @property
     def policy_bundle_snapshot(self) -> PolicyBundleSnapshot:
@@ -165,25 +177,32 @@ async def _replace_blocked_keys(
                     actor_user_id=actor_user_id,
                 )
 
-            async with session_scope() as session:
-                current = await get_policy_bundle_state(session)
-                if resource_kind == CatalogResourceKind.COMPONENT:
-                    current_keys = current.blocked_component_keys
-                    components = desired
-                    templates = current.blocked_template_keys
-                else:
-                    current_keys = current.blocked_template_keys
-                    components = current.blocked_component_keys
-                    templates = desired
-                committed = await replace_policy_bundle_state(
-                    session,
-                    expected_revision=current.revision,
-                    approved_provider_ids=current.approved_provider_ids,
-                    blocked_component_keys=components,
-                    blocked_template_keys=templates,
-                    actor_user_id=actor_user_id,
-                    reason=f"Replace blocked {resource_kind.value} catalog keys",
-                )
+            for attempt in range(_CATALOG_POLICY_CAS_ATTEMPTS):
+                async with session_scope() as session:
+                    current = await get_policy_bundle_state(session)
+                    if resource_kind == CatalogResourceKind.COMPONENT:
+                        current_keys = current.blocked_component_keys
+                        components = desired
+                        templates = current.blocked_template_keys
+                    else:
+                        current_keys = current.blocked_template_keys
+                        components = current.blocked_component_keys
+                        templates = desired
+                    try:
+                        committed = await replace_policy_bundle_state(
+                            session,
+                            expected_revision=current.revision,
+                            approved_provider_ids=current.approved_provider_ids,
+                            blocked_component_keys=components,
+                            blocked_template_keys=templates,
+                            actor_user_id=actor_user_id,
+                            reason=f"Replace blocked {resource_kind.value} catalog keys",
+                        )
+                    except PolicyBundleRevisionConflictError:
+                        if attempt == _CATALOG_POLICY_CAS_ATTEMPTS - 1:
+                            raise
+                        continue
+                    break
 
             apply_policy_bundle_state(committed)
             return CatalogPolicyUpdate(