modal 1.4.4.dev8__tar.gz → 1.4.4.dev10__tar.gz

{modal-1.4.4.dev8 → modal-1.4.4.dev10}/PKG-INFO

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

{modal-1.4.4.dev8 → modal-1.4.4.dev10}/modal/_utils/task_command_router_client.py

@@ -9,7 +9,7 @@ import typing
 import urllib.parse
 import weakref
 from contextlib import suppress
-from typing import AsyncGenerator, Optional
+from typing import AsyncGenerator, Callable, Optional, TypeVar
 
 import grpclib.client
 import grpclib.config
@@ -146,6 +146,10 @@ async def call_with_retries_on_transient_errors(
             await sleep_and_advance(e)
 
 
+_StdioReq = TypeVar("_StdioReq")
+_StdioResp = TypeVar("_StdioResp", sr_pb2.TaskExecStdioReadResponse, sr_pb2.SandboxStdioReadV2Response)
+
+
 async def fetch_command_router_access(server_client, task_id: str) -> api_pb2.TaskGetCommandRouterAccessResponse:
     """Fetch direct command router access info from Modal server."""
     return await server_client.stub.TaskGetCommandRouterAccess(
@@ -383,7 +387,7 @@ class TaskCommandRouterClient:
         file_descriptor: "api_pb2.FileDescriptor.ValueType",
         deadline: Optional[float] = None,
     ) -> AsyncGenerator[sr_pb2.TaskExecStdioReadResponse, None]:
-        """Stream stdout/stderr batches from the task, properly retrying on transient errors.
+        """Stream stdout/stderr batches from an exec'd command, retrying on transient errors.
 
         Args:
             task_id: The task ID of the task running the exec'd command.
@@ -413,6 +417,39 @@ class TaskCommandRouterClient:
                 async for item in stream:
                     yield item
 
+    async def sandbox_stdio_read(
+        self,
+        task_id: str,
+        # Quotes around the type required for protobuf 3.19.
+        file_descriptor: "api_pb2.FileDescriptor.ValueType",
+    ) -> AsyncGenerator[sr_pb2.SandboxStdioReadV2Response, None]:
+        """Stream stdout/stderr batches from a V2 sandbox.
+
+        Serves both live reads and post-exit reads.
+
+        Args:
+            task_id: The task ID hosting the V2 sandbox.
+            file_descriptor: The file descriptor to read from (stdout or stderr).
+        Returns:
+            AsyncGenerator[sr_pb2.SandboxStdioReadV2Response, None]: A stream of stdout/stderr batches.
+        Raises:
+            Errors: If retries are exhausted on transient errors or if there is
+              an error from the RPC itself.
+        """
+        if file_descriptor == api_pb2.FILE_DESCRIPTOR_STDOUT:
+            sr_fd = sr_pb2.SANDBOX_STDIO_FILE_DESCRIPTOR_STDOUT
+        elif file_descriptor == api_pb2.FILE_DESCRIPTOR_STDERR:
+            sr_fd = sr_pb2.SANDBOX_STDIO_FILE_DESCRIPTOR_STDERR
+        elif file_descriptor == api_pb2.FILE_DESCRIPTOR_INFO or file_descriptor == api_pb2.FILE_DESCRIPTOR_UNSPECIFIED:
+            raise ValueError(f"Unsupported file descriptor: {file_descriptor}")
+        else:
+            raise ValueError(f"Invalid file descriptor: {file_descriptor}")
+
+        with grpc_error_converter():
+            async with aclosing(self._stream_sandbox_stdio(task_id, sr_fd)) as stream:
+                async for item in stream:
+                    yield item
+
     async def exec_stdin_write(
         self, task_id: str, exec_id: str, offset: int, data: bytes, eof: bool
     ) -> sr_pb2.TaskExecStdinWriteResponse:
@@ -434,6 +471,25 @@ class TaskCommandRouterClient:
                 lambda: self._call_with_auth_retry(self._stub.TaskExecStdinWrite, request)
             )
 
