dex-python-sdk 1.20.0__tar.gz → 1.25.0__tar.gz

{dex_python_sdk-1.20.0 → dex_python_sdk-1.25.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dex-python-sdk
-Version: 1.20.0
+Version: 1.25.0
 Requires-Dist: grpcio>=1.83.0
 Requires-Dist: grpcio-status>=1.83.0
 Requires-Dist: protobuf>=7.35.1
@@ -95,14 +95,17 @@ 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.
 
-## Pull mode
+## Workers
 
-Set `WorkerOptions.pull` to run a `Worker` or `AsyncWorker` with outbound task
-streams. The Worker connects to the Dex server through `server_address` and
-does not open a listener. Advertise `worker.worker_target` through
-`ClientOptions.worker_target`. A single Flow can override the Client default
-with `StartFlowOptions(config_override=FlowConfig(worker_target=...))`; a
-SubFlow can use `SubFlowOptions` with the same `FlowConfig` override.
+Workers receive tasks through outbound pull streams. `Worker` and `AsyncWorker`
+need only an outbound connection to the Dex Server at `server_address`. Configure
+capacity and shutdown with `WorkerOptions.pull`; its default is
+`PullWorkerOptions()`. Clients and Flows use pull routing automatically.
+
+StartFlow and SubFlow always select pull Workers. Updating Flow configuration
+preserves its existing Worker target; an empty `FlowConfig()` changes nothing.
+Before upgrading an application with active push Flows, use an SDK version that
+supports Worker targets to migrate those Flows to pull routing.
 
 ```python
 import threading
@@ -118,10 +121,11 @@ worker = Worker(
 )
 worker_thread = threading.Thread(target=worker.start)
 worker_thread.start()
+worker.wait_pull_registered(timeout=30)
 client = Client(
     registry,
     cache,
-    ClientOptions(server_address="dex.example:8801", worker_target=worker.worker_target),
+    ClientOptions(server_address="dex.example:8801"),
 )
 try:
     client.start_flow(flow, flow_id, input_value)
@@ -135,13 +139,18 @@ The Worker polls its sticky queue and every registered Flow type queue. Flow
 type names must be unique within a Dex server. `max_concurrent_tasks` defaults
 to the CPU count, and `idle_streams_per_queue` defaults to one. At capacity,
 the Worker rejects an assignment so Dex can reassign it. Each stream runs one
-task at a time; Step heartbeat and stream writes use the same stream. After tasks finish, the Worker
-gracefully closes surplus idle streams while preserving readiness. After load drops,
+task at a time; Step heartbeat and stream writes use the same stream. After
+tasks finish, the Worker gracefully closes surplus idle streams while preserving
+readiness. The SDK sends an upstream Pong every 20 seconds while a task runs and
+stops before its final frame. This interval is internal and needs no Worker
+option. After load drops,
 the steady-state bound is `max_concurrent_tasks + (1 + number of Flow types) *
-idle_streams_per_queue` streams. The transient peak during a burst remains bounded
-by the matching caps. Size these options for `matching.maxConcurrentTasksCap` per
-app Worker ID per matching owner (256 by default) and `matching.maxStreams` per
-Matching instance (100,000 by default).
+idle_streams_per_queue` streams.
+Allow up to one more stream per queue for 30 seconds after reserve declines,
+or up to `1 + number of Flow types` additional streams. The transient peak
+during a burst remains bounded by the matching caps. Size these options for
+`matching.maxConcurrentTasksCap` per app Worker ID per matching owner (256 by
+default) and `matching.maxStreams` per Matching instance (100,000 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
@@ -155,10 +164,7 @@ 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.
 It must contain 1-63 ASCII letters, digits, dots, underscores, or hyphens and
-must not equal a registered Flow type. Pull mode cannot configure a bind address
-or an advertised address or headless target. Use `WorkerTarget(pull=True)` for
-Client defaults or Flow overrides when the Worker target is not directly in
-scope.
+must not equal a registered Flow type.
 
 Pull streams reconnect after failures. A server-directed reconnect waits for
 an active task to finish, then waits up to `stop_drain_timeout` for the server
@@ -181,7 +187,7 @@ Worker ID queue to register at least once. Timeout raises `TimeoutError`;
 async cancellation cancels only that wait. Startup failures reach waiters.
 A permanent rejection before readiness raises `PullRegistrationError`, whose
 `queue` and `error` identify the rejected queue and original transport error.
-Listening mode and stop before readiness raise `RuntimeError`.
+Stopping before readiness raises `RuntimeError`.
 
 Use `worker.pull_registration_state()` for subsequent readiness probes. Its
 immutable `queues` expose `queue`, `has_registered_stream`, and
@@ -206,11 +212,11 @@ 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
+Flows must stay responsive because RPCs share Worker task capacity with Steps.
+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.
 
@@ -222,15 +228,19 @@ 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.
+dispatch wait, then fall back according to their policy. A queue that declines
+a Step for reserved RPC capacity keeps one additional idle stream until 30
+seconds after its last decline. Other queues retain their configured count.
+The Worker does not open the extra stream while all task slots are busy. Set
+`idle_streams_per_queue` to at least two for heavy Step backlogs. 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.
+to distinguish reserve pressure from full-capacity rejections. The count excludes
+full-capacity rejections and unknown names. It counts decline events, including
+repeated declines of the same Step dispatch while Dex retries it. Read its rate
+to assess reserve pressure. It never resets and remains available after shutdown.
 
 ## Flow state I/O
 
@@ -388,7 +398,7 @@ explicit codec only for a custom encoding or a type Registry cannot derive.
 partitions them by definition type.
 
 `Worker` and `AsyncWorker` synchronize all registered Indexed Attributes with
-Dex Server before opening their listener. Existing indexes return immediately;
+Dex Server before opening outbound task streams. Existing indexes return immediately;
 failure or the default two-minute deadline aborts startup. An indexed
 `AttributeMap` must provide one fixed `index_key`.
 
@@ -784,8 +794,8 @@ environment:
 ./run-integration-tests.sh
 ```
 
-Set `DEX_INTEG_WORKER_MODE=pull` to run the suite with pull Workers instead of
-listening Workers. The default is `push`.
+The suite always uses pull Workers. Set `DEX_INTEG_ENGINE=zd` to test against
+server-zd, or leave it unset to use the Temporal-backed server.
 
 ### Measure integration coverage
 
@@ -858,9 +868,9 @@ See [CONTRIBUTING.md](../CONTRIBUTING.md#releases-monorepo-tags) for monorepo ta
 ## Server protocol compatibility
 
 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.
+common protocol, synchronize Attribute indexes, and only then open outbound
+task streams. The initial Python SDK interval is `[1,1]`. Missing information,
+invalid or disjoint intervals, and RPC failures stop startup before polling.
 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.

{dex_python_sdk-1.20.0 → dex_python_sdk-1.25.0}/README.md

@@ -79,14 +79,17 @@ 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.
 
-## Pull mode
+## Workers
 
-Set `WorkerOptions.pull` to run a `Worker` or `AsyncWorker` with outbound task
-streams. The Worker connects to the Dex server through `server_address` and
-does not open a listener. Advertise `worker.worker_target` through
-`ClientOptions.worker_target`. A single Flow can override the Client default
-with `StartFlowOptions(config_override=FlowConfig(worker_target=...))`; a
-SubFlow can use `SubFlowOptions` with the same `FlowConfig` override.
+Workers receive tasks through outbound pull streams. `Worker` and `AsyncWorker`
+need only an outbound connection to the Dex Server at `server_address`. Configure
+capacity and shutdown with `WorkerOptions.pull`; its default is
+`PullWorkerOptions()`. Clients and Flows use pull routing automatically.
+
+StartFlow and SubFlow always select pull Workers. Updating Flow configuration
+preserves its existing Worker target; an empty `FlowConfig()` changes nothing.
+Before upgrading an application with active push Flows, use an SDK version that
+supports Worker targets to migrate those Flows to pull routing.
 
 ```python
 import threading
@@ -102,10 +105,11 @@ worker = Worker(
 )
 worker_thread = threading.Thread(target=worker.start)
 worker_thread.start()
+worker.wait_pull_registered(timeout=30)
 client = Client(
     registry,
     cache,
-    ClientOptions(server_address="dex.example:8801", worker_target=worker.worker_target),
+    ClientOptions(server_address="dex.example:8801"),
 )
 try:
     client.start_flow(flow, flow_id, input_value)
@@ -119,13 +123,18 @@ The Worker polls its sticky queue and every registered Flow type queue. Flow
 type names must be unique within a Dex server. `max_concurrent_tasks` defaults
 to the CPU count, and `idle_streams_per_queue` defaults to one. At capacity,
 the Worker rejects an assignment so Dex can reassign it. Each stream runs one
-task at a time; Step heartbeat and stream writes use the same stream. After tasks finish, the Worker
-gracefully closes surplus idle streams while preserving readiness. After load drops,
+task at a time; Step heartbeat and stream writes use the same stream. After
+tasks finish, the Worker gracefully closes surplus idle streams while preserving
+readiness. The SDK sends an upstream Pong every 20 seconds while a task runs and
+stops before its final frame. This interval is internal and needs no Worker
+option. After load drops,
 the steady-state bound is `max_concurrent_tasks + (1 + number of Flow types) *
-idle_streams_per_queue` streams. The transient peak during a burst remains bounded
-by the matching caps. Size these options for `matching.maxConcurrentTasksCap` per
-app Worker ID per matching owner (256 by default) and `matching.maxStreams` per
-Matching instance (100,000 by default).
+idle_streams_per_queue` streams.
+Allow up to one more stream per queue for 30 seconds after reserve declines,
+or up to `1 + number of Flow types` additional streams. The transient peak
+during a burst remains bounded by the matching caps. Size these options for
+`matching.maxConcurrentTasksCap` per app Worker ID per matching owner (256 by
+default) and `matching.maxStreams` per Matching instance (100,000 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
@@ -139,10 +148,7 @@ 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.
 It must contain 1-63 ASCII letters, digits, dots, underscores, or hyphens and
-must not equal a registered Flow type. Pull mode cannot configure a bind address
-or an advertised address or headless target. Use `WorkerTarget(pull=True)` for
-Client defaults or Flow overrides when the Worker target is not directly in
-scope.
+must not equal a registered Flow type.
 
 Pull streams reconnect after failures. A server-directed reconnect waits for
 an active task to finish, then waits up to `stop_drain_timeout` for the server
@@ -165,7 +171,7 @@ Worker ID queue to register at least once. Timeout raises `TimeoutError`;
 async cancellation cancels only that wait. Startup failures reach waiters.
 A permanent rejection before readiness raises `PullRegistrationError`, whose
 `queue` and `error` identify the rejected queue and original transport error.
-Listening mode and stop before readiness raise `RuntimeError`.
+Stopping before readiness raises `RuntimeError`.
 
 Use `worker.pull_registration_state()` for subsequent readiness probes. Its
 immutable `queues` expose `queue`, `has_registered_stream`, and
@@ -190,11 +196,11 @@ 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
+Flows must stay responsive because RPCs share Worker task capacity with Steps.
+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.
 
@@ -206,15 +212,19 @@ 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.
+dispatch wait, then fall back according to their policy. A queue that declines
+a Step for reserved RPC capacity keeps one additional idle stream until 30
+seconds after its last decline. Other queues retain their configured count.
+The Worker does not open the extra stream while all task slots are busy. Set
+`idle_streams_per_queue` to at least two for heavy Step backlogs. 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.
+to distinguish reserve pressure from full-capacity rejections. The count excludes
+full-capacity rejections and unknown names. It counts decline events, including
+repeated declines of the same Step dispatch while Dex retries it. Read its rate
+to assess reserve pressure. It never resets and remains available after shutdown.
 
 ## Flow state I/O
 
@@ -372,7 +382,7 @@ explicit codec only for a custom encoding or a type Registry cannot derive.
 partitions them by definition type.
 
 `Worker` and `AsyncWorker` synchronize all registered Indexed Attributes with
-Dex Server before opening their listener. Existing indexes return immediately;
+Dex Server before opening outbound task streams. Existing indexes return immediately;
 failure or the default two-minute deadline aborts startup. An indexed
 `AttributeMap` must provide one fixed `index_key`.
 
@@ -768,8 +778,8 @@ environment:
 ./run-integration-tests.sh
 ```
 
-Set `DEX_INTEG_WORKER_MODE=pull` to run the suite with pull Workers instead of
-listening Workers. The default is `push`.
+The suite always uses pull Workers. Set `DEX_INTEG_ENGINE=zd` to test against
+server-zd, or leave it unset to use the Temporal-backed server.
 
 ### Measure integration coverage
 
@@ -842,9 +852,9 @@ See [CONTRIBUTING.md](../CONTRIBUTING.md#releases-monorepo-tags) for monorepo ta
 ## Server protocol compatibility
 
 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.
+common protocol, synchronize Attribute indexes, and only then open outbound
+task streams. The initial Python SDK interval is `[1,1]`. Missing information,
+invalid or disjoint intervals, and RPC failures stop startup before polling.
 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.

{dex_python_sdk-1.20.0 → dex_python_sdk-1.25.0}/dex/__init__.py

@@ -124,7 +124,7 @@ from dex.timer import Timer
 from dex.wait import Wait
 from dex.wait_options import WaitForAttributeOptions, WaitForStepCompletionOptions
 from dex.worker import Worker
-from dex.worker_options import PullWorkerOptions, WorkerOptions, WorkerTarget
+from dex.worker_options import PullWorkerOptions, WorkerOptions
 
 __all__ = [
     "BOOL",
@@ -226,7 +226,6 @@ __all__ = [
     "Worker",
     "WorkerInvocationError",
     "WorkerOptions",
-    "WorkerTarget",
     "PullWorkerOptions",
     "PullRegistrationError",
     "PullRegistrationState",

{dex_python_sdk-1.20.0 → dex_python_sdk-1.25.0}/dex/_grpc_errors.py

@@ -10,7 +10,7 @@ from __future__ import annotations
 
 import traceback
 from dataclasses import dataclass
-from typing import Any, Literal, cast
+from typing import Literal, cast
 
 import grpc
 from google.protobuf import any_pb2
@@ -165,23 +165,6 @@ def _worker_code(details: pb.ServiceErrorResponse | None) -> grpc.StatusCode | N
     )
 
 
-def abort_worker_error(
-    context: grpc.ServicerContext,
-    error: BaseException,
-) -> None:
-    context.abort_with_status(rpc_status.to_status(_worker_error_status(error)))
-
-
-async def async_abort_worker_error(
-    context: grpc.aio.ServicerContext[Any, Any],
-    error: BaseException,
-) -> None:
-    # types-grpcio on 3.11 omits aio abort_with_status; runtime provides it.
-    await cast(Any, context).abort_with_status(
-        rpc_status.to_status(_worker_error_status(error))
-    )
-
-
 def _worker_error_status(error: BaseException) -> status_pb2.Status:
     retry_after_error: RetryAfterError | None = None
     reported = error

{dex_python_sdk-1.20.0 → dex_python_sdk-1.25.0}/dex/_pull_worker.py

@@ -16,6 +16,7 @@ import random
 import re
 import socket
 import threading
+import time
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
 from typing import Any
@@ -37,6 +38,7 @@ from dex.worker_options import PullWorkerOptions
 
 _LOGGER = logging.getLogger(__name__)
 _GRPC_AIO_EOF = getattr(grpc.aio, "EOF")
+_BUSY_PULL_STREAM_KEEPALIVE_SECONDS = 20
 _PERMANENT_CODES = {
     grpc.StatusCode.UNAUTHENTICATED,
     grpc.StatusCode.PERMISSION_DENIED,
@@ -86,6 +88,7 @@ class PullWorkerRuntime:
         executor: ThreadPoolExecutor | None,
         app_worker_id: str | None = None,
         credential_runner: AsyncCredentialRunner | None = None,
+        reserve_pressure_linger: float = 30.0,
     ) -> None:
         self._registry = registry
         self.app_worker_id = app_worker_id or resolve_app_worker_id(registry, options)
@@ -107,6 +110,9 @@ class PullWorkerRuntime:
             options.reserved_rpc_tasks is None and self.max_concurrent_tasks == 1
         )
         self.reserved_rpc_declines = 0
+        self._last_reserve_decline: dict[str, float] = {}
+        self._reserve_pressure_timers: dict[str, asyncio.TimerHandle] = {}
+        self._reserve_pressure_linger = reserve_pressure_linger
         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
@@ -196,12 +202,26 @@ class PullWorkerRuntime:
                 and not stream.is_half_closing
                 for stream in self._streams
             )
-            for _ in range(self.idle_streams_per_queue - idle):
+            desired_idle = self.desired_idle_streams_per_queue(queue)
+            creation_target = (
+                self.idle_streams_per_queue
+                if self._busy_tasks == self.max_concurrent_tasks
+                else desired_idle
+            )
+            for _ in range(creation_target - idle):
                 stream = PullStream(self, stub, queue)
                 self._streams.add(stream)
                 stream.run_task = asyncio.create_task(stream.run())
                 stream.run_task.add_done_callback(stream.on_done)
 
+    def desired_idle_streams_per_queue(self, queue: str) -> int:
+        last_decline = self._last_reserve_decline.get(queue)
+        return self.idle_streams_per_queue + int(
+            self.reserved_rpc_tasks > 0
+            and last_decline is not None
+            and time.monotonic() - last_decline < self._reserve_pressure_linger
+        )
+
     def _stream_done(self, stream: PullStream, completed: asyncio.Task[None]) -> None:
         self._streams.discard(stream)
         if not completed.cancelled():
@@ -210,7 +230,9 @@ class PullWorkerRuntime:
                 _LOGGER.error("Worker task stream failed", exc_info=failure)
         self.signal()
 
-    def accept_task(self, task: pb.PollWorkerTask | None = None) -> bool:
+    def accept_task(
+        self, task: pb.PollWorkerTask | None = None, queue: str | None = None
+    ) -> bool:
         if self._stopping is None or self._stopping.is_set():
             return False
         if self._busy_tasks >= self.max_concurrent_tasks:
@@ -223,6 +245,18 @@ class PullWorkerRuntime:
             and self.unknown_task_reason(task) is None
         ):
             self.reserved_rpc_declines += 1
+            queue_name = queue or self.app_worker_id
+            self._last_reserve_decline[queue_name] = time.monotonic()
+            existing_timer = self._reserve_pressure_timers.pop(queue_name, None)
+            if existing_timer is not None:
+                existing_timer.cancel()
+            if self._loop is not None:
+                self._reserve_pressure_timers[queue_name] = self._loop.call_later(
+                    self._reserve_pressure_linger,
+                    self._expire_reserve_pressure,
+                    queue_name,
+                )
+            self.signal()
             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
@@ -232,6 +266,21 @@ class PullWorkerRuntime:
         self.signal()
         return True
 
+    def _expire_reserve_pressure(self, queue: str) -> None:
+        last_decline = self._last_reserve_decline.get(queue)
+        if last_decline is None:
+            return
+        remaining = self._reserve_pressure_linger - (time.monotonic() - last_decline)
+        if remaining > 0:
+            assert self._loop is not None
+            self._reserve_pressure_timers[queue] = self._loop.call_later(
+                remaining, self._expire_reserve_pressure, queue
+            )
+            return
+        self._last_reserve_decline.pop(queue, None)
+        self._reserve_pressure_timers.pop(queue, None)
+        self.signal()
+
     def finish_task(self) -> None:
         self._busy_tasks -= 1
         self.signal()
@@ -250,6 +299,9 @@ class PullWorkerRuntime:
             asyncio.get_running_loop().time() + self.stop_drain_timeout
         )
         self._stopping.set()
