modal 0.72.37__py3-none-any.whl → 0.72.39__py3-none-any.whl

modal/partial_function.pyi

@@ -25,7 +25,7 @@ ReturnType = typing.TypeVar("ReturnType", covariant=True)
 OriginalReturnType = typing.TypeVar("OriginalReturnType", covariant=True)
 
 class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
-    raw_f: typing.Callable[P, ReturnType]
+    raw_f: collections.abc.Callable[P, ReturnType]
     flags: _PartialFunctionFlags
     webhook_config: typing.Optional[modal_proto.api_pb2.WebhookConfig]
     is_generator: bool
@@ -38,7 +38,7 @@ class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
 
     def __init__(
         self,
-        raw_f: typing.Callable[P, ReturnType],
+        raw_f: collections.abc.Callable[P, ReturnType],
         flags: _PartialFunctionFlags,
         webhook_config: typing.Optional[modal_proto.api_pb2.WebhookConfig] = None,
         is_generator: typing.Optional[bool] = None,
@@ -49,13 +49,13 @@ class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
         force_build: bool = False,
         build_timeout: typing.Optional[int] = None,
     ): ...
-    def _get_raw_f(self) -> typing.Callable[P, ReturnType]: ...
+    def _get_raw_f(self) -> collections.abc.Callable[P, ReturnType]: ...
     def __get__(self, obj, objtype=None) -> modal.functions._Function[P, ReturnType, OriginalReturnType]: ...
     def __del__(self): ...
     def add_flags(self, flags) -> _PartialFunction: ...
 
 class PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
-    raw_f: typing.Callable[P, ReturnType]
+    raw_f: collections.abc.Callable[P, ReturnType]
     flags: _PartialFunctionFlags
     webhook_config: typing.Optional[modal_proto.api_pb2.WebhookConfig]
     is_generator: bool
@@ -68,7 +68,7 @@ class PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
 
     def __init__(
         self,
-        raw_f: typing.Callable[P, ReturnType],
+        raw_f: collections.abc.Callable[P, ReturnType],
         flags: _PartialFunctionFlags,
         webhook_config: typing.Optional[modal_proto.api_pb2.WebhookConfig] = None,
         is_generator: typing.Optional[bool] = None,
@@ -79,13 +79,15 @@ class PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
         force_build: bool = False,
         build_timeout: typing.Optional[int] = None,
     ): ...
-    def _get_raw_f(self) -> typing.Callable[P, ReturnType]: ...
+    def _get_raw_f(self) -> collections.abc.Callable[P, ReturnType]: ...
     def __get__(self, obj, objtype=None) -> modal.functions.Function[P, ReturnType, OriginalReturnType]: ...
     def __del__(self): ...
     def add_flags(self, flags) -> PartialFunction: ...
 
 def _find_partial_methods_for_user_cls(user_cls: type[typing.Any], flags: int) -> dict[str, _PartialFunction]: ...
-def _find_callables_for_obj(user_obj: typing.Any, flags: int) -> dict[str, typing.Callable[..., typing.Any]]: ...
+def _find_callables_for_obj(
+    user_obj: typing.Any, flags: int
+) -> dict[str, collections.abc.Callable[..., typing.Any]]: ...
 
 class _MethodDecoratorType:
     @typing.overload
