fix(framework): Reuse Runtime CA for AgentApp SDKs

framework/py/flwr/supercore/cli/flwr_agentapp.py

@@ -16,7 +16,9 @@
 
 
 import argparse
+import os
 from logging import DEBUG, INFO
+from pathlib import Path
 from queue import Queue
 
 from flwr.common.args import add_args_flwr_app_common, try_obtain_flwr_app_token
@@ -42,14 +44,20 @@ def flwr_agentapp() -> None:
         "`flwr-agentapp` will attempt to connect to SuperLink's Runtime API at %s",
         args.runtime_api_address,
     )
+    certificates = validate_and_resolve_root_certificates(
+        args.root_certificates, args.insecure
+    )
+    # Set the custom CA path for OpenAI SDK clients based on httpx
+    if args.root_certificates is not None:
+        os.environ["SSL_CERT_FILE"] = str(
+            Path(args.root_certificates).expanduser().resolve()
+        )
     run_agentapp(
         runtime_api_address=args.runtime_api_address,
         log_queue=log_queue,
         token=token,
         insecure=args.insecure,
-        certificates=validate_and_resolve_root_certificates(
-            args.root_certificates, args.insecure
-        ),
+        certificates=certificates,
         parent_pid=args.parent_pid,
         runtime_dependency_install=args.runtime_dependency_install,
     )

framework/py/flwr/supercore/cli/flwr_agentapp_test.py

@@ -16,6 +16,8 @@
 
 
 import importlib
+import os
+from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
@@ -128,3 +130,37 @@ def parse_args(self) -> SimpleNamespace:
     assert kwargs["certificates"] is None
     assert kwargs["parent_pid"] == 321
     assert kwargs["runtime_dependency_install"] is True
