dex-python-sdk 1.10.1__tar.gz → 1.12.0__tar.gz

{dex_python_sdk-1.10.1 → dex_python_sdk-1.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dex-python-sdk
-Version: 1.10.1
+Version: 1.12.0
 Requires-Dist: grpcio>=1.83.0
 Requires-Dist: grpcio-status>=1.83.0
 Requires-Dist: protobuf>=7.35.1
@@ -38,10 +38,19 @@ startup synchronization, blob hydration, and pull WorkerTaskService streams. A
 pull stream reads the token when it opens and again on every reconnect. For a
 custom credential, implement Credential.token() to return the current token;
 the method may be called concurrently, and an exception fails the call before
-it contacts Dex. AsyncClient and AsyncWorker also accept an awaitable token.
-Synchronous tokens run in an executor for async calls and on the calling thread
-for sync calls. Async calls bound token lookup by the RPC deadline when set;
-pull stream opens use a 30-second lookup budget. Sync token lookup must return
+it contacts Dex. Client and AsyncClient calls raise the credential's exception
+unchanged, including TimeoutError when the RPC deadline has not expired.
+AsyncClient and AsyncWorker also accept an awaitable token. Synchronous tokens
+run in a dedicated four-thread executor per async Client or Worker, created
+lazily and shut down without waiting for a blocked token call on close. Sync
+calls run token lookup on the calling thread. Async calls bound token lookup
+by the RPC deadline when set; pull stream opens use a 30-second lookup budget.
+Declare token with `async def` to invoke it on the event loop. A plain token
+function runs once in the SDK executor; if it returns a coroutine, that
+coroutine runs on the event loop. A loop-bound Future from another loop raises
+TypeError; use `async def` to create it on the right loop. A blocked token keeps
+one executor thread busy until it returns, and an unclosed Client or Worker may
+retain up to four idle credential threads. Sync token lookup must return
 promptly. By default, no authorization metadata is sent. Dex uses plaintext
 gRPC, so protect the network path that carries the token.
 
@@ -91,6 +100,14 @@ max_concurrent_tasks)` streams because busy streams are replaced with idle
 streams and retained after tasks finish. Size these options for the server's
 per-Worker stream cap (256 by default).
 
+The Worker rejects an unregistered Step or RPC by name. Dex retries that task
+on another pull Worker without consuming a Step attempt, including when the
+Flow was sticky to the rejecting Worker. If no capable Worker accepts it within
+the server's dispatch wait (10 seconds by default, or the local Step attempt),
+the Step follows its failure and retry policy; an RPC caller receives the named
+error. This reassignment requires a Dex Server that supports it. Against older
+Servers, the Worker returns the not-registered error as before.
+
 The generated `app_worker_id` uses the first nonempty `DEX_WORKER_INSTANCE`,
 `POD_NAME`, `HOSTNAME`, or system hostname, plus eight random base32 characters.
 An override must uniquely identify one Worker process and change after restart.
@@ -297,6 +314,39 @@ failures do not roll back Flow Attributes. `None` preserves current targets;
 an explicit empty list disables future synchronization while retaining protocol
 presence.
 
+Declare short-lived secrets, such as API keys, with `hard_delete_on_null=True`:
+
+```python
+api_key = dex.Attribute("api-key", str, hard_delete_on_null=True)
+provider_tokens = dex.AttributeMap("provider-token", str, hard_delete_on_null=True)
+```
+
+Dex stores these values only in its blob store, never in Flow history. Deleting
+the Attribute or setting it to `None` permanently deletes the stored value, and a
+deleted value reads as absent. Only string and object values, including
+`bytes`, are supported; `int`, `bool`, and `float` types are rejected. The
+option cannot be combined with an index or `sync_to_attribute_store`, and
+`wait_for_attribute_match` rejects it. Handlers receive the value inline, and
+the SDK never stores it in the local blob cache.
+
+Hard-delete writes require a Server with an enabled blob store and Server
+protocol version 3 or later. Workers still start against older Servers.
+There, a Step, RPC, or timeout handler result that writes a hard-delete
+Attribute fails that handler attempt with `RuntimeError`, as does a `wait_for`
+result whose SubFlow start initializes one. A Client checks the Server protocol
+once, before its first `start_flow` that initializes one.
+
+Each Flow keeps one stored copy per Attribute key, overwritten in place. A Step
+or RPC attempt writes or deletes the stored value before its result commits, so
+retries, other Steps, and RPCs can observe a write or deletion from an attempt
+that later failed. Concurrent writers of one key may leave a value other than
+the last committed write. A failed `start_flow` or SubFlow start for a running
+Flow ID still overwrites, or for `None` deletes, the running Flow's value.
+Executions that reuse a Flow ID share the value, and time travel does not
+restore a deleted value. Values remain until deleted or until retention and blob
+cleanup remove the Flow. Blob store backups or bucket versioning can retain
+copies.
+
 ```
 pip install dex-python-sdk==0.1.0
 ```
@@ -705,6 +755,11 @@ Both `Worker` and `AsyncWorker` call `GetServerInfo`, negotiate the highest
 common protocol, synchronize Attribute indexes, and only then bind
 WorkerService. The initial Python SDK interval is `[1,1]`. Missing information,
 invalid or disjoint intervals, and RPC failures stop startup before binding.
+The current Python SDK interval is `[1,3]`. Protocol 3 adds hard-delete
+Attributes. Workers record the negotiated version at startup and reject each
+handler result that contains a hard-delete Attribute write when it is below 3.
+Clients negotiate lazily, once, before the first `start_flow` that initializes
+a hard-delete Attribute.
 
 The diagnostic SDK version comes from Python distribution metadata for
 `dex-python-sdk` and falls back to `dev` when the distribution is unavailable.

{dex_python_sdk-1.10.1 → dex_python_sdk-1.12.0}/README.md

@@ -22,10 +22,19 @@ startup synchronization, blob hydration, and pull WorkerTaskService streams. A
 pull stream reads the token when it opens and again on every reconnect. For a
 custom credential, implement Credential.token() to return the current token;
 the method may be called concurrently, and an exception fails the call before
-it contacts Dex. AsyncClient and AsyncWorker also accept an awaitable token.
-Synchronous tokens run in an executor for async calls and on the calling thread
-for sync calls. Async calls bound token lookup by the RPC deadline when set;
-pull stream opens use a 30-second lookup budget. Sync token lookup must return
+it contacts Dex. Client and AsyncClient calls raise the credential's exception
+unchanged, including TimeoutError when the RPC deadline has not expired.
+AsyncClient and AsyncWorker also accept an awaitable token. Synchronous tokens
+run in a dedicated four-thread executor per async Client or Worker, created
+lazily and shut down without waiting for a blocked token call on close. Sync
+calls run token lookup on the calling thread. Async calls bound token lookup
+by the RPC deadline when set; pull stream opens use a 30-second lookup budget.
+Declare token with `async def` to invoke it on the event loop. A plain token
+function runs once in the SDK executor; if it returns a coroutine, that
+coroutine runs on the event loop. A loop-bound Future from another loop raises
+TypeError; use `async def` to create it on the right loop. A blocked token keeps
+one executor thread busy until it returns, and an unclosed Client or Worker may
+retain up to four idle credential threads. Sync token lookup must return
 promptly. By default, no authorization metadata is sent. Dex uses plaintext
 gRPC, so protect the network path that carries the token.
 
@@ -75,6 +84,14 @@ max_concurrent_tasks)` streams because busy streams are replaced with idle
 streams and retained after tasks finish. Size these options for the server's
 per-Worker stream cap (256 by default).
 