@@ -95,13 +97,13 @@ class _MethodDecoratorType:
     @typing.overload
     def __call__(
         self,
-        func: typing.Callable[
+        func: collections.abc.Callable[
             typing_extensions.Concatenate[typing.Any, P], collections.abc.Coroutine[typing.Any, typing.Any, ReturnType]
         ],
     ) -> PartialFunction[P, ReturnType, collections.abc.Coroutine[typing.Any, typing.Any, ReturnType]]: ...
     @typing.overload
     def __call__(
-        self, func: typing.Callable[typing_extensions.Concatenate[typing.Any, P], ReturnType]
+        self, func: collections.abc.Callable[typing_extensions.Concatenate[typing.Any, P], ReturnType]
     ) -> PartialFunction[P, ReturnType, ReturnType]: ...
 
 def _method(
@@ -122,7 +124,9 @@ def _web_endpoint(
     custom_domains: typing.Optional[collections.abc.Iterable[str]] = None,
     requires_proxy_auth: bool = False,
     wait_for_response: bool = True,
-) -> typing.Callable[[typing.Callable[P, ReturnType]], _PartialFunction[P, ReturnType, ReturnType]]: ...
+) -> collections.abc.Callable[
+    [collections.abc.Callable[P, ReturnType]], _PartialFunction[P, ReturnType, ReturnType]
+]: ...
 def _asgi_app(
     _warn_parentheses_missing=None,
     *,
@@ -130,7 +134,7 @@ def _asgi_app(
     custom_domains: typing.Optional[collections.abc.Iterable[str]] = None,
     requires_proxy_auth: bool = False,
     wait_for_response: bool = True,
-) -> typing.Callable[[typing.Callable[..., typing.Any]], _PartialFunction]: ...
+) -> collections.abc.Callable[[collections.abc.Callable[..., typing.Any]], _PartialFunction]: ...
 def _wsgi_app(
     _warn_parentheses_missing=None,
     *,
@@ -138,7 +142,7 @@ def _wsgi_app(
     custom_domains: typing.Optional[collections.abc.Iterable[str]] = None,
     requires_proxy_auth: bool = False,
     wait_for_response: bool = True,
-) -> typing.Callable[[typing.Callable[..., typing.Any]], _PartialFunction]: ...
+) -> collections.abc.Callable[[collections.abc.Callable[..., typing.Any]], _PartialFunction]: ...
 def _web_server(
     port: int,
     *,
@@ -146,31 +150,35 @@ def _web_server(
     label: typing.Optional[str] = None,
     custom_domains: typing.Optional[collections.abc.Iterable[str]] = None,
     requires_proxy_auth: bool = False,
-) -> typing.Callable[[typing.Callable[..., typing.Any]], _PartialFunction]: ...
+) -> collections.abc.Callable[[collections.abc.Callable[..., typing.Any]], _PartialFunction]: ...
 def _disallow_wrapping_method(f: _PartialFunction, wrapper: str) -> None: ...
 def _build(
     _warn_parentheses_missing=None, *, force: bool = False, timeout: int = 86400
-) -> typing.Callable[[typing.Union[typing.Callable[[typing.Any], typing.Any], _PartialFunction]], _PartialFunction]: ...
+) -> collections.abc.Callable[
+    [typing.Union[collections.abc.Callable[[typing.Any], typing.Any], _PartialFunction]], _PartialFunction
+]: ...
 def _enter(
     _warn_parentheses_missing=None, *, snap: bool = False
-) -> typing.Callable[[typing.Union[typing.Callable[[typing.Any], typing.Any], _PartialFunction]], _PartialFunction]: ...
+) -> collections.abc.Callable[
+    [typing.Union[collections.abc.Callable[[typing.Any], typing.Any], _PartialFunction]], _PartialFunction
+]: ...
 def _exit(
     _warn_parentheses_missing=None,
-) -> typing.Callable[
+) -> collections.abc.Callable[
     [
         typing.Union[
-            typing.Callable[
+            collections.abc.Callable[
                 [typing.Any, typing.Optional[type[BaseException]], typing.Optional[BaseException], typing.Any],
                 typing.Any,
             ],
-            typing.Callable[[typing.Any], typing.Any],
+            collections.abc.Callable[[typing.Any], typing.Any],
         ]
     ],
     _PartialFunction,
 ]: ...
 def _batched(
     _warn_parentheses_missing=None, *, max_batch_size: int, wait_ms: int
-) -> typing.Callable[[typing.Callable[..., typing.Any]], _PartialFunction]: ...
+) -> collections.abc.Callable[[collections.abc.Callable[..., typing.Any]], _PartialFunction]: ...
 def method(
     _warn_parentheses_missing=None,
     *,
@@ -186,7 +194,9 @@ def web_endpoint(
     custom_domains: typing.Optional[collections.abc.Iterable[str]] = None,
     requires_proxy_auth: bool = False,
     wait_for_response: bool = True,
-) -> typing.Callable[[typing.Callable[P, ReturnType]], PartialFunction[P, ReturnType, ReturnType]]: ...
+) -> collections.abc.Callable[
+    [collections.abc.Callable[P, ReturnType]], PartialFunction[P, ReturnType, ReturnType]
+]: ...
 def asgi_app(
     _warn_parentheses_missing=None,
     *,
@@ -194,7 +204,7 @@ def asgi_app(
     custom_domains: typing.Optional[collections.abc.Iterable[str]] = None,
     requires_proxy_auth: bool = False,
     wait_for_response: bool = True,
-) -> typing.Callable[[typing.Callable[..., typing.Any]], PartialFunction]: ...
+) -> collections.abc.Callable[[collections.abc.Callable[..., typing.Any]], PartialFunction]: ...
 def wsgi_app(
     _warn_parentheses_missing=None,
     *,
@@ -202,7 +212,7 @@ def wsgi_app(
     custom_domains: typing.Optional[collections.abc.Iterable[str]] = None,
     requires_proxy_auth: bool = False,
     wait_for_response: bool = True,
-) -> typing.Callable[[typing.Callable[..., typing.Any]], PartialFunction]: ...
+) -> collections.abc.Callable[[collections.abc.Callable[..., typing.Any]], PartialFunction]: ...
 def web_server(
     port: int,
     *,
@@ -210,27 +220,31 @@ def web_server(
     label: typing.Optional[str] = None,
     custom_domains: typing.Optional[collections.abc.Iterable[str]] = None,
     requires_proxy_auth: bool = False,
-) -> typing.Callable[[typing.Callable[..., typing.Any]], PartialFunction]: ...
+) -> collections.abc.Callable[[collections.abc.Callable[..., typing.Any]], PartialFunction]: ...
 def build(
     _warn_parentheses_missing=None, *, force: bool = False, timeout: int = 86400
-) -> typing.Callable[[typing.Union[typing.Callable[[typing.Any], typing.Any], PartialFunction]], PartialFunction]: ...
+) -> collections.abc.Callable[
+    [typing.Union[collections.abc.Callable[[typing.Any], typing.Any], PartialFunction]], PartialFunction
+]: ...
 def enter(
     _warn_parentheses_missing=None, *, snap: bool = False
-) -> typing.Callable[[typing.Union[typing.Callable[[typing.Any], typing.Any], PartialFunction]], PartialFunction]: ...
+) -> collections.abc.Callable[
+    [typing.Union[collections.abc.Callable[[typing.Any], typing.Any], PartialFunction]], PartialFunction
+]: ...
 def exit(
     _warn_parentheses_missing=None,
-) -> typing.Callable[
+) -> collections.abc.Callable[
     [
         typing.Union[
-            typing.Callable[
+            collections.abc.Callable[
                 [typing.Any, typing.Optional[type[BaseException]], typing.Optional[BaseException], typing.Any],
                 typing.Any,
             ],
-            typing.Callable[[typing.Any], typing.Any],
+            collections.abc.Callable[[typing.Any], typing.Any],
         ]
     ],
     PartialFunction,
 ]: ...
 def batched(
     _warn_parentheses_missing=None, *, max_batch_size: int, wait_ms: int
-) -> typing.Callable[[typing.Callable[..., typing.Any]], PartialFunction]: ...
+) -> collections.abc.Callable[[collections.abc.Callable[..., typing.Any]], PartialFunction]: ...

