From 9c49248f0f84438befb37ee34b9308931137f9a2 Mon Sep 17 00:00:00 2001 From: "balogh.adam@icloud.com" Date: Tue, 31 Mar 2026 20:24:30 -0400 Subject: [PATCH 1/2] select random tee --- src/opengradient/client/tee_registry.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/opengradient/client/tee_registry.py b/src/opengradient/client/tee_registry.py index 571e371..9e6c802 100644 --- a/src/opengradient/client/tee_registry.py +++ b/src/opengradient/client/tee_registry.py @@ -1,6 +1,7 @@ """TEE Registry client for fetching verified TEE endpoints and TLS certificates.""" import logging +import random import ssl from dataclasses import dataclass from typing import List, NamedTuple, Optional @@ -109,7 +110,7 @@ def get_active_tees_by_type(self, tee_type: int) -> List[TEEEndpoint]: def get_llm_tee(self) -> Optional[TEEEndpoint]: """ - Return the first active LLM proxy TEE from the registry. + Return a random active LLM proxy TEE from the registry. Returns: TEEEndpoint for an active LLM proxy TEE, or None if none are available. @@ -119,7 +120,7 @@ def get_llm_tee(self) -> Optional[TEEEndpoint]: logger.warning("No active LLM proxy TEEs found in registry") return None - return tees[0] + return random.choice(tees) def build_ssl_context_from_der(der_cert: bytes) -> ssl.SSLContext: From e59e6a8ba4ce727ae21eee867ca1a24e91b2735f Mon Sep 17 00:00:00 2001 From: "balogh.adam@icloud.com" Date: Tue, 31 Mar 2026 20:25:30 -0400 Subject: [PATCH 2/2] update test --- tests/tee_registry_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/tee_registry_test.py b/tests/tee_registry_test.py index faa1420..fa37e2c 100644 --- a/tests/tee_registry_test.py +++ b/tests/tee_registry_test.py @@ -144,7 +144,7 @@ def test_validator_type_label(self, mock_contract): class TestGetLlmTee: - def test_returns_first_active_tee(self, mock_contract): + def test_returns_random_active_tee(self, mock_contract): registry, contract = mock_contract contract.functions.getActiveTEEs.return_value.call.return_value = [ @@ -155,7 +155,7 @@ def test_returns_first_active_tee(self, mock_contract): result = registry.get_llm_tee() assert result is not None - assert result.endpoint == "https://tee-1.example.com" + assert result.endpoint in ("https://tee-1.example.com", "https://tee-2.example.com") def test_returns_none_when_no_tees(self, mock_contract): registry, contract = mock_contract