dex-python-sdk 1.14.1__tar.gz → 1.16.0__tar.gz

{dex_python_sdk-1.14.1 → dex_python_sdk-1.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dex-python-sdk
-Version: 1.14.1
+Version: 1.16.0
 Requires-Dist: grpcio>=1.83.0
 Requires-Dist: grpcio-status>=1.83.0
 Requires-Dist: protobuf>=7.35.1
@@ -25,6 +25,47 @@ checks its mtime and size at most once per second, and reloads it after rotation
 Whitespace around the token is removed. An unreadable or empty file fails calls
 until it is restored; the SDK does not send an old token.
 
+Use `ApiKeyCredential` when the application holds a long-lived API key and a
+token endpoint issues short-lived access tokens scoped to a server ID. Supply
+the token URL, server ID, and API key explicitly; the server ID is never derived
+from the endpoint host. The endpoint must use HTTPS, except loopback HTTP on
+`localhost`, `127.0.0.0/8`, or `::1` for local development.
+Loopback HTTP bypasses environment proxies so the API key stays on the host.
+HTTPS uses configured proxies with TLS through CONNECT. Token endpoint redirects
+are refused, including redirects that preserve the POST method.
+
+```python
+from dex import ApiKeyCredential, Client, ClientOptions
+
+credential = ApiKeyCredential(
+    "https://auth.example.com/token", "dex-server", "application-api-key"
+)
+client = Client(registry, cache, ClientOptions(credential=credential))
+try:
+    client.health_check()
+finally:
+    client.close()
+    credential.close()
+```
+
+The first call exchanges the key for a token. After 80% of its lifetime, a
+single background refresh begins while calls continue using the valid token.
+Callers waiting for an expired token share the exchange. Failed refreshes retry
+no more often than every five seconds while the previous token remains valid.
+Each exchange has a 10-second deadline by default, configurable through
+`exchange_timeout`. It covers connecting, sending, reading, and parsing. A Dex
+`UNAUTHENTICATED` response discards the rejected token
+and retries a unary call once with a fresh token. Pull Worker streams reconnect
+once immediately with a fresh token; a second rejection uses the normal backoff.
+The credential is thread-safe.
+It redacts the API key, issued access tokens, and JWT-shaped strings from
+exchange errors, including endpoint details and transport causes. Client
+calls raise `ApiKeyExchangeError` unchanged; pull Workers log it and retry.
+Closing the credential is optional. After all Clients and Workers sharing it
+are closed, `close()` prevents further token calls and waits at most until the
+current exchange deadline. An uninterruptible network operation may continue
+in its daemon thread after that deadline.
+
 ```python
 from dex import Client, ClientOptions, FileCredential, Worker, WorkerOptions
 
@@ -157,6 +198,40 @@ applies before and after readiness. Another registered stream for that queue
 prevents shutdown. The default is False and retries indefinitely. Always call
 `stop()` or `close()` afterward; both remain safe after `start()` returns.
 
+### Reserved RPC capacity
+
+Set `reserved_rpc_tasks` when synchronous RPC latency must stay responsive
+behind long-running Steps. When unset, it reserves one slot if the resolved task
+limit is at least two, otherwise zero. With the CPU-count default, Steps get one
+fewer slot than before. Set it explicitly to zero to opt out. With one slot and
+no explicit reserve, Worker start logs one warning: RPCs may wait behind
+long-running Steps; use a limit of at least two or set the reserve explicitly.
+Flows must stay responsive; Dex cannot prioritize synchronous RPCs like workflow
+engines prioritize reads. The value must be nonnegative and less than
+`max_concurrent_tasks` after its CPU-count default is resolved; invalid values
+fail construction. RPC and Step handlers count together: WaitFor and Execute
+tasks, including local attempts, are declined with `NO_CAPACITY` when the busy
+count reaches the task limit minus the reserve. API-direct and transactional
+InvokeWorkerRPC tasks can use the full limit.
+
+```python
+pull = PullWorkerOptions(max_concurrent_tasks=8, reserved_rpc_tasks=2)
+```
+
+A decline lets Dex try another Worker without consuming an attempt. If all
+Workers remain at their Step limit, a Step can wait until the dispatch wait
+expires, then receive `NO_WORKER_AVAILABLE`, consume an attempt, and follow its
+retry policy, as when all Workers are full. Local attempts wait the local
+dispatch wait, then fall back according to their policy. Each queue still keeps
+its configured idle streams; a reserve decline does not close a stream or change
+registration readiness.
+
+Each reserve decline records the Flow type through the dex._pull_worker logger
+at DEBUG level. Export the cumulative `pull_reserved_rpc_declines` Worker value
+to distinguish reserve pressure from full-capacity rejections. The count
+excludes full-capacity rejections and unknown names, never resets, and remains
+available after shutdown. Listening Workers return zero.
+
 ## Flow state I/O
 
 Applications read and write Flow state through typed RPCs. Client and AsyncClient

{dex_python_sdk-1.14.1 → dex_python_sdk-1.16.0}/README.md

@@ -9,6 +9,47 @@ checks its mtime and size at most once per second, and reloads it after rotation
 Whitespace around the token is removed. An unreadable or empty file fails calls
 until it is restored; the SDK does not send an old token.
 
+Use `ApiKeyCredential` when the application holds a long-lived API key and a
+token endpoint issues short-lived access tokens scoped to a server ID. Supply
+the token URL, server ID, and API key explicitly; the server ID is never derived
+from the endpoint host. The endpoint must use HTTPS, except loopback HTTP on
+`localhost`, `127.0.0.0/8`, or `::1` for local development.
+Loopback HTTP bypasses environment proxies so the API key stays on the host.
+HTTPS uses configured proxies with TLS through CONNECT. Token endpoint redirects
+are refused, including redirects that preserve the POST method.
+
+```python
+from dex import ApiKeyCredential, Client, ClientOptions
+
+credential = ApiKeyCredential(
+    "https://auth.example.com/token", "dex-server", "application-api-key"
+)
+client = Client(registry, cache, ClientOptions(credential=credential))
+try:
+    client.health_check()
+finally:
+    client.close()
+    credential.close()
+```
+
+The first call exchanges the key for a token. After 80% of its lifetime, a
+single background refresh begins while calls continue using the valid token.
+Callers waiting for an expired token share the exchange. Failed refreshes retry
+no more often than every five seconds while the previous token remains valid.
+Each exchange has a 10-second deadline by default, configurable through
+`exchange_timeout`. It covers connecting, sending, reading, and parsing. A Dex
+`UNAUTHENTICATED` response discards the rejected token
+and retries a unary call once with a fresh token. Pull Worker streams reconnect
+once immediately with a fresh token; a second rejection uses the normal backoff.
+The credential is thread-safe.
+It redacts the API key, issued access tokens, and JWT-shaped strings from
+exchange errors, including endpoint details and transport causes. Client
+calls raise `ApiKeyExchangeError` unchanged; pull Workers log it and retry.
+Closing the credential is optional. After all Clients and Workers sharing it
+are closed, `close()` prevents further token calls and waits at most until the
+current exchange deadline. An uninterruptible network operation may continue
+in its daemon thread after that deadline.
+
 ```python
 from dex import Client, ClientOptions, FileCredential, Worker, WorkerOptions
 
@@ -141,6 +182,40 @@ applies before and after readiness. Another registered stream for that queue
 prevents shutdown. The default is False and retries indefinitely. Always call
 `stop()` or `close()` afterward; both remain safe after `start()` returns.
 
+### Reserved RPC capacity
+
+Set `reserved_rpc_tasks` when synchronous RPC latency must stay responsive
+behind long-running Steps. When unset, it reserves one slot if the resolved task
+limit is at least two, otherwise zero. With the CPU-count default, Steps get one
+fewer slot than before. Set it explicitly to zero to opt out. With one slot and
+no explicit reserve, Worker start logs one warning: RPCs may wait behind
+long-running Steps; use a limit of at least two or set the reserve explicitly.
+Flows must stay responsive; Dex cannot prioritize synchronous RPCs like workflow
+engines prioritize reads. The value must be nonnegative and less than
+`max_concurrent_tasks` after its CPU-count default is resolved; invalid values
+fail construction. RPC and Step handlers count together: WaitFor and Execute
+tasks, including local attempts, are declined with `NO_CAPACITY` when the busy
+count reaches the task limit minus the reserve. API-direct and transactional
+InvokeWorkerRPC tasks can use the full limit.
+
+```python
+pull = PullWorkerOptions(max_concurrent_tasks=8, reserved_rpc_tasks=2)
+```
+
+A decline lets Dex try another Worker without consuming an attempt. If all
+Workers remain at their Step limit, a Step can wait until the dispatch wait
+expires, then receive `NO_WORKER_AVAILABLE`, consume an attempt, and follow its
+retry policy, as when all Workers are full. Local attempts wait the local
+dispatch wait, then fall back according to their policy. Each queue still keeps
+its configured idle streams; a reserve decline does not close a stream or change
+registration readiness.
+
+Each reserve decline records the Flow type through the dex._pull_worker logger
+at DEBUG level. Export the cumulative `pull_reserved_rpc_declines` Worker value
+to distinguish reserve pressure from full-capacity rejections. The count
+excludes full-capacity rejections and unknown names, never resets, and remains
+available after shutdown. Listening Workers return zero.
+
 ## Flow state I/O
 
 Applications read and write Flow state through typed RPCs. Client and AsyncClient

{dex_python_sdk-1.14.1 → dex_python_sdk-1.16.0}/dex/__init__.py

@@ -24,7 +24,6 @@ from dex.blob_cache import BlobCache, BlobCacheConfig, open_blob_cache
 from dex.channel import Channel, ChannelMap, ChannelMapLoad, ChannelMessage
 from dex.client import Client
 from dex.client_options import ClientOptions
-from dex.credential import Credential, FileCredential
 from dex.codec import (
     BOOL,
     BYTES,
@@ -39,6 +38,12 @@ from dex.codec import (
 )
 from dex.condition import ConditionCombination
 from dex.context import AsyncContext, Context, RecoveryErrorInfo
+from dex.credential import (
+    ApiKeyCredential,
+    ApiKeyExchangeError,
+    Credential,
+    FileCredential,
+)
 from dex.flow import Flow, PersistenceSchema, Registry, RPCResult, rpc
 from dex.flow_config import ActiveStepSearchMode, FlowConfig
 from dex.flow_info import (
@@ -62,6 +67,12 @@ from dex.flow_options import (
     TimeTravelType,
 )
 from dex.flow_result import FlowResult, StepCompletion
+from dex.pull_registration import (
+    PullQueueRegistrationState,
+    PullRegistrationError,
+    PullRegistrationState,
+)
+from dex.rpc_invoke_options import RPCInvokeOptions
 from dex.runtime_errors import (
     AttributeMapNotLoadedError,
     ChannelMessageNotFoundError,
@@ -80,7 +91,6 @@ from dex.runtime_errors import (
     ValueMappingError,
     WorkerInvocationError,
 )
-from dex.rpc_invoke_options import RPCInvokeOptions
 from dex.step import (
     RetryPolicy,
     Step,
@@ -113,11 +123,6 @@ from dex.timer import Timer
 from dex.wait import Wait
 from dex.wait_options import WaitForAttributeOptions, WaitForStepCompletionOptions
 from dex.worker import Worker
-from dex.pull_registration import (
-    PullRegistrationError,
-    PullRegistrationState,
-    PullQueueRegistrationState,
-)
 from dex.worker_options import PullWorkerOptions, WorkerOptions, WorkerTarget
 
 __all__ = [
@@ -138,6 +143,8 @@ __all__ = [
     "AsyncClient",
     "AsyncBufferedTextStream",
     "AsyncWorker",
+    "ApiKeyCredential",
+    "ApiKeyExchangeError",
     "BlobCache",
     "BlobCacheConfig",
     "BufferedTextStream",

{dex_python_sdk-1.14.1 → dex_python_sdk-1.16.0}/dex/_credential_transport.py

@@ -16,6 +16,7 @@ import asyncio
 import inspect
 import queue
 import threading
+import time
 from collections.abc import Callable
 from concurrent.futures import Future
 from typing import Any, NamedTuple, cast
@@ -168,6 +169,38 @@ def credential_deadline_error() -> grpc.aio.AioRpcError:
     )
 
 
+class _SyncCredentialDeadlineError(grpc.RpcError):
+    def code(self) -> grpc.StatusCode:
+        return grpc.StatusCode.DEADLINE_EXCEEDED
+
+    def details(self) -> str:
+        return "credential lookup exceeded RPC deadline"
+
+    def trailing_metadata(self) -> tuple[Any, ...]:
+        return ()
+
+
+def _bounded_credential_metadata(
+    credential: Credential, metadata: Any, timeout: float | None
+) -> Any:
+    if timeout is not None and timeout <= 0:
+        raise _SyncCredentialDeadlineError()
+    bounded_token = getattr(credential, "_token_with_timeout", None)
+    if not callable(bounded_token) or timeout is None:
+        return credential_metadata(credential, metadata)
+    try:
+        token = bounded_token(timeout)
+    except TimeoutError:
+        raise _SyncCredentialDeadlineError() from None
+    return tuple(metadata or ()) + (("authorization", "Bearer " + token),)
+
+
+def _reject_credential_token(credential: Credential, metadata: Any) -> None:
+    reject_token = getattr(credential, "_reject_token", None)
+    if callable(reject_token):
+        reject_token(metadata[-1][1][7:])
+
+
 class _SyncCredentialInterceptor(grpc.UnaryUnaryClientInterceptor):
     def __init__(self, credential: Credential) -> None:
         self._credential = credential
@@ -178,15 +211,49 @@ class _SyncCredentialInterceptor(grpc.UnaryUnaryClientInterceptor):
         client_call_details: grpc.ClientCallDetails,
         request: Any,
     ) -> Any:
+        started_at = time.monotonic()
+        timeout = client_call_details.timeout
+        metadata = _bounded_credential_metadata(
+            self._credential, client_call_details.metadata, timeout
+        )
+        if timeout is not None:
+            timeout -= time.monotonic() - started_at
+            if timeout <= 0:
+                raise _SyncCredentialDeadlineError()
         details = _SyncCallDetails(
             client_call_details.method,
-            client_call_details.timeout,
-            credential_metadata(self._credential, client_call_details.metadata),
+            timeout,
+            metadata,
             client_call_details.credentials,
             client_call_details.wait_for_ready,
             client_call_details.compression,
         )
-        return continuation(details, request)
+        try:
+            response = continuation(details, request)
+        except grpc.RpcError as error:
+            if error.code() != grpc.StatusCode.UNAUTHENTICATED:
+                raise
+            if not callable(getattr(self._credential, "_reject_token", None)):
+                raise
+        else:
+            if response.code() != grpc.StatusCode.UNAUTHENTICATED:
+                return response
+            if not callable(getattr(self._credential, "_reject_token", None)):
+                return response
+        _reject_credential_token(self._credential, details.metadata)
+        if client_call_details.timeout is not None:
+            timeout = client_call_details.timeout - (time.monotonic() - started_at)
+        retry_details = details._replace(
+            metadata=_bounded_credential_metadata(
+                self._credential, client_call_details.metadata, timeout
+            )
+        )
+        if client_call_details.timeout is not None:
+            timeout = client_call_details.timeout - (time.monotonic() - started_at)
+            if timeout <= 0:
+                raise _SyncCredentialDeadlineError()
+            retry_details = retry_details._replace(timeout=timeout)
+        return continuation(retry_details, request)
 
 
 class _AsyncCredentialInterceptor(grpc.aio.UnaryUnaryClientInterceptor):
@@ -216,7 +283,39 @@ class _AsyncCredentialInterceptor(grpc.aio.UnaryUnaryClientInterceptor):
             client_call_details.credentials,
             client_call_details.wait_for_ready,
         )
-        return await continuation(details, request)
+        call = await continuation(details, request)
+        try:
+            await call
+        except grpc.aio.AioRpcError as error:
+            if error.code() != grpc.StatusCode.UNAUTHENTICATED:
+                raise
+            if not callable(getattr(self._runner._credential, "_reject_token", None)):
+                raise
+            _reject_credential_token(self._runner._credential, metadata)
+            if client_call_details.timeout is not None:
+                remaining_timeout = client_call_details.timeout - (
+                    loop.time() - started_at
+                )
+                if remaining_timeout <= 0:
+                    raise credential_deadline_error() from None
+            retry_metadata = await async_credential_metadata(
+                self._runner, client_call_details.metadata, remaining_timeout
+            )
+            if client_call_details.timeout is not None:
+                remaining_timeout = client_call_details.timeout - (
+                    loop.time() - started_at
+                )
+                if remaining_timeout <= 0:
+                    raise credential_deadline_error() from None
+            retry_details = grpc.aio.ClientCallDetails(
+                client_call_details.method,
+                remaining_timeout,
+                retry_metadata,
+                client_call_details.credentials,
+                client_call_details.wait_for_ready,
+            )
+            return await continuation(retry_details, request)
+        return call
 
 
 def sync_channel(server_address: str, credential: Credential | None) -> grpc.Channel:

{dex_python_sdk-1.14.1 → dex_python_sdk-1.16.0}/dex/_pull_worker.py

@@ -22,7 +22,11 @@ from typing import Any
 
 import grpc
 
-from dex._credential_transport import AsyncCredentialRunner, async_credential_metadata
+from dex._credential_transport import (
+    AsyncCredentialRunner,
+    _reject_credential_token,
+    async_credential_metadata,
+)
 from dex._grpc_errors import _worker_error_status
 from dex._server_protocol import sdk_version
 from dex.dexpb import dex_pb2 as pb
@@ -94,6 +98,15 @@ class PullWorkerRuntime:
         self._permanent_error: PullRegistrationError | None = None
         self._rejection_stop_task: asyncio.Task[None] | None = None
         self.max_concurrent_tasks = options.max_concurrent_tasks or os.cpu_count() or 1
+        self.reserved_rpc_tasks = (
+            options.reserved_rpc_tasks
+            if options.reserved_rpc_tasks is not None
+            else int(self.max_concurrent_tasks >= 2)
+        )
+        self._should_warn_rpc_capacity = (
+            options.reserved_rpc_tasks is None and self.max_concurrent_tasks == 1
+        )
+        self.reserved_rpc_declines = 0
         self.idle_streams_per_queue = options.idle_streams_per_queue or 1
         self.stop_drain_timeout = options.stop_drain_timeout.total_seconds() or 30
         self._dispatcher = dispatcher
@@ -132,6 +145,10 @@ class PullWorkerRuntime:
         return None
 
     async def run(self) -> None:
+        if self._should_warn_rpc_capacity:
+            _LOGGER.warning(
+                "RPCs may wait behind long-running Steps; use max_concurrent_tasks >= 2 or set reserved_rpc_tasks explicitly"
+            )
         self._loop = asyncio.get_running_loop()
         self._loop_thread_id = threading.get_ident()
         self._wake = asyncio.Event()
@@ -193,11 +210,24 @@ class PullWorkerRuntime:
                 _LOGGER.error("Worker task stream failed", exc_info=failure)
         self.signal()
 
-    def accept_task(self) -> bool:
+    def accept_task(self, task: pb.PollWorkerTask | None = None) -> bool:
         if self._stopping is None or self._stopping.is_set():
             return False
         if self._busy_tasks >= self.max_concurrent_tasks:
             return False
+        if (
+            task is not None
+            and task.WhichOneof("request") in ("wait_for", "execute")
+            and self.reserved_rpc_tasks > 0
+            and self._busy_tasks >= self.max_concurrent_tasks - self.reserved_rpc_tasks
+            and self.unknown_task_reason(task) is None
+        ):
+            self.reserved_rpc_declines += 1
+            request = task.wait_for if task.HasField("wait_for") else task.execute
+            _LOGGER.debug(
+                "decline Step task for reserved RPC capacity: %s", request.flow_type
+            )
+            return False
         self._busy_tasks += 1
         self.signal()
         return True
@@ -288,6 +318,7 @@ class PullStream:
         self._call: Any = None
         self._send_lock = asyncio.Lock()
         self._is_reconnecting = False
+        self._stream_metadata: Any = None
         self._deadline_handle: asyncio.TimerHandle | None = None
         self._supports_unknown_name_rejection = False
 
@@ -296,6 +327,7 @@ class PullStream:
 
     async def run(self) -> None:
         backoff = 0.1
+        has_retried_unauthenticated = False
         while not self.runtime.is_stopping:
             self.is_draining = False
             was_registered = False
@@ -310,12 +342,35 @@ class PullStream:
                 )
             finally:
                 code = error.code() if error is not None else None
+                if self._is_registered:
+                    has_retried_unauthenticated = False
+                credential = (
+                    self.runtime._credential_runner._credential
+                    if self.runtime._credential_runner is not None
+                    else None
+                )
+                can_reject_token = (
+                    code is grpc.StatusCode.UNAUTHENTICATED
+                    and self._stream_metadata is not None
+                    and credential is not None
+                    and callable(getattr(credential, "_reject_token", None))
+                )
+                should_retry_unauthenticated = (
+                    can_reject_token and not has_retried_unauthenticated
+                )
+                if can_reject_token and credential is not None:
+                    _reject_credential_token(credential, self._stream_metadata)
                 rejection = self.runtime.registration.record_stream_result(
                     self.queue,
                     self._is_registered,
-                    error if code in _PERMANENT_CODES else None,
+                    (
+                        error
+                        if code in _PERMANENT_CODES and not should_retry_unauthenticated
+                        else None
+                    ),
                 )
                 self._is_registered = False
+                self._stream_metadata = None
                 if (
                     rejection is not None
                     and self.runtime._stop_on_permanent_stream_rejection
@@ -331,6 +386,11 @@ class PullStream:
                 or self._is_retiring
             ):
                 break
+            if should_retry_unauthenticated:
+                has_retried_unauthenticated = True
+                continue
+            if code is grpc.StatusCode.UNAUTHENTICATED:
+                has_retried_unauthenticated = False
             if self._is_reconnecting:
                 self._is_reconnecting = False
                 backoff = 0.1
@@ -364,6 +424,7 @@ class PullStream:
         if self.runtime.is_stopping:
             return False
         self._supports_unknown_name_rejection = False
+        self._stream_metadata = None
         if self.runtime._credential_runner is None:
             self._call = self.stub.PollWorkerTask()
         else:
@@ -372,6 +433,7 @@ class PullStream:
             )
             if self.runtime.is_stopping:
                 return False
+            self._stream_metadata = metadata
             self._call = self.stub.PollWorkerTask(metadata=metadata)
         await self._send(
             pb.PollWorkerTaskRequest(
@@ -524,7 +586,7 @@ class PullStream:
         if unknown_reason is not None:
             await self._reject_task(task, unknown_reason)
             return
-        if self.is_busy or self.is_draining or not self.runtime.accept_task():
+        if self.is_busy or self.is_draining or not self.runtime.accept_task(task):
             await self._reject_task(task)
             return
         self.is_busy = True

{dex_python_sdk-1.14.1 → dex_python_sdk-1.16.0}/dex/async_worker.py

@@ -221,6 +221,23 @@ class AsyncWorker:
             raise RuntimeError("wait_pull_registered requires a pull Worker")
         await self._pull.registration.wait_async(timeout)
 
+    @property
+    def pull_reserved_rpc_declines(self) -> int:
+        """Return cumulative Step declines caused by reserved RPC capacity.
+
+        Full-capacity rejections and unknown names are excluded. The count never
+        resets during this Worker's lifetime and stays available after stop.
+        Reading it alongside handlers is safe. Listening Workers return zero.
+        Export this value to a monitoring counter to identify reserve pressure.
+
+        Returns:
+            The cumulative reserve-decline count, or zero for a listening Worker.
+
+        Examples:
+            >>> declines = worker.pull_reserved_rpc_declines
+        """
+        return 0 if self._pull is None else self._pull.reserved_rpc_declines
+
     def pull_registration_state(self) -> PullRegistrationState:
         """Return an immutable snapshot of pull queue registrations.