modal/queue.pyi

@@ -85,6 +85,8 @@ class _Queue(modal._object._Object):
         self, *, partition: typing.Optional[str] = None, item_poll_timeout: float = 0.0
     ) -> collections.abc.AsyncGenerator[typing.Any, None]: ...
 
+SUPERSELF = typing.TypeVar("SUPERSELF", covariant=True)
+
 class Queue(modal.object.Object):
     def __init__(self): ...
     @staticmethod
@@ -139,13 +141,13 @@ class Queue(modal.object.Object):
 
     delete: __delete_spec
 
-    class ___get_nonblocking_spec(typing_extensions.Protocol):
+    class ___get_nonblocking_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, partition: typing.Optional[str], n_values: int) -> list[typing.Any]: ...
         async def aio(self, partition: typing.Optional[str], n_values: int) -> list[typing.Any]: ...
 
-    _get_nonblocking: ___get_nonblocking_spec
+    _get_nonblocking: ___get_nonblocking_spec[typing_extensions.Self]
 
-    class ___get_blocking_spec(typing_extensions.Protocol):
+    class ___get_blocking_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, partition: typing.Optional[str], timeout: typing.Optional[float], n_values: int
         ) -> list[typing.Any]: ...
@@ -153,15 +155,15 @@ class Queue(modal.object.Object):
             self, partition: typing.Optional[str], timeout: typing.Optional[float], n_values: int
         ) -> list[typing.Any]: ...
 
