feat: support region

Author: fioncat

tests/test_connection_config.py

@@ -3,9 +3,10 @@
 
 def test_api_url_defaults_correctly(monkeypatch):
     monkeypatch.setenv("AGENTBOX_DOMAIN", "")
+    monkeypatch.delenv("AGENTBOX_REGION", raising=False)
 
     config = ConnectionConfig()
-    assert config.api_url == "https://api.sandbox.ucloudai.com"
+    assert config.api_url == "https://api.cn-wlcb.sandbox.ucloudai.com"
 
 
 def test_api_url_in_args():
@@ -25,3 +26,45 @@ def test_api_url_has_correct_priority(monkeypatch):
 
     config = ConnectionConfig(api_url="http://localhost:8080")
     assert config.api_url == "http://localhost:8080"
+
+
+def test_region_cn_wlcb(monkeypatch):
+    monkeypatch.delenv("AGENTBOX_DOMAIN", raising=False)
+    monkeypatch.setenv("AGENTBOX_REGION", "cn-wlcb")
+
+    config = ConnectionConfig()
+    assert config.domain == "cn-wlcb.sandbox.ucloudai.com"
+    assert config.api_url == "https://api.cn-wlcb.sandbox.ucloudai.com"
+
+
+def test_region_us_ca(monkeypatch):
+    monkeypatch.delenv("AGENTBOX_DOMAIN", raising=False)
+    monkeypatch.setenv("AGENTBOX_REGION", "us-ca")
+
+    config = ConnectionConfig()
+    assert config.domain == "us-ca.sandbox.ucloudai.com"
+    assert config.api_url == "https://api.us-ca.sandbox.ucloudai.com"
+
+
+def test_domain_takes_priority_over_region(monkeypatch):
+    monkeypatch.setenv("AGENTBOX_DOMAIN", "custom.example.com")
+    monkeypatch.setenv("AGENTBOX_REGION", "us-ca")
+
+    config = ConnectionConfig()
+    assert config.domain == "custom.example.com"
+    assert config.api_url == "https://api.custom.example.com"
+
+
+def test_domain_arg_takes_priority_over_region(monkeypatch):
+    monkeypatch.setenv("AGENTBOX_REGION", "us-ca")
+
+    config = ConnectionConfig(domain="override.example.com")
+    assert config.domain == "override.example.com"
+
+
+def test_default_region_is_cn_wlcb(monkeypatch):
+    monkeypatch.delenv("AGENTBOX_DOMAIN", raising=False)
+    monkeypatch.delenv("AGENTBOX_REGION", raising=False)
+
+    config = ConnectionConfig()
+    assert config.domain == "cn-wlcb.sandbox.ucloudai.com"

ucloud_sandbox/connection_config.py

@@ -12,6 +12,13 @@
 KEEPALIVE_PING_INTERVAL_SEC = 50  # 50 seconds
 KEEPALIVE_PING_HEADER = "Keepalive-Ping-Interval"
 
+REGION_DOMAIN_MAP = {
+    "cn-wlcb": "cn-wlcb.sandbox.ucloudai.com",
+    "us-ca": "us-ca.sandbox.ucloudai.com",
+}
+
+DEFAULT_REGION = "cn-wlcb"
+
 
 class ApiParams(TypedDict, total=False):
     """
@@ -30,7 +37,7 @@ class ApiParams(TypedDict, total=False):
     """AgentBox API Key to use for authentication, defaults to `AGENTBOX_API_KEY` environment variable."""
 
     domain: Optional[str]
-    """AgentBox domain to use, defaults to `AGENTBOX_DOMAIN` environment variable."""
+    """AgentBox domain to use, defaults to `AGENTBOX_DOMAIN` environment variable. If not set, derived from `AGENTBOX_REGION`."""
 
     api_url: Optional[str]
     """URL to use for the API, defaults to `https://api.<domain>`. For internal use only."""
@@ -52,9 +59,17 @@ class ConnectionConfig:
 
     envd_port = 49983
 
+    @staticmethod
+    def _region():
+        return os.getenv("AGENTBOX_REGION") or DEFAULT_REGION
+
     @staticmethod
     def _domain():
-        return os.getenv("AGENTBOX_DOMAIN") or "sandbox.ucloudai.com"
+        explicit_domain = os.getenv("AGENTBOX_DOMAIN")
+        if explicit_domain:
+            return explicit_domain
+        region = ConnectionConfig._region()
+        return REGION_DOMAIN_MAP.get(region, REGION_DOMAIN_MAP[DEFAULT_REGION])
 
     @staticmethod
     def _debug():

ucloud_sandbox/template_async/main.py

@@ -219,7 +219,6 @@ async def build(
                     )
                 )
 
-            domain = domain or os.environ.get("AGENTBOX_DOMAIN", "sandbox.ucloudai.com")
             config = ConnectionConfig(
                 domain=domain, api_key=api_key or os.environ.get("AGENTBOX_API_KEY")
             )
@@ -309,7 +308,6 @@ async def build_in_background(
         )
         ```
         """
-        domain = domain or os.environ.get("AGENTBOX_DOMAIN", "sandbox.ucloudai.com")
         config = ConnectionConfig(
             domain=domain, api_key=api_key or os.environ.get("AGENTBOX_API_KEY")
         )
@@ -353,7 +351,6 @@ async def get_build_status(
         status = await AsyncTemplate.get_build_status(build_info, logs_offset=0)
         ```
         """
-        domain = domain or os.environ.get("AGENTBOX_DOMAIN", "sandbox.ucloudai.com")
         config = ConnectionConfig(
             domain=domain, api_key=api_key or os.environ.get("AGENTBOX_API_KEY")
         )

ucloud_sandbox/template_sync/main.py

@@ -219,7 +219,6 @@ def build(
                     )
                 )
 
-            domain = domain or os.environ.get("AGENTBOX_DOMAIN", "sandbox.ucloudai.com")
             config = ConnectionConfig(
                 domain=domain, api_key=api_key or os.environ.get("AGENTBOX_API_KEY")
             )
@@ -309,7 +308,6 @@ def build_in_background(
         )
         ```
         """
-        domain = domain or os.environ.get("AGENTBOX_DOMAIN", "sandbox.ucloudai.com")
         config = ConnectionConfig(
             domain=domain, api_key=api_key or os.environ.get("AGENTBOX_API_KEY")
         )
@@ -353,7 +351,6 @@ def get_build_status(
         status = Template.get_build_status(build_info, logs_offset=0)
         ```
         """
-        domain = domain or os.environ.get("AGENTBOX_DOMAIN", "sandbox.ucloudai.com")
         config = ConnectionConfig(
             domain=domain, api_key=api_key or os.environ.get("AGENTBOX_API_KEY")
         )