feat(sdk/python): add not_before/not_after/with_app_info TLS key options

Extend get_tls_key with the three GetTlsKeyArgs proto fields (6/7/8)
added on dstack OS >= 0.5.7. New params are kw-only so positional call
sites keep working unchanged. When any new option is provided, the SDK
probes the Version RPC first and raises a clear error on older OS images
that lack it (same gating pattern as ed25519 key derivation).

Mirrors the JS SDK change in dc5fc21.

Author: Leechael

sdk/python/src/dstack_sdk/dstack_client.py

@@ -401,6 +401,19 @@ async def _ensure_algorithm_supported(self, algorithm: str) -> None:
                 "OS version too old (Version RPC unavailable)"
             )
 
+    async def _ensure_tls_key_options_supported(
+        self, feature_names: List[str]
+    ) -> None:
+        """Check OS version when 0.5.7+ TLS key options are requested."""
+        try:
+            await self.version()
+        except Exception:
+            features = ", ".join(feature_names)
+            raise RuntimeError(
+                f"TLS key options [{features}] are not supported: "
+                "OS version too old (Version RPC unavailable)"
+            )
+
     async def get_key(
         self,
         path: str | None = None,
@@ -477,8 +490,28 @@ async def get_tls_key(
         usage_ra_tls: bool = False,
         usage_server_auth: bool = True,
         usage_client_auth: bool = False,
+        *,
+        not_before: Optional[int] = None,
+        not_after: Optional[int] = None,
+        with_app_info: Optional[bool] = None,
     ) -> GetTlsKeyResponse:
-        """Request a TLS key from the service with optional parameters."""
+        """Request a TLS key from the service with optional parameters.
+
+        ``not_before`` / ``not_after`` (seconds since UNIX epoch) and
+        ``with_app_info`` require dstack OS >= 0.5.7. When any of them is set,
+        the SDK probes the guest agent ``Version`` RPC first and raises a
+        clear error on older OS images.
+        """
+        new_features: List[str] = []
+        if not_before is not None:
+            new_features.append("not_before")
+        if not_after is not None:
+            new_features.append("not_after")
+        if with_app_info is not None:
+            new_features.append("with_app_info")
+        if new_features:
+            await self._ensure_tls_key_options_supported(new_features)
+
         data: Dict[str, Any] = {
             "subject": subject or "",
             "usage_ra_tls": usage_ra_tls,
@@ -487,6 +520,12 @@ async def get_tls_key(
         }
         if alt_names:
             data["alt_names"] = list(alt_names)
+        if not_before is not None:
+            data["not_before"] = not_before
+        if not_after is not None:
+            data["not_after"] = not_after
+        if with_app_info is not None:
+            data["with_app_info"] = with_app_info
 
         result = await self._send_rpc_request("GetTlsKey", data)
         return GetTlsKeyResponse(**result)
@@ -604,6 +643,10 @@ def get_tls_key(
         usage_ra_tls: bool = False,
         usage_server_auth: bool = True,
         usage_client_auth: bool = False,
+        *,
+        not_before: Optional[int] = None,
+        not_after: Optional[int] = None,
+        with_app_info: Optional[bool] = None,
     ) -> GetTlsKeyResponse:
         """Request a TLS key from the service with optional parameters."""
         raise NotImplementedError

sdk/python/tests/test_client.py

@@ -649,3 +649,67 @@ async def test_mixed_sync_async_calls():
     # Both should work and return valid results
     assert len(sync_result.decode_key()) == 32
     assert len(async_result.decode_key()) == 32
+
+
+@pytest.mark.asyncio
+async def test_get_tls_key_new_options_payload(monkeypatch):
+    """0.5.7+ TLS options reach the payload and trigger the Version probe."""
+    calls: list = []
+
+    async def fake_send(self, method, payload):
+        calls.append((method, payload))
+        if method == "Version":
+            return {"version": "0.5.7", "rev": "test"}
+        return {"key": "k", "certificate_chain": []}
+
+    monkeypatch.setenv("DSTACK_SIMULATOR_ENDPOINT", "http://localhost:0")
+    monkeypatch.setattr(AsyncDstackClient, "_send_rpc_request", fake_send)
+    client = AsyncDstackClient()
+    result = await client.get_tls_key(
+        subject="api.example.com",
+        not_before=1_700_000_000,
+        not_after=1_800_000_000,
+        with_app_info=True,
+    )
+    assert isinstance(result, GetTlsKeyResponse)
+    assert [c[0] for c in calls] == ["Version", "GetTlsKey"]
+    payload = calls[1][1]
+    assert payload["not_before"] == 1_700_000_000
+    assert payload["not_after"] == 1_800_000_000
+    assert payload["with_app_info"] is True
+
+
+@pytest.mark.asyncio
+async def test_get_tls_key_legacy_options_skip_version_probe(monkeypatch):
+    """Calls without the new options must NOT probe Version (backward compat)."""
+    calls: list = []
+
+    async def fake_send(self, method, payload):
+        calls.append((method, payload))
+        return {"key": "k", "certificate_chain": []}
+
+    monkeypatch.setenv("DSTACK_SIMULATOR_ENDPOINT", "http://localhost:0")
+    monkeypatch.setattr(AsyncDstackClient, "_send_rpc_request", fake_send)
+    client = AsyncDstackClient()
+    await client.get_tls_key(subject="api.example.com")
+    assert [c[0] for c in calls] == ["GetTlsKey"]
+    payload = calls[0][1]
+    assert "not_before" not in payload
+    assert "not_after" not in payload
+    assert "with_app_info" not in payload
+
+
+@pytest.mark.asyncio
+async def test_get_tls_key_new_options_require_version(monkeypatch):
+    """Raise a clear error when Version RPC is unavailable on older OS."""
+
+    async def fake_send(self, method, payload):
+        if method == "Version":
+            raise RuntimeError("Version not implemented")
+        return {"key": "k", "certificate_chain": []}
+
+    monkeypatch.setenv("DSTACK_SIMULATOR_ENDPOINT", "http://localhost:0")
+    monkeypatch.setattr(AsyncDstackClient, "_send_rpc_request", fake_send)
+    client = AsyncDstackClient()
+    with pytest.raises(RuntimeError, match="TLS key options"):
+        await client.get_tls_key(with_app_info=False)