-    _get_blocking: ___get_blocking_spec
+    _get_blocking: ___get_blocking_spec[typing_extensions.Self]
 
-    class __clear_spec(typing_extensions.Protocol):
+    class __clear_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, *, partition: typing.Optional[str] = None, all: bool = False) -> None: ...
         async def aio(self, *, partition: typing.Optional[str] = None, all: bool = False) -> None: ...
 
-    clear: __clear_spec
+    clear: __clear_spec[typing_extensions.Self]
 
-    class __get_spec(typing_extensions.Protocol):
+    class __get_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, block: bool = True, timeout: typing.Optional[float] = None, *, partition: typing.Optional[str] = None
         ) -> typing.Optional[typing.Any]: ...
@@ -169,9 +171,9 @@ class Queue(modal.object.Object):
             self, block: bool = True, timeout: typing.Optional[float] = None, *, partition: typing.Optional[str] = None
         ) -> typing.Optional[typing.Any]: ...
 
-    get: __get_spec
+    get: __get_spec[typing_extensions.Self]
 
-    class __get_many_spec(typing_extensions.Protocol):
+    class __get_many_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self,
             n_values: int,
@@ -189,9 +191,9 @@ class Queue(modal.object.Object):
             partition: typing.Optional[str] = None,
         ) -> list[typing.Any]: ...
 
-    get_many: __get_many_spec
+    get_many: __get_many_spec[typing_extensions.Self]
 
-    class __put_spec(typing_extensions.Protocol):
+    class __put_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self,
             v: typing.Any,
@@ -211,9 +213,9 @@ class Queue(modal.object.Object):
             partition_ttl: int = 86400,
         ) -> None: ...
 
-    put: __put_spec
+    put: __put_spec[typing_extensions.Self]
 
-    class __put_many_spec(typing_extensions.Protocol):
+    class __put_many_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self,
             vs: list[typing.Any],
@@ -233,9 +235,9 @@ class Queue(modal.object.Object):
             partition_ttl: int = 86400,
         ) -> None: ...
 
-    put_many: __put_many_spec
+    put_many: __put_many_spec[typing_extensions.Self]
 
-    class ___put_many_blocking_spec(typing_extensions.Protocol):
+    class ___put_many_blocking_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self,
             partition: typing.Optional[str],
@@ -251,21 +253,21 @@ class Queue(modal.object.Object):
             timeout: typing.Optional[float] = None,
         ): ...
 
-    _put_many_blocking: ___put_many_blocking_spec
+    _put_many_blocking: ___put_many_blocking_spec[typing_extensions.Self]
 
-    class ___put_many_nonblocking_spec(typing_extensions.Protocol):
+    class ___put_many_nonblocking_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, partition: typing.Optional[str], partition_ttl: int, vs: list[typing.Any]): ...
         async def aio(self, partition: typing.Optional[str], partition_ttl: int, vs: list[typing.Any]): ...
 
-    _put_many_nonblocking: ___put_many_nonblocking_spec
+    _put_many_nonblocking: ___put_many_nonblocking_spec[typing_extensions.Self]
 
-    class __len_spec(typing_extensions.Protocol):
+    class __len_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, *, partition: typing.Optional[str] = None, total: bool = False) -> int: ...
         async def aio(self, *, partition: typing.Optional[str] = None, total: bool = False) -> int: ...
 