+The Worker rejects an unregistered Step or RPC by name. Dex retries that task
+on another pull Worker without consuming a Step attempt, including when the
+Flow was sticky to the rejecting Worker. If no capable Worker accepts it within
+the server's dispatch wait (10 seconds by default, or the local Step attempt),
+the Step follows its failure and retry policy; an RPC caller receives the named
+error. This reassignment requires a Dex Server that supports it. Against older
+Servers, the Worker returns the not-registered error as before.
+
 The generated `app_worker_id` uses the first nonempty `DEX_WORKER_INSTANCE`,
 `POD_NAME`, `HOSTNAME`, or system hostname, plus eight random base32 characters.
 An override must uniquely identify one Worker process and change after restart.
@@ -281,6 +298,39 @@ failures do not roll back Flow Attributes. `None` preserves current targets;
 an explicit empty list disables future synchronization while retaining protocol
 presence.
 
+Declare short-lived secrets, such as API keys, with `hard_delete_on_null=True`:
+
+```python
+api_key = dex.Attribute("api-key", str, hard_delete_on_null=True)
+provider_tokens = dex.AttributeMap("provider-token", str, hard_delete_on_null=True)
+```
+
+Dex stores these values only in its blob store, never in Flow history. Deleting
+the Attribute or setting it to `None` permanently deletes the stored value, and a
+deleted value reads as absent. Only string and object values, including
+`bytes`, are supported; `int`, `bool`, and `float` types are rejected. The
+option cannot be combined with an index or `sync_to_attribute_store`, and
+`wait_for_attribute_match` rejects it. Handlers receive the value inline, and
+the SDK never stores it in the local blob cache.
+
+Hard-delete writes require a Server with an enabled blob store and Server
+protocol version 3 or later. Workers still start against older Servers.
+There, a Step, RPC, or timeout handler result that writes a hard-delete
+Attribute fails that handler attempt with `RuntimeError`, as does a `wait_for`
+result whose SubFlow start initializes one. A Client checks the Server protocol
+once, before its first `start_flow` that initializes one.
+
+Each Flow keeps one stored copy per Attribute key, overwritten in place. A Step
+or RPC attempt writes or deletes the stored value before its result commits, so
+retries, other Steps, and RPCs can observe a write or deletion from an attempt
+that later failed. Concurrent writers of one key may leave a value other than
+the last committed write. A failed `start_flow` or SubFlow start for a running
+Flow ID still overwrites, or for `None` deletes, the running Flow's value.
+Executions that reuse a Flow ID share the value, and time travel does not
+restore a deleted value. Values remain until deleted or until retention and blob
+cleanup remove the Flow. Blob store backups or bucket versioning can retain
+copies.
+
 ```
 pip install dex-python-sdk==0.1.0
 ```