+
+
+def test_flwr_agentapp_exposes_explicit_root_certificates(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    """Expose the Runtime root certificate file to SDK clients."""
+    monkeypatch.setenv("SSL_CERT_FILE", "inherited-ca.pem")
+    certificate_path = tmp_path / "runtime-ca.pem"
+    certificate_path.write_bytes(b"root-certificates")
+    monkeypatch.chdir(tmp_path)
+    args = SimpleNamespace(
+        insecure=False,
+        runtime_api_address="runtime.example:9092",
+        token="test-token",
+        root_certificates=certificate_path.name,
+        parent_pid=None,
+        runtime_dependency_install=False,
+    )
+
+    class _Parser:
+        def parse_args(self) -> SimpleNamespace:
+            """Return a fixed namespace with explicit root certificates."""
+            return args
+
+    with (
+        patch.object(flwr_agentapp_module, "_parse_args_run_flwr_agentapp", _Parser),
+        patch.object(flwr_agentapp_module, "mirror_output_to_queue"),
+        patch.object(flwr_agentapp_module, "restore_output"),
+        patch.object(flwr_agentapp_module, "run_agentapp") as run_agentapp,
+    ):
+        flwr_agentapp_module.flwr_agentapp()
+
+    assert os.environ["SSL_CERT_FILE"] == str(certificate_path.resolve())
+    assert run_agentapp.call_args.kwargs["certificates"] == b"root-certificates"

framework/py/flwr/supercore/task_process/agent/run_agentapp.py

@@ -16,11 +16,9 @@
 
 
 import os
-import ssl
 from logging import DEBUG, ERROR
 from pathlib import Path
 from queue import Queue
-from tempfile import NamedTemporaryFile
 
 import httpx
 
@@ -71,8 +69,6 @@
 _AGENT_INPUT_KEY = "agent.input"
 _RUNTIME_API_KEY_ENV = "FLWR_RUNTIME_API_KEY"
 _RUNTIME_BASE_URL_ENV = "FLWR_RUNTIME_BASE_URL"
-_SSL_CERT_FILE_ENV = "SSL_CERT_FILE"
-_SSL_CERT_DIR_ENV = "SSL_CERT_DIR"
 
 
 def run_agentapp(  # pylint: disable=R0912, R0913, R0914, R0915, R0917, W0212
@@ -104,7 +100,6 @@ def run_agentapp(  # pylint: disable=R0912, R0913, R0914, R0915, R0917, W0212
     heartbeat_sender = None
     context: Context | None = None
     runtime_env_dir: Path | None = None
-    runtime_root_certificates_path: Path | None = None
     exit_code = ExitCode.SUCCESS
 
     def on_exit() -> None:
@@ -134,8 +129,6 @@ def on_exit() -> None:
         grid.close()
 
         cleanup_app_runtime_environment(runtime_env_dir)
-        if runtime_root_certificates_path is not None:
-            runtime_root_certificates_path.unlink(missing_ok=True)
 
     register_signal_handlers(
         event_type=EventType.FLWR_AGENTAPP_RUN_LEAVE,
@@ -247,9 +240,7 @@ def on_exit() -> None:
         connectors = RuntimeAgentConnectors(responses)
         agent = RuntimeAgentSession(responses=responses, connectors=connectors)
 
-        runtime_root_certificates_path = _set_runtime_environment(
-            runtime_api_address, token, insecure, certificates
-        )
+        _set_runtime_environment(runtime_api_address, token, insecure)
 
         # Load and run the AgentApp
         agent_app = load_app(agent_app_attr, LoadAgentAppError, app_path)
@@ -291,61 +282,9 @@ def _set_runtime_environment(
     runtime_api_address: str,
     token: str,
     insecure: bool,
-    certificates: bytes | None = None,
-) -> Path | None:
+) -> None:
     """Expose the Open Responses-compatible Runtime endpoint to the AgentApp."""
     scheme = "http" if insecure else "https"
     address = runtime_api_address.rstrip("/")
     os.environ[_RUNTIME_BASE_URL_ENV] = f"{scheme}://{address}/v1/runtime"
     os.environ[_RUNTIME_API_KEY_ENV] = token
-
-    if certificates is None:
-        return None
-
-    # OpenAI/httpx reads custom CAs from SSL_CERT_FILE, which requires a path.
-    # Extend the public and inherited roots with the Runtime CA for this process.
-    public_ca_context = httpx.create_ssl_context(trust_env=False)
-    trusted_certificates = [
-        *public_ca_context.get_ca_certs(binary_form=True),
-        *_load_inherited_ca_certificates(),
-    ]
-    with NamedTemporaryFile(
-        mode="wb", prefix="flwr-runtime-ca-", suffix=".pem", delete=False
-    ) as certificate_file:
-        for trusted_certificate in dict.fromkeys(trusted_certificates):
-            certificate_file.write(
-                ssl.DER_cert_to_PEM_cert(trusted_certificate).encode("ascii")
-            )
-        certificate_file.write(b"\n")
-        certificate_file.write(certificates)
-    certificate_path = Path(certificate_file.name)
-    os.environ[_SSL_CERT_FILE_ENV] = str(certificate_path)
-    return certificate_path
-
-
-def _load_inherited_ca_certificates() -> list[bytes]:
-    """Load roots configured through the standard OpenSSL environment."""
-    ca_file = os.environ.get(_SSL_CERT_FILE_ENV)
-    ca_dir = os.environ.get(_SSL_CERT_DIR_ENV)
-    if not ca_file and not ca_dir:
-        return []
-
-    context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
-    if ca_file:
-        context.load_verify_locations(cafile=ca_file)
-    if ca_dir:
-        for directory in ca_dir.split(os.pathsep):
-            if not directory:
-                continue
-            try:
-                ca_paths = sorted(Path(directory).iterdir())
-            except OSError:
-                continue
-            for ca_path in ca_paths:
-                if not ca_path.is_file():
-                    continue
-                try:
-                    context.load_verify_locations(cafile=ca_path)
-                except OSError:
-                    continue
-    return context.get_ca_certs(binary_form=True)

framework/py/flwr/supercore/task_process/agent/run_agentapp_test.py

@@ -15,11 +15,7 @@
 """Tests for the AgentApp process environment."""
 
 import os
-import ssl
-from pathlib import Path
-from unittest.mock import Mock
 
-import httpx
 import pytest
 
 from .run_agentapp import _set_runtime_environment
@@ -32,90 +28,9 @@ def test_set_runtime_environment(
     """Expose the Runtime Responses base URL and AgentApp task token."""
     monkeypatch.delenv("FLWR_RUNTIME_BASE_URL", raising=False)
     monkeypatch.delenv("FLWR_RUNTIME_API_KEY", raising=False)
-    monkeypatch.delenv("SSL_CERT_FILE", raising=False)
-    certificate_path = _set_runtime_environment(
-        "runtime.example:9092", "task-token", insecure=insecure
-    )
+    _set_runtime_environment("runtime.example:9092", "task-token", insecure=insecure)
 
     assert os.environ["FLWR_RUNTIME_BASE_URL"] == (
         f"{scheme}://runtime.example:9092/v1/runtime"
     )
     assert os.environ["FLWR_RUNTIME_API_KEY"] == "task-token"
-    assert certificate_path is None
-    assert "SSL_CERT_FILE" not in os.environ
-
-
-def test_set_runtime_environment_exposes_root_certificates(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    """Add custom Runtime root certificates to the public trust bundle."""
-    monkeypatch.delenv("SSL_CERT_FILE", raising=False)
-    monkeypatch.delenv("SSL_CERT_DIR", raising=False)
-
-    certificate_path = _set_runtime_environment(
-        "runtime.example:9092",
-        "task-token",
-        insecure=False,
-        certificates=b"root-certificates",
-    )
-
-    assert certificate_path is not None
-    try:
-        assert os.environ["SSL_CERT_FILE"] == str(certificate_path)
-        certificate_bundle = certificate_path.read_bytes()
-        assert certificate_bundle.startswith(b"-----BEGIN CERTIFICATE-----")
-        assert certificate_bundle.endswith(b"\nroot-certificates")
-    finally:
-        certificate_path.unlink(missing_ok=True)
-
-
-@pytest.mark.parametrize("ca_env", ["SSL_CERT_FILE", "SSL_CERT_DIR"])
-def test_set_runtime_environment_preserves_inherited_root_certificates(
-    monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ca_env: str
-) -> None:
-    """Keep inherited custom roots when adding the Runtime root certificate."""
-    inherited_certificates = httpx.create_ssl_context(trust_env=False).get_ca_certs(
-        binary_form=True
-    )[:2]
-    inherited_certificates_pem = [
-        ssl.DER_cert_to_PEM_cert(certificate).encode("ascii")
-        for certificate in inherited_certificates
-    ]
-    if ca_env == "SSL_CERT_FILE":
-        inherited_ca_path = tmp_path / "inherited-ca.pem"
-        inherited_ca_path.write_bytes(inherited_certificates_pem[0])
-        ca_env_value = str(inherited_ca_path)
-        expected_certificates = inherited_certificates_pem[:1]
-    else:
-        inherited_ca_directories = []
-        for index, certificate in enumerate(inherited_certificates_pem):
-            inherited_ca_directory = tmp_path / f"inherited-cas-{index}"
-            inherited_ca_directory.mkdir()
-            (inherited_ca_directory / "inherited-ca.pem").write_bytes(certificate)
-            inherited_ca_directories.append(str(inherited_ca_directory))
-        ca_env_value = os.pathsep.join(inherited_ca_directories)
-        expected_certificates = inherited_certificates_pem
-    monkeypatch.setenv(ca_env, ca_env_value)
-    monkeypatch.delenv(
-        "SSL_CERT_DIR" if ca_env == "SSL_CERT_FILE" else "SSL_CERT_FILE",
-        raising=False,
-    )
-    public_ca_context = Mock()
-    public_ca_context.get_ca_certs.return_value = []
-    monkeypatch.setattr(httpx, "create_ssl_context", lambda **_: public_ca_context)
-
-    certificate_path = _set_runtime_environment(
-        "runtime.example:9092",
-        "task-token",
-        insecure=False,
-        certificates=b"runtime-root-certificate",
-    )
-
-    assert certificate_path is not None
-    try:
-        certificate_bundle = certificate_path.read_bytes()
-        for certificate in expected_certificates:
-            assert certificate in certificate_bundle
-        assert certificate_bundle.endswith(b"\nruntime-root-certificate")
-    finally:
-        certificate_path.unlink(missing_ok=True)