feat(sdk): validate algorithm in getKey by checking OS version

When getKey is called with a non-secp256k1 algorithm (e.g. ed25519),
the SDK now calls the Version RPC to check if the OS supports it.
If the Version RPC is unavailable (OS <= 0.5.6), the call is rejected
with a clear error instead of silently returning the wrong key type.
The version check result is cached for the lifetime of the client.

Author: kvinwang

sdk/go/dstack/client.go

@@ -19,6 +19,7 @@ import (
 	"net/http"
 	"os"
 	"strings"
+	"sync"
 )
 
 // Represents the response from a TLS key derivation request.
@@ -174,6 +175,9 @@ type DstackClient struct {
 	baseURL    string
 	httpClient *http.Client
 	logger     *slog.Logger
+
+	versionOnce   sync.Once
+	versionSupported bool
 }
 
 // Functional option for configuring a DstackClient.
@@ -369,8 +373,37 @@ func (c *DstackClient) GetTlsKey(
 	return &response, nil
 }
 
+// requiresVersionCheck returns true for algorithms that need OS >= 0.5.7.
+func requiresVersionCheck(algorithm string) bool {
+	switch algorithm {
+	case "secp256k1", "secp256k1_prehashed", "k256", "":
+		return false
+	default:
+		return true
+	}
+}
+
+// ensureAlgorithmSupported checks the OS version when a non-secp256k1 algorithm is requested.
+// On old OS (no Version RPC), it returns an error to prevent silent key type mismatch.
+func (c *DstackClient) ensureAlgorithmSupported(ctx context.Context, algorithm string) error {
+	if !requiresVersionCheck(algorithm) {
+		return nil
+	}
+	c.versionOnce.Do(func() {
+		_, err := c.GetVersion(ctx)
+		c.versionSupported = err == nil
+	})
+	if !c.versionSupported {
+		return fmt.Errorf("algorithm %q is not supported: OS version too old (Version RPC unavailable)", algorithm)
+	}
+	return nil
+}
+
 // Gets a key from the dstack service.
 func (c *DstackClient) GetKey(ctx context.Context, path string, purpose string, algorithm string) (*GetKeyResponse, error) {
+	if err := c.ensureAlgorithmSupported(ctx, algorithm); err != nil {
+		return nil, err
+	}
 	payload := map[string]interface{}{
 		"path":      path,
 		"purpose":   purpose,

sdk/js/src/index.ts

@@ -177,8 +177,11 @@ export interface TlsKeyOptions {
   usageClientAuth?: boolean;
 }
 
+const SECP256K1_ALGORITHMS = new Set(['secp256k1', 'secp256k1_prehashed', 'k256', ''])
+
 export class DstackClient<T extends TcbInfo = TcbInfoV05x> {
   protected endpoint: string
+  private _versionChecked: boolean | undefined = undefined
 
   constructor(endpoint: string | undefined = undefined) {
     if (endpoint === undefined) {
@@ -202,7 +205,23 @@ export class DstackClient<T extends TcbInfo = TcbInfoV05x> {
     this.endpoint = endpoint
   }
 
+  private async ensureAlgorithmSupported(algorithm: string): Promise<void> {
+    if (SECP256K1_ALGORITHMS.has(algorithm)) return
+    if (this._versionChecked === undefined) {
+      try {
+        await this.version()
+        this._versionChecked = true
+      } catch {
+        this._versionChecked = false
+      }
+    }
+    if (!this._versionChecked) {
+      throw new Error(`algorithm "${algorithm}" is not supported: OS version too old (Version RPC unavailable)`)
+    }
+  }
+
   async getKey(path: string, purpose: string = '', algorithm: string = 'secp256k1'): Promise<GetKeyResponse> {
+    await this.ensureAlgorithmSupported(algorithm)
     const payload = JSON.stringify({
       path: path,
       purpose: purpose,

sdk/python/src/dstack_sdk/dstack_client.py

@@ -303,6 +303,7 @@ def __init__(
         self._sync_client: Optional[httpx.Client] = None
         self._client_ref_count = 0
         self._timeout = timeout
+        self._version_checked: Optional[bool] = None
 
         if endpoint.startswith("http://") or endpoint.startswith("https://"):
             self.async_transport = httpx.AsyncHTTPTransport()
@@ -380,13 +381,30 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
                 self._sync_client.close()
                 self._sync_client = None
 
+    async def _ensure_algorithm_supported(self, algorithm: str) -> None:
+        """Check OS version when a non-secp256k1 algorithm is requested."""
+        if algorithm in ("secp256k1", "secp256k1_prehashed", "k256", ""):
+            return
+        if self._version_checked is None:
+            try:
+                await self.version()
+                self._version_checked = True
+            except Exception:
+                self._version_checked = False
+        if not self._version_checked:
+            raise RuntimeError(
+                f'algorithm "{algorithm}" is not supported: '
+                "OS version too old (Version RPC unavailable)"
+            )
+
     async def get_key(
         self,
         path: str | None = None,
         purpose: str | None = None,
         algorithm: str = "secp256k1",
     ) -> GetKeyResponse:
         """Derive a key from the given path, purpose, and algorithm."""
+        await self._ensure_algorithm_supported(algorithm)
         data: Dict[str, Any] = {
             "path": path or "",
             "purpose": purpose or "",

sdk/rust/examples/dstack_client_usage.rs

@@ -34,7 +34,7 @@ async fn main() -> anyhow::Result<()> {
 
     // 2. Derive a key
     let response = client
-        .get_key(Some("my-app".to_string()), Some("encryption".to_string()))
+        .get_key(Some("my-app".to_string()), Some("encryption".to_string()), None)
         .await?;
     println!("Key derived successfully!");
     println!("  Key length: {}", response.key.len());
@@ -93,7 +93,7 @@ async fn main() -> anyhow::Result<()> {
     );
 
     // 6. Get a simple key without purpose
-    let response = client.get_key(Some("simple-key".to_string()), None).await?;
+    let response = client.get_key(Some("simple-key".to_string()), None, None).await?;
     println!("Simple key derived successfully!");
     println!("  Key: {}", response.key);
 

sdk/rust/src/dstack_client.rs

@@ -11,6 +11,7 @@ use reqwest::Client;
 use serde::{de::DeserializeOwned, Serialize};
 use serde_json::{json, Value};
 use std::env;
+use std::sync::Mutex;
 
 pub use dstack_sdk_types::dstack::*;
 
@@ -60,6 +61,14 @@ pub enum ClientKind {
 
 pub trait BaseClient {}
 
+/// Algorithms that require OS >= 0.5.7 (i.e., Version RPC support).
+fn requires_version_check(algorithm: &str) -> bool {
+    !matches!(
+        algorithm,
+        "secp256k1" | "secp256k1_prehashed" | "k256" | ""
+    )
+}
+
 /// The main client for interacting with the dstack service
 pub struct DstackClient {
     /// The base URL for HTTP requests
@@ -68,6 +77,8 @@ pub struct DstackClient {
     endpoint: String,
     /// The type of client (HTTP or Unix domain socket)
     client: ClientKind,
+    /// Cached result of version check: None = not checked, Some(true) = new OS, Some(false) = old OS
+    version_checked: Mutex<Option<bool>>,
 }
 
 impl BaseClient for DstackClient {}
@@ -86,6 +97,7 @@ impl DstackClient {
             base_url,
             endpoint,
             client,
+            version_checked: Mutex::new(None),
         }
     }
 
@@ -126,15 +138,43 @@ impl DstackClient {
         }
     }
 
+    /// Check if the OS supports the given algorithm, querying version if needed.
+    async fn ensure_algorithm_supported(&self, algorithm: &str) -> Result<()> {
+        if !requires_version_check(algorithm) {
+            return Ok(());
+        }
+        let already_checked = self.version_checked.lock().unwrap().clone();
+        match already_checked {
+            Some(true) => Ok(()),
+            Some(false) => anyhow::bail!(
+                "algorithm \"{algorithm}\" is not supported: OS version too old (Version RPC unavailable)"
+            ),
+            None => {
+                let supported = self.version().await.is_ok();
+                *self.version_checked.lock().unwrap() = Some(supported);
+                if supported {
+                    Ok(())
+                } else {
+                    anyhow::bail!(
+                        "algorithm \"{algorithm}\" is not supported: OS version too old (Version RPC unavailable)"
+                    )
+                }
+            }
+        }
+    }
+
     pub async fn get_key(
         &self,
         path: Option<String>,
         purpose: Option<String>,
+        algorithm: Option<String>,
     ) -> Result<GetKeyResponse> {
+        let algorithm = algorithm.unwrap_or_else(|| "secp256k1".to_string());
+        self.ensure_algorithm_supported(&algorithm).await?;
         let data = json!({
             "path": path.unwrap_or_default(),
             "purpose": purpose.unwrap_or_default(),
-            "algorithm": "secp256k1", // Default or specify as needed
+            "algorithm": algorithm,
         });
         let response = self.send_rpc_request("/GetKey", &data).await?;
         let response = serde_json::from_value::<GetKeyResponse>(response)?;

sdk/rust/tests/test_client.rs

@@ -12,7 +12,7 @@ use sha2::{Digest, Sha256};
 #[tokio::test]
 async fn test_async_client_get_key() {
     let client = AsyncDstackClient::new(None);
-    let result = client.get_key(None, None).await.unwrap();
+    let result = client.get_key(None, None, None).await.unwrap();
     assert!(!result.key.is_empty());
     assert_eq!(result.decode_key().unwrap().len(), 32);
 }

sdk/rust/tests/test_eth.rs

@@ -12,7 +12,7 @@ use dstack_sdk_types::dstack::GetKeyResponse;
 async fn test_async_to_keypair() {
     let client = DstackClient::new(None);
     let result = client
-        .get_key(Some("test".to_string()), None)
+        .get_key(Some("test".to_string()), None, None)
         .await
         .expect("get_key failed");