@@ -689,6 +739,11 @@ Both `Worker` and `AsyncWorker` call `GetServerInfo`, negotiate the highest
 common protocol, synchronize Attribute indexes, and only then bind
 WorkerService. The initial Python SDK interval is `[1,1]`. Missing information,
 invalid or disjoint intervals, and RPC failures stop startup before binding.
+The current Python SDK interval is `[1,3]`. Protocol 3 adds hard-delete
+Attributes. Workers record the negotiated version at startup and reject each
+handler result that contains a hard-delete Attribute write when it is below 3.
+Clients negotiate lazily, once, before the first `start_flow` that initializes
+a hard-delete Attribute.
 
 The diagnostic SDK version comes from Python distribution metadata for
 `dex-python-sdk` and falls back to `dev` when the distribution is unavailable.

{dex_python_sdk-1.10.1 → dex_python_sdk-1.12.0}/dex/_async_worker_dispatcher.py

@@ -16,6 +16,7 @@ from typing import Any, Callable
 
 from dex._async_value_hydrator import AsyncValueHydrator
 from dex._invocation_context import InvocationContext, InvocationMethod
+from dex._server_protocol import NegotiatedServerProtocol
 from dex._value_mapper import ValueMapper
 from dex._worker_dispatcher import _TIMEOUT_HANDLER_STEP_TYPE, WorkerDispatcher
 from dex.dexpb import dex_pb2 as pb
@@ -64,10 +65,16 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
         registry: Registry,
         values: ValueMapper,
         hydrator: AsyncValueHydrator,
+        server_protocol: NegotiatedServerProtocol,
     ) -> None:
