diff --git a/framework/e2e/test_control_api.sh b/framework/e2e/test_control_api.sh index 125c49ed690f..be54588b041f 100755 --- a/framework/e2e/test_control_api.sh +++ b/framework/e2e/test_control_api.sh @@ -77,6 +77,9 @@ background_pids=() cleanup() { if [ "${#background_pids[@]}" -gt 0 ]; then kill "${background_pids[@]}" 2>/dev/null || true + sleep 1 + # SuperNodes can keep retrying after the SuperLink exits. + kill -KILL "${background_pids[@]}" 2>/dev/null || true wait "${background_pids[@]}" 2>/dev/null || true fi } diff --git a/framework/py/flwr/superlink/dependencies/run_source.py b/framework/py/flwr/superlink/dependencies/run_source.py new file mode 100644 index 000000000000..ecc2bf977475 --- /dev/null +++ b/framework/py/flwr/superlink/dependencies/run_source.py @@ -0,0 +1,35 @@ +# Copyright 2026 Flower Labs GmbH. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""FastAPI dependency for Control API run-source attribution.""" + +from typing import Annotated + +from fastapi import Depends, Header + +from flwr.superlink.run_source import ( + RUN_SOURCE_METADATA_KEY, + RunStartSource, + resolve_run_start_source, +) + + +def get_run_source( + run_source: Annotated[str | None, Header(alias=RUN_SOURCE_METADATA_KEY)] = None, +) -> RunStartSource: + """Return the normalized run source from the request header.""" + return resolve_run_start_source(run_source) + + +RunSourceDependency = Annotated[RunStartSource, Depends(get_run_source)] diff --git a/framework/py/flwr/superlink/dependencies/run_source_test.py b/framework/py/flwr/superlink/dependencies/run_source_test.py new file mode 100644 index 000000000000..15570f778aa9 --- /dev/null +++ b/framework/py/flwr/superlink/dependencies/run_source_test.py @@ -0,0 +1,27 @@ +# Copyright 2026 Flower Labs GmbH. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for the Control API run-source dependency.""" + +from .run_source import get_run_source + + +def test_get_run_source_defaults_to_unknown() -> None: + """Return unknown when the request has no source header.""" + assert get_run_source() == "unknown" + + +def test_get_run_source_normalizes_header_value() -> None: + """Normalize a caller-provided source header.""" + assert get_run_source("web_ui") == "web_ui" diff --git a/framework/py/flwr/superlink/extensions.py b/framework/py/flwr/superlink/extensions.py index 60f957fb2f1e..09e29d898854 100644 --- a/framework/py/flwr/superlink/extensions.py +++ b/framework/py/flwr/superlink/extensions.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""SuperLink FastAPI extension hooks.""" +"""SuperLink extension hooks.""" from collections.abc import Callable, Mapping from contextlib import AbstractAsyncContextManager @@ -27,11 +27,11 @@ from flwr.common.logger import log from flwr.supercore.run import Run +from flwr.superlink.run_source import RunStartSource SuperLinkLifespanContext = Callable[ [FastAPI], AbstractAsyncContextManager[Mapping[str, Any] | None] ] -RunStartSource = Literal["cli", "web_ui", "automation", "unknown"] ResultDeliveryChannel = Literal["logs", "chat"] _SGXT_MODULE = "flwr.ee.superlink.extensions" @@ -102,7 +102,8 @@ def notify_run_started(run: Run, source: RunStartSource) -> None: non-blocking and best effort; the Flower framework does not create a background thread or event loop for it. The run snapshot is copied before handing it to the extension so the callback cannot mutate the object used - to build the successful StartRun response. + to build the successful StartRun response. The source is also best-effort + caller attribution and must not be used for authorization decisions. """ try: sgxt = _try_import_sgxt() diff --git a/framework/py/flwr/superlink/routers/control/router.py b/framework/py/flwr/superlink/routers/control/router.py index 660eaded0177..1b4572dc44f1 100644 --- a/framework/py/flwr/superlink/routers/control/router.py +++ b/framework/py/flwr/superlink/routers/control/router.py @@ -80,6 +80,7 @@ from flwr.supercore.protobuf.translation import get_protobuf_request from flwr.superlink.dependencies.account import get_account from flwr.superlink.dependencies.linkstate import get_linkstate +from flwr.superlink.dependencies.run_source import RunSourceDependency from flwr.superlink.servicer.control import control_handlers router = APIRouter(prefix="/v1/control", tags=["Control"], route_class=ProtobufRoute) @@ -93,10 +94,17 @@ def start_run( request: Annotated[StartRunRequest, Depends(get_protobuf_request)], linkstate: LinkStateDependency, account: AccountDependency, + run_source: RunSourceDependency, ) -> StartRunResponse: """Start a run.""" # Temporary: pass an empty Fleet API type - return control_handlers.start_run(request, account, linkstate, "") + return control_handlers.start_run( + request, + account, + linkstate, + "", + source=run_source, + ) @router.post("/list-runs") diff --git a/framework/py/flwr/superlink/routers/control/router_test.py b/framework/py/flwr/superlink/routers/control/router_test.py index e98f359d6d3a..2ce5bb7b56a1 100644 --- a/framework/py/flwr/superlink/routers/control/router_test.py +++ b/framework/py/flwr/superlink/routers/control/router_test.py @@ -16,7 +16,7 @@ from datetime import datetime -from unittest.mock import Mock +from unittest.mock import Mock, patch from fastapi import FastAPI, HTTPException, Request, Response, status from fastapi.routing import APIRoute @@ -26,6 +26,8 @@ from flwr.proto.control_pb2 import ( # pylint: disable=E0611 ListRunsRequest, ListRunsResponse, + StartRunRequest, + StartRunResponse, ) from flwr.server.superlink.linkstate import LinkState from flwr.supercore.auth.typing import AccountInfo @@ -41,6 +43,7 @@ from flwr.superlink.dependencies.linkstate import get_linkstate from flwr.superlink.routers.control.middlewares import ControlAuthenticationMiddleware from flwr.superlink.routers.control.router import router +from flwr.superlink.routers.control.router import start_run as start_run_route _ACCOUNT = AccountInfo(flwr_aid=NOOP_FLWR_AID, account_name="account") @@ -75,6 +78,43 @@ def test_all_control_routes_have_protobuf_request_types() -> None: assert route_keys == control_request_types +def test_start_run_forwards_resolved_source() -> None: + """Forward the normalized source to the control handler.""" + request = StartRunRequest() + linkstate = Mock() + expected = StartRunResponse(run_id=1) + + with patch( + "flwr.superlink.routers.control.router.control_handlers.start_run", + return_value=expected, + ) as start_run: + response = start_run_route(request, linkstate, _ACCOUNT, "unknown") + + assert response is expected + start_run.assert_called_once_with( + request, + _ACCOUNT, + linkstate, + "", + source="unknown", + ) + + +def test_start_run_forwards_caller_provided_source() -> None: + """Forward a caller-provided analytics source without treating it as auth.""" + request = StartRunRequest() + linkstate = Mock() + expected = StartRunResponse(run_id=1) + + with patch( + "flwr.superlink.routers.control.router.control_handlers.start_run", + return_value=expected, + ) as start_run: + start_run_route(request, linkstate, _ACCOUNT, run_source="cli") + + assert start_run.call_args.kwargs["source"] == "cli" + + def test_protobuf_request_without_handler_response_returns_internal_error() -> None: """A configured protobuf route must store its handler response in request state.""" app = FastAPI() diff --git a/framework/py/flwr/superlink/run_source.py b/framework/py/flwr/superlink/run_source.py new file mode 100644 index 000000000000..5c8450817a38 --- /dev/null +++ b/framework/py/flwr/superlink/run_source.py @@ -0,0 +1,36 @@ +# Copyright 2026 Flower Labs GmbH. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Run-start source attribution utilities.""" + +from typing import Literal, cast, get_args + +RunStartSource = Literal["cli", "web_ui", "automation", "unknown"] +RUN_SOURCE_METADATA_KEY = "x-flwr-run-source" +_RUN_START_SOURCES = frozenset(get_args(RunStartSource)) + + +def resolve_run_start_source(value: str | None) -> RunStartSource: + """Normalize a caller-provided source label for analytics. + + Source attribution is intentionally best effort. Callers can only affect + the analytics label for their own request, so recognized values are + trusted and invalid values fall back to ``unknown``. This value is not a + security or authorization signal. + """ + if value is None: + return "unknown" + if value not in _RUN_START_SOURCES: + return "unknown" + return cast(RunStartSource, value) diff --git a/framework/py/flwr/superlink/run_source_test.py b/framework/py/flwr/superlink/run_source_test.py new file mode 100644 index 000000000000..1edff1c80342 --- /dev/null +++ b/framework/py/flwr/superlink/run_source_test.py @@ -0,0 +1,33 @@ +# Copyright 2026 Flower Labs GmbH. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for run-start source attribution utilities.""" + +from pytest import mark + +from .run_source import RunStartSource, resolve_run_start_source + + +@mark.parametrize( + ("value", "expected"), + [ + (None, "unknown"), + ("web_ui", "web_ui"), + ("automation", "automation"), + ("not-a-source", "unknown"), + ], +) +def test_resolve_run_start_source(value: str | None, expected: RunStartSource) -> None: + """Normalize best-effort caller attribution without treating it as auth.""" + assert resolve_run_start_source(value) == expected diff --git a/framework/py/flwr/superlink/servicer/control/control_handlers.py b/framework/py/flwr/superlink/servicer/control/control_handlers.py index d3c6eda7cd8b..34733f0f3434 100644 --- a/framework/py/flwr/superlink/servicer/control/control_handlers.py +++ b/framework/py/flwr/superlink/servicer/control/control_handlers.py @@ -168,6 +168,7 @@ from flwr.superlink.artifact_provider import ArtifactProvider from flwr.superlink.auth_plugin import ControlAuthnPlugin from flwr.superlink.federation.noop_federation_manager import NoOpFederationManager +from flwr.superlink.run_source import RunStartSource class InvalidConnectorRequestError(FlowerError): @@ -467,7 +468,8 @@ def start_run( # pylint: disable=too-many-branches,too-many-locals,too-many-sta account: AccountInfo, state: LinkState, fleet_api_type: str | None, - source: extensions.RunStartSource = "unknown", + *, + source: RunStartSource = "unknown", ) -> StartRunResponse: """Create run ID.""" log(INFO, "ControlServicer.StartRun") diff --git a/framework/py/flwr/superlink/servicer/control/control_handlers_test.py b/framework/py/flwr/superlink/servicer/control/control_handlers_test.py index a876edb735db..f7100acf684a 100644 --- a/framework/py/flwr/superlink/servicer/control/control_handlers_test.py +++ b/framework/py/flwr/superlink/servicer/control/control_handlers_test.py @@ -140,13 +140,19 @@ def test_start_run_notifies_extension_after_persisting_run(self) -> None: request = StartRunRequest(federation=NOOP_FEDERATION_ID) request.app_spec = "@flwr/demo==0.0.1" request.fab.hash_str = fab_hash - response = start_run(request, self.account, self.state, None) + response = start_run( + request, + self.account, + self.state, + None, + source="web_ui", + ) run = self.state.get_run_info(run_ids=[response.run_id])[0] notify_run_started.assert_called_once() notified_run, source = notify_run_started.call_args.args self.assertEqual(notified_run.run_id, run.run_id) - self.assertEqual(source, "unknown") + self.assertEqual(source, "web_ui") def test_start_run_rejects_unknown_fab_hash(self) -> None: """Test StartRun rejects an unknown FAB hash without an app spec.""" diff --git a/framework/py/flwr/superlink/servicer/control/control_servicer.py b/framework/py/flwr/superlink/servicer/control/control_servicer.py index 0d3705daebf5..4a3a86967328 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer.py @@ -97,8 +97,10 @@ from flwr.supercore.auth.typing import AccountInfo from flwr.supercore.error import ApiErrorCode, FlowerError from flwr.supercore.object_store import ObjectStoreFactory +from flwr.supercore.utils import get_metadata_str from flwr.superlink.artifact_provider import ArtifactProvider from flwr.superlink.auth_plugin import ControlAuthnPlugin +from flwr.superlink.run_source import RUN_SOURCE_METADATA_KEY, resolve_run_start_source from . import control_handlers from .control_account_auth_interceptor import get_current_account_info @@ -127,8 +129,17 @@ def StartRun( self, request: StartRunRequest, context: grpc.ServicerContext ) -> StartRunResponse: """Create run ID.""" + # Best-effort analytics attribution only; + # trust caller-provided labels and default missing sources to unknown. + metadata = context.invocation_metadata() + run_source = get_metadata_str(metadata, RUN_SOURCE_METADATA_KEY) + return control_handlers.start_run( - request, _get_account(), self.linkstate_factory.state(), self.fleet_api_type + request, + _get_account(), + self.linkstate_factory.state(), + self.fleet_api_type, + source=resolve_run_start_source(run_source), ) def StreamLogs( # pylint: disable=C0103 diff --git a/framework/py/flwr/superlink/servicer/control/control_servicer_test.py b/framework/py/flwr/superlink/servicer/control/control_servicer_test.py index cdc9c7f11505..aaaaa2e8d79c 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer_test.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer_test.py @@ -67,6 +67,7 @@ ShowFederationRequest, ShowFederationResponse, StartRunRequest, + StartRunResponse, StopRunRequest, StreamLogsRequest, StreamLogsResponse, @@ -102,6 +103,7 @@ ) from flwr.superlink.auth_plugin import NoOpControlAuthnPlugin from flwr.superlink.federation import NoOpFederationManager +from flwr.superlink.run_source import RUN_SOURCE_METADATA_KEY from flwr.superlink.servicer.control.control_account_auth_interceptor import ( shared_account_info, ) @@ -195,6 +197,12 @@ def setUp(self) -> None: shared_account_info.set(account_info) self.state = self.servicer.linkstate_factory.state() + def _make_start_run_context(self) -> MagicMock: + """Create a gRPC context with empty invocation metadata.""" + context = MagicMock(spec=grpc.ServicerContext) + context.invocation_metadata.return_value = () + return context + def _create_dummy_run(self, flwr_aid: str | None) -> int: return self.state.create_run( "flwr/demo", @@ -392,7 +400,7 @@ def test_start_run(self) -> None: ) as mock_get_metadata_from_config, ): mock_get_metadata_from_config.return_value = (fab_id, fab_version) - response = self.servicer.StartRun(request, Mock()) + response = self.servicer.StartRun(request, self._make_start_run_context()) runs = self.state.get_run_info(run_ids=[response.run_id]) run_info = runs[0] if runs else None @@ -412,6 +420,37 @@ def test_start_run(self) -> None: self.assertEqual(run_context.run_id, response.run_id) self.assertEqual(run_context.series_id, response.series_id) + def test_start_run_forwards_best_effort_source_metadata(self) -> None: + """Forward caller-provided source metadata for analytics attribution.""" + context = Mock() + context.invocation_metadata.return_value = ( + (RUN_SOURCE_METADATA_KEY, "web_ui"), + ) + expected = StartRunResponse(run_id=42) + + with patch( + "flwr.superlink.servicer.control.control_servicer.control_handlers.start_run", + return_value=expected, + ) as start_run: + response = self.servicer.StartRun(StartRunRequest(), context) + + self.assertIs(response, expected) + self.assertEqual(start_run.call_args.kwargs["source"], "web_ui") + + def test_start_run_defaults_to_unknown_source_without_metadata(self) -> None: + """Do not infer a caller from the gRPC transport.""" + context = Mock() + context.invocation_metadata.return_value = () + expected = StartRunResponse(run_id=42) + + with patch( + "flwr.superlink.servicer.control.control_servicer.control_handlers.start_run", + return_value=expected, + ) as start_run: + self.servicer.StartRun(StartRunRequest(), context) + + self.assertEqual(start_run.call_args.kwargs["source"], "unknown") + def test_start_run_validates_and_binds_oauth_connectors(self) -> None: """StartRun should bind canonical connected OAuth connector refs.""" flow = _OAuthFlow() @@ -443,7 +482,7 @@ def test_start_run_validates_and_binds_oauth_connectors(self) -> None: return_value=("flwr/demo", "1.0.0"), ), ): - response = self.servicer.StartRun(request, Mock()) + response = self.servicer.StartRun(request, self._make_start_run_context()) self.assertEqual( list(self.state.get_run_connector_refs(run_id=response.run_id)), @@ -490,7 +529,7 @@ def test_start_run_rejects_connectors_for_capable_federation( ), self.assertRaises(FlowerError) as error, ): - self.servicer.StartRun(request, Mock()) + self.servicer.StartRun(request, self._make_start_run_context()) self.assertEqual(error.exception.code, ApiErrorCode.INVALID_CONNECTOR_REQUEST) self.assertEqual( @@ -531,7 +570,7 @@ def test_start_run_rejects_unavailable_oauth_connector( ), self.assertRaises(FlowerError) as error, ): - self.servicer.StartRun(request, Mock()) + self.servicer.StartRun(request, self._make_start_run_context()) self.assertEqual(error.exception.code, expected_code) self.assertEqual(list(self.state.get_run_info()), []) @@ -545,7 +584,7 @@ def test_start_run_defaults_to_account_simulation_federation(self) -> None: self.servicer.linkstate_factory.state_instance = None with self.assertRaises(RuntimeError): - self.servicer.StartRun(StartRunRequest(), Mock()) + self.servicer.StartRun(StartRunRequest(), self._make_start_run_context()) federation_manager.exists.assert_called_once_with(expected_federation_id) @@ -578,7 +617,7 @@ def test_start_run_uses_existing_series_id(self) -> None: ): mock_get_fab_config.return_value = {"tool": {"flwr": {"app": {}}}} mock_get_metadata_from_config.return_value = ("flwr/demo", "v1.0.0") - response = self.servicer.StartRun(request, Mock()) + response = self.servicer.StartRun(request, self._make_start_run_context()) run = self.state.get_run_info(run_ids=[response.run_id])[0] run_context = self.state.get_run_series_context(series_id) @@ -631,7 +670,7 @@ def test_start_run_creates_task_with_matching_type( "tool": {"flwr": {"app": {"config": {"train": {"lr": 0.1}}}}} } mock_get_metadata_from_config.return_value = ("flwr/demo", "v1.0.0") - response = self.servicer.StartRun(request, Mock()) + response = self.servicer.StartRun(request, self._make_start_run_context()) runs = self.state.get_run_info(run_ids=[response.run_id]) tasks = self.state.get_tasks() @@ -679,7 +718,7 @@ def test_start_run_creates_agentapp_run_from_local_fab(self) -> None: } } mock_get_metadata_from_config.return_value = ("flwr/agent", "0.1.0") - response = self.servicer.StartRun(request, Mock()) + response = self.servicer.StartRun(request, self._make_start_run_context()) runs = self.state.get_run_info(run_ids=[response.run_id]) tasks = self.state.get_tasks() @@ -710,7 +749,7 @@ def test_start_run_raises_if_create_run_fails(self) -> None: request.fab.hash_str = hashlib.sha256(fab_content).hexdigest() request.fab.content = fab_content request.federation = NOOP_FEDERATION_ID - context = Mock() + context = self._make_start_run_context() with ( patch( @@ -759,7 +798,7 @@ def test_start_run_returns_note_for_remote_app(self) -> None: "anne-dev/simple-legacy-127", "0.1.0", ) - response = self.servicer.StartRun(request, Mock()) + response = self.servicer.StartRun(request, self._make_start_run_context()) assert response.HasField("note") assert response.note @@ -804,7 +843,7 @@ def test_start_run_accepts_valid_nested_override_keys(self) -> None: } } mock_get_metadata_from_config.return_value = ("flwr/demo", "v1.0.0") - response = self.servicer.StartRun(request, Mock()) + response = self.servicer.StartRun(request, self._make_start_run_context()) runs = self.state.get_run_info(run_ids=[response.run_id]) run_info = runs[0] if runs else None @@ -824,7 +863,7 @@ def test_start_run_rejects_unknown_override_keys(self) -> None: request.federation = NOOP_FEDERATION_ID for key, value in user_config_to_proto({"unknown.key": 10}).items(): request.override_config[key].CopyFrom(value) - context = Mock() + context = self._make_start_run_context() # Execute/Assert with ( @@ -851,7 +890,7 @@ def test_start_run_denied_when_not_entitled(self) -> None: request.fab.content = b"test FAB content" request.federation = NOOP_FEDERATION_ID - context = Mock() + context = self._make_start_run_context() with ( patch( @@ -917,7 +956,7 @@ def test_start_run_calls_can_execute_with_expected_args( "tool": {"flwr": {"app": {"config": {"train": {"lr": 0.1}}}}} } mock_get_metadata_from_config.return_value = ("flwr/demo", "v1.0.0") - _ = self.servicer.StartRun(request, Mock()) + _ = self.servicer.StartRun(request, self._make_start_run_context()) mock_can_execute.assert_called_once_with( self.aid,