From 977c32ee768c50f39eb056322b0b5d89884cc62d Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Mon, 24 Aug 2026 15:10:09 +0200 Subject: [PATCH 01/18] refactor(framework): Simplify run-started extension hook --- framework/py/flwr/superlink/extensions.py | 38 ++++++++++- .../py/flwr/superlink/extensions_test.py | 63 +++++++++++++++++++ .../servicer/control/control_handlers.py | 7 ++- .../servicer/control/control_handlers_test.py | 38 +++++++++++ 4 files changed, 143 insertions(+), 3 deletions(-) create mode 100644 framework/py/flwr/superlink/extensions_test.py diff --git a/framework/py/flwr/superlink/extensions.py b/framework/py/flwr/superlink/extensions.py index b051a9a7953a..49d09c16ba13 100644 --- a/framework/py/flwr/superlink/extensions.py +++ b/framework/py/flwr/superlink/extensions.py @@ -14,19 +14,24 @@ # ============================================================================== """SuperLink FastAPI extension hooks.""" - from collections.abc import Callable, Mapping from contextlib import AbstractAsyncContextManager +from copy import deepcopy from importlib import import_module +from logging import WARNING from types import ModuleType -from typing import Any, cast +from typing import Any, Literal, cast from fastapi import FastAPI from starlette.middleware import Middleware +from flwr.common.logger import log +from flwr.supercore.run import Run + SuperLinkLifespanContext = Callable[ [FastAPI], AbstractAsyncContextManager[Mapping[str, Any] | None] ] +RunStartSource = Literal["cli", "web_ui", "automation", "unknown"] _SGXT_MODULE = "flwr.ee.superlink.extensions" @@ -87,3 +92,32 @@ def get_lifespan_contexts() -> tuple[SuperLinkLifespanContext, ...]: if get_sgxt_lifespan_contexts is None: return () return get_sgxt_lifespan_contexts() + + +def notify_run_started(run: Run, source: RunStartSource) -> None: + """Notify an optional extension after a run has been persisted. + + The callback is synchronous by design. Extensions must keep this hook + 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. + """ + try: + sgxt = _try_import_sgxt() + if sgxt is None: + return + + on_run_started = cast( + Callable[[Run, RunStartSource], None] | None, + getattr(sgxt, "on_run_started", None), + ) + if on_run_started is not None: + on_run_started(deepcopy(run), source) + except Exception as exc: # pylint: disable=broad-exception-caught + log( + WARNING, + "Run-start extension notification failed: %s.", + type(exc).__name__, + exc_info=exc, + ) diff --git a/framework/py/flwr/superlink/extensions_test.py b/framework/py/flwr/superlink/extensions_test.py new file mode 100644 index 000000000000..93e7a447c60c --- /dev/null +++ b/framework/py/flwr/superlink/extensions_test.py @@ -0,0 +1,63 @@ +# 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 SuperLink extension notifications.""" + +from types import ModuleType +from unittest.mock import Mock + +from pytest import MonkeyPatch + +from flwr.supercore.run import Run + +from . import extensions + + +def test_notify_run_started_passes_a_snapshot_to_the_extension( + monkeypatch: MonkeyPatch, +) -> None: + """Pass a copy of the persisted run to the optional extension.""" + callback = Mock() + module = ModuleType("flwr.ee.superlink.extensions") + module.on_run_started = callback # type: ignore[attr-defined] + monkeypatch.setattr(extensions, "_try_import_sgxt", lambda: module) + run = Run.create_empty(123) + + extensions.notify_run_started(run, "unknown") + + callback.assert_called_once() + notified_run, source = callback.call_args.args + assert notified_run == run + assert notified_run is not run + assert source == "unknown" + + +def test_notify_run_started_skips_missing_extension(monkeypatch: MonkeyPatch) -> None: + """Do nothing when the optional extension package is absent.""" + monkeypatch.setattr(extensions, "_try_import_sgxt", lambda: None) + + extensions.notify_run_started(Run.create_empty(123), "unknown") + + +def test_notify_run_started_isolates_extension_import_failure( + monkeypatch: MonkeyPatch, +) -> None: + """Keep a persisted run successful when extension discovery fails.""" + + def fail_import() -> ModuleType | None: + raise RuntimeError("extension import failed") + + monkeypatch.setattr(extensions, "_try_import_sgxt", fail_import) + + extensions.notify_run_started(Run.create_empty(123), "unknown") diff --git a/framework/py/flwr/superlink/servicer/control/control_handlers.py b/framework/py/flwr/superlink/servicer/control/control_handlers.py index b1272089fb30..af7219a2def6 100644 --- a/framework/py/flwr/superlink/servicer/control/control_handlers.py +++ b/framework/py/flwr/superlink/servicer/control/control_handlers.py @@ -164,6 +164,7 @@ resolve_account_ids, strict_json_dumps, ) +from flwr.superlink import extensions from flwr.superlink.artifact_provider import ArtifactProvider from flwr.superlink.auth_plugin import ControlAuthnPlugin from flwr.superlink.federation.noop_federation_manager import NoOpFederationManager @@ -466,6 +467,7 @@ 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", ) -> StartRunResponse: """Create run ID.""" log(INFO, "ControlServicer.StartRun") @@ -651,9 +653,11 @@ def start_run( # pylint: disable=too-many-branches,too-many-locals,too-many-sta log_msg = f"Created run {run_id} in federation {run.federation_id}" log(INFO, log_msg) - return StartRunResponse( + response = StartRunResponse( run_id=run_id, note=note, series_id=series_id, federation=run.federation_id ) + extensions.notify_run_started(run, source) + return response def stream_logs( @@ -901,6 +905,7 @@ def dispatch_automation( AccountInfo(flwr_aid=flwr_aid, account_name=""), state, None, + source="automation", ) except Exception as exc: # pylint: disable=broad-exception-caught state.finish_automation( 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 9f873713c899..01344ba482c3 100644 --- a/framework/py/flwr/superlink/servicer/control/control_handlers_test.py +++ b/framework/py/flwr/superlink/servicer/control/control_handlers_test.py @@ -108,6 +108,44 @@ def test_start_run_reuses_fab_by_hash(self) -> None: [("@flwr/demo", fab_hash, TaskType.SERVER_APP)], ) + def test_start_run_notifies_extension_after_persisting_run(self) -> None: + """Notify the optional extension with the persisted run snapshot.""" + fab_content = b"stored FAB" + fab_hash = hashlib.sha256(fab_content).hexdigest() + self.state.store_app( + fab=Fab(fab_hash, fab_content, {}), + federation_id=NOOP_FEDERATION_ID, + app_id="@flwr/demo", + app_type=TaskType.SERVER_APP, + added_by=self.account.flwr_aid, + ) + + with ( + patch( + "flwr.superlink.servicer.control.control_handlers.get_fab_config", + return_value={"tool": {"flwr": {"app": {}}}}, + ), + patch( + "flwr.superlink.servicer.control.control_handlers" + ".get_metadata_from_config", + return_value=("flwr/demo", "v0.0.1"), + ), + patch( + "flwr.superlink.servicer.control.control_handlers" + ".extensions.notify_run_started" + ) as notify_run_started, + ): + 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) + + 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") + def test_start_run_rejects_unknown_fab_hash(self) -> None: """Test StartRun rejects an unknown FAB hash without an app spec.""" request = StartRunRequest(federation=NOOP_FEDERATION_ID) From 0a3209d37e23c9477c2052862edb02ed28383705 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Tue, 25 Aug 2026 13:54:38 +0200 Subject: [PATCH 02/18] feat(framework): Propagate run source attribution --- framework/py/flwr/superlink/extensions.py | 27 +++++++++++- .../py/flwr/superlink/extensions_test.py | 19 ++++++++- .../flwr/superlink/routers/control/router.py | 14 ++++++- .../superlink/routers/control/router_test.py | 42 ++++++++++++++++++- .../servicer/control/control_handlers.py | 1 + .../servicer/control/control_handlers_test.py | 10 ++++- .../servicer/control/control_servicer.py | 23 +++++++++- .../servicer/control/control_servicer_test.py | 33 +++++++++++++++ 8 files changed, 161 insertions(+), 8 deletions(-) diff --git a/framework/py/flwr/superlink/extensions.py b/framework/py/flwr/superlink/extensions.py index 49d09c16ba13..5e8e22c9bd03 100644 --- a/framework/py/flwr/superlink/extensions.py +++ b/framework/py/flwr/superlink/extensions.py @@ -32,9 +32,33 @@ [FastAPI], AbstractAsyncContextManager[Mapping[str, Any] | None] ] RunStartSource = Literal["cli", "web_ui", "automation", "unknown"] +RUN_SOURCE_METADATA_KEY = "x-flwr-run-source" +_RUN_START_SOURCES = frozenset({"cli", "web_ui", "automation", "unknown"}) _SGXT_MODULE = "flwr.ee.superlink.extensions" +def resolve_run_start_source( + value: str | bytes | None, *, default: RunStartSource +) -> 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 default + if isinstance(value, bytes): + try: + value = value.decode("ascii") + except UnicodeDecodeError: + return "unknown" + if value not in _RUN_START_SOURCES: + return "unknown" + return cast(RunStartSource, value) + + def _try_import_sgxt() -> ModuleType | None: """Return the SuperGrid Extensions module when it is installed.""" try: @@ -101,7 +125,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/extensions_test.py b/framework/py/flwr/superlink/extensions_test.py index 93e7a447c60c..c4b29f3803fd 100644 --- a/framework/py/flwr/superlink/extensions_test.py +++ b/framework/py/flwr/superlink/extensions_test.py @@ -17,13 +17,30 @@ from types import ModuleType from unittest.mock import Mock -from pytest import MonkeyPatch +from pytest import MonkeyPatch, mark from flwr.supercore.run import Run from . import extensions +@mark.parametrize( + ("value", "default", "expected"), + [ + (None, "cli", "cli"), + ("web_ui", "cli", "web_ui"), + (b"automation", "cli", "automation"), + ("not-a-source", "cli", "unknown"), + (b"\\xff", "cli", "unknown"), + ], +) +def test_resolve_run_start_source( + value: str | bytes | None, default: extensions.RunStartSource, expected: str +) -> None: + """Normalize best-effort caller attribution without treating it as auth.""" + assert extensions.resolve_run_start_source(value, default=default) == expected + + def test_notify_run_started_passes_a_snapshot_to_the_extension( monkeypatch: MonkeyPatch, ) -> None: diff --git a/framework/py/flwr/superlink/routers/control/router.py b/framework/py/flwr/superlink/routers/control/router.py index 660eaded0177..c91dfa3aba1c 100644 --- a/framework/py/flwr/superlink/routers/control/router.py +++ b/framework/py/flwr/superlink/routers/control/router.py @@ -16,7 +16,7 @@ from typing import Annotated -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, Header from flwr.proto.control_pb2 import ( # pylint: disable=E0611 AcceptInvitationRequest, @@ -78,6 +78,7 @@ from flwr.supercore.auth.typing import AccountInfo from flwr.supercore.protobuf.routing import ProtobufRoute from flwr.supercore.protobuf.translation import get_protobuf_request +from flwr.superlink import extensions from flwr.superlink.dependencies.account import get_account from flwr.superlink.dependencies.linkstate import get_linkstate from flwr.superlink.servicer.control import control_handlers @@ -93,10 +94,19 @@ def start_run( request: Annotated[StartRunRequest, Depends(get_protobuf_request)], linkstate: LinkStateDependency, account: AccountDependency, + run_source: Annotated[ + str | None, Header(alias=extensions.RUN_SOURCE_METADATA_KEY) + ] = None, ) -> 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=extensions.resolve_run_start_source(run_source, default="web_ui"), + ) @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..c221a778210b 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_defaults_to_best_effort_web_ui_source() -> None: + """Treat the HTTP control route as a Web UI producer by default.""" + 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) + + assert response is expected + start_run.assert_called_once_with( + request, + _ACCOUNT, + linkstate, + "", + source="web_ui", + ) + + +def test_start_run_accepts_best_effort_source_header() -> 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/servicer/control/control_handlers.py b/framework/py/flwr/superlink/servicer/control/control_handlers.py index af7219a2def6..4ed7eb5cb7d9 100644 --- a/framework/py/flwr/superlink/servicer/control/control_handlers.py +++ b/framework/py/flwr/superlink/servicer/control/control_handlers.py @@ -467,6 +467,7 @@ 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", ) -> StartRunResponse: """Create run ID.""" 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 01344ba482c3..d8d3a0d0b6e8 100644 --- a/framework/py/flwr/superlink/servicer/control/control_handlers_test.py +++ b/framework/py/flwr/superlink/servicer/control/control_handlers_test.py @@ -138,13 +138,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..1de630a39dc1 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer.py @@ -97,6 +97,7 @@ from flwr.supercore.auth.typing import AccountInfo from flwr.supercore.error import ApiErrorCode, FlowerError from flwr.supercore.object_store import ObjectStoreFactory +from flwr.superlink import extensions from flwr.superlink.artifact_provider import ArtifactProvider from flwr.superlink.auth_plugin import ControlAuthnPlugin @@ -127,8 +128,28 @@ def StartRun( self, request: StartRunRequest, context: grpc.ServicerContext ) -> StartRunResponse: """Create run ID.""" + metadata = context.invocation_metadata() + # This is best-effort analytics attribution, not authentication. We + # trust callers to label their own requests and default direct gRPC + # callers to the CLI source when no label is provided. + run_source = ( + next( + ( + value + for key, value in metadata + if key == extensions.RUN_SOURCE_METADATA_KEY + ), + None, + ) + if isinstance(metadata, (tuple, list)) + else None + ) 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=extensions.resolve_run_start_source(run_source, default="cli"), ) 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 444abb1af58f..3dfa64b3ae62 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, @@ -100,6 +101,7 @@ RegisterSupernodeContext, StartRunContext, ) +from flwr.superlink import extensions from flwr.superlink.auth_plugin import NoOpControlAuthnPlugin from flwr.superlink.federation import NoOpFederationManager from flwr.superlink.servicer.control.control_account_auth_interceptor import ( @@ -412,6 +414,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 = ( + (extensions.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_cli_source_without_metadata(self) -> None: + """Default direct gRPC callers to the CLI analytics source.""" + 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"], "cli") + def test_start_run_validates_and_binds_oauth_connectors(self) -> None: """StartRun should bind canonical connected OAuth connector refs.""" flow = _OAuthFlow() From d4aa97d2c56710e37c29b3f74afe4b9777b2f1b3 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Tue, 25 Aug 2026 14:41:38 +0200 Subject: [PATCH 03/18] fix(framework): Avoid inferring HTTP run source --- framework/py/flwr/superlink/routers/control/router.py | 2 +- framework/py/flwr/superlink/routers/control/router_test.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/framework/py/flwr/superlink/routers/control/router.py b/framework/py/flwr/superlink/routers/control/router.py index c91dfa3aba1c..1a5ec574bec8 100644 --- a/framework/py/flwr/superlink/routers/control/router.py +++ b/framework/py/flwr/superlink/routers/control/router.py @@ -105,7 +105,7 @@ def start_run( account, linkstate, "", - source=extensions.resolve_run_start_source(run_source, default="web_ui"), + source=extensions.resolve_run_start_source(run_source, default="unknown"), ) diff --git a/framework/py/flwr/superlink/routers/control/router_test.py b/framework/py/flwr/superlink/routers/control/router_test.py index c221a778210b..bf2192314732 100644 --- a/framework/py/flwr/superlink/routers/control/router_test.py +++ b/framework/py/flwr/superlink/routers/control/router_test.py @@ -78,8 +78,8 @@ def test_all_control_routes_have_protobuf_request_types() -> None: assert route_keys == control_request_types -def test_start_run_defaults_to_best_effort_web_ui_source() -> None: - """Treat the HTTP control route as a Web UI producer by default.""" +def test_start_run_defaults_to_unknown_source_without_header() -> None: + """Avoid inferring a source from the HTTP transport.""" request = StartRunRequest() linkstate = Mock() expected = StartRunResponse(run_id=1) @@ -96,7 +96,7 @@ def test_start_run_defaults_to_best_effort_web_ui_source() -> None: _ACCOUNT, linkstate, "", - source="web_ui", + source="unknown", ) From 83c3bf26a497a73d14526671e70125059cac4bc7 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Tue, 25 Aug 2026 14:49:51 +0200 Subject: [PATCH 04/18] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- framework/py/flwr/superlink/extensions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/framework/py/flwr/superlink/extensions.py b/framework/py/flwr/superlink/extensions.py index 59a8500d5f99..94b99d8d2ab7 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 From 6153c404496bb5c358b155321058775d27c21a2b Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Tue, 25 Aug 2026 14:51:06 +0200 Subject: [PATCH 05/18] fix(framework): Derive run sources from literal --- framework/py/flwr/superlink/extensions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/framework/py/flwr/superlink/extensions.py b/framework/py/flwr/superlink/extensions.py index 59a8500d5f99..bfc8813896c7 100644 --- a/framework/py/flwr/superlink/extensions.py +++ b/framework/py/flwr/superlink/extensions.py @@ -20,7 +20,7 @@ from importlib import import_module from logging import WARNING from types import ModuleType -from typing import Any, Literal, cast +from typing import Any, Literal, cast, get_args from fastapi import FastAPI from starlette.middleware import Middleware @@ -34,7 +34,7 @@ RunStartSource = Literal["cli", "web_ui", "automation", "unknown"] ResultDeliveryChannel = Literal["logs", "chat"] RUN_SOURCE_METADATA_KEY = "x-flwr-run-source" -_RUN_START_SOURCES = frozenset({"cli", "web_ui", "automation", "unknown"}) +_RUN_START_SOURCES = frozenset(get_args(RunStartSource)) _SGXT_MODULE = "flwr.ee.superlink.extensions" From 03eee2e623b8ca7f4ef3bd44a2e090c723603148 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Tue, 25 Aug 2026 15:07:00 +0200 Subject: [PATCH 06/18] fix(framework): Do not infer gRPC run source --- .../py/flwr/superlink/servicer/control/control_servicer.py | 6 +++--- .../superlink/servicer/control/control_servicer_test.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/framework/py/flwr/superlink/servicer/control/control_servicer.py b/framework/py/flwr/superlink/servicer/control/control_servicer.py index 1de630a39dc1..6f89709d7db2 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer.py @@ -130,8 +130,8 @@ def StartRun( """Create run ID.""" metadata = context.invocation_metadata() # This is best-effort analytics attribution, not authentication. We - # trust callers to label their own requests and default direct gRPC - # callers to the CLI source when no label is provided. + # trust callers to label their own requests and default unlabeled + # requests to unknown when no source is provided. run_source = ( next( ( @@ -149,7 +149,7 @@ def StartRun( _get_account(), self.linkstate_factory.state(), self.fleet_api_type, - source=extensions.resolve_run_start_source(run_source, default="cli"), + source=extensions.resolve_run_start_source(run_source, default="unknown"), ) 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 c7c0f9773165..85371bc9499e 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer_test.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer_test.py @@ -431,8 +431,8 @@ def test_start_run_forwards_best_effort_source_metadata(self) -> None: self.assertIs(response, expected) self.assertEqual(start_run.call_args.kwargs["source"], "web_ui") - def test_start_run_defaults_to_cli_source_without_metadata(self) -> None: - """Default direct gRPC callers to the CLI analytics source.""" + 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) @@ -443,7 +443,7 @@ def test_start_run_defaults_to_cli_source_without_metadata(self) -> None: ) as start_run: self.servicer.StartRun(StartRunRequest(), context) - self.assertEqual(start_run.call_args.kwargs["source"], "cli") + 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.""" From 6be54607a2f3a1654a62a960a4c46cd44dda9951 Mon Sep 17 00:00:00 2001 From: Heng Pan Date: Tue, 25 Aug 2026 14:31:10 +0100 Subject: [PATCH 07/18] Apply suggestion from @panh99 --- .../superlink/servicer/control/control_servicer.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/framework/py/flwr/superlink/servicer/control/control_servicer.py b/framework/py/flwr/superlink/servicer/control/control_servicer.py index 6f89709d7db2..d476b5f38cac 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer.py @@ -132,18 +132,7 @@ def StartRun( # This is best-effort analytics attribution, not authentication. We # trust callers to label their own requests and default unlabeled # requests to unknown when no source is provided. - run_source = ( - next( - ( - value - for key, value in metadata - if key == extensions.RUN_SOURCE_METADATA_KEY - ), - None, - ) - if isinstance(metadata, (tuple, list)) - else None - ) + run_source = dict(metadata).get(extensions.RUN_SOURCE_METADATA_KEY) return control_handlers.start_run( request, _get_account(), From 57146213fcd13aee1318d8e52f2fffacb80e5611 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Tue, 25 Aug 2026 19:00:48 +0200 Subject: [PATCH 08/18] Simplify run source resolution fallback --- framework/py/flwr/superlink/extensions.py | 6 ++---- framework/py/flwr/superlink/extensions_test.py | 16 ++++++++-------- .../py/flwr/superlink/routers/control/router.py | 2 +- .../servicer/control/control_servicer.py | 15 +++++++++++++-- 4 files changed, 24 insertions(+), 15 deletions(-) diff --git a/framework/py/flwr/superlink/extensions.py b/framework/py/flwr/superlink/extensions.py index ebe264954c5e..0f19cf613e92 100644 --- a/framework/py/flwr/superlink/extensions.py +++ b/framework/py/flwr/superlink/extensions.py @@ -38,9 +38,7 @@ _SGXT_MODULE = "flwr.ee.superlink.extensions" -def resolve_run_start_source( - value: str | bytes | None, *, default: RunStartSource -) -> RunStartSource: +def resolve_run_start_source(value: str | bytes | None) -> RunStartSource: """Normalize a caller-provided source label for analytics. Source attribution is intentionally best effort. Callers can only affect @@ -49,7 +47,7 @@ def resolve_run_start_source( security or authorization signal. """ if value is None: - return default + return "unknown" if isinstance(value, bytes): try: value = value.decode("ascii") diff --git a/framework/py/flwr/superlink/extensions_test.py b/framework/py/flwr/superlink/extensions_test.py index 6265ba1f6bf9..fea92c53b269 100644 --- a/framework/py/flwr/superlink/extensions_test.py +++ b/framework/py/flwr/superlink/extensions_test.py @@ -25,20 +25,20 @@ @mark.parametrize( - ("value", "default", "expected"), + ("value", "expected"), [ - (None, "cli", "cli"), - ("web_ui", "cli", "web_ui"), - (b"automation", "cli", "automation"), - ("not-a-source", "cli", "unknown"), - (b"\\xff", "cli", "unknown"), + (None, "unknown"), + ("web_ui", "web_ui"), + (b"automation", "automation"), + ("not-a-source", "unknown"), + (b"\\xff", "unknown"), ], ) def test_resolve_run_start_source( - value: str | bytes | None, default: extensions.RunStartSource, expected: str + value: str | bytes | None, expected: extensions.RunStartSource ) -> None: """Normalize best-effort caller attribution without treating it as auth.""" - assert extensions.resolve_run_start_source(value, default=default) == expected + assert extensions.resolve_run_start_source(value) == expected def test_notify_run_started_passes_a_snapshot_to_the_extension( diff --git a/framework/py/flwr/superlink/routers/control/router.py b/framework/py/flwr/superlink/routers/control/router.py index 1a5ec574bec8..0489bb0ccdc1 100644 --- a/framework/py/flwr/superlink/routers/control/router.py +++ b/framework/py/flwr/superlink/routers/control/router.py @@ -105,7 +105,7 @@ def start_run( account, linkstate, "", - source=extensions.resolve_run_start_source(run_source, default="unknown"), + source=extensions.resolve_run_start_source(run_source), ) diff --git a/framework/py/flwr/superlink/servicer/control/control_servicer.py b/framework/py/flwr/superlink/servicer/control/control_servicer.py index d476b5f38cac..357fa6e13ff0 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer.py @@ -132,13 +132,24 @@ def StartRun( # This is best-effort analytics attribution, not authentication. We # trust callers to label their own requests and default unlabeled # requests to unknown when no source is provided. - run_source = dict(metadata).get(extensions.RUN_SOURCE_METADATA_KEY) + run_source = ( + next( + ( + value + for key, value in metadata + if key == extensions.RUN_SOURCE_METADATA_KEY + ), + None, + ) + if isinstance(metadata, (tuple, list)) + else None + ) return control_handlers.start_run( request, _get_account(), self.linkstate_factory.state(), self.fleet_api_type, - source=extensions.resolve_run_start_source(run_source, default="unknown"), + source=extensions.resolve_run_start_source(run_source), ) def StreamLogs( # pylint: disable=C0103 From c1900f563e6b947744bc37d1cdc2a4f42e26d563 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Tue, 25 Aug 2026 19:02:11 +0200 Subject: [PATCH 09/18] Extract run source request dependency --- .../flwr/superlink/dependencies/run_source.py | 35 +++++++++++++++++++ .../superlink/dependencies/run_source_test.py | 27 ++++++++++++++ .../flwr/superlink/routers/control/router.py | 10 +++--- .../superlink/routers/control/router_test.py | 8 ++--- 4 files changed, 70 insertions(+), 10 deletions(-) create mode 100644 framework/py/flwr/superlink/dependencies/run_source.py create mode 100644 framework/py/flwr/superlink/dependencies/run_source_test.py 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..cd4a918513c5 --- /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.extensions 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/routers/control/router.py b/framework/py/flwr/superlink/routers/control/router.py index 0489bb0ccdc1..1b4572dc44f1 100644 --- a/framework/py/flwr/superlink/routers/control/router.py +++ b/framework/py/flwr/superlink/routers/control/router.py @@ -16,7 +16,7 @@ from typing import Annotated -from fastapi import APIRouter, Depends, Header +from fastapi import APIRouter, Depends from flwr.proto.control_pb2 import ( # pylint: disable=E0611 AcceptInvitationRequest, @@ -78,9 +78,9 @@ from flwr.supercore.auth.typing import AccountInfo from flwr.supercore.protobuf.routing import ProtobufRoute from flwr.supercore.protobuf.translation import get_protobuf_request -from flwr.superlink import extensions 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) @@ -94,9 +94,7 @@ def start_run( request: Annotated[StartRunRequest, Depends(get_protobuf_request)], linkstate: LinkStateDependency, account: AccountDependency, - run_source: Annotated[ - str | None, Header(alias=extensions.RUN_SOURCE_METADATA_KEY) - ] = None, + run_source: RunSourceDependency, ) -> StartRunResponse: """Start a run.""" # Temporary: pass an empty Fleet API type @@ -105,7 +103,7 @@ def start_run( account, linkstate, "", - source=extensions.resolve_run_start_source(run_source), + source=run_source, ) diff --git a/framework/py/flwr/superlink/routers/control/router_test.py b/framework/py/flwr/superlink/routers/control/router_test.py index bf2192314732..2ce5bb7b56a1 100644 --- a/framework/py/flwr/superlink/routers/control/router_test.py +++ b/framework/py/flwr/superlink/routers/control/router_test.py @@ -78,8 +78,8 @@ def test_all_control_routes_have_protobuf_request_types() -> None: assert route_keys == control_request_types -def test_start_run_defaults_to_unknown_source_without_header() -> None: - """Avoid inferring a source from the HTTP transport.""" +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) @@ -88,7 +88,7 @@ def test_start_run_defaults_to_unknown_source_without_header() -> None: "flwr.superlink.routers.control.router.control_handlers.start_run", return_value=expected, ) as start_run: - response = start_run_route(request, linkstate, _ACCOUNT) + response = start_run_route(request, linkstate, _ACCOUNT, "unknown") assert response is expected start_run.assert_called_once_with( @@ -100,7 +100,7 @@ def test_start_run_defaults_to_unknown_source_without_header() -> None: ) -def test_start_run_accepts_best_effort_source_header() -> None: +def test_start_run_forwards_caller_provided_source() -> None: """Forward a caller-provided analytics source without treating it as auth.""" request = StartRunRequest() linkstate = Mock() From d9fed81e553c5b98fd1e191a56932ef5eddfd841 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Wed, 26 Aug 2026 10:48:22 +0200 Subject: [PATCH 10/18] Update framework/py/flwr/superlink/servicer/control/control_servicer.py Co-authored-by: Heng Pan --- framework/py/flwr/superlink/servicer/control/control_servicer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/framework/py/flwr/superlink/servicer/control/control_servicer.py b/framework/py/flwr/superlink/servicer/control/control_servicer.py index 357fa6e13ff0..e65d5dc4ad26 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer.py @@ -97,6 +97,7 @@ 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 import extensions from flwr.superlink.artifact_provider import ArtifactProvider from flwr.superlink.auth_plugin import ControlAuthnPlugin From 2db4dfb901643c4cf2f5efd3da333ccea3772d0f Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Wed, 26 Aug 2026 10:48:29 +0200 Subject: [PATCH 11/18] Update framework/py/flwr/superlink/servicer/control/control_servicer.py Co-authored-by: Heng Pan --- .../servicer/control/control_servicer.py | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/framework/py/flwr/superlink/servicer/control/control_servicer.py b/framework/py/flwr/superlink/servicer/control/control_servicer.py index e65d5dc4ad26..fde6e56317b0 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer.py @@ -129,22 +129,11 @@ 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() - # This is best-effort analytics attribution, not authentication. We - # trust callers to label their own requests and default unlabeled - # requests to unknown when no source is provided. - run_source = ( - next( - ( - value - for key, value in metadata - if key == extensions.RUN_SOURCE_METADATA_KEY - ), - None, - ) - if isinstance(metadata, (tuple, list)) - else None - ) + run_source = get_metadata_str(metadata, extensions.RUN_SOURCE_METADATA_KEY) + return control_handlers.start_run( request, _get_account(), From 20353c4b88f7163a550af802e6714e977c344b6e Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Wed, 26 Aug 2026 11:08:24 +0200 Subject: [PATCH 12/18] Refactor run-source attribution utilities --- .../flwr/superlink/dependencies/run_source.py | 2 +- framework/py/flwr/superlink/extensions.py | 26 +---------- .../py/flwr/superlink/extensions_test.py | 19 +------- framework/py/flwr/superlink/run_source.py | 43 +++++++++++++++++++ .../py/flwr/superlink/run_source_test.py | 36 ++++++++++++++++ .../servicer/control/control_handlers.py | 3 +- .../servicer/control/control_servicer.py | 11 +++-- .../servicer/control/control_servicer_test.py | 4 +- 8 files changed, 94 insertions(+), 50 deletions(-) create mode 100644 framework/py/flwr/superlink/run_source.py create mode 100644 framework/py/flwr/superlink/run_source_test.py diff --git a/framework/py/flwr/superlink/dependencies/run_source.py b/framework/py/flwr/superlink/dependencies/run_source.py index cd4a918513c5..ecc2bf977475 100644 --- a/framework/py/flwr/superlink/dependencies/run_source.py +++ b/framework/py/flwr/superlink/dependencies/run_source.py @@ -18,7 +18,7 @@ from fastapi import Depends, Header -from flwr.superlink.extensions import ( +from flwr.superlink.run_source import ( RUN_SOURCE_METADATA_KEY, RunStartSource, resolve_run_start_source, diff --git a/framework/py/flwr/superlink/extensions.py b/framework/py/flwr/superlink/extensions.py index 0f19cf613e92..09e29d898854 100644 --- a/framework/py/flwr/superlink/extensions.py +++ b/framework/py/flwr/superlink/extensions.py @@ -20,44 +20,22 @@ from importlib import import_module from logging import WARNING from types import ModuleType -from typing import Any, Literal, cast, get_args +from typing import Any, Literal, cast from fastapi import FastAPI from starlette.middleware import Middleware 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"] -RUN_SOURCE_METADATA_KEY = "x-flwr-run-source" -_RUN_START_SOURCES = frozenset(get_args(RunStartSource)) _SGXT_MODULE = "flwr.ee.superlink.extensions" -def resolve_run_start_source(value: str | bytes | 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 isinstance(value, bytes): - try: - value = value.decode("ascii") - except UnicodeDecodeError: - return "unknown" - if value not in _RUN_START_SOURCES: - return "unknown" - return cast(RunStartSource, value) - - def _try_import_sgxt() -> ModuleType | None: """Return the SuperGrid Extensions module when it is installed.""" try: diff --git a/framework/py/flwr/superlink/extensions_test.py b/framework/py/flwr/superlink/extensions_test.py index fea92c53b269..cf2aefe0f25e 100644 --- a/framework/py/flwr/superlink/extensions_test.py +++ b/framework/py/flwr/superlink/extensions_test.py @@ -17,30 +17,13 @@ from types import ModuleType from unittest.mock import Mock -from pytest import MonkeyPatch, mark +from pytest import MonkeyPatch from flwr.supercore.run import Run from . import extensions -@mark.parametrize( - ("value", "expected"), - [ - (None, "unknown"), - ("web_ui", "web_ui"), - (b"automation", "automation"), - ("not-a-source", "unknown"), - (b"\\xff", "unknown"), - ], -) -def test_resolve_run_start_source( - value: str | bytes | None, expected: extensions.RunStartSource -) -> None: - """Normalize best-effort caller attribution without treating it as auth.""" - assert extensions.resolve_run_start_source(value) == expected - - def test_notify_run_started_passes_a_snapshot_to_the_extension( monkeypatch: MonkeyPatch, ) -> None: diff --git a/framework/py/flwr/superlink/run_source.py b/framework/py/flwr/superlink/run_source.py new file mode 100644 index 000000000000..3be2062062ba --- /dev/null +++ b/framework/py/flwr/superlink/run_source.py @@ -0,0 +1,43 @@ +# 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 | bytes | 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. The bytes form keeps the normalizer + tolerant of callers passing raw gRPC metadata values; only ASCII labels + are accepted. + """ + if value is None: + return "unknown" + if isinstance(value, bytes): + try: + value = value.decode("ascii") + except UnicodeDecodeError: + 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..edd4b33d72cc --- /dev/null +++ b/framework/py/flwr/superlink/run_source_test.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. +# ============================================================================== +"""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"), + (b"automation", "automation"), + ("not-a-source", "unknown"), + (b"\\xff", "unknown"), + ], +) +def test_resolve_run_start_source( + value: str | bytes | 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 896c24cb0d4e..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): @@ -468,7 +469,7 @@ def start_run( # pylint: disable=too-many-branches,too-many-locals,too-many-sta 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_servicer.py b/framework/py/flwr/superlink/servicer/control/control_servicer.py index fde6e56317b0..5f1f37b492ee 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer.py @@ -98,9 +98,12 @@ 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 import extensions 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 @@ -129,17 +132,17 @@ def StartRun( self, request: StartRunRequest, context: grpc.ServicerContext ) -> StartRunResponse: """Create run ID.""" - # Best-effort analytics attribution only; + # 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, extensions.RUN_SOURCE_METADATA_KEY) + 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, - source=extensions.resolve_run_start_source(run_source), + 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 85371bc9499e..557bb732e1c3 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer_test.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer_test.py @@ -101,9 +101,9 @@ RegisterSupernodeContext, StartRunContext, ) -from flwr.superlink import extensions 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, ) @@ -418,7 +418,7 @@ 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 = ( - (extensions.RUN_SOURCE_METADATA_KEY, "web_ui"), + (RUN_SOURCE_METADATA_KEY, "web_ui"), ) expected = StartRunResponse(run_id=42) From f01b57c8faf636b7d45a542863cb3344cdb02b99 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Wed, 26 Aug 2026 11:13:14 +0200 Subject: [PATCH 13/18] Fix run-source import ordering --- .../py/flwr/superlink/servicer/control/control_servicer.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/framework/py/flwr/superlink/servicer/control/control_servicer.py b/framework/py/flwr/superlink/servicer/control/control_servicer.py index 5f1f37b492ee..4a3a86967328 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer.py @@ -100,10 +100,7 @@ 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 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 From 7864abc9a74d6a09b8fd3ca94b94c82a2c61a88d Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Wed, 26 Aug 2026 11:17:19 +0200 Subject: [PATCH 14/18] Configure StartRun test metadata --- .../servicer/control/control_servicer_test.py | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) 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 557bb732e1c3..e9af98ea81fe 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer_test.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer_test.py @@ -197,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", @@ -394,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 @@ -476,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)), @@ -523,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( @@ -564,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()), []) @@ -578,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) @@ -611,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) @@ -664,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() @@ -712,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() @@ -792,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 @@ -837,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 @@ -950,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, From a340d11e19974ac505f6bdaa70ba3bee217a4835 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Wed, 26 Aug 2026 11:21:12 +0200 Subject: [PATCH 15/18] Configure remaining StartRun test contexts --- .../superlink/servicer/control/control_servicer_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 e9af98ea81fe..aaaaa2e8d79c 100644 --- a/framework/py/flwr/superlink/servicer/control/control_servicer_test.py +++ b/framework/py/flwr/superlink/servicer/control/control_servicer_test.py @@ -749,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( @@ -863,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 ( @@ -890,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( From 6554dcecf8a76224a383451f031a00c62a7158e7 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Wed, 26 Aug 2026 11:39:01 +0200 Subject: [PATCH 16/18] Bound E2E service cleanup --- framework/e2e/test_control_api.sh | 3 +++ 1 file changed, 3 insertions(+) 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 } From 073d50b84b6d5a73c70d378a51f3840f79eafbd8 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Wed, 26 Aug 2026 11:55:32 +0200 Subject: [PATCH 17/18] Narrow run-source input to strings --- framework/py/flwr/superlink/run_source.py | 11 ++--------- framework/py/flwr/superlink/run_source_test.py | 5 ++--- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/framework/py/flwr/superlink/run_source.py b/framework/py/flwr/superlink/run_source.py index 3be2062062ba..5c8450817a38 100644 --- a/framework/py/flwr/superlink/run_source.py +++ b/framework/py/flwr/superlink/run_source.py @@ -21,23 +21,16 @@ _RUN_START_SOURCES = frozenset(get_args(RunStartSource)) -def resolve_run_start_source(value: str | bytes | None) -> 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. The bytes form keeps the normalizer - tolerant of callers passing raw gRPC metadata values; only ASCII labels - are accepted. + security or authorization signal. """ if value is None: return "unknown" - if isinstance(value, bytes): - try: - value = value.decode("ascii") - except UnicodeDecodeError: - 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 index edd4b33d72cc..5244d9179db2 100644 --- a/framework/py/flwr/superlink/run_source_test.py +++ b/framework/py/flwr/superlink/run_source_test.py @@ -24,13 +24,12 @@ [ (None, "unknown"), ("web_ui", "web_ui"), - (b"automation", "automation"), + ("automation", "automation"), ("not-a-source", "unknown"), - (b"\\xff", "unknown"), ], ) def test_resolve_run_start_source( - value: str | bytes | None, expected: RunStartSource + value: str | None, expected: RunStartSource ) -> None: """Normalize best-effort caller attribution without treating it as auth.""" assert resolve_run_start_source(value) == expected From 7f635e50a3f181bcc7c85c8a76f023fba7d01c0b Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Wed, 26 Aug 2026 11:58:09 +0200 Subject: [PATCH 18/18] Format run-source test signature --- framework/py/flwr/superlink/run_source_test.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/framework/py/flwr/superlink/run_source_test.py b/framework/py/flwr/superlink/run_source_test.py index 5244d9179db2..1edff1c80342 100644 --- a/framework/py/flwr/superlink/run_source_test.py +++ b/framework/py/flwr/superlink/run_source_test.py @@ -28,8 +28,6 @@ ("not-a-source", "unknown"), ], ) -def test_resolve_run_start_source( - value: str | None, expected: RunStartSource -) -> None: +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