Support API tokens in SDK (lmstudio-ai#163)

ryan-the-crayon · ncoghlan · web-flow · commit 3900323bfb22 · 2025-10-18T00:33:20.000+07:00
* Pass api_token to Client or AsyncClient constructor
* Pass api_token to configure_default_client() in the convenience API
* Set LMSTUDIO_API_TOKEN in the process environment

---------

Co-authored-by: Nick Coghlan &lt;ncoghlan@gmail.com&gt;
diff --git a/src/lmstudio/async_api.py b/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()
diff --git a/src/lmstudio/json_api.py b/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,12 @@
 
 DEFAULT_TTL = 60 * 60  # By default, leaves idle models loaded for an hour
 
+# lmstudio-js and lmstudio-python use the same API token environment variable
+_ENV_API_TOKEN = "LMSTUDIO_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]]
 
@@ -401,7 +409,9 @@ def from_details(message: str, details: DictObject) -> "LMStudioServerError":
         if display_data:
             specific_error: LMStudioServerError | None = None
             match display_data:
-                case {"code": "generic.noModelMatchingQuery"}:
+                case {"code": "generic.noModelMatchingQuery"} | {
+                    "code": "generic.pathNotFound"
+                }:
                     specific_error = LMStudioModelNotFoundError(str(default_error))
                 case {"code": "generic.presetNotFound"}:
                     specific_error = LMStudioPresetNotFoundError(str(default_error))
@@ -2041,10 +2051,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,25 +2099,57 @@ 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:
+    @classmethod
+    def _create_auth_from_token(cls, api_token: str | None) -> DictObject:
+        """Create an LM Studio websocket auth message from an API token.
+
+        If no token is given, and none is set in the environment,
+        falls back to generating a client scoped guest identifier
+        """
+        if api_token is None:
+            api_token = os.environ.get(_ENV_API_TOKEN, None)
+        if api_token:  # Accept empty string as equivalent to None
+            match = _LMS_API_TOKEN_REGEX.match(api_token.strip())
+            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 cls._format_auth_message(client_identifier, client_passkey)
+
+        return cls._format_auth_message()
+
+    def _create_auth_message(self, api_token: str | None = None) -> DictObject:
         """Create an LM Studio websocket authentication message."""
-        return self._format_auth_message()
+        # This is an instance method purely so subclasses may override it
+        return self._create_auth_from_token(api_token)
 
 
 TClient = TypeVar("TClient", bound=ClientBase)
diff --git a/src/lmstudio/plugin/runner.py b/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)
 
diff --git a/src/lmstudio/sync_api.py b/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()
 
diff --git a/tests/test_sessions.py b/tests/test_sessions.py
@@ -1,21 +1,26 @@
 """Test common client session behaviour."""
 
 import logging
+import os
+
 from typing import Generator
+from unittest import mock
 
 import pytest
 from pytest import LogCaptureFixture as LogCap
 
 from lmstudio import (
     AsyncClient,
     Client,
+    LMStudioValueError,
     LMStudioWebsocketError,
 )
 from lmstudio.async_api import (
     _AsyncLMStudioWebsocket,
     _AsyncSession,
     _AsyncSessionSystem,
 )
