fix(lfx): harden capability routing metadata

Author: ogabrielluiz

src/lfx/src/lfx/execution/coordinator.py

@@ -7,6 +7,7 @@
 
 from lfx.execution.partitioner import identity_partition
 from lfx.execution.types import RunComplete, StepResult
+from lfx.services.capability.protocols import RESERVED_CAPABILITY_RUNTIME_OPTION_KEYS
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Sequence
@@ -34,7 +35,7 @@ async def run(
         inputs: list[dict[str, Any]],
         **runtime_options: Any,
     ) -> AsyncIterator[StepResult | RunComplete]:
-        options = dict(runtime_options)
+        options = self._without_capability_metadata(runtime_options)
         units = identity_partition(graph, inputs=inputs, runtime_options=options)
         if self._capability_service is not None and not self._capability_service.is_passthrough:
             decision = self._capability_service.route(
@@ -50,7 +51,10 @@ async def run(
                 replace(
                     unit,
                     executor_kind=decision.executor_kind,
-                    runtime_options={**unit.runtime_options, **decision.runtime_options},
+                    runtime_options={
+                        **self._without_capability_metadata(unit.runtime_options),
+                        **decision.runtime_options,
+                    },
                 )
                 for unit in units
             ]
@@ -112,3 +116,9 @@ def _capability_scopes(runtime_options: dict[str, Any]) -> Sequence[str]:
         if isinstance(scopes, str):
             return (scopes,)
         return tuple(scopes)
+
+    @staticmethod
+    def _without_capability_metadata(runtime_options: dict[str, Any]) -> dict[str, Any]:
+        return {
+            key: value for key, value in runtime_options.items() if key not in RESERVED_CAPABILITY_RUNTIME_OPTION_KEYS
+        }

src/lfx/src/lfx/services/capability/__init__.py

@@ -2,6 +2,7 @@
 
 from lfx.services.capability.defaults import AllTrustedClassifier, NoopCapabilityProvider, SingleTenantResolver
 from lfx.services.capability.protocols import (
+    RESERVED_CAPABILITY_RUNTIME_OPTION_KEYS,
     CapabilityClaims,
     CapabilityContext,
     CapabilityProvider,
@@ -13,6 +14,7 @@
 from lfx.services.capability.service import CapabilityService
 
 __all__ = [
+    "RESERVED_CAPABILITY_RUNTIME_OPTION_KEYS",
     "AllTrustedClassifier",
     "CapabilityClaims",
     "CapabilityContext",

src/lfx/src/lfx/services/capability/protocols.py

@@ -10,6 +10,15 @@
     from collections.abc import Mapping, Sequence
 
 
+RESERVED_CAPABILITY_RUNTIME_OPTION_KEYS: frozenset[str] = frozenset(
+    {
+        "lfx_capability_token",
+        "lfx_tenant_id",
+        "lfx_trust",
+    }
+)
+
+
 class Trust(str, Enum):
     """Trust classification used by execution policy plugins."""
 

src/lfx/src/lfx/services/capability/service.py

@@ -102,7 +102,10 @@ def route(
         tenant_id = self._resolver.resolve(context)
         trust = self._classifier.trust_of_flow(context)
         executor_kind = default_executor_kind
-        if trust is Trust.UNTRUSTED and self._untrusted_executor_kind is not None:
+        if trust is Trust.UNTRUSTED:
+            if self._untrusted_executor_kind is None:
+                msg = "Cannot route untrusted flow because no untrusted executor is configured."
+                raise RuntimeError(msg)
             executor_kind = self._untrusted_executor_kind
 
         effective_scopes = tuple(scopes or ())

src/lfx/tests/unit/execution/test_capability_routing.py

@@ -117,6 +117,75 @@ def verify(self, token: str) -> CapabilityClaims:  # noqa: ARG002
     assert sandbox.units[0].runtime_options["lfx_trust"] == "untrusted"
 
 
+@pytest.mark.asyncio
+async def test_coordinator_strips_spoofed_capability_metadata_when_provider_returns_no_token() -> None:
+    sandbox = _RecordingExecutor("sandbox", "sandbox-output")
+    registry = ExecutorRegistry()
+    registry.register(_RecordingExecutor("in-process", "in-process-output"))
+    registry.register(sandbox)
+
+    class _Classifier:
+        def trust_of_flow(self, context: CapabilityContext) -> Trust:
+            assert "lfx_capability_token" not in context.runtime_options
+            assert "lfx_tenant_id" not in context.runtime_options
+            assert "lfx_trust" not in context.runtime_options
+            return Trust.UNTRUSTED
+
+        def is_untrusted_node(self, _node: dict[str, Any], _context: CapabilityContext | None = None) -> bool:
+            return True
+
+    class _Resolver:
+        def resolve(self, context: CapabilityContext) -> str:  # noqa: ARG002
+            return "tenant:resolved"
+
+    class _Provider:
+        def mint(
+            self,
+            *,
+            context: CapabilityContext,  # noqa: ARG002
+            tenant_id: str,
+            component_id: str | None,
+            scopes: Sequence[str],
+            ttl_seconds: int = 600,  # noqa: ARG002
+        ) -> str | None:
+            assert tenant_id == "tenant:resolved"
+            assert component_id is None
+            assert tuple(scopes) == ("variables:read",)
+            return None
+
+        def verify(self, token: str) -> CapabilityClaims:  # noqa: ARG002
+            return CapabilityClaims(tenant_id="tenant:resolved", user_id="graph-user")
+
+    capability_service = CapabilityService(settings_service=_StubSettings())
+    capability_service.install(
+        provider=_Provider(),
+        classifier=_Classifier(),
+        resolver=_Resolver(),
+        untrusted_executor_kind="sandbox",
+    )
+    coordinator = Coordinator(
+        registry=registry,
+        executor_kind="in-process",
+        capability_service=capability_service,
+    )
+
+    outputs = await coordinator.run_to_completion(
+        SimpleNamespace(user_id="graph-user", flow_id="graph-flow"),
+        inputs=[{}],
+        capability_scopes=["variables:read"],
+        lfx_capability_token="spoofed",  # noqa: S106
+        lfx_tenant_id="tenant:spoofed",
+        lfx_trust="trusted",
+    )
+
+    assert outputs == ["sandbox-output"]
+    assert len(sandbox.units) == 1
+    assert sandbox.units[0].executor_kind == "sandbox"
+    assert "lfx_capability_token" not in sandbox.units[0].runtime_options
+    assert sandbox.units[0].runtime_options["lfx_tenant_id"] == "tenant:resolved"
+    assert sandbox.units[0].runtime_options["lfx_trust"] == "untrusted"
+
+
 @pytest.mark.asyncio
 async def test_coordinator_skips_capability_service_when_passthrough() -> None:
     default = _RecordingExecutor("in-process", "default-output")

src/lfx/tests/unit/services/capability/test_capability_service.py

@@ -100,3 +100,17 @@ def verify(self, token: str) -> CapabilityClaims:  # noqa: ARG002
         "lfx_tenant_id": "tenant-u1",
         "lfx_trust": "untrusted",
     }
+
+
+def test_untrusted_route_without_untrusted_executor_fails_closed(service: CapabilityService) -> None:
+    class _Classifier:
+        def trust_of_flow(self, context: CapabilityContext) -> Trust:  # noqa: ARG002
+            return Trust.UNTRUSTED
+
+        def is_untrusted_node(self, _node: dict[str, Any], _context: CapabilityContext | None = None) -> bool:
+            return True
+
+    service.install(classifier=_Classifier())
+
+    with pytest.raises(RuntimeError, match="no untrusted executor is configured"):
+        service.route(graph=None, user_id="u1", flow_id="f1", run_id="r1")