-    len: __len_spec
+    len: __len_spec[typing_extensions.Self]
 
-    class __iterate_spec(typing_extensions.Protocol):
+    class __iterate_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, *, partition: typing.Optional[str] = None, item_poll_timeout: float = 0.0
         ) -> typing.Generator[typing.Any, None, None]: ...
@@ -273,4 +275,4 @@ class Queue(modal.object.Object):
             self, *, partition: typing.Optional[str] = None, item_poll_timeout: float = 0.0
         ) -> collections.abc.AsyncGenerator[typing.Any, None]: ...
 
-    iterate: __iterate_spec
+    iterate: __iterate_spec[typing_extensions.Self]

modal/sandbox.py

@@ -517,8 +517,9 @@ class _Sandbox(_Object, type_prefix="sb"):
             raise InvalidError(f"workdir must be an absolute path, got: {workdir}")
 
         # Force secret resolution so we can pass the secret IDs to the backend.
+        # TODO should we parallelize this?
         for secret in secrets:
-            await secret.resolve(client=self._client)
+            await secret.hydrate(client=self._client)
 
         task_id = await self._get_task_id()
         req = api_pb2.ContainerExecRequest(

modal/sandbox.pyi

@@ -155,6 +155,8 @@ class _Sandbox(modal._object._Object):
         client: typing.Optional[modal.client._Client] = None,
     ) -> collections.abc.AsyncGenerator[_Sandbox, None]: ...
 
+SUPERSELF = typing.TypeVar("SUPERSELF", covariant=True)
+
 class Sandbox(modal.object.Object):
     _result: typing.Optional[modal_proto.api_pb2.GenericResult]
     _stdout: modal.io_streams.StreamReader[str]
@@ -264,49 +266,49 @@ class Sandbox(modal.object.Object):
 
     from_id: __from_id_spec
 
-    class __set_tags_spec(typing_extensions.Protocol):
+    class __set_tags_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, tags: dict[str, str], *, client: typing.Optional[modal.client.Client] = None): ...
         async def aio(self, tags: dict[str, str], *, client: typing.Optional[modal.client.Client] = None): ...
 
-    set_tags: __set_tags_spec
+    set_tags: __set_tags_spec[typing_extensions.Self]
 
-    class __snapshot_filesystem_spec(typing_extensions.Protocol):
+    class __snapshot_filesystem_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, timeout: int = 55) -> modal.image.Image: ...
         async def aio(self, timeout: int = 55) -> modal.image.Image: ...
 
-    snapshot_filesystem: __snapshot_filesystem_spec
+    snapshot_filesystem: __snapshot_filesystem_spec[typing_extensions.Self]
 
-    class __wait_spec(typing_extensions.Protocol):
+    class __wait_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, raise_on_termination: bool = True): ...
         async def aio(self, raise_on_termination: bool = True): ...
 
-    wait: __wait_spec
+    wait: __wait_spec[typing_extensions.Self]
 
-    class __tunnels_spec(typing_extensions.Protocol):
+    class __tunnels_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, timeout: int = 50) -> dict[int, modal._tunnel.Tunnel]: ...
         async def aio(self, timeout: int = 50) -> dict[int, modal._tunnel.Tunnel]: ...
 
-    tunnels: __tunnels_spec
+    tunnels: __tunnels_spec[typing_extensions.Self]
 
-    class __terminate_spec(typing_extensions.Protocol):
+    class __terminate_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self): ...
         async def aio(self): ...
 
-    terminate: __terminate_spec
+    terminate: __terminate_spec[typing_extensions.Self]
 
-    class __poll_spec(typing_extensions.Protocol):
+    class __poll_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self) -> typing.Optional[int]: ...
         async def aio(self) -> typing.Optional[int]: ...
 
-    poll: __poll_spec
+    poll: __poll_spec[typing_extensions.Self]
 
-    class ___get_task_id_spec(typing_extensions.Protocol):
+    class ___get_task_id_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self): ...
         async def aio(self): ...
 
-    _get_task_id: ___get_task_id_spec
+    _get_task_id: ___get_task_id_spec[typing_extensions.Self]
 
