API token in SDK

src/lmstudio/async_api.py

@@ -1492,9 +1492,11 @@ async def embed(
 class AsyncClient(ClientBase):
     """Async SDK client interface."""
 
-    def __init__(self, api_host: str | None = None) -> None:
+    def __init__(
+        self, api_host: str | None = None, api_token: str | None = None
+    ) -> None:
         """Initialize API client."""
-        super().__init__(api_host)
+        super().__init__(api_host, api_token)
         self._resources = AsyncExitStack()
         self._sessions: dict[str, _AsyncSession] = {}
         self._task_manager = AsyncTaskManager()

src/lmstudio/json_api.py

@@ -13,6 +13,8 @@
 import copy
 import inspect
 import json
+import os
+import re
 import sys
 import uuid
 import warnings
@@ -197,6 +199,13 @@
 
 DEFAULT_TTL = 60 * 60  # By default, leaves idle models loaded for an hour
 
+# API token environment variable name currently has an underscore prefix
+# The prefix will be removed once the interface has been set by lmstudio-js
+_ENV_API_TOKEN = "_LMS_SDK_API_TOKEN"
+_LMS_API_TOKEN_REGEX = re.compile(
+    r"^sk-lm-(?P<clientIdentifier>[A-Za-z0-9]{8}):(?P<clientPasskey>[A-Za-z0-9]{20})$"
+)
+
 # Require a coroutine (not just any awaitable) for run_coroutine_threadsafe compatibility
 SendMessageAsync: TypeAlias = Callable[[DictObject], Coroutine[Any, Any, None]]
 
@@ -2041,10 +2050,12 @@ def _ensure_connected(self, usage: str) -> None | NoReturn:
 class ClientBase:
     """Common base class for SDK client interfaces."""
 
-    def __init__(self, api_host: str | None = None) -> None:
+    def __init__(
+        self, api_host: str | None = None, api_token: str | None = None
+    ) -> None:
         """Initialize API client."""
         self._api_host = api_host
-        self._auth_details = self._create_auth_message()
+        self._auth_details = self._create_auth_message(api_token)
 
     @property
     def api_host(self) -> str:
@@ -2087,24 +2098,46 @@ def _format_auth_message(
         client_id: str | None = None, client_key: str | None = None
     ) -> DictObject:
         """Create an LM Studio websocket authentication message."""
-        # Note: authentication (in its current form) is primarily a cooperative
+        # Note: the authentication fields are used for two distinct purposes.
+        # When extracted from an API token (see _create_auth_message below),
+        # they are an actual authentication & authorisation mechanism.
+        # When generated internally by the SDK, they are instead a cooperative
         # resource management mechanism that allows the server to appropriately
         # manage client-scoped resources (such as temporary file handles).
-        # As such, the client ID and client passkey are currently more a two part
-        # client identifier than they are an adversarial security measure. This is
-        # sufficient to prevent accidental conflicts and, in combination with secure
-        # websocket support, would be sufficient to ensure that access to the running
-        # client was required to extract the auth details.
-        client_identifier = client_id if client_id is not None else str(uuid.uuid4())
+        # As such, when the API host isn't configured to require API tokens,
+        # the client ID and client key are more a two part client
+        # identifier than they are an adversarial security measure.
+        client_identifier = (
+            client_id if client_id is not None else f"guest:{str(uuid.uuid4())}"
+        )
         client_passkey = client_key if client_key is not None else str(uuid.uuid4())
         return {
             "authVersion": 1,
             "clientIdentifier": client_identifier,
             "clientPasskey": client_passkey,
         }
 
-    def _create_auth_message(self) -> DictObject:
+    def _create_auth_message(self, api_token: str | None = None) -> DictObject:
         """Create an LM Studio websocket authentication message."""
+        if api_token is None:
+            api_token = os.getenv(_ENV_API_TOKEN, None)
+        if api_token is not None:
+            match = _LMS_API_TOKEN_REGEX.match(api_token)
+            if match is None:
+                raise LMStudioValueError(
+                    "The api_token argument does not look like a valid LM Studio API token.\n\n"
+                    "LM Studio API tokens are obtained from LM Studio, and they look like this:\n"
+                    "sk-lm-xxxxxxxxxxxxxxxxxxxxxxxxxxxxx."
+                )
+            groups = match.groupdict()
+            client_identifier = groups.get("clientIdentifier")
+            client_passkey = groups.get("clientPasskey")
+            if client_identifier is None or client_passkey is None:
+                raise LMStudioValueError(
+                    "Unexpected error parsing api_token: required token fields were not detected."
+                )
+            return self._format_auth_message(client_identifier, client_passkey)
+
         return self._format_auth_message()
 
 

src/lmstudio/plugin/runner.py

@@ -103,10 +103,10 @@ def __init__(
         _AsyncSessionPlugins,
     )
 
-    def _create_auth_message(self) -> DictObject:
+    def _create_auth_message(self, api_token: str | None = None) -> DictObject:
         """Create an LM Studio websocket authentication message."""
         if self._client_id is None or self._client_key is None:
-            return super()._create_auth_message()
+            return super()._create_auth_message(api_token)
         # Use plugin credentials to unlock the full plugin client API
         return self._format_auth_message(self._client_id, self._client_key)
 

src/lmstudio/sync_api.py

@@ -1540,9 +1540,11 @@ def embed(
 class Client(ClientBase):
     """Synchronous SDK client interface."""
 
-    def __init__(self, api_host: str | None = None) -> None:
+    def __init__(
+        self, api_host: str | None = None, api_token: str | None = None
+    ) -> None:
         """Initialize API client."""
-        super().__init__(api_host)
+        super().__init__(api_host, api_token)
         self._resources = rm = ExitStack()
         self._ws_thread = ws_thread = AsyncWebsocketThread(dict(client=repr(self)))
         ws_thread.start()
@@ -1699,40 +1701,43 @@ def list_loaded_models(
 
 
 # Convenience API
-_default_api_host = None
+_default_api_host: str | None = None
+_default_api_token: str | None = None
 _default_client: Client | None = None
 
 
 @sdk_public_api()
-def configure_default_client(api_host: str) -> None:
+def configure_default_client(api_host: str, api_token: str | None = None) -> None:
     """Set the server API host for the default global client (without creating the client)."""
     global _default_api_host
     if _default_client is not None:
         raise LMStudioClientError(
-            "Default client is already created, cannot set its API host."
+            "Default client is already created, cannot set its API host or token."
         )
     _default_api_host = api_host
+    _default_api_token = api_token
 
 
 @sdk_public_api()
 def get_default_client(api_host: str | None = None) -> Client:
     """Get the default global client (creating it if necessary)."""
+    # Note: call configure_default_client() explicitly to set the API token
     global _default_client
     if api_host is not None:
         # This will raise an exception if the client already exists
         configure_default_client(api_host)
     if _default_client is None:
-        _default_client = Client(_default_api_host)
+        _default_client = Client(_default_api_host, _default_api_token)
         _default_client._ensure_api_host_is_valid()
     return _default_client
 
 
 def _reset_default_client() -> None:
     # Allow the test suite to reset the client without
     # having to poke directly at the module's internals
-    global _default_api_host, _default_client
+    global _default_api_host, _default_api_token, _default_client
     previous_client = _default_client
-    _default_api_host = _default_client = None
+    _default_api_host = _default_api_token = _default_client = None
     if previous_client is not None:
         previous_client.close()