+from lmstudio.json_api import ClientBase
 from lmstudio.sync_api import (
     SyncLMStudioWebsocket,
     _SyncSession,
@@ -24,6 +29,93 @@
 from lmstudio._ws_impl import AsyncTaskManager
 from lmstudio._ws_thread import AsyncWebsocketThread
 
+# This API token is structurally valid
+_VALID_API_TOKEN = "sk-lm-abcDEF78:abcDEF7890abcDEF7890"
+
+
+@pytest.mark.parametrize("client_cls", [AsyncClient, Client])
+def test_auth_message_default(client_cls: ClientBase) -> None:
+    with mock.patch.dict(os.environ) as env:
+        env.pop("LMSTUDIO_API_TOKEN", None)
+        auth_message = client_cls._create_auth_from_token(None)
+        assert auth_message["authVersion"] == 1
+        assert auth_message["clientIdentifier"].startswith("guest:")
+        client_key = auth_message["clientPasskey"]
+        assert client_key != ""
+        assert isinstance(client_key, str)
+
+
+@pytest.mark.parametrize("client_cls", [AsyncClient, Client])
+def test_auth_message_empty_token(client_cls: ClientBase) -> None:
+    with mock.patch.dict(os.environ) as env:
+        # Set a valid token in the env to ensure it is ignored
+        env["LMSTUDIO_API_TOKEN"] = _VALID_API_TOKEN
+        auth_message = client_cls._create_auth_from_token("")
+        assert auth_message["authVersion"] == 1
+        assert auth_message["clientIdentifier"].startswith("guest:")
+        client_key = auth_message["clientPasskey"]
+        assert client_key != ""
+        assert isinstance(client_key, str)
+
+
+@pytest.mark.parametrize("client_cls", [AsyncClient, Client])
+def test_auth_message_empty_token_from_env(client_cls: ClientBase) -> None:
+    with mock.patch.dict(os.environ) as env:
+        env["LMSTUDIO_API_TOKEN"] = ""
+        auth_message = client_cls._create_auth_from_token(None)
+        assert auth_message["authVersion"] == 1
+        assert auth_message["clientIdentifier"].startswith("guest:")
+        client_key = auth_message["clientPasskey"]
+        assert client_key != ""
+        assert isinstance(client_key, str)
+
+
+@pytest.mark.parametrize("client_cls", [AsyncClient, Client])
+def test_auth_message_valid_token(client_cls: ClientBase) -> None:
+    auth_message = client_cls._create_auth_from_token(_VALID_API_TOKEN)
+    assert auth_message["authVersion"] == 1
+    assert auth_message["clientIdentifier"] == "abcDEF78"
+    assert auth_message["clientPasskey"] == "abcDEF7890abcDEF7890"
+
+
+@pytest.mark.parametrize("client_cls", [AsyncClient, Client])
+def test_auth_message_valid_token_from_env(client_cls: ClientBase) -> None:
+    with mock.patch.dict(os.environ) as env:
+        env["LMSTUDIO_API_TOKEN"] = _VALID_API_TOKEN
+        auth_message = client_cls._create_auth_from_token(None)
+        assert auth_message["authVersion"] == 1
+        assert auth_message["clientIdentifier"] == "abcDEF78"
+        assert auth_message["clientPasskey"] == "abcDEF7890abcDEF7890"
+
+
+_INVALID_TOKENS = [
+    "missing-token-prefix",
+    "sk-lm-missing-id-and-key-separator",
+    "sk-lm-invalid_id:invalid_key",
+    "sk-lm-idtoolong:abcDEF7890abcDEF7890",
+    "sk-lm-abcDEF78:keytooshort",
+]
+
+
+@pytest.mark.parametrize("client_cls", [AsyncClient, Client])
+@pytest.mark.parametrize("api_token", _INVALID_TOKENS)
+def test_auth_message_invalid_token(client_cls: ClientBase, api_token: str) -> None:
+    with mock.patch.dict(os.environ) as env:
+        env["LMSTUDIO_API_TOKEN"] = _VALID_API_TOKEN
+        with pytest.raises(LMStudioValueError):
+            client_cls._create_auth_from_token(api_token)
+
+
+@pytest.mark.parametrize("client_cls", [AsyncClient, Client])
+@pytest.mark.parametrize("api_token", _INVALID_TOKENS)
+def test_auth_message_invalid_token_from_env(
+    client_cls: ClientBase, api_token: str
+) -> None:
+    with mock.patch.dict(os.environ) as env:
+        env["LMSTUDIO_API_TOKEN"] = api_token
+        with pytest.raises(LMStudioValueError):
+            client_cls._create_auth_from_token(None)
+
 
 async def check_connected_async_session(session: _AsyncSession) -> None:
     assert session.connected
@@ -160,7 +252,7 @@ def test_implicit_reconnection_sync(caplog: LogCap) -> None:
 async def test_websocket_cm_async(caplog: LogCap) -> None:
     caplog.set_level(logging.DEBUG)
     api_host = await AsyncClient.find_default_local_api_host()
-    auth_details = AsyncClient._format_auth_message()
+    auth_details = AsyncClient._create_auth_from_token(None)
     tm = AsyncTaskManager(on_activation=None)
     lmsws = _AsyncLMStudioWebsocket(tm, f"http://{api_host}/system", auth_details)
     # SDK client websockets start out disconnected
@@ -200,7 +292,7 @@ def ws_thread() -> Generator[AsyncWebsocketThread, None, None]:
 def test_websocket_cm_sync(ws_thread: AsyncWebsocketThread, caplog: LogCap) -> None:
     caplog.set_level(logging.DEBUG)
     api_host = Client.find_default_local_api_host()
-    auth_details = Client._format_auth_message()
+    auth_details = Client._create_auth_from_token(None)
     lmsws = SyncLMStudioWebsocket(ws_thread, f"http://{api_host}/system", auth_details)
     # SDK client websockets start out disconnected
     assert not lmsws.connected
diff --git a/tox.ini b/tox.ini
@@ -19,6 +19,7 @@ allowlist_externals = pytest
 passenv =
     CI
     LMS_*
+    LMSTUDIO_*
 commands =
     # Even the "slow" tests aren't absurdly slow, so default to running them
     pytest {posargs} tests/
