modal 1.5.6.dev12__tar.gz → 1.5.6.dev13__tar.gz

{modal-1.5.6.dev12 → modal-1.5.6.dev13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.5.6.dev12
+Version: 1.5.6.dev13
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License-Expression: Apache-2.0

{modal-1.5.6.dev12 → modal-1.5.6.dev13}/modal/_functions.py

@@ -64,7 +64,7 @@ from ._utils.function_utils import (
     parse_gpu_config,
     validate_target_concurrency,
 )
-from ._utils.grpc_utils import Retry, RetryWarningMessage
+from ._utils.grpc_utils import Retry, RetryWarningMessage, is_class_function_lookup_error
 from ._utils.mount_utils import validate_network_file_systems, validate_volumes, validate_volumes_by_object_id
 from .call_graph import InputInfo, _reconstruct_call_graph
 from .client import _Client
@@ -1361,9 +1361,13 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 env_context = (
                     f" (in the '{load_context.environment_name}' environment)" if load_context.environment_name else ""
                 )
-                raise NotFoundError(
-                    f"Lookup failed for Function '{name}' from the '{app_name}' app{env_context}: {exc}."
-                ) from None
+                msg = f"Lookup failed for Function '{name}' from the '{app_name}' app{env_context}: {exc}."
+                if is_class_function_lookup_error(exc):
+                    msg = (
+                        f"Function '{name}' from the '{app_name}' app{env_context} is a modal.Cls. "
+                        f"Use `modal.Cls.from_name` instead."
+                    )
+                raise NotFoundError(msg) from None
 
             print_server_warnings(response.server_warnings)
 

{modal-1.5.6.dev12 → modal-1.5.6.dev13}/modal/_output/status.py

@@ -46,16 +46,17 @@ class FunctionCreationStatus:
 
         if not self.response:
             self.status_row.finish(f"Unknown error when creating function {self.tag}")
+            return
 
-        elif web_url := self.response.handle_metadata.web_url:
+        for warning in self.response.server_warnings:
+            self.status_row.warn(warning)
+
+        if web_url := self.response.handle_metadata.web_url:
             url_info = self.response.function.web_url_info
             requires_proxy_auth = self.response.function.webhook_config.requires_proxy_auth
             proxy_auth_suffix = " 🔑" if requires_proxy_auth else ""
             # Ensure terms used here match terms used in modal.com/docs/guide/webhook-urls doc.
             suffix = _get_suffix_from_web_url_info(url_info)
-            # TODO: this is only printed when we're showing progress. Maybe move this somewhere else.
-            for warning in self.response.server_warnings:
-                self.status_row.warn(warning)
             self.status_row.finish(
                 f"Created web function {self.tag} => [magenta underline]{web_url}[/magenta underline]"
                 f"{proxy_auth_suffix}{suffix}"
@@ -78,8 +79,6 @@ class FunctionCreationStatus:
                 )
 
         else:
-            for warning in self.response.server_warnings:
-                self.status_row.warn(warning)
             self.status_row.finish(f"Created function {self.tag}.")
             if self.response.function.method_definitions_set:
                 for method_definition in self.response.function.method_definitions.values():

{modal-1.5.6.dev12 → modal-1.5.6.dev13}/modal/_utils/auth_token_manager.py

@@ -4,6 +4,7 @@ import base64
 import json
 import random
 import time
+from collections.abc import Awaitable, Callable
 from typing import Any
 
 from modal._utils.grpc_utils import DEFAULT_MAX_RETRIES, Retry
@@ -15,8 +16,15 @@ from .logger import logger
 AUTH_DENIED_EXCEPTIONS = (AuthError, PermissionDeniedError)
 
 
+TokenFetcher = Callable[[Retry], Awaitable[str]]
+
+
 class _AuthTokenManager:
-    """Handles fetching and refreshing of the input plane auth token."""
+    """Handles fetching and refreshing of a short-lived JWT.
+
+    Defaults to the input plane auth token (`AuthTokenGet`); pass `fetch` to manage a different token with the
+    same caching and refresh behavior.
+    """
 
     # Fraction of the token's lifetime to use before refreshing.
     REFRESH_FRACTION = 0.5
@@ -26,7 +34,7 @@ class _AuthTokenManager:
     # only ever delays a refresh, so the margin before expiry stays at no less than
     # 1 - REFRESH_FRACTION - REFRESH_JITTER of the lifetime.
     REFRESH_JITTER = 0.1
-    # Bound each AuthTokenGet attempt so a refresh can't hang.
+    # Bound each fetch attempt so a refresh can't hang.
     AUTH_TOKEN_TIMEOUT = 5.0
     AUTH_TOKEN_RETRY_TOTAL_TIMEOUT = 3 * AUTH_TOKEN_TIMEOUT
     # After a failed refresh, wait before hitting the server again, growing exponentially with
@@ -36,8 +44,9 @@ class _AuthTokenManager:
     # If the token doesn't have an expiry field, default to current time plus this value (not expected).
     DEFAULT_EXPIRY_OFFSET = 20 * 60
 
-    def __init__(self, stub: "modal_api_grpc.ModalClientModal"):
+    def __init__(self, stub: "modal_api_grpc.ModalClientModal", fetch: TokenFetcher | None = None):
         self._stub = stub
+        self._fetch = fetch or self._fetch_auth_token
         self._token = ""
         self._expiry = 0.0
         self._refresh_at = 0.0
@@ -105,31 +114,33 @@ class _AuthTokenManager:
             # Retrieve the exception, if any, so a failure isn't logged as unhandled when no caller is waiting.
             task.exception()
 
+    async def _fetch_auth_token(self, retry: Retry) -> str:
+        resp: api_pb2.AuthTokenGetResponse = await self._stub.AuthTokenGet(api_pb2.AuthTokenGetRequest(), retry=retry)
+        return resp.token
+
     async def _fetch_token(self):
-        """Make the AuthTokenGet request and cache its token, or arm the failure backoff."""
+        """Fetch a new token and cache it, or arm the failure backoff."""
         try:
-            resp: api_pb2.AuthTokenGetResponse = await self._stub.AuthTokenGet(
-                api_pb2.AuthTokenGetRequest(),
-                # No cached token to fall back on, so a failure is user-visible: retry transient errors.
-                # Otherwise one attempt, and _retry_after handles the cooldown.
-                retry=Retry(
-                    attempt_timeout=self.AUTH_TOKEN_TIMEOUT,
-                    total_timeout=self.AUTH_TOKEN_RETRY_TOTAL_TIMEOUT,
-                    max_retries=0 if self._token else DEFAULT_MAX_RETRIES,
-                ),
+            # No cached token to fall back on, so a failure is user-visible: retry transient errors.
+            # Otherwise one attempt, and _retry_after handles the cooldown.
+            retry = Retry(
+                attempt_timeout=self.AUTH_TOKEN_TIMEOUT,
+                total_timeout=self.AUTH_TOKEN_RETRY_TOTAL_TIMEOUT,
+                max_retries=0 if self._token else DEFAULT_MAX_RETRIES,
             )
-            if not resp.token:
+            token = await self._fetch(retry)
+            if not token:
                 # Not expected
                 raise ExecutionError(
                     "Internal error: Did not receive auth token from server. Please contact Modal support."
                 )
-            if exp := self._decode_jwt(resp.token).get("exp"):
+            if exp := self._decode_jwt(token).get("exp"):
                 expiry = float(exp)
             else:
                 # This should never happen.
-                logger.warning("x-modal-auth-token does not contain exp field")
+                logger.warning("Auth token does not contain exp field")
                 expiry = time.time() + self.DEFAULT_EXPIRY_OFFSET
-            self._token = resp.token
+            self._token = token
             self._set_expiry(expiry)
             self._retry_after = 0.0
             self._backoff = self.FAILURE_BACKOFF_BASE

{modal-1.5.6.dev12 → modal-1.5.6.dev13}/modal/_utils/grpc_utils.py

@@ -25,7 +25,7 @@ from grpclib.encoding.base import CodecBase, StatusDetailsCodecBase
 from grpclib.exceptions import StreamTerminatedError
 from grpclib.protocol import H2Protocol
 
-from modal.exception import ClientClosed, ConnectionError
+from modal.exception import ClientClosed, ConnectionError, NotFoundError
 from modal_proto import api_pb2
 from modal_version import __version__
 
@@ -121,12 +121,12 @@ class ModalChannel(grpclib.client.Channel):
         self.__target_host = host
         self.__target_port = port
         self.__ssl_context = ssl_context
-        self.__closed = False
-        self.__closed_error_message = closed_error_message
+        self._permanently_closed = False
+        self._closed_error_message = closed_error_message
 
     async def __connect__(self):
-        if self.__closed and self.__closed_error_message is not None:
-            raise ClientClosed(self.__closed_error_message)
+        if self._permanently_closed and self._closed_error_message is not None:
+            raise ClientClosed(self._closed_error_message)
         return await super().__connect__()
 
     async def _create_connection(self) -> H2Protocol:
@@ -161,8 +161,13 @@ class ModalChannel(grpclib.client.Channel):
             and not self._protocol.connection.is_closing()
         )
 
+    def release_connection(self):
+        """Release the connection while leaving the channel reusable."""
+        return super().close()
+
     def close(self):
-        self.__closed = True
+        """Close the channel permanently."""
+        self._permanently_closed = True
         return super().close()
 
 
@@ -450,6 +455,15 @@ async def retry_transient_errors(
     return await _retry_transient_errors(fn, req, retry=Retry(max_retries=max_retries))
 
 
+def is_class_function_lookup_error(exc: NotFoundError) -> bool:
+    """Whether a function lookup failed because the requested name identifies a class."""
+    return any(
+        isinstance(detail, api_pb2.FunctionLookupError)
+        and detail.reason == api_pb2.FunctionLookupError.REASON_CLASS_NAME_USED
+        for detail in (getattr(exc, "_grpc_details", None) or [])
+    )
+
+
 def get_server_retry_policy(exc: Exception) -> api_pb2.RPCRetryPolicy | None:
     """Get server retry policy."""
     if not isinstance(exc, GRPCError) or not exc.details:

modal-1.5.6.dev13/modal/_utils/idle_countdown.py

@@ -0,0 +1,64 @@
+# Copyright Modal Labs 2025
+import asyncio
+from collections.abc import Callable
+
+
+class IdleCountdown:
+    """An extendable idle timeout tracker."""
+
+    def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
+        self._loop = loop
+        self._handle: asyncio.TimerHandle | None = None
+        self._retired = False
+        self._expire_on_idle = False
+        self._release: Callable[[], None] | None = None
+
+    def _assert_owning_loop(self) -> None:
+        assert asyncio.get_running_loop() is self._loop, "Idle bookkeeping must run on its owning event loop"
+
+    def expire_when_idle(self) -> None:
+        """Release at the current or next idle boundary, even if the timeout is disabled.
+
+        Must run on the owning event loop without yielding control.
+        """
+        self._assert_owning_loop()
+        self._expire_on_idle = True
+        if self._release is not None and not self._retired:
+            release = self._release
+            self.stop()
+            release()
+
+    def stop(self) -> None:
+        """Cancel the countdown on the owning event loop without yielding control."""
+        self._assert_owning_loop()
+        self._release = None
+        if self._handle is not None:
+            self._handle.cancel()
+            self._handle = None
+
+    def retire(self) -> None:
+        """Stop the countdown and refuse every subsequent arm.
+
+        Must run on the owning event loop without yielding control.
+        """
+        self._assert_owning_loop()
+        self._retired = True
+        self.stop()
+
+    def arm(self, timeout_secs: float, release: Callable[[], None]) -> None:
+        """Start the countdown again, calling `release` when it elapses.
+
+        Must run on the owning event loop without yielding control.
+        A timeout of zero or less never fires unless `expire_when_idle` was
+        requested. `release` runs as a plain callback on the event loop, so it
+        must not block.
+        """
+        self.stop()
+        if self._retired:
+            return
+        self._release = release
+        if self._expire_on_idle:
+            timeout_secs = 0
+        elif timeout_secs <= 0:
+            return
+        self._handle = self._loop.call_later(timeout_secs, release)

{modal-1.5.6.dev12 → modal-1.5.6.dev13}/modal/_utils/task_command_router_client.py

@@ -1,6 +1,7 @@
 # Copyright Modal Labs 2025
 import asyncio
 import base64
+import contextlib
 import io
 import json
 import socket
@@ -18,8 +19,15 @@ import grpclib.events
 from grpclib import GRPCError, Status
 from grpclib.exceptions import StreamTerminatedError
 
-from modal.config import logger
-from modal.exception import AuthError, ExecTimeoutError, InternalError, ServiceError, TimeoutError as ModalTimeoutError
+from modal.config import config, logger
+from modal.exception import (
+    AuthError,
+    ClientClosed,
+    ExecTimeoutError,
+    InternalError,
+    ServiceError,
+    TimeoutError as ModalTimeoutError,
+)
 from modal_proto import api_pb2, task_command_router_pb2 as sr_pb2
 from modal_proto.task_command_router_grpc import TaskCommandRouterStub
 
@@ -27,6 +35,7 @@ from .._grpc_client import grpc_error_converter
 from .._utils.grpc_utils import ModalChannel, create_channel_config
 from .async_utils import aclosing, retry
 from .grpc_utils import RETRYABLE_GRPC_STATUS_CODES
+from .idle_countdown import IdleCountdown
 
 STREAMING_STDIN_CHUNK_SIZE = 256 * 1024
 
@@ -271,7 +280,10 @@ class TaskCommandRouterClient:
         loop = asyncio.get_running_loop()
         jwt_refresh_lock = asyncio.Lock()
 
-        return cls(server_client, task_id, url, jwt, channel, loop, jwt_refresh_lock, sandbox_id=sandbox_id)
+        client = cls(server_client, task_id, url, jwt, channel, loop, jwt_refresh_lock, sandbox_id=sandbox_id)
+        # An unused client releases its connection too.
+        client._arm_idle_release()
+        return client
 
     @classmethod
     async def init(
@@ -323,7 +335,7 @@ class TaskCommandRouterClient:
         task_id: str,
         server_url: str,
         jwt: str,
-        channel: grpclib.client.Channel,
+        channel: ModalChannel,
         loop: asyncio.AbstractEventLoop,
         jwt_refresh_lock: asyncio.Lock,
         *,
@@ -337,7 +349,7 @@ class TaskCommandRouterClient:
         # even if finalization happens from a different thread (e.g. via synchronicity).
         self._loop = loop
 
-        # Attach bearer token on all requests to the worker-side router service.
+        # Attach the bearer token to every request.
         self._server_client = server_client
         self._task_id = task_id
         self._sandbox_id = sandbox_id
@@ -356,6 +368,11 @@ class TaskCommandRouterClient:
 
         self._closed = False
 
+        self.idle_timeout_secs: float = config["sandbox_channel_idle_timeout"]
+        self._inflight = 0
+        self._idle = IdleCountdown(loop)
+        self._release_seq = 0
+
         self._channel_finalizer = weakref.finalize(
             self,
             _finalize_channel,
@@ -369,18 +386,96 @@ class TaskCommandRouterClient:
     def _is_v2_sandbox(self) -> bool:
         return self._sandbox_id is not None or _is_v2_task_id(self._task_id)
 
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise ClientClosed("Unable to perform operation on a detached sandbox")
+
     def _get_metadata(self):
+        self._ensure_open()
         return {"authorization": f"Bearer {self._jwt}"}
 
-    async def close(self) -> None:
-        """Close the client."""
+    @contextlib.asynccontextmanager
+    async def _lease(self):
+        """Keep the connection active for the duration of an operation."""
+        self._ensure_open()
+        self._take_lease()
+        try:
+            yield
+        finally:
+            self._drop_lease()
+
+    def _assert_owning_loop(self) -> None:
+        assert asyncio.get_running_loop() is self._loop, "Channel bookkeeping must run on its owning event loop"
+
+    def _take_lease(self) -> None:
+        """Acquire a lease on the owning event loop without yielding control."""
+        self._idle.stop()
+        self._inflight += 1
+
+    def _drop_lease(self) -> None:
+        """Release a lease on the owning event loop without yielding control."""
+        self._assert_owning_loop()
+        self._inflight -= 1
+        if self._inflight == 0:
+            self._arm_idle_release()
+
+    def _arm_idle_release(self) -> None:
+        """Arm idle release on the owning event loop without yielding control."""
+        self._idle.arm(self.idle_timeout_secs, self._release_if_idle)
+
+    def _release_if_idle(self) -> None:
+        """Release the inactive connection, permanently if the client is closed.
+
+        Must run on the owning event loop without yielding control.
+        """
+        self._assert_owning_loop()
+        if self._inflight:
+            return
         if self._closed:
+            self._close_channel()
+            return
+        if not self._channel._connected:
             return
+        logger.debug(f"Releasing idle command router channel for task {self._task_id}")
+        self._release_seq += 1
+        self._channel.release_connection()
+
+    @contextlib.asynccontextmanager
+    async def _lease_suspended(self):
+        """Temporarily suspend the caller's lease."""
+        self._drop_lease()
+        try:
+            yield
+        finally:
+            self._take_lease()
+
+    def _was_released_since(self, release_seq: int) -> bool:
+        """Whether an idle release invalidated the current stream."""
+        return self._release_seq != release_seq
+
+    async def close_when_idle(self) -> None:
+        """Request the "close soon" behavior used by Sandbox.detach().
 
+        Reject new operations and close the channel as soon as active leases
+        finish, without waiting for or interrupting those operations.
+        """
+        self._assert_owning_loop()
         self._closed = True
-        self._channel.close()
+        self._idle.expire_when_idle()
+        self._release_if_idle()
+
+    async def close(self) -> None:
+        """Close the client permanently; later use raises `ClientClosed`."""
+        self._assert_owning_loop()
+        self._closed = True
+        self._idle.retire()
+        self._close_channel()
+
+    def _close_channel(self) -> None:
+        """Close on the owning event loop without yielding, including from an idle timer."""
+        self._assert_owning_loop()
         if self._channel_finalizer.alive:
-            # skip the finalizer if we've closed the channel anyway
+            self._channel.close()
             self._channel_finalizer.detach()
 
     async def exec_start(self, request: sr_pb2.TaskExecStartRequest) -> sr_pb2.TaskExecStartResponse:
@@ -590,27 +685,28 @@ class TaskCommandRouterClient:
         Does not retry; `exec_stdin_write_stream` owns resume.
         """
         with grpc_error_converter():
-            stream = self._stub.TaskExecStdinWriteStream.open(metadata=self._get_metadata())
-            async with stream as s:
-                start = sr_pb2.TaskExecStdinWriteStreamRequest(
-                    start=sr_pb2.TaskExecStdinWriteStreamStart(
-                        task_id=task_id,
-                        exec_id=exec_id,
-                        offset=start_offset,
-                    ),
-                )
-                await s.send_message(start)
-                async for data in chunks:
-                    if not data:
-                        continue
-                    await s.send_message(sr_pb2.TaskExecStdinWriteStreamRequest(data=data))
-                # The server closes stdin only on this explicit End message. A
-                # stream that breaks before it leaves stdin open for resume.
-                await s.send_message(
-                    sr_pb2.TaskExecStdinWriteStreamRequest(end=sr_pb2.TaskExecStdinWriteStreamEnd()),
-                    end=True,
-                )
-                return await s.recv_message()
+            async with self._lease():
+                stream = self._stub.TaskExecStdinWriteStream.open(metadata=self._get_metadata())
+                async with stream as s:
+                    start = sr_pb2.TaskExecStdinWriteStreamRequest(
+                        start=sr_pb2.TaskExecStdinWriteStreamStart(
+                            task_id=task_id,
+                            exec_id=exec_id,
+                            offset=start_offset,
+                        ),
+                    )
+                    await s.send_message(start)
+                    async for data in chunks:
+                        if not data:
+                            continue
+                        await s.send_message(sr_pb2.TaskExecStdinWriteStreamRequest(data=data))
+                    # Only an explicit End marks stdin as closed. A failed write
+                    # stream remains resumable.
+                    await s.send_message(
+                        sr_pb2.TaskExecStdinWriteStreamRequest(end=sr_pb2.TaskExecStdinWriteStreamEnd()),
+                        end=True,
+                    )
+                    return await s.recv_message()
 
     async def exec_stdin_status(self, task_id: str, exec_id: str) -> sr_pb2.TaskExecStdinStatusResponse:
         """Read the current stdin write status for an exec'd command, to support retries from the right offset.
@@ -778,14 +874,15 @@ class TaskCommandRouterClient:
             self._jwt_exp = _parse_jwt_expiration(jwt)
 
     async def _call_with_auth_retry(self, func, *args, **kwargs):
-        try:
-            return await func(*args, **kwargs, metadata=self._get_metadata())
-        except GRPCError as exc:
-            if exc.status == Status.UNAUTHENTICATED:
-                await self._refresh_jwt()
-                # Retry with the original arguments preserved
+        async with self._lease():
+            try:
                 return await func(*args, **kwargs, metadata=self._get_metadata())
-            raise
+            except GRPCError as exc:
+                if exc.status == Status.UNAUTHENTICATED:
+                    await self._refresh_jwt()
+                    # Retry with the original arguments preserved
+                    return await func(*args, **kwargs, metadata=self._get_metadata())
+                raise
 
     async def _stream_stdio_with_retries(
         self,
@@ -795,17 +892,14 @@ class TaskCommandRouterClient:
         deadline_label: str,
         deadline: float | None = None,
     ) -> AsyncGenerator[_StdioResp, None]:
-        """Drive a streaming-stdio RPC with offset bookkeeping, transient-error
-        retries, and JWT-refresh auth retries.
-
-        Shared by [`_stream_stdio`] (exec stdio) and [`_stream_sandbox_stdio`]
-        (V2 sandbox top-level stdio); both response types have a ``bytes data``
-        field that this helper uses to advance the offset. For V2 sandbox
-        responses (which carry ``starting_offset``), the offset is rebased off
-        the first chunk of each attempt so transient reconnects don't miss
-        bytes.
+        """Read stdio with offset tracking, authentication, and transient retries.
+
+        The connection is kept active while waiting for output and may be
+        released while the caller holds a chunk.
         """
         offset = 0
+        # Set when an idle release invalidates the stream while its lease is suspended.
+        reopen_after_release = False
         delay_secs = self.stream_stdio_retry_delay_secs
         delay_factor = self.stream_stdio_retry_delay_factor
         num_retries_remaining = self.stream_stdio_max_retries
@@ -815,6 +909,7 @@ class TaskCommandRouterClient:
 
         async def sleep_and_update_delay_and_num_retries_remaining(e: Exception):
             nonlocal delay_secs, num_retries_remaining
+            self._ensure_open()
             logger.debug(f"Retrying stdio read with delay {delay_secs}s due to error: {e}")
             if deadline is not None and deadline - time.monotonic() <= delay_secs:
                 raise ExecTimeoutError(f"Deadline exceeded while streaming stdio for {deadline_label}")
@@ -826,41 +921,54 @@ class TaskCommandRouterClient:
         while True:
             timeout = max(0, deadline - time.monotonic()) if deadline is not None else None
             try:
-                stream = stub_method.open(timeout=timeout, metadata=self._get_metadata())
-                async with stream as s:
-                    req = request_factory(offset)
-
-                    # Auth retry is scoped to a single refresh per streaming attempt. While auth metadata is
-                    # sent on request start, UNAUTHENTICATED may sometimes surface during iteration,
-                    # so we handle it at both send and receive boundaries.
-                    is_first_chunk_of_attempt = True
-                    try:
-                        await s.send_message(req, end=True)
-                        async for item in s:
-                            # We successfully authenticated after a JWT refresh, reset the auth retry flag.
-                            if did_auth_retry:
-                                did_auth_retry = False
-                            # Any received chunk is progress: reset the backoff and refill the
-                            # retry budget so it bounds consecutive failures, not failures over
-                            # the stream's lifetime.
-                            delay_secs = self.stream_stdio_retry_delay_secs
-                            num_retries_remaining = self.stream_stdio_max_retries
-                            # Track it so transient reconnects request the
-                            # correct next byte.
-                            if is_first_chunk_of_attempt and isinstance(item, sr_pb2.SandboxStdioReadV2Response):
-                                offset = item.starting_offset
-                            is_first_chunk_of_attempt = False
-                            offset += len(item.data)
-                            yield item
-                    except GRPCError as exc:
-                        if exc.status == Status.UNAUTHENTICATED and not did_auth_retry:
-                            await self._refresh_jwt()
-                            # Mark that we've retried authentication for this streaming attempt, to
-                            # prevent subsequent retries.
-                            did_auth_retry = True
-                            continue
-                        raise
-
+                async with self._lease():
+                    release_seq_at_open = self._release_seq
+                    stream = stub_method.open(timeout=timeout, metadata=self._get_metadata())
+                    async with stream as s:
+                        req = request_factory(offset)
+
+                        # Auth retry is scoped to a single refresh per streaming attempt. While auth metadata is
+                        # sent on request start, UNAUTHENTICATED may sometimes surface during iteration,
+                        # so we handle it at both send and receive boundaries.
+                        is_first_chunk_of_attempt = True
+                        try:
+                            await s.send_message(req, end=True)
+                            async for item in s:
+                                # We successfully authenticated after a JWT refresh, reset the auth retry flag.
+                                if did_auth_retry:
+                                    did_auth_retry = False
+                                # A successful chunk ends the current run of transient failures.
+                                delay_secs = self.stream_stdio_retry_delay_secs
+                                num_retries_remaining = self.stream_stdio_max_retries
+                                # Track it so transient reconnects request the
+                                # correct next byte.
+                                if is_first_chunk_of_attempt and isinstance(item, sr_pb2.SandboxStdioReadV2Response):
+                                    offset = item.starting_offset
+                                is_first_chunk_of_attempt = False
+                                offset += len(item.data)
+                                async with self._lease_suspended():
+                                    yield item
+                                # Reject output buffered while the consumer held the chunk
+                                # if the Sandbox was detached during the yield.
+                                self._ensure_open()
+                                # Resume from the current offset if an idle release made
+                                # this stream stale while the lease was suspended.
+                                if self._was_released_since(release_seq_at_open):
+                                    reopen_after_release = True
+                                    break
+                        except GRPCError as exc:
+                            if exc.status == Status.UNAUTHENTICATED and not did_auth_retry:
+                                await self._refresh_jwt()
+                                # Mark that we've retried authentication for this streaming attempt, to
+                                # prevent subsequent retries.
+                                did_auth_retry = True
+                                continue
+                            raise
+
+                if reopen_after_release:
+                    # Idle release does not consume the transient-error retry budget.
+                    reopen_after_release = False
+                    continue
                 # We successfully streamed all output.
                 return
             except GRPCError as e:
@@ -872,7 +980,9 @@ class TaskCommandRouterClient:
                 # StreamTerminatedError are not properly raised in grpclib<=0.4.7
                 # fixed in https://github.com/vmagamedov/grpclib/issues/185
                 # TODO: update to newer version (>=0.4.8) once stable
-                if num_retries_remaining > 0 and "_write_appdata" in str(e):
+                if "_write_appdata" not in str(e):
+                    raise e
+                if num_retries_remaining > 0:
                     await sleep_and_update_delay_and_num_retries_remaining(e)
                 else:
                     raise e