+        if server_protocol is None:
+            raise ValueError(
+                "AsyncWorkerDispatcher requires a NegotiatedServerProtocol"
+            )
         self._registry = registry
         self._values = values
         self._async_hydrator = hydrator
+        self._server_protocol = server_protocol
 
     async def invoke_wait_for(  # type: ignore[override]
         self,
@@ -127,7 +134,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             if not isinstance(wait, Wait):
                 raise TypeError("wait_for must return Wait")
             response = pb.InvokeWaitForMethodResponse(
-                upsert_attributes=list(context.attribute_writes.values()),
+                upsert_attributes=self._result_attribute_writes(context),
                 upsert_step_exe_locals=list(context.local_writes.values()),
                 record_events=context.events,
                 publish_to_channel=context.publications,
@@ -229,7 +236,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
                 raise TypeError("execute must return StepDecision")
             response = pb.InvokeExecuteMethodResponse(
                 step_decision=self._map_decision(flow, decision),
-                upsert_attributes=list(context.attribute_writes.values()),
+                upsert_attributes=self._result_attribute_writes(context),
                 record_events=context.events,
                 upsert_step_exe_locals=list(context.local_writes.values()),
                 publish_to_channel=context.publications,
@@ -309,7 +316,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
                 raise TypeError("handle_timeout must return StepDecision")
             response = pb.InvokeExecuteMethodResponse(
                 step_decision=self._map_decision(flow, decision),
-                upsert_attributes=list(context.attribute_writes.values()),
+                upsert_attributes=self._result_attribute_writes(context),
                 record_events=context.events,
                 upsert_step_exe_locals=list(context.local_writes.values()),
                 publish_to_channel=context.publications,
@@ -361,7 +368,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             returned = await returned
         try:
             response = pb.InvokeWorkerRPCResponse(
-                upsert_attributes=list(context.attribute_writes.values()),
+                upsert_attributes=self._result_attribute_writes(context),
                 record_events=context.events,
                 delete_from_channel=context.channel_deletions,
                 publish_to_channel=context.publications,

dex_python_sdk-1.12.0/dex/_credential_transport.py

@@ -0,0 +1,242 @@
+# Copyright (c) 2026 Super Durable, Inc.
+#
+# Licensed under the Sustainable Use License 1.0.
+# You may not use this file except in compliance with the License.
+# See the LICENSE file in the repository root.
+#
+# SPDX-License-Identifier: LicenseRef-Sustainable-Use-1.0
+
+"""Attach bearer metadata to outbound gRPC calls."""
+
+# mypy: disable-error-code=type-arg
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import queue
+import threading
+from collections.abc import Callable
+from concurrent.futures import Future
+from typing import Any, NamedTuple, cast
+
+import grpc
+
+from dex.credential import Credential
+
+_TOKEN_WORKERS = 4
+
+
+class AsyncCredentialRunner:
+    def __init__(self, credential: Credential) -> None:
+        self._credential = credential
+        self._lock = threading.Lock()
+        self._jobs: queue.SimpleQueue[Future[Any] | None] = queue.SimpleQueue()
+        self._threads: list[threading.Thread] = []
+        self._pending: set[Future[Any]] = set()
+        self._closed = False
+
+    async def token(self) -> str:
+        token_method = self._credential.token
+        if inspect.iscoroutinefunction(token_method) or inspect.iscoroutinefunction(
+            getattr(token_method, "__call__", None)
+        ):
+            token = token_method()
+            if inspect.isawaitable(token):
+                return await token
+            return token
+        token = await asyncio.wrap_future(self._submit())
+        if inspect.isawaitable(token):
+            if isinstance(token, asyncio.Future):
+                try:
+                    token_loop = token.get_loop()
+                except RuntimeError as error:
+                    raise TypeError(
+                        "loop-bound credentials must use async def token()"
+                    ) from error
+                if token_loop is not asyncio.get_running_loop():
+                    raise TypeError("loop-bound credentials must use async def token()")
+            return await token
+        return token
+
+    def _submit(self) -> Future[Any]:
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("credential runner is closed")
+            future: Future[Any] = Future()
+            self._pending.add(future)
+            if len(self._threads) < _TOKEN_WORKERS:
+                thread = threading.Thread(
+                    target=self._run,
+                    name="dex-python-credential",
+                    daemon=True,
+                )
+                self._threads.append(thread)
+                thread.start()
+            self._jobs.put(future)
+            return future
+
+    def _run(self) -> None:
+        while True:
+            future = self._jobs.get()
+            if future is None:
+                return
+            with self._lock:
+                if future.done() or not future.set_running_or_notify_cancel():
+                    self._pending.discard(future)
+                    continue
+            try:
+                token = self._credential.token()
+                failure = None
+            except BaseException as error:
+                token = None
+                failure = error
+            with self._lock:
+                if not future.done():
+                    if failure is None:
+                        future.set_result(token)
+                    else:
+                        future.set_exception(failure)
+                self._pending.discard(future)
+
+    def close(self) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            for future in self._pending:
+                if not future.done():
+                    future.set_exception(RuntimeError("credential runner is closed"))
+            self._pending.clear()
+            while True:
+                try:
+                    self._jobs.get_nowait()
+                except queue.Empty:
+                    break
+            for _ in self._threads:
+                self._jobs.put(None)
+
+
+class _SyncCallDetails(NamedTuple):
+    method: str
+    timeout: float | None
+    metadata: Any
+    credentials: grpc.CallCredentials | None
+    wait_for_ready: bool | None
+    compression: grpc.Compression | None
+
+
+def credential_metadata(credential: Credential | None, metadata: Any = None) -> Any:
+    if credential is None:
+        return metadata
+    token = credential.token()
+    if inspect.isawaitable(token):
+        if inspect.iscoroutine(token):
+            token.close()
+        raise TypeError("synchronous calls require a synchronous credential")
+    return tuple(metadata or ()) + (("authorization", "Bearer " + token),)
+
+
+async def async_credential_metadata(
+    runner: AsyncCredentialRunner, metadata: Any = None, timeout: float | None = None
+) -> Any:
+    if timeout is None:
+        token = await runner.token()
+    else:
+        lookup = asyncio.create_task(runner.token())
+        try:
+            completed, _ = await asyncio.wait({lookup}, timeout=timeout)
+            if not completed:
+                lookup.cancel()
+                await asyncio.gather(lookup, return_exceptions=True)
+                raise credential_deadline_error()
+            token = lookup.result()
+        except BaseException:
+            lookup.cancel()
+            await asyncio.gather(lookup, return_exceptions=True)
+            raise
+    return tuple(metadata or ()) + (("authorization", "Bearer " + token),)
+
+
+def credential_deadline_error() -> grpc.aio.AioRpcError:
+    return grpc.aio.AioRpcError(
+        grpc.StatusCode.DEADLINE_EXCEEDED,
+        grpc.aio.Metadata(),
+        grpc.aio.Metadata(),
+        "credential lookup exceeded RPC deadline",
+        None,
+    )
+
+
+class _SyncCredentialInterceptor(grpc.UnaryUnaryClientInterceptor):
+    def __init__(self, credential: Credential) -> None:
+        self._credential = credential
+
+    def intercept_unary_unary(
+        self,
+        continuation: Callable[..., Any],
+        client_call_details: grpc.ClientCallDetails,
+        request: Any,
+    ) -> Any:
+        details = _SyncCallDetails(
+            client_call_details.method,
+            client_call_details.timeout,
+            credential_metadata(self._credential, client_call_details.metadata),
+            client_call_details.credentials,
+            client_call_details.wait_for_ready,
+            client_call_details.compression,
+        )
+        return continuation(details, request)
+
+
+class _AsyncCredentialInterceptor(grpc.aio.UnaryUnaryClientInterceptor):
+    def __init__(self, runner: AsyncCredentialRunner) -> None:
+        self._runner = runner
+
+    async def intercept_unary_unary(
+        self,
+        continuation: Callable[..., Any],
+        client_call_details: grpc.aio.ClientCallDetails,
+        request: Any,
+    ) -> Any:
+        loop = asyncio.get_running_loop()
+        started_at = loop.time()
+        metadata = await async_credential_metadata(
+            self._runner, client_call_details.metadata, client_call_details.timeout
+        )
+        remaining_timeout = client_call_details.timeout
+        if remaining_timeout is not None:
+            remaining_timeout -= loop.time() - started_at
+            if remaining_timeout <= 0:
+                raise credential_deadline_error()
+        details = grpc.aio.ClientCallDetails(
+            client_call_details.method,
+            remaining_timeout,
+            metadata,
+            client_call_details.credentials,
+            client_call_details.wait_for_ready,
+        )
+        return await continuation(details, request)
+
+
+def sync_channel(server_address: str, credential: Credential | None) -> grpc.Channel:
+    channel = grpc.insecure_channel(server_address)
+    if credential is None:
+        return channel
+    return grpc.intercept_channel(channel, _SyncCredentialInterceptor(credential))
+
+
+def async_channel(
+    server_address: str, credential_runner: AsyncCredentialRunner | None
+) -> grpc.aio.Channel:
+    if credential_runner is None:
+        return grpc.aio.insecure_channel(server_address)
+    return grpc.aio.insecure_channel(
+        server_address,
+        interceptors=[
+            cast(
+                grpc.aio.ClientInterceptor,
+                _AsyncCredentialInterceptor(credential_runner),
+            )
+        ],
+    )

{dex_python_sdk-1.10.1 → dex_python_sdk-1.12.0}/dex/_invocation_context.py

@@ -15,7 +15,12 @@ from urllib.parse import unquote
 
 from dex._utils import require_map_instance, require_name
 from dex._value_mapper import ValueMapper
-from dex.attribute import Attribute, AttributeMap, _apply_attribute_store_sync
+from dex.attribute import (
+    Attribute,
+    AttributeMap,
+    _apply_attribute_store_sync,
+    _apply_hard_delete_on_null,
+)
 from dex.channel import Channel, ChannelMap, ChannelMessage
 from dex.codec import Codec
 from dex.context import RecoveryErrorInfo
@@ -349,6 +354,7 @@ class InvocationContext:
         if index is not None:
             write.index_config.CopyFrom(index)
         _apply_attribute_store_sync(write, definition)
+        _apply_hard_delete_on_null(write, definition)
         self.attribute_writes[key] = write
 
     def _delete_attribute(
@@ -366,6 +372,7 @@ class InvocationContext:
         if index is not None:
             write.index_config.CopyFrom(index)
         _apply_attribute_store_sync(write, definition)
+        _apply_hard_delete_on_null(write, definition)
         self.attribute_writes[key] = write
 
     def _publish_channel(

{dex_python_sdk-1.10.1 → dex_python_sdk-1.12.0}/dex/_pull_worker.py

@@ -22,13 +22,12 @@ from typing import Any
 
 import grpc
 
-from dex._credential_transport import async_credential_metadata
+from dex._credential_transport import AsyncCredentialRunner, 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
 from dex.dexpb import dex_pb2_grpc
 from dex.flow import Registry
-from dex.credential import Credential
 from dex.worker_options import PullWorkerOptions
 
 _LOGGER = logging.getLogger(__name__)
@@ -81,8 +80,9 @@ class PullWorkerRuntime:
         server_address: str,
         executor: ThreadPoolExecutor | None,
         app_worker_id: str | None = None,
-        credential: Credential | None = None,
+        credential_runner: AsyncCredentialRunner | None = None,
     ) -> None:
+        self._registry = registry
         self.app_worker_id = app_worker_id or resolve_app_worker_id(registry, options)
         self.flow_types = sorted(registry._registered_flows)
         self.queues = (self.app_worker_id, *self.flow_types)
@@ -92,7 +92,7 @@ class PullWorkerRuntime:
         self._dispatcher = dispatcher
         self._server_address = server_address
         self._executor = executor
-        self._credential = credential
+        self._credential_runner = credential_runner
         self._loop: asyncio.AbstractEventLoop | None = None
         self._loop_thread_id: int | None = None
         self._started = threading.Event()
@@ -104,6 +104,26 @@ class PullWorkerRuntime:
         self._busy_tasks = 0
         self._unsupported_logged = False
 
+    def unknown_task_reason(
+        self, task: pb.PollWorkerTask
+    ) -> pb.PollWorkerTaskRejectReason | None:
+        kind = task.WhichOneof("request")
+        if kind == "rpc":
+            flow = self._registry._registered_flows.get(task.rpc.flow_type)
+            if flow is None or task.rpc.rpc_name not in flow.rpcs:
+                return pb.POLL_WORKER_TASK_REJECT_REASON_UNKNOWN_RPC
+        elif kind in ("wait_for", "execute"):
+            request = task.wait_for if kind == "wait_for" else task.execute
+            flow = self._registry._registered_flows.get(request.flow_type)
+            if flow is None:
+                return pb.POLL_WORKER_TASK_REJECT_REASON_UNKNOWN_STEP
+            if kind == "execute" and request.step_type == "sys:timeout_handler":
+                if not flow.has_timeout_handler:
+                    return pb.POLL_WORKER_TASK_REJECT_REASON_UNKNOWN_STEP
+            elif request.step_type not in flow.steps:
+                return pb.POLL_WORKER_TASK_REJECT_REASON_UNKNOWN_STEP
+        return None
+
     async def run(self) -> None:
         self._loop = asyncio.get_running_loop()
         self._loop_thread_id = threading.get_ident()
@@ -250,6 +270,7 @@ class PullStream:
         self._send_lock = asyncio.Lock()
         self._is_reconnecting = False
         self._deadline_handle: asyncio.TimerHandle | None = None
+        self._supports_unknown_name_rejection = False
 
     def on_done(self, completed: asyncio.Task[None]) -> None:
         self.runtime._stream_done(self, completed)
@@ -300,12 +321,17 @@ class PullStream:
                 backoff = min(backoff * 2, 5)
 
     async def _poll(self) -> bool:
-        if self.runtime._credential is None:
+        if self.runtime.is_stopping:
+            return False
+        self._supports_unknown_name_rejection = False
+        if self.runtime._credential_runner is None:
             self._call = self.stub.PollWorkerTask()
         else:
             metadata = await async_credential_metadata(
-                self.runtime._credential, timeout=30.0
+                self.runtime._credential_runner, timeout=30.0
             )
+            if self.runtime.is_stopping:
+                return False
             self._call = self.stub.PollWorkerTask(metadata=metadata)
         await self._send(
             pb.PollWorkerTaskRequest(
@@ -361,6 +387,9 @@ class PullStream:
                         return False
                     if kind != "registered":
                         raise RuntimeError("Worker task stream did not register")
+                    self._supports_unknown_name_rejection = (
+                        frame.registered.supports_unknown_name_rejection
+                    )
                     was_registered = True
                     continue
                 if kind == "task":
@@ -440,6 +469,14 @@ class PullStream:
                     pass
 
     async def _start_task(self, task: pb.PollWorkerTask) -> None:
+        unknown_reason = (
+            self.runtime.unknown_task_reason(task)
+            if self._supports_unknown_name_rejection
+            else None
+        )
+        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():
             await self._reject_task(task)
             return
@@ -455,12 +492,16 @@ class PullStream:
             )
         self.runtime.signal()
 
-    async def _reject_task(self, task: pb.PollWorkerTask) -> None:
+    async def _reject_task(
+        self,
+        task: pb.PollWorkerTask,
+        reason: pb.PollWorkerTaskRejectReason = pb.POLL_WORKER_TASK_REJECT_REASON_NO_CAPACITY,
+    ) -> None:
         await self._send(
             pb.PollWorkerTaskRequest(
                 reject=pb.PollWorkerTaskReject(
                     task_id=task.task_id,
-                    reason=pb.POLL_WORKER_TASK_REJECT_REASON_NO_CAPACITY,
+                    reason=reason,
                 )
             )
         )