From 91ef11e05078868539f5161af1140a678a9033f1 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Fri, 23 Oct 2026 22:31:29 -0700
Subject: [PATCH] Return bool from is_pathways_backend_used and validate proxy
 connection timeout.

PiperOrigin-RevId: 1006105622
---
 pathwaysutils/_initialize.py             |  4 +++-
 pathwaysutils/proxy_backend.py           | 14 +++++++++++++-
 pathwaysutils/test/initialize_test.py    |  6 ++++--
 pathwaysutils/test/proxy_backend_test.py | 12 ++++++++++++
 4 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/pathwaysutils/_initialize.py b/pathwaysutils/_initialize.py
index de7bb46..7a133eb 100644
--- a/pathwaysutils/_initialize.py
+++ b/pathwaysutils/_initialize.py
@@ -40,7 +40,9 @@ def is_pathways_backend_used() -> bool:
   is not necessarily which backend is ultimately selected or there may be more
   than one platform specified and another may have higher priority.
   """
-  return jax.config.jax_platforms and "proxy" in jax.config.jax_platforms
+  return bool(
+      jax.config.jax_platforms and "proxy" in jax.config.jax_platforms
+  )
 
 
 def _is_persistence_enabled() -> bool:
diff --git a/pathwaysutils/proxy_backend.py b/pathwaysutils/proxy_backend.py
index 2507a8a..0f01e98 100644
--- a/pathwaysutils/proxy_backend.py
+++ b/pathwaysutils/proxy_backend.py
@@ -26,7 +26,19 @@ def make_client():
     options = ifrt_proxy.ClientConnectionOptions()
     timeout_secs = os.environ.get("PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS")
     if timeout_secs:
-      options.connection_timeout_in_seconds = int(timeout_secs)
+      try:
+        timeout_val = int(timeout_secs)
+      except ValueError as e:
+        raise ValueError(
+            "PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS must be an integer, got:"
+            f" {timeout_secs}"
+        ) from e
+      if timeout_val <= 0:
+        raise ValueError(
+            "PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS must be positive, got:"
+            f" {timeout_secs}"
+        )
+      options.connection_timeout_in_seconds = timeout_val
     return ifrt_proxy.get_client(
         jax.config.read("jax_backend_target"),
         options,
diff --git a/pathwaysutils/test/initialize_test.py b/pathwaysutils/test/initialize_test.py
index 9228a49..e2e33f1 100644
--- a/pathwaysutils/test/initialize_test.py
+++ b/pathwaysutils/test/initialize_test.py
@@ -68,15 +68,17 @@ def test_initialize_more_than_once(self, initialization_count):
     )
 
   @parameterized.named_parameters(
+      ("none", None),
       ("empty", ""),
       ("cpu", "cpu"),
       ("tpu", "tpu"),
       ("gpu", "gpu"),
       ("cpu,tpu,gpu", "cpu,tpu,gpu"),
   )
-  def test_not_is_pathways_backend_used(self, platform: str):
+  def test_not_is_pathways_backend_used(self, platform: str | None):
     jax.config.update("jax_platforms", platform)
-    self.assertFalse(_initialize.is_pathways_backend_used())
+    result = _initialize.is_pathways_backend_used()
+    self.assertIs(result, False)
 
   @parameterized.named_parameters(
       ("proxy", "proxy"),
diff --git a/pathwaysutils/test/proxy_backend_test.py b/pathwaysutils/test/proxy_backend_test.py
index 13fce79..2c53010 100644
--- a/pathwaysutils/test/proxy_backend_test.py
+++ b/pathwaysutils/test/proxy_backend_test.py
@@ -95,6 +95,18 @@ def test_proxy_backend_registration_without_timeout(self):
     options = args[1]
     self.assertNotEqual(options.connection_timeout_in_seconds, 42)
 
+  def test_proxy_backend_registration_with_invalid_timeout(self):
+    self.enter_context(
+        mock.patch.dict(
+            os.environ, {"PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS": "-5"}
+        )
+    )
+    proxy_backend.register_backend_factory()
+    with self.assertRaisesRegex(
+        RuntimeError, "PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS must be positive"
+    ):
+      backend.backends()
+
 
 if __name__ == "__main__":
   absltest.main()