-    class __exec_spec(typing_extensions.Protocol):
+    class __exec_spec(typing_extensions.Protocol[SUPERSELF]):
         @typing.overload
         def __call__(
             self,
@@ -364,9 +366,9 @@ class Sandbox(modal.object.Object):
             _pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
         ) -> modal.container_process.ContainerProcess[bytes]: ...
 
-    exec: __exec_spec
+    exec: __exec_spec[typing_extensions.Self]
 
-    class __open_spec(typing_extensions.Protocol):
+    class __open_spec(typing_extensions.Protocol[SUPERSELF]):
         @typing.overload
         def __call__(self, path: str, mode: _typeshed.OpenTextMode) -> modal.file_io.FileIO[str]: ...
         @typing.overload
@@ -376,27 +378,27 @@ class Sandbox(modal.object.Object):
         @typing.overload
         async def aio(self, path: str, mode: _typeshed.OpenBinaryMode) -> modal.file_io.FileIO[bytes]: ...
 
-    open: __open_spec
+    open: __open_spec[typing_extensions.Self]
 
-    class __ls_spec(typing_extensions.Protocol):
+    class __ls_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, path: str) -> list[str]: ...
         async def aio(self, path: str) -> list[str]: ...
 
-    ls: __ls_spec
+    ls: __ls_spec[typing_extensions.Self]
 
-    class __mkdir_spec(typing_extensions.Protocol):
+    class __mkdir_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, path: str, parents: bool = False) -> None: ...
         async def aio(self, path: str, parents: bool = False) -> None: ...
 
-    mkdir: __mkdir_spec
+    mkdir: __mkdir_spec[typing_extensions.Self]
 
-    class __rm_spec(typing_extensions.Protocol):
+    class __rm_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, path: str, recursive: bool = False) -> None: ...
         async def aio(self, path: str, recursive: bool = False) -> None: ...
 
-    rm: __rm_spec
+    rm: __rm_spec[typing_extensions.Self]
 
-    class __watch_spec(typing_extensions.Protocol):
+    class __watch_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self,
             path: str,
@@ -412,7 +414,7 @@ class Sandbox(modal.object.Object):
             timeout: typing.Optional[int] = None,
         ) -> typing.AsyncIterator[modal.file_io.FileWatchEvent]: ...
 
-    watch: __watch_spec
+    watch: __watch_spec[typing_extensions.Self]
 
     @property
     def stdout(self) -> modal.io_streams.StreamReader[str]: ...

modal/token_flow.pyi

@@ -13,10 +13,12 @@ class _TokenFlow:
         self, timeout: float = 40.0, grpc_extra_timeout: float = 5.0
     ) -> typing.Optional[modal_proto.api_pb2.TokenFlowWaitResponse]: ...
 
+SUPERSELF = typing.TypeVar("SUPERSELF", covariant=True)
+
 class TokenFlow:
     def __init__(self, client: modal.client.Client): ...
 
-    class __start_spec(typing_extensions.Protocol):
+    class __start_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, utm_source: typing.Optional[str] = None, next_url: typing.Optional[str] = None
         ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[tuple[str, str, str]]: ...
@@ -24,9 +26,9 @@ class TokenFlow:
             self, utm_source: typing.Optional[str] = None, next_url: typing.Optional[str] = None
         ) -> typing.AsyncContextManager[tuple[str, str, str]]: ...
 
-    start: __start_spec
+    start: __start_spec[typing_extensions.Self]
 
-    class __finish_spec(typing_extensions.Protocol):
+    class __finish_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, timeout: float = 40.0, grpc_extra_timeout: float = 5.0
         ) -> typing.Optional[modal_proto.api_pb2.TokenFlowWaitResponse]: ...
@@ -34,7 +36,7 @@ class TokenFlow:
             self, timeout: float = 40.0, grpc_extra_timeout: float = 5.0
         ) -> typing.Optional[modal_proto.api_pb2.TokenFlowWaitResponse]: ...
 
-    finish: __finish_spec
+    finish: __finish_spec[typing_extensions.Self]
 
 async def _new_token(
     *,