+        for timer in self._reserve_pressure_timers.values():
+            timer.cancel()
+        self._reserve_pressure_timers.clear()
         try:
             for stream in tuple(self._streams):
                 stream.drain()
@@ -317,6 +369,8 @@ class PullStream:
         self._active = threading.Event()
         self._call: Any = None
         self._send_lock = asyncio.Lock()
+        self._keepalive_task: asyncio.Task[None] | None = None
+        self._is_final_frame_sent = False
         self._is_reconnecting = False
         self._stream_metadata: Any = None
         self._deadline_handle: asyncio.TimerHandle | None = None
@@ -424,6 +478,7 @@ class PullStream:
         if self.runtime.is_stopping:
             return False
         self._supports_unknown_name_rejection = False
+        self._is_final_frame_sent = False
         self._stream_metadata = None
         if self.runtime._credential_runner is None:
             self._call = self.stub.PollWorkerTask()
@@ -471,9 +526,11 @@ class PullStream:
                             and frame.WhichOneof("frame") == "task"
                         ):
                             await self._reject_task(frame.task)
-                    self.is_half_closing = True
-                    self.runtime.signal()
-                    await self._call.done_writing()
+                    async with self._send_lock:
+                        self._stop_keepalive()
+                        self.is_half_closing = True
+                        self.runtime.signal()
+                        await self._call.done_writing()
                     await self._await_stream_eof(reader)
                     reader = None
                     self._is_reconnecting = not self._is_retiring
@@ -515,6 +572,8 @@ class PullStream:
                         self._is_reconnecting = True
                         return was_registered
         finally:
+            async with self._send_lock:
+                self._stop_keepalive()
             retirement.cancel()
             await asyncio.gather(retirement, return_exceptions=True)
             self.is_half_closing = False
@@ -586,12 +645,18 @@ 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(task):
+        if (
+            self.is_busy
+            or self.is_draining
+            or not self.runtime.accept_task(task, self.queue)
+        ):
             await self._reject_task(task)
             return
         self.is_busy = True
         self.task_id = task.task_id
         self._active.set()
+        self._is_final_frame_sent = False
+        self._keepalive_task = asyncio.create_task(self._send_busy_keepalives())
         self.handler_task = asyncio.create_task(self._execute(task))
         if task.HasField("deadline"):
             deadline = task.deadline.ToDatetime(tzinfo=timezone.utc)
@@ -716,14 +781,68 @@ class PullStream:
     async def _send_frame(self, frame: pb.PollWorkerTaskFrame) -> None:
         await self._send(pb.PollWorkerTaskRequest(task_frame=frame))
 
+    async def _send_busy_keepalives(self) -> None:
+        try:
+            while True:
+                await asyncio.sleep(_BUSY_PULL_STREAM_KEEPALIVE_SECONDS)
+                async with self._send_lock:
+                    if self._is_final_frame_sent or self.is_half_closing:
+                        return
+                    await self._call.write(
+                        pb.PollWorkerTaskRequest(pong=pb.PollWorkerTaskPong())
+                    )
+        except asyncio.CancelledError:
+            return
+        except Exception:
+            _LOGGER.exception("send Worker keepalive Pong failed")
+
+    def _stop_keepalive(self) -> None:
+        if self._keepalive_task is not None:
+            self._keepalive_task.cancel()
+            self._keepalive_task = None
+
     async def _send(self, request: pb.PollWorkerTaskRequest) -> None:
         async with self._send_lock:
+            if request.HasField("task_frame"):
+                frame = request.task_frame
+                is_progress_frame = (
+                    frame.HasField("wait_for")
+                    and frame.wait_for.WhichOneof("output")
+                    in ("heartbeat", "stream_write")
+                    or frame.HasField("execute")
+                    and frame.execute.WhichOneof("output")
+                    in ("heartbeat", "stream_write")
+                )
+                if (
+                    frame.task_id != self.task_id
+                    or self._is_final_frame_sent
+                    or not self._active.is_set()
+                    and is_progress_frame
+                ):
+                    return
+                if (
+                    frame.HasField("rpc")
+                    or frame.HasField("worker_error")
+                    or frame.HasField("wait_for")
+                    and frame.wait_for.WhichOneof("output") == "result"
+                    or frame.HasField("execute")
+                    and frame.execute.WhichOneof("output") == "result"
+                ):
+                    self._stop_keepalive()
+                    self._is_final_frame_sent = True
+            if request.HasField("pong") and (
+                self._is_final_frame_sent or self.is_half_closing
+            ):
+                return
             await self._call.write(request)
 
     async def _finish_handler(self) -> None:
         if self.handler_task is None:
             return
         await asyncio.gather(self.handler_task, return_exceptions=True)
+        async with self._send_lock:
+            self._stop_keepalive()
+            self._is_final_frame_sent = False
         if self._deadline_handle is not None:
             self._deadline_handle.cancel()
             self._deadline_handle = None
@@ -744,7 +863,7 @@ class PullStream:
             and not candidate.is_draining
             for candidate in self.runtime._streams
         )
-        if idle_replacements >= self.runtime.idle_streams_per_queue:
+        if idle_replacements >= self.runtime.desired_idle_streams_per_queue(self.queue):
             self.is_draining = True
             self._is_retiring = True
             self._retirement_requested.set()