+    async def sandbox_stdin_write_v2(
+        self, task_id: str, offset: int, data: bytes, eof: bool
+    ) -> sr_pb2.SandboxStdinWriteV2Response:
+        """Write to the stdin stream of a V2 sandbox's entrypoint process.
+
+        Args:
+            task_id: The task ID of the V2 sandbox.
+            offset: The offset to start writing to.
+            eof: Whether to close the stdin stream after writing the data.
+        Raises:
+            Other errors: If retries are exhausted on transient errors or if there's an error
+              from the RPC itself.
+        """
+        request = sr_pb2.SandboxStdinWriteV2Request(task_id=task_id, offset=offset, data=data, eof=eof)
+        with grpc_error_converter():
+            return await call_with_retries_on_transient_errors(
+                lambda: self._call_with_auth_retry(self._stub.SandboxStdinWriteV2, request)
+            )
+
     async def exec_poll(
         self, task_id: str, exec_id: str, deadline: Optional[float] = None
     ) -> sr_pb2.TaskExecPollResponse:
@@ -558,16 +614,23 @@ class TaskCommandRouterClient:
                 return await func(*args, **kwargs, metadata=self._get_metadata())
             raise
 
-    async def _stream_stdio(
+    async def _stream_stdio_with_retries(
         self,
-        task_id: str,
-        exec_id: str,
-        # Quotes around the type required for protobuf 3.19.
-        file_descriptor: "sr_pb2.TaskExecStdioFileDescriptor.ValueType",
+        *,
+        stub_method: "grpclib.client.UnaryStreamMethod[_StdioReq, _StdioResp]",
+        request_factory: Callable[[int], _StdioReq],
+        deadline_label: str,
         deadline: Optional[float] = None,
-    ) -> AsyncGenerator[sr_pb2.TaskExecStdioReadResponse, None]:
-        """Stream stdio from the task, properly updating the offset and retrying on transient errors.
-        Raises ExecTimeoutError if the deadline is exceeded.
+    ) -> 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.
         """
         offset = 0
         delay_secs = self.stream_stdio_retry_delay_secs
@@ -581,7 +644,7 @@ class TaskCommandRouterClient:
             nonlocal delay_secs, num_retries_remaining
             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 exec {exec_id}")
+                raise ExecTimeoutError(f"Deadline exceeded while streaming stdio for {deadline_label}")
 
             await asyncio.sleep(delay_secs)
             delay_secs *= delay_factor
@@ -590,18 +653,14 @@ class TaskCommandRouterClient:
         while True:
             timeout = max(0, deadline - time.monotonic()) if deadline is not None else None
             try:
-                stream = self._stub.TaskExecStdioRead.open(timeout=timeout, metadata=self._get_metadata())
+                stream = stub_method.open(timeout=timeout, metadata=self._get_metadata())
                 async with stream as s:
-                    req = sr_pb2.TaskExecStdioReadRequest(
-                        task_id=task_id,
-                        exec_id=exec_id,
-                        offset=offset,
-                        file_descriptor=file_descriptor,
-                    )
+                    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:
@@ -610,6 +669,11 @@ class TaskCommandRouterClient:
                                 did_auth_retry = False
                             # Reset retry backoff after any successful chunk.
                             delay_secs = self.stream_stdio_retry_delay_secs
+                            # 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:
@@ -652,6 +716,63 @@ class TaskCommandRouterClient:
                 else:
                     raise ConnectionError(str(e))
 
+    async def _stream_stdio(
+        self,
+        task_id: str,
+        exec_id: str,
+        # Quotes around the type required for protobuf 3.19.
+        file_descriptor: "sr_pb2.TaskExecStdioFileDescriptor.ValueType",
+        deadline: Optional[float] = None,
+    ) -> AsyncGenerator[sr_pb2.TaskExecStdioReadResponse, None]:
+        """Stream exec stdio from the task, retrying on transient errors.
+        Raises ExecTimeoutError if the deadline is exceeded.
+        """
+
+        def request_factory(offset: int) -> sr_pb2.TaskExecStdioReadRequest:
+            return sr_pb2.TaskExecStdioReadRequest(
+                task_id=task_id,
+                exec_id=exec_id,
+                offset=offset,
+                file_descriptor=file_descriptor,
+            )
+
+        async with aclosing(
+            self._stream_stdio_with_retries(
+                stub_method=self._stub.TaskExecStdioRead,
+                request_factory=request_factory,
+                deadline_label=f"exec {exec_id}",
+                deadline=deadline,
+            )
+        ) as stream:
+            async for item in stream:
+                yield item
+
+    async def _stream_sandbox_stdio(
+        self,
+        task_id: str,
+        # Quotes around the type required for protobuf 3.19.
+        file_descriptor: "sr_pb2.SandboxStdioFileDescriptor.ValueType",
+    ) -> AsyncGenerator[sr_pb2.SandboxStdioReadV2Response, None]:
+        """Stream V2 sandbox top-level stdio from the task, retrying on transient errors."""
+
+        def request_factory(offset: int) -> sr_pb2.SandboxStdioReadV2Request:
+            return sr_pb2.SandboxStdioReadV2Request(
+                task_id=task_id,
+                offset=offset,
+                file_descriptor=file_descriptor,
+            )
+
+        async with aclosing(
+            self._stream_stdio_with_retries(
+                stub_method=self._stub.SandboxStdioReadV2,
+                request_factory=request_factory,
+                deadline_label=f"sandbox {task_id}",
+                deadline=None,
+            )
+        ) as stream:
+            async for item in stream:
+                yield item
+
     async def mount_image(self, request: sr_pb2.TaskMountDirectoryRequest):
         with grpc_error_converter():
             return await call_with_retries_on_transient_errors(
@@ -664,21 +785,10 @@ class TaskCommandRouterClient:
                 lambda: self._call_with_auth_retry(self._stub.TaskUnmountDirectory, request)
             )
 
-    async def snapshot_directory(
-        self, request: sr_pb2.TaskSnapshotDirectoryRequest, **kwargs
-    ) -> sr_pb2.TaskSnapshotDirectoryResponse:
-        with grpc_error_converter():
-            return await call_with_retries_on_transient_errors(
-                lambda: self._call_with_auth_retry(self._stub.TaskSnapshotDirectory, request, **kwargs)
-            )
-
-    async def snapshot_filesystem(
-        self, request: sr_pb2.TaskSnapshotFilesystemRequest, *, timeout: float, **kwargs
-    ) -> sr_pb2.TaskSnapshotFilesystemResponse:
-        # Compute the overall deadline once; each retry attempt passes the
-        # remaining budget as the per-call gRPC timeout so we honor the
-        # caller-specified `timeout` across retries instead of giving each
-        # attempt a fresh full window.
+    async def _snapshot_with_deadline(self, rpc, request, *, timeout: float, **kwargs):
+        # helper method for snapshot_directory and snapshot_filesystem to handle grpc
+        # deadlines in a consistent way, converting any error to TimeoutError after passing
+        # the total deadline budget
         timeout_deadline = time.monotonic() + timeout
 
         def call():
@@ -687,12 +797,8 @@ class TaskCommandRouterClient:
                 # doesn't matter which exception type this is
                 # as it will be caught by the catch all below
                 raise ModalTimeoutError("Timeout expired")
+            return self._call_with_auth_retry(rpc, request, timeout=call_timeout, **kwargs)
 
-            return self._call_with_auth_retry(
-                self._stub.TaskSnapshotFilesystem, request, timeout=call_timeout, **kwargs
-            )
-
-        # Any failure observed at or after the deadline is treated as a timeout
         try:
             with grpc_error_converter():
                 return await call_with_retries_on_transient_errors(
@@ -704,3 +810,13 @@ class TaskCommandRouterClient:
             if time.monotonic() >= timeout_deadline:
                 raise ModalTimeoutError("Timeout expired")
             raise
+
+    async def snapshot_directory(
+        self, request: sr_pb2.TaskSnapshotDirectoryRequest, *, timeout: float, **kwargs
+    ) -> sr_pb2.TaskSnapshotDirectoryResponse:
+        return await self._snapshot_with_deadline(self._stub.TaskSnapshotDirectory, request, timeout=timeout, **kwargs)
+
+    async def snapshot_filesystem(
+        self, request: sr_pb2.TaskSnapshotFilesystemRequest, *, timeout: float, **kwargs
+    ) -> sr_pb2.TaskSnapshotFilesystemResponse:
+        return await self._snapshot_with_deadline(self._stub.TaskSnapshotFilesystem, request, timeout=timeout, **kwargs)

{modal-1.4.4.dev8 → modal-1.4.4.dev10}/modal/app.pyi

@@ -55,7 +55,7 @@ def check_sequence(items: typing.Sequence[typing.Any], item_type: type[typing.An
 
 CLS_T = typing.TypeVar("CLS_T", bound="type[typing.Any]")
 
-P = typing_extensions.ParamSpec("P")
+P = typing.ParamSpec("P")
 
 ReturnType = typing.TypeVar("ReturnType")
 

{modal-1.4.4.dev8 → modal-1.4.4.dev10}/modal/client.pyi

@@ -35,7 +35,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.4.4.dev8",
+        version: str = "1.4.4.dev10",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.
@@ -175,7 +175,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.4.4.dev8",
+        version: str = "1.4.4.dev10",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

{modal-1.4.4.dev8 → modal-1.4.4.dev10}/modal/container_process.py

@@ -13,9 +13,9 @@ from .config import logger
 from .exception import ExecTimeoutError, InteractiveTimeoutError, InvalidError
 from .io_streams import (
     _StreamReader,
-    _StreamReaderThroughCommandRouterParams,
+    _StreamReaderThroughSandboxExecCommandRouterParams,
     _StreamWriter,
-    _StreamWriterThroughCommandRouterParams,
+    _StreamWriterThroughCommandRouterSandboxExecParams,
 )
 from .stream_type import StreamType
 
@@ -58,7 +58,7 @@ class _ContainerProcess(Generic[T]):
         self._by_line = by_line
         self._task_id = task_id
         self._stdout = _StreamReader[T](
-            _StreamReaderThroughCommandRouterParams(
+            _StreamReaderThroughSandboxExecCommandRouterParams(
                 file_descriptor=api_pb2.FILE_DESCRIPTOR_STDOUT,
                 task_id=self._task_id,
                 object_id=process_id,
@@ -70,7 +70,7 @@ class _ContainerProcess(Generic[T]):
             by_line=by_line,
         )
         self._stderr = _StreamReader[T](
-            _StreamReaderThroughCommandRouterParams(
+            _StreamReaderThroughSandboxExecCommandRouterParams(
                 file_descriptor=api_pb2.FILE_DESCRIPTOR_STDERR,
                 task_id=self._task_id,
                 object_id=process_id,
@@ -82,7 +82,7 @@ class _ContainerProcess(Generic[T]):
             by_line=by_line,
         )
         self._stdin = _StreamWriter(
-            _StreamWriterThroughCommandRouterParams(
+            _StreamWriterThroughCommandRouterSandboxExecParams(
                 task_id=self._task_id,
                 object_id=process_id,
                 command_router_client=self._command_router_client,

{modal-1.4.4.dev8 → modal-1.4.4.dev10}/modal/functions.pyi

@@ -27,7 +27,7 @@ import typing_extensions
 
 ReturnType_INNER = typing.TypeVar("ReturnType_INNER", covariant=True)
 
-P_INNER = typing_extensions.ParamSpec("P_INNER")
+P_INNER = typing.ParamSpec("P_INNER")
 
 class Function(
     typing.Generic[modal._functions.P, modal._functions.ReturnType, modal._functions.OriginalReturnType],

{modal-1.4.4.dev8 → modal-1.4.4.dev10}/modal/image.pyi

@@ -29,7 +29,7 @@ class _AutoDockerIgnoreSentinel:
 
 AUTO_DOCKERIGNORE: _AutoDockerIgnoreSentinel
 
-P = typing_extensions.ParamSpec("P")
+P = typing.ParamSpec("P")
 
 def _validate_python_version(
     python_version: typing.Optional[str],

{modal-1.4.4.dev8 → modal-1.4.4.dev10}/modal/io_streams.py

@@ -4,7 +4,8 @@ import codecs
 import contextlib
 import io
 import sys
-from collections.abc import AsyncGenerator, AsyncIterator
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
 from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
@@ -231,9 +232,9 @@ class _StreamReaderThroughServerParams:
 
 
 @dataclass(frozen=True)
-class _StreamReaderThroughCommandRouterParams:
-    """Parameters for a ``_StreamReader`` that reads container-process stdio
-    directly from the worker via the task command router."""
+class _StreamReaderThroughSandboxExecCommandRouterParams:
+    """Parameters for a ``_StreamReader`` that reads sandbox-exec stdio
+    directly from the worker via the task command router (``exec_stdio_read``)."""
 
     file_descriptor: "api_pb2.FileDescriptor.ValueType"
     task_id: str
@@ -242,10 +243,50 @@ class _StreamReaderThroughCommandRouterParams:
     deadline: Optional[float]
 
 
-async def _stdio_stream_from_command_router(
-    params: _StreamReaderThroughCommandRouterParams,
+@dataclass(frozen=True)
+class _StreamReaderThroughSandboxCommandRouterParams:
+    """Parameters for a ``_StreamReader`` that reads a V2 sandbox's
+    stdio directly from the worker via the task command router
+    (``sandbox_stdio_read``)."""
+
+    file_descriptor: "api_pb2.FileDescriptor.ValueType"
+    sandbox_id: str
+    # Lazily fetches ``(task_id, command_router_client)`` the first time the
+    # stream is iterated. Captures the sandbox handle so we only mint a JWT
+    # and open a connection to the worker when stdio is actually read.
+    resolve_router: Callable[[], Awaitable[tuple[str, TaskCommandRouterClient]]]
+
+
+_StreamReaderThroughCommandRouterParams = (
+    _StreamReaderThroughSandboxExecCommandRouterParams | _StreamReaderThroughSandboxCommandRouterParams
+)
+
+
+async def _stdio_stream_from_sandbox_command_router(
+    params: _StreamReaderThroughSandboxCommandRouterParams,
+) -> AsyncGenerator[bytes, None]:
+    """Stream raw bytes from a V2 sandbox's primary stdio via ``sandbox_stdio_read``."""
+    task_id, command_router_client = await params.resolve_router()
+    first_chunk = True
+    async with aclosing(command_router_client.sandbox_stdio_read(task_id, params.file_descriptor)) as stream:
+        async for item in stream:
+            if len(item.data) == 0:
+                raise ValueError("Received empty message streaming stdio from sandbox.")
+            if first_chunk:
+                first_chunk = False
+                if item.starting_offset > 0:
+                    logger.warning(
+                        f"V2 sandbox {params.sandbox_id} stdio: dropped first "
+                        f"{item.starting_offset} bytes; only the most recent portion "
+                        f"of output is retained."
+                    )
+            yield item.data
+
+
+async def _stdio_stream_from_sandbox_exec_command_router(
+    params: _StreamReaderThroughSandboxExecCommandRouterParams,
 ) -> AsyncGenerator[bytes, None]:
-    """Stream raw bytes from the router client."""
+    """Stream raw bytes from a V2 sandbox-exec'd process via ``exec_stdio_read``."""
     async with aclosing(
         params.command_router_client.exec_stdio_read(
             params.task_id, params.object_id, params.file_descriptor, params.deadline
@@ -254,9 +295,7 @@ async def _stdio_stream_from_command_router(
         try:
             async for item in stream:
                 if len(item.data) == 0:
-                    # This is an error.
                     raise ValueError("Received empty message streaming stdio from sandbox.")
-
                 yield item.data
         except ExecTimeoutError:
             logger.debug(f"Deadline exceeded while streaming stdio for exec {params.object_id}")
@@ -265,20 +304,22 @@ async def _stdio_stream_from_command_router(
             return
 
 
-class _BytesStreamReaderThroughCommandRouter:
-    """
-    StreamReader implementation that will read directly from the worker that
-    hosts the sandbox.
+def _stdio_stream_from_command_router(
+    params: _StreamReaderThroughCommandRouterParams,
+) -> AsyncGenerator[bytes, None]:
+    """Dispatch between the V2-sandbox primary stdio and the V2-sandbox-exec
+    stdio streams, both of which yield raw bytes."""
+    if isinstance(params, _StreamReaderThroughSandboxCommandRouterParams):
+        return _stdio_stream_from_sandbox_command_router(params)
+    return _stdio_stream_from_sandbox_exec_command_router(params)
 
-    This implementation is used for non-text streams.
-    """
 
-    def __init__(
-        self,
-        params: _StreamReaderThroughCommandRouterParams,
-    ) -> None:
+class _BytesStreamReaderThroughCommandRouter:
+    """StreamReader that yields raw bytes from the router-backed stdio source
+    (either V2 sandbox top-level stdio or V2 sandbox-exec stdio)."""
+
+    def __init__(self, params: _StreamReaderThroughCommandRouterParams) -> None:
         self._params = params
-        self._stream = None
 
     @property
     def file_descriptor(self) -> int:
@@ -300,18 +341,10 @@ class _BytesStreamReaderThroughCommandRouter:
 
 
 class _TextStreamReaderThroughCommandRouter:
-    """
-    StreamReader implementation that will read directly from the worker
-    that hosts the sandbox.
+    """StreamReader that yields UTF-8-decoded text from the router-backed
+    stdio source."""
 
-    This implementation is used for text streams.
-    """
-
-    def __init__(
-        self,
-        params: _StreamReaderThroughCommandRouterParams,
-        by_line: bool,
-    ) -> None:
+    def __init__(self, params: _StreamReaderThroughCommandRouterParams, by_line: bool) -> None:
         self._params = params
         self._by_line = by_line
 
@@ -513,15 +546,28 @@ class _StreamWriterThroughServerParams:
 
 
 @dataclass(frozen=True)
-class _StreamWriterThroughCommandRouterParams:
-    """Parameters for a ``_StreamWriter`` that writes container-process stdin
-    directly to the worker via the task command router."""
+class _StreamWriterThroughCommandRouterSandboxExecParams:
+    """Parameters for a ``_StreamWriter`` that writes the stdin of a process
+    spawned via ``sb.exec(...)`` directly to the worker via the task command
+    router."""
 
     task_id: str
     object_id: str
     command_router_client: TaskCommandRouterClient
 
 
+@dataclass(frozen=True)
+class _StreamWriterThroughCommandRouterSandboxParams:
+    """Parameters for a ``_StreamWriter`` that writes a V2 sandbox entrypoint's
+    stdin directly to the worker via the task command router.
+    """
+
+    # Lazily fetches ``(task_id, command_router_client)`` the first time the
+    # writer drains. Captures the sandbox handle so we only mint a JWT and
+    # open a connection to the worker when stdin is actually written.
+    resolve_router: Callable[[], Awaitable[tuple[str, TaskCommandRouterClient]]]
+
+
 class _StreamWriterThroughServer:
     """Provides an interface to buffer and write to a sandbox stream (`stdin`) via the server."""
 
@@ -584,14 +630,18 @@ class _StreamWriterThroughServer:
             raise ValueError(str(exc))
 
 
-class _StreamWriterThroughCommandRouter:
-    def __init__(self, params: _StreamWriterThroughCommandRouterParams) -> None:
-        self._object_id = params.object_id
-        self._command_router_client = params.command_router_client
-        self._task_id = params.task_id
-        self._is_closed = False
-        self._buffer = bytearray()
-        self._offset = 0
+class _StreamWriterThroughCommandRouterBuffer(ABC):
+    """Buffering/draining logic for stream writers that flush data
+    to the task command router."""
+
+    def __init__(self) -> None:
+        self._buffer: bytearray = bytearray()
+        self._is_closed: bool = False
+        self._offset: int = 0
+
+    @abstractmethod
+    async def stdin_write(self, data: bytes, eof: bool) -> None:
+        """Write the given chunk (with optional EOF) to the command router."""
 
     def write(self, data: Union[bytes, bytearray, memoryview, str]) -> None:
         if self._is_closed:
@@ -613,29 +663,58 @@ class _StreamWriterThroughCommandRouter:
         # NB: There's no need to prevent writing eof twice, because the command router will ignore the second EOF.
         if self._buffer or eof:
             data = bytes(self._buffer)
-            await self._command_router_client.exec_stdin_write(
-                task_id=self._task_id, exec_id=self._object_id, offset=self._offset, data=data, eof=eof
-            )
+            await self.stdin_write(data, eof)
             # Only clear the buffer after writing the data to the command router is successful.
             # This allows the client to retry drain() in the event of an exception (though
-            # exec_stdin_write already retries on transient errors, so most users will probably
-            # not do this).
+            # the underlying write call already retries on transient errors, so most users will
+            # probably not do this).
             self._buffer.clear()
             self._offset += len(data)
 
 
+class _StreamWriterThroughCommandRouterSandboxExec(_StreamWriterThroughCommandRouterBuffer):
+    def __init__(self, params: _StreamWriterThroughCommandRouterSandboxExecParams) -> None:
+        super().__init__()
+        self._object_id = params.object_id
+        self._command_router_client = params.command_router_client
+        self._task_id = params.task_id
+
+    async def stdin_write(self, data: bytes, eof: bool) -> None:
+        await self._command_router_client.exec_stdin_write(
+            task_id=self._task_id, exec_id=self._object_id, offset=self._offset, data=data, eof=eof
+        )
+
+
+class _StreamWriterThroughCommandRouterSandbox(_StreamWriterThroughCommandRouterBuffer):
+    """Write a V2 sandbox entrypoint's stdin directly to the worker
+    via the task command router."""
+
+    def __init__(self, params: _StreamWriterThroughCommandRouterSandboxParams) -> None:
+        super().__init__()
+        self._resolve_router = params.resolve_router
+
+    async def stdin_write(self, data: bytes, eof: bool) -> None:
+        task_id, client = await self._resolve_router()
+        await client.sandbox_stdin_write_v2(task_id=task_id, offset=self._offset, data=data, eof=eof)
+
+
 class _StreamWriter:
     """Provides an interface to buffer and write logs to a sandbox or container process stream (`stdin`)."""
 
     def __init__(
         self,
-        params: Union[_StreamWriterThroughServerParams, _StreamWriterThroughCommandRouterParams],
+        params: Union[
+            _StreamWriterThroughServerParams,
+            _StreamWriterThroughCommandRouterSandboxExecParams,
+            _StreamWriterThroughCommandRouterSandboxParams,
+        ],
     ) -> None:
         """mdmd:hidden"""
-        if isinstance(params, _StreamWriterThroughCommandRouterParams):
-            self._impl = _StreamWriterThroughCommandRouter(params)
+        if isinstance(params, _StreamWriterThroughCommandRouterSandboxExecParams):
+            self._impl = _StreamWriterThroughCommandRouterSandboxExec(params)
+        elif isinstance(params, _StreamWriterThroughCommandRouterSandboxParams):
+            self._impl = _StreamWriterThroughCommandRouterSandbox(params)
         else:
-            # Sandbox stdin is written via the server.
             self._impl = _StreamWriterThroughServer(params)
 
     def write(self, data: Union[bytes, bytearray, memoryview, str]) -> None: