diff --git a/.changeset/young-onions-cut.md b/.changeset/young-onions-cut.md new file mode 100644 index 0000000000..8a2e1f3878 --- /dev/null +++ b/.changeset/young-onions-cut.md @@ -0,0 +1,5 @@ +--- +'@e2b/python-sdk': minor +--- + +Changes the connectrpc implementation to use the official plugin instead diff --git a/.gitattributes b/.gitattributes index 8a1dff1bdb..f61d388004 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,2 +1,4 @@ packages/python-sdk/e2b/api/client/** linguist-generated=true +packages/python-sdk/e2b/envd/filesystem/** linguist-generated=true +packages/python-sdk/e2b/envd/process/** linguist-generated=true **/*.gen.ts linguist-generated=true diff --git a/codegen.Dockerfile b/codegen.Dockerfile index 34ebf5e92c..f38e6b2d26 100644 --- a/codegen.Dockerfile +++ b/codegen.Dockerfile @@ -5,10 +5,6 @@ RUN go install github.com/bufbuild/buf/cmd/buf@v1.50.1 && \ go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.28.1 && \ go install connectrpc.com/connect/cmd/protoc-gen-connect-go@v1.18.1 -# Install our custom protoc plugin, connect-python -COPY ./packages/connect-python /packages/connect-python -RUN cd /packages/connect-python && make bin/protoc-gen-connect-python - FROM python:3.10 diff --git a/packages/python-sdk/Makefile b/packages/python-sdk/Makefile index 572c0fbd99..958e4beb9b 100644 --- a/packages/python-sdk/Makefile +++ b/packages/python-sdk/Makefile @@ -16,10 +16,6 @@ generate-volume-api: ruff format . generate-envd: - if [ ! -f "/go/bin/protoc-gen-connect-python" ]; then \ - $(MAKE) -C $(ROOT_DIR)/packages/connect-python build; \ - fi - cd $(ROOT_DIR)/spec/envd && pwd && buf generate --template buf-python.gen.yaml ./scripts/fix-python-pb.sh diff --git a/packages/python-sdk/e2b/envd/filesystem/filesystem_connect.py b/packages/python-sdk/e2b/envd/filesystem/filesystem_connect.py index e995585e88..effaab697d 100644 --- a/packages/python-sdk/e2b/envd/filesystem/filesystem_connect.py +++ b/packages/python-sdk/e2b/envd/filesystem/filesystem_connect.py @@ -1,193 +1,734 @@ -# Code generated by protoc-gen-connect-python 0.1.0.dev2, DO NOT EDIT. -from typing import Any, Generator, Coroutine, AsyncGenerator, Optional -from httpcore import ConnectionPool, AsyncConnectionPool +# -*- coding: utf-8 -*- +# Generated by https://github.com/connectrpc/connect-python. DO NOT EDIT! +# source: filesystem/filesystem.proto -import e2b_connect as connect +from collections.abc import AsyncGenerator, AsyncIterator, Iterable, Iterator, Mapping +from typing import Protocol +from connectrpc.client import ConnectClient, ConnectClientSync +from connectrpc.code import Code +from connectrpc.codec import Codec +from connectrpc.compression import Compression +from connectrpc.errors import ConnectError +from connectrpc.interceptor import Interceptor, InterceptorSync +from connectrpc.method import IdempotencyLevel, MethodInfo +from connectrpc.request import Headers, RequestContext +from connectrpc.server import ( + ConnectASGIApplication, + ConnectWSGIApplication, + Endpoint, + EndpointSync, +) from e2b.envd.filesystem import filesystem_pb2 as filesystem_dot_filesystem__pb2 -FilesystemName = "filesystem.Filesystem" +class Filesystem(Protocol): + async def stat( + self, request: filesystem_dot_filesystem__pb2.StatRequest, ctx: RequestContext + ) -> filesystem_dot_filesystem__pb2.StatResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def make_dir( + self, + request: filesystem_dot_filesystem__pb2.MakeDirRequest, + ctx: RequestContext, + ) -> filesystem_dot_filesystem__pb2.MakeDirResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def move( + self, request: filesystem_dot_filesystem__pb2.MoveRequest, ctx: RequestContext + ) -> filesystem_dot_filesystem__pb2.MoveResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") -class FilesystemClient: + async def list_dir( + self, + request: filesystem_dot_filesystem__pb2.ListDirRequest, + ctx: RequestContext, + ) -> filesystem_dot_filesystem__pb2.ListDirResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def remove( + self, request: filesystem_dot_filesystem__pb2.RemoveRequest, ctx: RequestContext + ) -> filesystem_dot_filesystem__pb2.RemoveResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + def watch_dir( + self, + request: filesystem_dot_filesystem__pb2.WatchDirRequest, + ctx: RequestContext, + ) -> AsyncIterator[filesystem_dot_filesystem__pb2.WatchDirResponse]: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def create_watcher( + self, + request: filesystem_dot_filesystem__pb2.CreateWatcherRequest, + ctx: RequestContext, + ) -> filesystem_dot_filesystem__pb2.CreateWatcherResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def get_watcher_events( + self, + request: filesystem_dot_filesystem__pb2.GetWatcherEventsRequest, + ctx: RequestContext, + ) -> filesystem_dot_filesystem__pb2.GetWatcherEventsResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def remove_watcher( + self, + request: filesystem_dot_filesystem__pb2.RemoveWatcherRequest, + ctx: RequestContext, + ) -> filesystem_dot_filesystem__pb2.RemoveWatcherResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + +class FilesystemASGIApplication(ConnectASGIApplication[Filesystem]): def __init__( self, - base_url: str, - *, - pool: Optional[ConnectionPool] = None, - async_pool: Optional[AsyncConnectionPool] = None, - compressor=None, - json=False, - **opts, - ): - self._stat = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{FilesystemName}/Stat", - response_type=filesystem_dot_filesystem__pb2.StatResponse, - compressor=compressor, - json=json, - **opts, - ) - self._make_dir = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{FilesystemName}/MakeDir", - response_type=filesystem_dot_filesystem__pb2.MakeDirResponse, - compressor=compressor, - json=json, - **opts, - ) - self._move = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{FilesystemName}/Move", - response_type=filesystem_dot_filesystem__pb2.MoveResponse, - compressor=compressor, - json=json, - **opts, - ) - self._list_dir = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{FilesystemName}/ListDir", - response_type=filesystem_dot_filesystem__pb2.ListDirResponse, - compressor=compressor, - json=json, - **opts, - ) - self._remove = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{FilesystemName}/Remove", - response_type=filesystem_dot_filesystem__pb2.RemoveResponse, - compressor=compressor, - json=json, - **opts, - ) - self._watch_dir = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{FilesystemName}/WatchDir", - response_type=filesystem_dot_filesystem__pb2.WatchDirResponse, - compressor=compressor, - json=json, - **opts, - ) - self._create_watcher = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{FilesystemName}/CreateWatcher", - response_type=filesystem_dot_filesystem__pb2.CreateWatcherResponse, - compressor=compressor, - json=json, - **opts, - ) - self._get_watcher_events = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{FilesystemName}/GetWatcherEvents", - response_type=filesystem_dot_filesystem__pb2.GetWatcherEventsResponse, - compressor=compressor, - json=json, - **opts, - ) - self._remove_watcher = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{FilesystemName}/RemoveWatcher", - response_type=filesystem_dot_filesystem__pb2.RemoveWatcherResponse, - compressor=compressor, - json=json, - **opts, + service: Filesystem | AsyncGenerator[Filesystem], + *, + interceptors: Iterable[Interceptor] = (), + read_max_bytes: int | None = None, + compressions: Iterable[Compression] | None = None, + codecs: Iterable[Codec] | None = None, + ) -> None: + super().__init__( + service=service, + endpoints=lambda svc: { + "/filesystem.Filesystem/Stat": Endpoint.unary( + method=MethodInfo( + name="Stat", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.StatRequest, + output=filesystem_dot_filesystem__pb2.StatResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.stat, + ), + "/filesystem.Filesystem/MakeDir": Endpoint.unary( + method=MethodInfo( + name="MakeDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.MakeDirRequest, + output=filesystem_dot_filesystem__pb2.MakeDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.make_dir, + ), + "/filesystem.Filesystem/Move": Endpoint.unary( + method=MethodInfo( + name="Move", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.MoveRequest, + output=filesystem_dot_filesystem__pb2.MoveResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.move, + ), + "/filesystem.Filesystem/ListDir": Endpoint.unary( + method=MethodInfo( + name="ListDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.ListDirRequest, + output=filesystem_dot_filesystem__pb2.ListDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.list_dir, + ), + "/filesystem.Filesystem/Remove": Endpoint.unary( + method=MethodInfo( + name="Remove", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.RemoveRequest, + output=filesystem_dot_filesystem__pb2.RemoveResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.remove, + ), + "/filesystem.Filesystem/WatchDir": Endpoint.server_stream( + method=MethodInfo( + name="WatchDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.WatchDirRequest, + output=filesystem_dot_filesystem__pb2.WatchDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.watch_dir, + ), + "/filesystem.Filesystem/CreateWatcher": Endpoint.unary( + method=MethodInfo( + name="CreateWatcher", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.CreateWatcherRequest, + output=filesystem_dot_filesystem__pb2.CreateWatcherResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.create_watcher, + ), + "/filesystem.Filesystem/GetWatcherEvents": Endpoint.unary( + method=MethodInfo( + name="GetWatcherEvents", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.GetWatcherEventsRequest, + output=filesystem_dot_filesystem__pb2.GetWatcherEventsResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.get_watcher_events, + ), + "/filesystem.Filesystem/RemoveWatcher": Endpoint.unary( + method=MethodInfo( + name="RemoveWatcher", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.RemoveWatcherRequest, + output=filesystem_dot_filesystem__pb2.RemoveWatcherResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.remove_watcher, + ), + }, + interceptors=interceptors, + read_max_bytes=read_max_bytes, + compressions=compressions, + codecs=codecs, ) - def stat( - self, req: filesystem_dot_filesystem__pb2.StatRequest, **opts + @property + def path(self) -> str: + """Returns the URL path to mount the application to when serving multiple applications.""" + return "/filesystem.Filesystem" + + +class FilesystemClient(ConnectClient): + async def stat( + self, + request: filesystem_dot_filesystem__pb2.StatRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, ) -> filesystem_dot_filesystem__pb2.StatResponse: - return self._stat.call_unary(req, **opts) + return await self.execute_unary( + request=request, + method=MethodInfo( + name="Stat", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.StatRequest, + output=filesystem_dot_filesystem__pb2.StatResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) - def astat( - self, req: filesystem_dot_filesystem__pb2.StatRequest, **opts - ) -> Coroutine[Any, Any, filesystem_dot_filesystem__pb2.StatResponse]: - return self._stat.acall_unary(req, **opts) + async def make_dir( + self, + request: filesystem_dot_filesystem__pb2.MakeDirRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.MakeDirResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="MakeDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.MakeDirRequest, + output=filesystem_dot_filesystem__pb2.MakeDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + async def move( + self, + request: filesystem_dot_filesystem__pb2.MoveRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.MoveResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="Move", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.MoveRequest, + output=filesystem_dot_filesystem__pb2.MoveResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + async def list_dir( + self, + request: filesystem_dot_filesystem__pb2.ListDirRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.ListDirResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="ListDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.ListDirRequest, + output=filesystem_dot_filesystem__pb2.ListDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + async def remove( + self, + request: filesystem_dot_filesystem__pb2.RemoveRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.RemoveResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="Remove", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.RemoveRequest, + output=filesystem_dot_filesystem__pb2.RemoveResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def watch_dir( + self, + request: filesystem_dot_filesystem__pb2.WatchDirRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> AsyncIterator[filesystem_dot_filesystem__pb2.WatchDirResponse]: + return self.execute_server_stream( + request=request, + method=MethodInfo( + name="WatchDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.WatchDirRequest, + output=filesystem_dot_filesystem__pb2.WatchDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + async def create_watcher( + self, + request: filesystem_dot_filesystem__pb2.CreateWatcherRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.CreateWatcherResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="CreateWatcher", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.CreateWatcherRequest, + output=filesystem_dot_filesystem__pb2.CreateWatcherResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + async def get_watcher_events( + self, + request: filesystem_dot_filesystem__pb2.GetWatcherEventsRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.GetWatcherEventsResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="GetWatcherEvents", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.GetWatcherEventsRequest, + output=filesystem_dot_filesystem__pb2.GetWatcherEventsResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + async def remove_watcher( + self, + request: filesystem_dot_filesystem__pb2.RemoveWatcherRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.RemoveWatcherResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="RemoveWatcher", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.RemoveWatcherRequest, + output=filesystem_dot_filesystem__pb2.RemoveWatcherResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + +class FilesystemSync(Protocol): + def stat( + self, request: filesystem_dot_filesystem__pb2.StatRequest, ctx: RequestContext + ) -> filesystem_dot_filesystem__pb2.StatResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def make_dir( - self, req: filesystem_dot_filesystem__pb2.MakeDirRequest, **opts + self, + request: filesystem_dot_filesystem__pb2.MakeDirRequest, + ctx: RequestContext, ) -> filesystem_dot_filesystem__pb2.MakeDirResponse: - return self._make_dir.call_unary(req, **opts) - - def amake_dir( - self, req: filesystem_dot_filesystem__pb2.MakeDirRequest, **opts - ) -> Coroutine[Any, Any, filesystem_dot_filesystem__pb2.MakeDirResponse]: - return self._make_dir.acall_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def move( - self, req: filesystem_dot_filesystem__pb2.MoveRequest, **opts + self, request: filesystem_dot_filesystem__pb2.MoveRequest, ctx: RequestContext ) -> filesystem_dot_filesystem__pb2.MoveResponse: - return self._move.call_unary(req, **opts) - - def amove( - self, req: filesystem_dot_filesystem__pb2.MoveRequest, **opts - ) -> Coroutine[Any, Any, filesystem_dot_filesystem__pb2.MoveResponse]: - return self._move.acall_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def list_dir( - self, req: filesystem_dot_filesystem__pb2.ListDirRequest, **opts + self, + request: filesystem_dot_filesystem__pb2.ListDirRequest, + ctx: RequestContext, ) -> filesystem_dot_filesystem__pb2.ListDirResponse: - return self._list_dir.call_unary(req, **opts) - - def alist_dir( - self, req: filesystem_dot_filesystem__pb2.ListDirRequest, **opts - ) -> Coroutine[Any, Any, filesystem_dot_filesystem__pb2.ListDirResponse]: - return self._list_dir.acall_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def remove( - self, req: filesystem_dot_filesystem__pb2.RemoveRequest, **opts + self, request: filesystem_dot_filesystem__pb2.RemoveRequest, ctx: RequestContext ) -> filesystem_dot_filesystem__pb2.RemoveResponse: - return self._remove.call_unary(req, **opts) - - def aremove( - self, req: filesystem_dot_filesystem__pb2.RemoveRequest, **opts - ) -> Coroutine[Any, Any, filesystem_dot_filesystem__pb2.RemoveResponse]: - return self._remove.acall_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def watch_dir( - self, req: filesystem_dot_filesystem__pb2.WatchDirRequest, **opts - ) -> Generator[filesystem_dot_filesystem__pb2.WatchDirResponse, Any, None]: - return self._watch_dir.call_server_stream(req, **opts) - - def awatch_dir( - self, req: filesystem_dot_filesystem__pb2.WatchDirRequest, **opts - ) -> AsyncGenerator[filesystem_dot_filesystem__pb2.WatchDirResponse, Any]: - return self._watch_dir.acall_server_stream(req, **opts) + self, + request: filesystem_dot_filesystem__pb2.WatchDirRequest, + ctx: RequestContext, + ) -> Iterator[filesystem_dot_filesystem__pb2.WatchDirResponse]: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def create_watcher( - self, req: filesystem_dot_filesystem__pb2.CreateWatcherRequest, **opts + self, + request: filesystem_dot_filesystem__pb2.CreateWatcherRequest, + ctx: RequestContext, ) -> filesystem_dot_filesystem__pb2.CreateWatcherResponse: - return self._create_watcher.call_unary(req, **opts) - - def acreate_watcher( - self, req: filesystem_dot_filesystem__pb2.CreateWatcherRequest, **opts - ) -> Coroutine[Any, Any, filesystem_dot_filesystem__pb2.CreateWatcherResponse]: - return self._create_watcher.acall_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def get_watcher_events( - self, req: filesystem_dot_filesystem__pb2.GetWatcherEventsRequest, **opts + self, + request: filesystem_dot_filesystem__pb2.GetWatcherEventsRequest, + ctx: RequestContext, ) -> filesystem_dot_filesystem__pb2.GetWatcherEventsResponse: - return self._get_watcher_events.call_unary(req, **opts) - - def aget_watcher_events( - self, req: filesystem_dot_filesystem__pb2.GetWatcherEventsRequest, **opts - ) -> Coroutine[Any, Any, filesystem_dot_filesystem__pb2.GetWatcherEventsResponse]: - return self._get_watcher_events.acall_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def remove_watcher( - self, req: filesystem_dot_filesystem__pb2.RemoveWatcherRequest, **opts + self, + request: filesystem_dot_filesystem__pb2.RemoveWatcherRequest, + ctx: RequestContext, ) -> filesystem_dot_filesystem__pb2.RemoveWatcherResponse: - return self._remove_watcher.call_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + +class FilesystemWSGIApplication(ConnectWSGIApplication): + def __init__( + self, + service: FilesystemSync, + interceptors: Iterable[InterceptorSync] = (), + read_max_bytes: int | None = None, + compressions: Iterable[Compression] | None = None, + codecs: Iterable[Codec] | None = None, + ) -> None: + super().__init__( + endpoints={ + "/filesystem.Filesystem/Stat": EndpointSync.unary( + method=MethodInfo( + name="Stat", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.StatRequest, + output=filesystem_dot_filesystem__pb2.StatResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.stat, + ), + "/filesystem.Filesystem/MakeDir": EndpointSync.unary( + method=MethodInfo( + name="MakeDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.MakeDirRequest, + output=filesystem_dot_filesystem__pb2.MakeDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.make_dir, + ), + "/filesystem.Filesystem/Move": EndpointSync.unary( + method=MethodInfo( + name="Move", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.MoveRequest, + output=filesystem_dot_filesystem__pb2.MoveResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.move, + ), + "/filesystem.Filesystem/ListDir": EndpointSync.unary( + method=MethodInfo( + name="ListDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.ListDirRequest, + output=filesystem_dot_filesystem__pb2.ListDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.list_dir, + ), + "/filesystem.Filesystem/Remove": EndpointSync.unary( + method=MethodInfo( + name="Remove", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.RemoveRequest, + output=filesystem_dot_filesystem__pb2.RemoveResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.remove, + ), + "/filesystem.Filesystem/WatchDir": EndpointSync.server_stream( + method=MethodInfo( + name="WatchDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.WatchDirRequest, + output=filesystem_dot_filesystem__pb2.WatchDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.watch_dir, + ), + "/filesystem.Filesystem/CreateWatcher": EndpointSync.unary( + method=MethodInfo( + name="CreateWatcher", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.CreateWatcherRequest, + output=filesystem_dot_filesystem__pb2.CreateWatcherResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.create_watcher, + ), + "/filesystem.Filesystem/GetWatcherEvents": EndpointSync.unary( + method=MethodInfo( + name="GetWatcherEvents", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.GetWatcherEventsRequest, + output=filesystem_dot_filesystem__pb2.GetWatcherEventsResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.get_watcher_events, + ), + "/filesystem.Filesystem/RemoveWatcher": EndpointSync.unary( + method=MethodInfo( + name="RemoveWatcher", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.RemoveWatcherRequest, + output=filesystem_dot_filesystem__pb2.RemoveWatcherResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.remove_watcher, + ), + }, + interceptors=interceptors, + read_max_bytes=read_max_bytes, + compressions=compressions, + codecs=codecs, + ) - def aremove_watcher( - self, req: filesystem_dot_filesystem__pb2.RemoveWatcherRequest, **opts - ) -> Coroutine[Any, Any, filesystem_dot_filesystem__pb2.RemoveWatcherResponse]: - return self._remove_watcher.acall_unary(req, **opts) + @property + def path(self) -> str: + """Returns the URL path to mount the application to when serving multiple applications.""" + return "/filesystem.Filesystem" + + +class FilesystemClientSync(ConnectClientSync): + def stat( + self, + request: filesystem_dot_filesystem__pb2.StatRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.StatResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="Stat", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.StatRequest, + output=filesystem_dot_filesystem__pb2.StatResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def make_dir( + self, + request: filesystem_dot_filesystem__pb2.MakeDirRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.MakeDirResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="MakeDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.MakeDirRequest, + output=filesystem_dot_filesystem__pb2.MakeDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def move( + self, + request: filesystem_dot_filesystem__pb2.MoveRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.MoveResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="Move", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.MoveRequest, + output=filesystem_dot_filesystem__pb2.MoveResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def list_dir( + self, + request: filesystem_dot_filesystem__pb2.ListDirRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.ListDirResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="ListDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.ListDirRequest, + output=filesystem_dot_filesystem__pb2.ListDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def remove( + self, + request: filesystem_dot_filesystem__pb2.RemoveRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.RemoveResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="Remove", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.RemoveRequest, + output=filesystem_dot_filesystem__pb2.RemoveResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def watch_dir( + self, + request: filesystem_dot_filesystem__pb2.WatchDirRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> Iterator[filesystem_dot_filesystem__pb2.WatchDirResponse]: + return self.execute_server_stream( + request=request, + method=MethodInfo( + name="WatchDir", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.WatchDirRequest, + output=filesystem_dot_filesystem__pb2.WatchDirResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def create_watcher( + self, + request: filesystem_dot_filesystem__pb2.CreateWatcherRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.CreateWatcherResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="CreateWatcher", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.CreateWatcherRequest, + output=filesystem_dot_filesystem__pb2.CreateWatcherResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def get_watcher_events( + self, + request: filesystem_dot_filesystem__pb2.GetWatcherEventsRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.GetWatcherEventsResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="GetWatcherEvents", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.GetWatcherEventsRequest, + output=filesystem_dot_filesystem__pb2.GetWatcherEventsResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def remove_watcher( + self, + request: filesystem_dot_filesystem__pb2.RemoveWatcherRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> filesystem_dot_filesystem__pb2.RemoveWatcherResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="RemoveWatcher", + service_name="filesystem.Filesystem", + input=filesystem_dot_filesystem__pb2.RemoveWatcherRequest, + output=filesystem_dot_filesystem__pb2.RemoveWatcherResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) diff --git a/packages/python-sdk/e2b/envd/process/process_connect.py b/packages/python-sdk/e2b/envd/process/process_connect.py index dc09fbe698..d0cf47aea1 100644 --- a/packages/python-sdk/e2b/envd/process/process_connect.py +++ b/packages/python-sdk/e2b/envd/process/process_connect.py @@ -1,174 +1,644 @@ -# Code generated by protoc-gen-connect-python 0.1.0.dev2, DO NOT EDIT. -from typing import Any, Generator, Coroutine, AsyncGenerator, Optional -from httpcore import ConnectionPool, AsyncConnectionPool +# -*- coding: utf-8 -*- +# Generated by https://github.com/connectrpc/connect-python. DO NOT EDIT! +# source: process/process.proto -import e2b_connect as connect +from collections.abc import AsyncGenerator, AsyncIterator, Iterable, Iterator, Mapping +from typing import Protocol +from connectrpc.client import ConnectClient, ConnectClientSync +from connectrpc.code import Code +from connectrpc.codec import Codec +from connectrpc.compression import Compression +from connectrpc.errors import ConnectError +from connectrpc.interceptor import Interceptor, InterceptorSync +from connectrpc.method import IdempotencyLevel, MethodInfo +from connectrpc.request import Headers, RequestContext +from connectrpc.server import ( + ConnectASGIApplication, + ConnectWSGIApplication, + Endpoint, + EndpointSync, +) from e2b.envd.process import process_pb2 as process_dot_process__pb2 -ProcessName = "process.Process" +class Process(Protocol): + async def list( + self, request: process_dot_process__pb2.ListRequest, ctx: RequestContext + ) -> process_dot_process__pb2.ListResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + def connect( + self, request: process_dot_process__pb2.ConnectRequest, ctx: RequestContext + ) -> AsyncIterator[process_dot_process__pb2.ConnectResponse]: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + def start( + self, request: process_dot_process__pb2.StartRequest, ctx: RequestContext + ) -> AsyncIterator[process_dot_process__pb2.StartResponse]: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def update( + self, request: process_dot_process__pb2.UpdateRequest, ctx: RequestContext + ) -> process_dot_process__pb2.UpdateResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def stream_input( + self, + request: AsyncIterator[process_dot_process__pb2.StreamInputRequest], + ctx: RequestContext, + ) -> process_dot_process__pb2.StreamInputResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def send_input( + self, request: process_dot_process__pb2.SendInputRequest, ctx: RequestContext + ) -> process_dot_process__pb2.SendInputResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def send_signal( + self, request: process_dot_process__pb2.SendSignalRequest, ctx: RequestContext + ) -> process_dot_process__pb2.SendSignalResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + async def close_stdin( + self, request: process_dot_process__pb2.CloseStdinRequest, ctx: RequestContext + ) -> process_dot_process__pb2.CloseStdinResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") -class ProcessClient: + +class ProcessASGIApplication(ConnectASGIApplication[Process]): def __init__( self, - base_url: str, - *, - pool: Optional[ConnectionPool] = None, - async_pool: Optional[AsyncConnectionPool] = None, - compressor=None, - json=False, - **opts, - ): - self._list = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{ProcessName}/List", - response_type=process_dot_process__pb2.ListResponse, - compressor=compressor, - json=json, - **opts, - ) - self._connect = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{ProcessName}/Connect", - response_type=process_dot_process__pb2.ConnectResponse, - compressor=compressor, - json=json, - **opts, - ) - self._start = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{ProcessName}/Start", - response_type=process_dot_process__pb2.StartResponse, - compressor=compressor, - json=json, - **opts, - ) - self._update = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{ProcessName}/Update", - response_type=process_dot_process__pb2.UpdateResponse, - compressor=compressor, - json=json, - **opts, - ) - self._stream_input = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{ProcessName}/StreamInput", - response_type=process_dot_process__pb2.StreamInputResponse, - compressor=compressor, - json=json, - **opts, - ) - self._send_input = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{ProcessName}/SendInput", - response_type=process_dot_process__pb2.SendInputResponse, - compressor=compressor, - json=json, - **opts, - ) - self._send_signal = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{ProcessName}/SendSignal", - response_type=process_dot_process__pb2.SendSignalResponse, - compressor=compressor, - json=json, - **opts, - ) - self._close_stdin = connect.Client( - pool=pool, - async_pool=async_pool, - url=f"{base_url}/{ProcessName}/CloseStdin", - response_type=process_dot_process__pb2.CloseStdinResponse, - compressor=compressor, - json=json, - **opts, + service: Process | AsyncGenerator[Process], + *, + interceptors: Iterable[Interceptor] = (), + read_max_bytes: int | None = None, + compressions: Iterable[Compression] | None = None, + codecs: Iterable[Codec] | None = None, + ) -> None: + super().__init__( + service=service, + endpoints=lambda svc: { + "/process.Process/List": Endpoint.unary( + method=MethodInfo( + name="List", + service_name="process.Process", + input=process_dot_process__pb2.ListRequest, + output=process_dot_process__pb2.ListResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.list, + ), + "/process.Process/Connect": Endpoint.server_stream( + method=MethodInfo( + name="Connect", + service_name="process.Process", + input=process_dot_process__pb2.ConnectRequest, + output=process_dot_process__pb2.ConnectResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.connect, + ), + "/process.Process/Start": Endpoint.server_stream( + method=MethodInfo( + name="Start", + service_name="process.Process", + input=process_dot_process__pb2.StartRequest, + output=process_dot_process__pb2.StartResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.start, + ), + "/process.Process/Update": Endpoint.unary( + method=MethodInfo( + name="Update", + service_name="process.Process", + input=process_dot_process__pb2.UpdateRequest, + output=process_dot_process__pb2.UpdateResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.update, + ), + "/process.Process/StreamInput": Endpoint.client_stream( + method=MethodInfo( + name="StreamInput", + service_name="process.Process", + input=process_dot_process__pb2.StreamInputRequest, + output=process_dot_process__pb2.StreamInputResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.stream_input, + ), + "/process.Process/SendInput": Endpoint.unary( + method=MethodInfo( + name="SendInput", + service_name="process.Process", + input=process_dot_process__pb2.SendInputRequest, + output=process_dot_process__pb2.SendInputResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.send_input, + ), + "/process.Process/SendSignal": Endpoint.unary( + method=MethodInfo( + name="SendSignal", + service_name="process.Process", + input=process_dot_process__pb2.SendSignalRequest, + output=process_dot_process__pb2.SendSignalResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.send_signal, + ), + "/process.Process/CloseStdin": Endpoint.unary( + method=MethodInfo( + name="CloseStdin", + service_name="process.Process", + input=process_dot_process__pb2.CloseStdinRequest, + output=process_dot_process__pb2.CloseStdinResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=svc.close_stdin, + ), + }, + interceptors=interceptors, + read_max_bytes=read_max_bytes, + compressions=compressions, + codecs=codecs, ) - def list( - self, req: process_dot_process__pb2.ListRequest, **opts - ) -> process_dot_process__pb2.ListResponse: - return self._list.call_unary(req, **opts) + @property + def path(self) -> str: + """Returns the URL path to mount the application to when serving multiple applications.""" + return "/process.Process" - def alist( - self, req: process_dot_process__pb2.ListRequest, **opts - ) -> Coroutine[Any, Any, process_dot_process__pb2.ListResponse]: - return self._list.acall_unary(req, **opts) - def connect( - self, req: process_dot_process__pb2.ConnectRequest, **opts - ) -> Generator[process_dot_process__pb2.ConnectResponse, Any, None]: - return self._connect.call_server_stream(req, **opts) +class ProcessClient(ConnectClient): + async def list( + self, + request: process_dot_process__pb2.ListRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.ListResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="List", + service_name="process.Process", + input=process_dot_process__pb2.ListRequest, + output=process_dot_process__pb2.ListResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) - def aconnect( - self, req: process_dot_process__pb2.ConnectRequest, **opts - ) -> AsyncGenerator[process_dot_process__pb2.ConnectResponse, Any]: - return self._connect.acall_server_stream(req, **opts) + def connect( + self, + request: process_dot_process__pb2.ConnectRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> AsyncIterator[process_dot_process__pb2.ConnectResponse]: + return self.execute_server_stream( + request=request, + method=MethodInfo( + name="Connect", + service_name="process.Process", + input=process_dot_process__pb2.ConnectRequest, + output=process_dot_process__pb2.ConnectResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) def start( - self, req: process_dot_process__pb2.StartRequest, **opts - ) -> Generator[process_dot_process__pb2.StartResponse, Any, None]: - return self._start.call_server_stream(req, **opts) + self, + request: process_dot_process__pb2.StartRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> AsyncIterator[process_dot_process__pb2.StartResponse]: + return self.execute_server_stream( + request=request, + method=MethodInfo( + name="Start", + service_name="process.Process", + input=process_dot_process__pb2.StartRequest, + output=process_dot_process__pb2.StartResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + async def update( + self, + request: process_dot_process__pb2.UpdateRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.UpdateResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="Update", + service_name="process.Process", + input=process_dot_process__pb2.UpdateRequest, + output=process_dot_process__pb2.UpdateResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + async def stream_input( + self, + request: AsyncIterator[process_dot_process__pb2.StreamInputRequest], + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.StreamInputResponse: + return await self.execute_client_stream( + request=request, + method=MethodInfo( + name="StreamInput", + service_name="process.Process", + input=process_dot_process__pb2.StreamInputRequest, + output=process_dot_process__pb2.StreamInputResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + async def send_input( + self, + request: process_dot_process__pb2.SendInputRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.SendInputResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="SendInput", + service_name="process.Process", + input=process_dot_process__pb2.SendInputRequest, + output=process_dot_process__pb2.SendInputResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) - def astart( - self, req: process_dot_process__pb2.StartRequest, **opts - ) -> AsyncGenerator[process_dot_process__pb2.StartResponse, Any]: - return self._start.acall_server_stream(req, **opts) + async def send_signal( + self, + request: process_dot_process__pb2.SendSignalRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.SendSignalResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="SendSignal", + service_name="process.Process", + input=process_dot_process__pb2.SendSignalRequest, + output=process_dot_process__pb2.SendSignalResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + async def close_stdin( + self, + request: process_dot_process__pb2.CloseStdinRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.CloseStdinResponse: + return await self.execute_unary( + request=request, + method=MethodInfo( + name="CloseStdin", + service_name="process.Process", + input=process_dot_process__pb2.CloseStdinRequest, + output=process_dot_process__pb2.CloseStdinResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + +class ProcessSync(Protocol): + def list( + self, request: process_dot_process__pb2.ListRequest, ctx: RequestContext + ) -> process_dot_process__pb2.ListResponse: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + def connect( + self, request: process_dot_process__pb2.ConnectRequest, ctx: RequestContext + ) -> Iterator[process_dot_process__pb2.ConnectResponse]: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + def start( + self, request: process_dot_process__pb2.StartRequest, ctx: RequestContext + ) -> Iterator[process_dot_process__pb2.StartResponse]: + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def update( - self, req: process_dot_process__pb2.UpdateRequest, **opts + self, request: process_dot_process__pb2.UpdateRequest, ctx: RequestContext ) -> process_dot_process__pb2.UpdateResponse: - return self._update.call_unary(req, **opts) - - def aupdate( - self, req: process_dot_process__pb2.UpdateRequest, **opts - ) -> Coroutine[Any, Any, process_dot_process__pb2.UpdateResponse]: - return self._update.acall_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def stream_input( - self, req: process_dot_process__pb2.StreamInputRequest, **opts + self, + request: Iterator[process_dot_process__pb2.StreamInputRequest], + ctx: RequestContext, ) -> process_dot_process__pb2.StreamInputResponse: - return self._stream_input.call_client_stream(req, **opts) - - def astream_input( - self, req: process_dot_process__pb2.StreamInputRequest, **opts - ) -> Coroutine[Any, Any, process_dot_process__pb2.StreamInputResponse]: - return self._stream_input.acall_client_stream(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def send_input( - self, req: process_dot_process__pb2.SendInputRequest, **opts + self, request: process_dot_process__pb2.SendInputRequest, ctx: RequestContext ) -> process_dot_process__pb2.SendInputResponse: - return self._send_input.call_unary(req, **opts) - - def asend_input( - self, req: process_dot_process__pb2.SendInputRequest, **opts - ) -> Coroutine[Any, Any, process_dot_process__pb2.SendInputResponse]: - return self._send_input.acall_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def send_signal( - self, req: process_dot_process__pb2.SendSignalRequest, **opts + self, request: process_dot_process__pb2.SendSignalRequest, ctx: RequestContext ) -> process_dot_process__pb2.SendSignalResponse: - return self._send_signal.call_unary(req, **opts) - - def asend_signal( - self, req: process_dot_process__pb2.SendSignalRequest, **opts - ) -> Coroutine[Any, Any, process_dot_process__pb2.SendSignalResponse]: - return self._send_signal.acall_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") def close_stdin( - self, req: process_dot_process__pb2.CloseStdinRequest, **opts + self, request: process_dot_process__pb2.CloseStdinRequest, ctx: RequestContext ) -> process_dot_process__pb2.CloseStdinResponse: - return self._close_stdin.call_unary(req, **opts) + raise ConnectError(Code.UNIMPLEMENTED, "Not implemented") + + +class ProcessWSGIApplication(ConnectWSGIApplication): + def __init__( + self, + service: ProcessSync, + interceptors: Iterable[InterceptorSync] = (), + read_max_bytes: int | None = None, + compressions: Iterable[Compression] | None = None, + codecs: Iterable[Codec] | None = None, + ) -> None: + super().__init__( + endpoints={ + "/process.Process/List": EndpointSync.unary( + method=MethodInfo( + name="List", + service_name="process.Process", + input=process_dot_process__pb2.ListRequest, + output=process_dot_process__pb2.ListResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.list, + ), + "/process.Process/Connect": EndpointSync.server_stream( + method=MethodInfo( + name="Connect", + service_name="process.Process", + input=process_dot_process__pb2.ConnectRequest, + output=process_dot_process__pb2.ConnectResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.connect, + ), + "/process.Process/Start": EndpointSync.server_stream( + method=MethodInfo( + name="Start", + service_name="process.Process", + input=process_dot_process__pb2.StartRequest, + output=process_dot_process__pb2.StartResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.start, + ), + "/process.Process/Update": EndpointSync.unary( + method=MethodInfo( + name="Update", + service_name="process.Process", + input=process_dot_process__pb2.UpdateRequest, + output=process_dot_process__pb2.UpdateResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.update, + ), + "/process.Process/StreamInput": EndpointSync.client_stream( + method=MethodInfo( + name="StreamInput", + service_name="process.Process", + input=process_dot_process__pb2.StreamInputRequest, + output=process_dot_process__pb2.StreamInputResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.stream_input, + ), + "/process.Process/SendInput": EndpointSync.unary( + method=MethodInfo( + name="SendInput", + service_name="process.Process", + input=process_dot_process__pb2.SendInputRequest, + output=process_dot_process__pb2.SendInputResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.send_input, + ), + "/process.Process/SendSignal": EndpointSync.unary( + method=MethodInfo( + name="SendSignal", + service_name="process.Process", + input=process_dot_process__pb2.SendSignalRequest, + output=process_dot_process__pb2.SendSignalResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.send_signal, + ), + "/process.Process/CloseStdin": EndpointSync.unary( + method=MethodInfo( + name="CloseStdin", + service_name="process.Process", + input=process_dot_process__pb2.CloseStdinRequest, + output=process_dot_process__pb2.CloseStdinResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + function=service.close_stdin, + ), + }, + interceptors=interceptors, + read_max_bytes=read_max_bytes, + compressions=compressions, + codecs=codecs, + ) + + @property + def path(self) -> str: + """Returns the URL path to mount the application to when serving multiple applications.""" + return "/process.Process" + + +class ProcessClientSync(ConnectClientSync): + def list( + self, + request: process_dot_process__pb2.ListRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.ListResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="List", + service_name="process.Process", + input=process_dot_process__pb2.ListRequest, + output=process_dot_process__pb2.ListResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def connect( + self, + request: process_dot_process__pb2.ConnectRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> Iterator[process_dot_process__pb2.ConnectResponse]: + return self.execute_server_stream( + request=request, + method=MethodInfo( + name="Connect", + service_name="process.Process", + input=process_dot_process__pb2.ConnectRequest, + output=process_dot_process__pb2.ConnectResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def start( + self, + request: process_dot_process__pb2.StartRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> Iterator[process_dot_process__pb2.StartResponse]: + return self.execute_server_stream( + request=request, + method=MethodInfo( + name="Start", + service_name="process.Process", + input=process_dot_process__pb2.StartRequest, + output=process_dot_process__pb2.StartResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def update( + self, + request: process_dot_process__pb2.UpdateRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.UpdateResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="Update", + service_name="process.Process", + input=process_dot_process__pb2.UpdateRequest, + output=process_dot_process__pb2.UpdateResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def stream_input( + self, + request: Iterator[process_dot_process__pb2.StreamInputRequest], + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.StreamInputResponse: + return self.execute_client_stream( + request=request, + method=MethodInfo( + name="StreamInput", + service_name="process.Process", + input=process_dot_process__pb2.StreamInputRequest, + output=process_dot_process__pb2.StreamInputResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def send_input( + self, + request: process_dot_process__pb2.SendInputRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.SendInputResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="SendInput", + service_name="process.Process", + input=process_dot_process__pb2.SendInputRequest, + output=process_dot_process__pb2.SendInputResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) - def aclose_stdin( - self, req: process_dot_process__pb2.CloseStdinRequest, **opts - ) -> Coroutine[Any, Any, process_dot_process__pb2.CloseStdinResponse]: - return self._close_stdin.acall_unary(req, **opts) + def send_signal( + self, + request: process_dot_process__pb2.SendSignalRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.SendSignalResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="SendSignal", + service_name="process.Process", + input=process_dot_process__pb2.SendSignalRequest, + output=process_dot_process__pb2.SendSignalResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) + + def close_stdin( + self, + request: process_dot_process__pb2.CloseStdinRequest, + *, + headers: Headers | Mapping[str, str] | None = None, + timeout_ms: int | None = None, + ) -> process_dot_process__pb2.CloseStdinResponse: + return self.execute_unary( + request=request, + method=MethodInfo( + name="CloseStdin", + service_name="process.Process", + input=process_dot_process__pb2.CloseStdinRequest, + output=process_dot_process__pb2.CloseStdinResponse, + idempotency_level=IdempotencyLevel.UNKNOWN, + ), + headers=headers, + timeout_ms=timeout_ms, + ) diff --git a/packages/python-sdk/e2b/envd/pyqwest_httpx_adapter.py b/packages/python-sdk/e2b/envd/pyqwest_httpx_adapter.py new file mode 100644 index 0000000000..10e13944e2 --- /dev/null +++ b/packages/python-sdk/e2b/envd/pyqwest_httpx_adapter.py @@ -0,0 +1,347 @@ +from contextlib import asynccontextmanager, contextmanager +from typing import Any, AsyncIterable, AsyncIterator, Iterable, Iterator, Mapping + +import httpx +from connectrpc.code import Code +from connectrpc.errors import ConnectError +from pyqwest import FullResponse +from pyqwest import Headers as HTTPHeaders + +from e2b.envd.rpc import STREAM_REQUEST_TIMEOUT_HEADER + +_STREAM_REQUEST_TIMEOUT_HEADER = STREAM_REQUEST_TIMEOUT_HEADER.lower() +_REMOTE_PROTOCOL_RETRIES = 3 + + +def _headers(headers: httpx.Headers) -> HTTPHeaders: + return HTTPHeaders(headers.multi_items()) + + +def _prepare_headers(headers) -> tuple[Any, float | None]: + if headers is None: + return None, None + + if hasattr(headers, "items"): + items = list(headers.items()) + else: + items = list(headers) + + request_timeout = None + filtered_headers = [] + for name, value in items: + if name.lower() == _STREAM_REQUEST_TIMEOUT_HEADER: + request_timeout = float(value) + continue + + filtered_headers.append((name, value)) + + return filtered_headers, request_timeout + + +def _timeout(timeout: float | None, request_timeout: float | None) -> Any: + if request_timeout is None or request_timeout == 0: + return timeout + + return httpx.Timeout( + timeout=None, + connect=request_timeout, + read=timeout, + write=request_timeout, + pool=request_timeout, + ) + + +def _stream_timeout(timeout: float | None, request_timeout: float | None) -> Any: + if request_timeout is None or request_timeout == 0: + return httpx.Timeout( + timeout=None, + connect=None, + read=timeout, + write=None, + pool=None, + ) + + return _timeout(timeout, request_timeout) + + +def _request_timeout_error(e: httpx.TimeoutException) -> ConnectError: + return ConnectError(Code.CANCELED, str(e) or "Request timed out") + + +def _stream_timeout_error(e: httpx.TimeoutException) -> ConnectError: + if isinstance(e, httpx.ReadTimeout): + return ConnectError(Code.DEADLINE_EXCEEDED, str(e) or "Stream timed out") + + return _request_timeout_error(e) + + +def _retry_remote_protocol(call): + for _ in range(_REMOTE_PROTOCOL_RETRIES): + try: + return call() + except httpx.RemoteProtocolError: + pass + + return call() + + +async def _aretry_remote_protocol(call): + for _ in range(_REMOTE_PROTOCOL_RETRIES): + try: + return await call() + except httpx.RemoteProtocolError: + pass + + return await call() + + +def _sync_content(content): + if isinstance(content, Iterable) and not isinstance( + content, (bytes, bytearray, str) + ): + return list(content) + + return content + + +async def _async_content(content): + if isinstance(content, AsyncIterable): + return b"".join([chunk async for chunk in content]) + + if isinstance(content, Iterable) and not isinstance( + content, (bytes, bytearray, str) + ): + return b"".join(content) + + return content + + +class _SyncStreamResponse: + def __init__(self, response: httpx.Response) -> None: + self.status = response.status_code + self.headers = _headers(response.headers) + self.trailers = HTTPHeaders() + self.content = _iter_stream_bytes(response) + + +class _AsyncStreamResponse: + def __init__(self, response: httpx.Response) -> None: + self.status = response.status_code + self.headers = _headers(response.headers) + self.trailers = HTTPHeaders() + self.content = _aiter_stream_bytes(response) + + +def _iter_stream_bytes(response: httpx.Response): + try: + yield from response.iter_bytes() + except httpx.TimeoutException as e: + raise _stream_timeout_error(e) from e + + +async def _aiter_stream_bytes(response: httpx.Response): + try: + async for chunk in response.aiter_bytes(): + yield chunk + except httpx.TimeoutException as e: + raise _stream_timeout_error(e) from e + + +@contextmanager +def _open_stream(open_stream): + try: + with open_stream() as response: + yield response + except httpx.TimeoutException as e: + raise _request_timeout_error(e) from e + + +@asynccontextmanager +async def _aopen_stream(open_stream): + try: + async with open_stream() as response: + yield response + except httpx.TimeoutException as e: + raise _request_timeout_error(e) from e + + +class PyqwestHTTPXAdapter: + def __init__(self, transport: httpx.BaseTransport) -> None: + self._client = httpx.Client(transport=transport, timeout=None) + + def get( + self, + url: str, + headers: Any = None, + *, + timeout: float | None = None, + params: Mapping[str, str] | None = None, + ) -> FullResponse: + headers, request_timeout = _prepare_headers(headers) + try: + response = _retry_remote_protocol( + lambda: self._client.get( + url, + headers=headers, + timeout=_timeout(timeout, request_timeout), + params=params, + ) + ) + except httpx.TimeoutException as e: + raise _request_timeout_error(e) from e + + return FullResponse( + response.status_code, + _headers(response.headers), + response.content, + HTTPHeaders(), + ) + + def post( + self, + url: str, + headers: Any = None, + content=None, + *, + timeout: float | None = None, + params: Mapping[str, str] | None = None, + ) -> FullResponse: + headers, request_timeout = _prepare_headers(headers) + content = _sync_content(content) + try: + response = _retry_remote_protocol( + lambda: self._client.post( + url, + headers=headers, + content=content, + timeout=_timeout(timeout, request_timeout), + params=params, + ) + ) + except httpx.TimeoutException as e: + raise _request_timeout_error(e) from e + + return FullResponse( + response.status_code, + _headers(response.headers), + response.content, + HTTPHeaders(), + ) + + @contextmanager + def stream( + self, + method: str, + url: str, + headers: Any = None, + content=None, + *, + timeout: float | None = None, + params: Mapping[str, str] | None = None, + ) -> Iterator[_SyncStreamResponse]: + headers, request_timeout = _prepare_headers(headers) + content = _sync_content(content) + + def open_stream(): + return self._client.stream( + method, + url, + headers=headers, + content=content, + timeout=_stream_timeout(timeout, request_timeout), + params=params, + ) + + with _open_stream(open_stream) as response: + yield _SyncStreamResponse(response) + + +class AsyncPyqwestHTTPXAdapter: + def __init__(self, transport: httpx.AsyncBaseTransport) -> None: + self._client = httpx.AsyncClient(transport=transport, timeout=None) + + async def get( + self, + url: str, + headers: Any = None, + *, + timeout: float | None = None, + params: Mapping[str, str] | None = None, + ) -> FullResponse: + headers, request_timeout = _prepare_headers(headers) + try: + response = await _aretry_remote_protocol( + lambda: self._client.get( + url, + headers=headers, + timeout=_timeout(timeout, request_timeout), + params=params, + ) + ) + except httpx.TimeoutException as e: + raise _request_timeout_error(e) from e + + return FullResponse( + response.status_code, + _headers(response.headers), + response.content, + HTTPHeaders(), + ) + + async def post( + self, + url: str, + headers: Any = None, + content=None, + *, + timeout: float | None = None, + params: Mapping[str, str] | None = None, + ) -> FullResponse: + headers, request_timeout = _prepare_headers(headers) + content = await _async_content(content) + try: + response = await _aretry_remote_protocol( + lambda: self._client.post( + url, + headers=headers, + content=content, + timeout=_timeout(timeout, request_timeout), + params=params, + ) + ) + except httpx.TimeoutException as e: + raise _request_timeout_error(e) from e + + return FullResponse( + response.status_code, + _headers(response.headers), + response.content, + HTTPHeaders(), + ) + + @asynccontextmanager + async def stream( + self, + method: str, + url: str, + headers: Any = None, + content=None, + *, + timeout: float | None = None, + params: Mapping[str, str] | None = None, + ) -> AsyncIterator[_AsyncStreamResponse]: + headers, request_timeout = _prepare_headers(headers) + content = await _async_content(content) + + def open_stream(): + return self._client.stream( + method, + url, + headers=headers, + content=content, + timeout=_stream_timeout(timeout, request_timeout), + params=params, + ) + + async with _aopen_stream(open_stream) as response: + yield _AsyncStreamResponse(response) diff --git a/packages/python-sdk/e2b/envd/rpc.py b/packages/python-sdk/e2b/envd/rpc.py index f8d263c5e8..2bc1f90121 100644 --- a/packages/python-sdk/e2b/envd/rpc.py +++ b/packages/python-sdk/e2b/envd/rpc.py @@ -2,7 +2,11 @@ from typing import Callable, Optional from packaging.version import Version -from e2b_connect.client import Code, ConnectException +from connectrpc.code import Code +from google.protobuf.json_format import MessageToJson, Parse +from google.protobuf.message import Message +from connectrpc.errors import ConnectError +from connectrpc.request import RequestContext from e2b.exceptions import ( SandboxException, @@ -16,18 +20,20 @@ from e2b.connection_config import Username, default_username from e2b.envd.versions import ENVD_DEFAULT_USER +STREAM_REQUEST_TIMEOUT_HEADER = "E2B-Stream-Request-Timeout" + _DEFAULT_RPC_ERROR_MAP: dict[Code, Callable[[str], Exception]] = { - Code.invalid_argument: InvalidArgumentException, - Code.unauthenticated: AuthenticationException, - Code.not_found: NotFoundException, - Code.unavailable: format_sandbox_timeout_exception, - Code.resource_exhausted: lambda message: RateLimitException( + Code.INVALID_ARGUMENT: InvalidArgumentException, + Code.UNAUTHENTICATED: AuthenticationException, + Code.NOT_FOUND: NotFoundException, + Code.UNAVAILABLE: format_sandbox_timeout_exception, + Code.RESOURCE_EXHAUSTED: lambda message: RateLimitException( f"{message}: Rate limit exceeded, please try again later." ), - Code.canceled: lambda message: TimeoutException( + Code.CANCELED: lambda message: TimeoutException( f"{message}: This error is likely due to exceeding 'request_timeout'. You can pass the request timeout value as an option when making the request." ), - Code.deadline_exceeded: lambda message: TimeoutException( + Code.DEADLINE_EXCEEDED: lambda message: TimeoutException( f"{message}: This error is likely due to exceeding 'timeout' — the total time a long running request (like process or directory watch) can be active. It can be modified by passing 'timeout' when making the request. Use '0' to disable the timeout." ), } @@ -39,22 +45,100 @@ def handle_rpc_exception( ): """Handle errors from envd RPC calls by mapping gRPC status codes to specific exception types. - :param e: The caught exception, expected to be a ``ConnectException``. + :param e: The caught exception, expected to be a ``ConnectError``. :param error_map: Optional map of gRPC codes to exception factories that override the defaults. - :return: The corresponding exception, or the original exception if not a ``ConnectException``. + :return: The corresponding exception, or the original exception if not a ``ConnectError``. """ - if isinstance(e, ConnectException): - if error_map and e.status in error_map: - return error_map[e.status](e.message) + if isinstance(e, ConnectError): + if error_map and e.code in error_map: + return error_map[e.code](e.message) - if e.status in _DEFAULT_RPC_ERROR_MAP: - return _DEFAULT_RPC_ERROR_MAP[e.status](e.message) + if e.code in _DEFAULT_RPC_ERROR_MAP: + return _DEFAULT_RPC_ERROR_MAP[e.code](e.message) - return SandboxException(f"{e.status}: {e.message}") + return SandboxException(f"{e.code}: {e.message}") else: return e +class ProtoJSONCodec: + def name(self) -> str: + return "json" + + def encode(self, message: Message) -> bytes: + return MessageToJson(message).encode() + + def decode(self, data: bytes | bytearray, message: Message): + Parse(data.decode(), message, ignore_unknown_fields=True) + return message + + +def request_timeout_ms(timeout: Optional[float]) -> Optional[int]: + if timeout is None: + return None + + return int(timeout * 1000) + + +def stream_timeout_ms( + timeout: Optional[float], +) -> Optional[int]: + if timeout == 0: + return None + + return request_timeout_ms(timeout) + + +def stream_request_headers( + headers: dict[str, str], + request_timeout: Optional[float], +) -> dict[str, str]: + if request_timeout is None or request_timeout == 0: + return headers + + return { + **headers, + STREAM_REQUEST_TIMEOUT_HEADER: str(request_timeout), + } + + +class SandboxHeadersInterceptor: + def __init__(self, headers: dict[str, str]) -> None: + self._headers = headers + + def _add_headers(self, ctx: RequestContext) -> None: + request_headers = ctx.request_headers() + for key, value in self._headers.items(): + if key not in request_headers: + request_headers[key] = value + + def on_start_sync(self, ctx: RequestContext) -> None: + self._add_headers(ctx) + + async def on_start(self, ctx: RequestContext) -> None: + self._add_headers(ctx) + + def on_end_sync( + self, token: None, ctx: RequestContext, error: Exception | None + ) -> None: + return None + + async def on_end( + self, token: None, ctx: RequestContext, error: Exception | None + ) -> None: + return None + + +def connect_client_kwargs(headers: dict[str, str], http_client): + return { + "codec": ProtoJSONCodec(), + "accept_compression": (), + "send_compression": None, + "interceptors": (SandboxHeadersInterceptor(headers),), + "http_client": http_client, + } + + def authentication_header( envd_version: Version, user: Optional[Username] = None ) -> dict[str, str]: diff --git a/packages/python-sdk/e2b/sandbox_async/commands/command.py b/packages/python-sdk/e2b/sandbox_async/commands/command.py index 32b75fd26b..c9deee4b02 100644 --- a/packages/python-sdk/e2b/sandbox_async/commands/command.py +++ b/packages/python-sdk/e2b/sandbox_async/commands/command.py @@ -1,7 +1,7 @@ from typing import Dict, List, Literal, Optional, Union, overload -import e2b_connect -import httpcore +from connectrpc.code import Code +from connectrpc.errors import ConnectError from packaging.version import Version from e2b.connection_config import ( ConnectionConfig, @@ -9,8 +9,16 @@ KEEPALIVE_PING_HEADER, KEEPALIVE_PING_INTERVAL_SEC, ) +from e2b.envd.pyqwest_httpx_adapter import AsyncPyqwestHTTPXAdapter from e2b.envd.process import process_connect, process_pb2 -from e2b.envd.rpc import authentication_header, handle_rpc_exception +from e2b.envd.rpc import ( + authentication_header, + connect_client_kwargs, + handle_rpc_exception, + request_timeout_ms, + stream_request_headers, + stream_timeout_ms, +) from e2b.envd.versions import ENVD_COMMANDS_STDIN from e2b.exceptions import SandboxException from e2b.sandbox.commands.main import ProcessInfo @@ -28,18 +36,14 @@ def __init__( self, envd_api_url: str, connection_config: ConnectionConfig, - pool: httpcore.AsyncConnectionPool, + rpc_client: AsyncPyqwestHTTPXAdapter, envd_version: Version, ) -> None: self._connection_config = connection_config self._envd_version = envd_version self._rpc = process_connect.ProcessClient( envd_api_url, - # TODO: Fix and enable compression again — the headers compression is not solved for streaming. - # compressor=e2b_connect.GzipCompressor, - async_pool=pool, - json=True, - headers=connection_config.sandbox_headers, + **connect_client_kwargs(connection_config.sandbox_headers, rpc_client), ) async def list( @@ -54,10 +58,10 @@ async def list( :return: List of running commands and PTY sessions """ try: - res = await self._rpc.alist( + res = await self._rpc.list( process_pb2.ListRequest(), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) return [ @@ -89,19 +93,19 @@ async def kill( :return: `True` if the command was killed, `False` if the command was not found """ try: - await self._rpc.asend_signal( + await self._rpc.send_signal( process_pb2.SendSignalRequest( process=process_pb2.ProcessSelector(pid=pid), signal=process_pb2.Signal.SIGNAL_SIGKILL, ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) return True except Exception as e: - if isinstance(e, e2b_connect.ConnectException): - if e.status == e2b_connect.Code.not_found: + if isinstance(e, ConnectError): + if e.code == Code.NOT_FOUND: return False raise handle_rpc_exception(e) @@ -119,15 +123,15 @@ async def send_stdin( :param request_timeout: Timeout for the request in **seconds** """ try: - await self._rpc.asend_input( + await self._rpc.send_input( process_pb2.SendInputRequest( process=process_pb2.ProcessSelector(pid=pid), input=process_pb2.ProcessInput( stdin=data.encode(), ), ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) except Exception as e: @@ -246,7 +250,7 @@ async def _start( on_stdout: Optional[OutputHandler[Stdout]], on_stderr: Optional[OutputHandler[Stderr]], ) -> AsyncCommandHandle: - events = self._rpc.astart( + events = self._rpc.start( process_pb2.StartRequest( process=process_pb2.ProcessConfig( cmd="/bin/bash", @@ -256,14 +260,14 @@ async def _start( ), stdin=stdin, ), - headers={ - **authentication_header(self._envd_version, user), - KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), - }, - timeout=timeout, - request_timeout=self._connection_config.get_request_timeout( - request_timeout + headers=stream_request_headers( + { + **authentication_header(self._envd_version, user), + KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), + }, + self._connection_config.get_request_timeout(request_timeout), ), + timeout_ms=stream_timeout_ms(timeout), ) try: @@ -304,17 +308,17 @@ async def connect( :return: `AsyncCommandHandle` handle to interact with the running command """ - events = self._rpc.aconnect( + events = self._rpc.connect( process_pb2.ConnectRequest( process=process_pb2.ProcessSelector(pid=pid), ), - timeout=timeout, - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=stream_timeout_ms(timeout), + headers=stream_request_headers( + { + KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), + }, + self._connection_config.get_request_timeout(request_timeout), ), - headers={ - KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), - }, ) try: diff --git a/packages/python-sdk/e2b/sandbox_async/commands/command_handle.py b/packages/python-sdk/e2b/sandbox_async/commands/command_handle.py index dfe34e8383..304b5772c3 100644 --- a/packages/python-sdk/e2b/sandbox_async/commands/command_handle.py +++ b/packages/python-sdk/e2b/sandbox_async/commands/command_handle.py @@ -5,6 +5,7 @@ Callable, Any, AsyncGenerator, + AsyncIterator, Union, Tuple, Coroutine, @@ -76,8 +77,8 @@ def __init__( self, pid: int, handle_kill: Callable[[], Coroutine[Any, Any, bool]], - events: AsyncGenerator[ - Union[process_pb2.StartResponse, process_pb2.ConnectResponse], Any + events: AsyncIterator[ + Union[process_pb2.StartResponse, process_pb2.ConnectResponse] ], on_stdout: Optional[OutputHandler[Stdout]] = None, on_stderr: Optional[OutputHandler[Stderr]] = None, diff --git a/packages/python-sdk/e2b/sandbox_async/commands/pty.py b/packages/python-sdk/e2b/sandbox_async/commands/pty.py index 3585b13246..94bc45588c 100644 --- a/packages/python-sdk/e2b/sandbox_async/commands/pty.py +++ b/packages/python-sdk/e2b/sandbox_async/commands/pty.py @@ -1,8 +1,7 @@ from typing import Dict, Optional -import e2b_connect -import httpcore - +from connectrpc.code import Code +from connectrpc.errors import ConnectError from packaging.version import Version from e2b.envd.process import process_connect, process_pb2 from e2b.connection_config import ( @@ -12,7 +11,15 @@ KEEPALIVE_PING_INTERVAL_SEC, ) from e2b.exceptions import SandboxException -from e2b.envd.rpc import authentication_header, handle_rpc_exception +from e2b.envd.pyqwest_httpx_adapter import AsyncPyqwestHTTPXAdapter +from e2b.envd.rpc import ( + authentication_header, + connect_client_kwargs, + handle_rpc_exception, + request_timeout_ms, + stream_request_headers, + stream_timeout_ms, +) from e2b.sandbox.commands.command_handle import PtySize from e2b.sandbox_async.commands.command_handle import ( AsyncCommandHandle, @@ -30,18 +37,14 @@ def __init__( self, envd_api_url: str, connection_config: ConnectionConfig, - pool: httpcore.AsyncConnectionPool, + rpc_client: AsyncPyqwestHTTPXAdapter, envd_version: Version, ) -> None: self._connection_config = connection_config self._envd_version = envd_version self._rpc = process_connect.ProcessClient( envd_api_url, - # TODO: Fix and enable compression again — the headers compression is not solved for streaming. - # compressor=e2b_connect.GzipCompressor, - async_pool=pool, - json=True, - headers=connection_config.sandbox_headers, + **connect_client_kwargs(connection_config.sandbox_headers, rpc_client), ) async def kill( @@ -58,19 +61,19 @@ async def kill( :return: `true` if the PTY was killed, `false` if the PTY was not found """ try: - await self._rpc.asend_signal( + await self._rpc.send_signal( process_pb2.SendSignalRequest( process=process_pb2.ProcessSelector(pid=pid), signal=process_pb2.Signal.SIGNAL_SIGKILL, ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) return True except Exception as e: - if isinstance(e, e2b_connect.ConnectException): - if e.status == e2b_connect.Code.not_found: + if isinstance(e, ConnectError): + if e.code == Code.NOT_FOUND: return False raise handle_rpc_exception(e) @@ -88,15 +91,15 @@ async def send_stdin( :param request_timeout: Timeout for the request in **seconds** """ try: - await self._rpc.asend_input( + await self._rpc.send_input( process_pb2.SendInputRequest( process=process_pb2.ProcessSelector(pid=pid), input=process_pb2.ProcessInput( pty=data, ), ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) except Exception as e: @@ -129,7 +132,7 @@ async def create( envs.setdefault("TERM", "xterm-256color") envs.setdefault("LANG", "C.UTF-8") envs.setdefault("LC_ALL", "C.UTF-8") - events = self._rpc.astart( + events = self._rpc.start( process_pb2.StartRequest( process=process_pb2.ProcessConfig( cmd="/bin/bash", @@ -141,14 +144,14 @@ async def create( size=process_pb2.PTY.Size(rows=size.rows, cols=size.cols) ), ), - headers={ - **authentication_header(self._envd_version, user), - KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), - }, - timeout=timeout, - request_timeout=self._connection_config.get_request_timeout( - request_timeout + headers=stream_request_headers( + { + **authentication_header(self._envd_version, user), + KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), + }, + self._connection_config.get_request_timeout(request_timeout), ), + timeout_ms=stream_timeout_ms(timeout), ) try: @@ -185,17 +188,17 @@ async def connect( :return: Handle to interact with the PTY """ - events = self._rpc.aconnect( + events = self._rpc.connect( process_pb2.ConnectRequest( process=process_pb2.ProcessSelector(pid=pid), ), - timeout=timeout, - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=stream_timeout_ms(timeout), + headers=stream_request_headers( + { + KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), + }, + self._connection_config.get_request_timeout(request_timeout), ), - headers={ - KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), - }, ) try: @@ -229,14 +232,14 @@ async def resize( :param size: New size of the PTY :param request_timeout: Timeout for the request in **seconds** """ - await self._rpc.aupdate( + await self._rpc.update( process_pb2.UpdateRequest( process=process_pb2.ProcessSelector(pid=pid), pty=process_pb2.PTY( size=process_pb2.PTY.Size(rows=size.rows, cols=size.cols), ), ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) diff --git a/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py b/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py index 309f4f6169..6f77910e6f 100644 --- a/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py +++ b/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py @@ -3,11 +3,11 @@ from typing import IO, AsyncIterator, List, Literal, Optional, Union, overload -import httpcore import httpx +from connectrpc.code import Code +from connectrpc.errors import ConnectError from packaging.version import Version -import e2b_connect as connect from e2b.connection_config import ( KEEPALIVE_PING_HEADER, KEEPALIVE_PING_INTERVAL_SEC, @@ -17,7 +17,15 @@ ) from e2b.envd.api import ENVD_API_FILES_ROUTE, ahandle_envd_api_exception from e2b.envd.filesystem import filesystem_connect, filesystem_pb2 -from e2b.envd.rpc import authentication_header, handle_rpc_exception +from e2b.envd.pyqwest_httpx_adapter import AsyncPyqwestHTTPXAdapter +from e2b.envd.rpc import ( + authentication_header, + connect_client_kwargs, + handle_rpc_exception, + request_timeout_ms, + stream_request_headers, + stream_timeout_ms, +) from e2b.envd.versions import ( ENVD_DEFAULT_USER, ENVD_OCTET_STREAM_UPLOAD, @@ -39,18 +47,30 @@ from e2b.sandbox.filesystem.watch_handle import FilesystemEvent from e2b.sandbox_async.filesystem.watch_handle import AsyncWatchHandle from e2b.sandbox_async.utils import OutputHandler -from e2b_connect.client import Code _FILESYSTEM_RPC_ERROR_MAP = { - Code.not_found: FileNotFoundException, + Code.NOT_FOUND: FileNotFoundException, } _FILESYSTEM_HTTP_ERROR_MAP = { 404: FileNotFoundException, } +_ENOENT_MESSAGE = "no such file or directory" + def _handle_filesystem_rpc_exception(e: Exception) -> Exception: + if isinstance(e, ConnectError): + if e.code == Code.NOT_FOUND: + return FileNotFoundException(e.message) + + if ( + e.code == Code.UNKNOWN + # Older envd versions returned ENOENT as UNKNOWN instead of NOT_FOUND. + and _ENOENT_MESSAGE in e.message.lower() + ): + return FileNotFoundException(e.message) + return handle_rpc_exception(e, _FILESYSTEM_RPC_ERROR_MAP) @@ -68,22 +88,17 @@ def __init__( envd_api_url: str, envd_version: Version, connection_config: ConnectionConfig, - pool: httpcore.AsyncConnectionPool, + rpc_client: AsyncPyqwestHTTPXAdapter, envd_api: httpx.AsyncClient, ) -> None: self._envd_api_url = envd_api_url self._envd_version = envd_version self._connection_config = connection_config - self._pool = pool self._envd_api = envd_api self._rpc = filesystem_connect.FilesystemClient( envd_api_url, - # TODO: Fix and enable compression again — the headers compression is not solved for streaming. - # compressor=e2b_connect.GzipCompressor, - async_pool=pool, - json=True, - headers=connection_config.sandbox_headers, + **connect_client_kwargs(connection_config.sandbox_headers, rpc_client), ) @overload @@ -372,10 +387,10 @@ async def list( raise InvalidArgumentException("depth should be at least 1") try: - res = await self._rpc.alist_dir( + res = await self._rpc.list_dir( filesystem_pb2.ListDirRequest(path=path, depth=depth), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) @@ -425,10 +440,10 @@ async def exists( :return: `True` if the file or directory exists, `False` otherwise """ try: - await self._rpc.astat( + await self._rpc.stat( filesystem_pb2.StatRequest(path=path), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) @@ -436,10 +451,10 @@ async def exists( return True except Exception as e: - if isinstance(e, connect.ConnectException): - if e.status == connect.Code.not_found: - return False - raise _handle_filesystem_rpc_exception(e) + err = _handle_filesystem_rpc_exception(e) + if isinstance(err, FileNotFoundException): + return False + raise err async def get_info( self, @@ -457,10 +472,10 @@ async def get_info( :return: Information about the file or directory like name, type, and path """ try: - r = await self._rpc.astat( + r = await self._rpc.stat( filesystem_pb2.StatRequest(path=path), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) @@ -498,10 +513,10 @@ async def remove( :param request_timeout: Timeout for the request in **seconds** """ try: - await self._rpc.aremove( + await self._rpc.remove( filesystem_pb2.RemoveRequest(path=path), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) @@ -526,13 +541,13 @@ async def rename( :return: Information about the renamed file or directory """ try: - r = await self._rpc.amove( + r = await self._rpc.move( filesystem_pb2.MoveRequest( source=old_path, destination=new_path, ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) @@ -573,18 +588,18 @@ async def make_dir( :return: `True` if the directory was created, `False` if the directory already exists """ try: - await self._rpc.amake_dir( + await self._rpc.make_dir( filesystem_pb2.MakeDirRequest(path=path), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) return True except Exception as e: - if isinstance(e, connect.ConnectException): - if e.status == connect.Code.already_exists: + if isinstance(e, ConnectError): + if e.code == Code.ALREADY_EXISTS: return False raise _handle_filesystem_rpc_exception(e) @@ -617,16 +632,16 @@ async def watch_dir( "You can do this by running `e2b template build` in the directory with the template." ) - events = self._rpc.awatch_dir( + events = self._rpc.watch_dir( filesystem_pb2.WatchDirRequest(path=path, recursive=recursive), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=stream_timeout_ms(timeout), + headers=stream_request_headers( + { + **authentication_header(self._envd_version, user), + KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), + }, + self._connection_config.get_request_timeout(request_timeout), ), - timeout=timeout, - headers={ - **authentication_header(self._envd_version, user), - KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), - }, ) try: diff --git a/packages/python-sdk/e2b/sandbox_async/filesystem/watch_handle.py b/packages/python-sdk/e2b/sandbox_async/filesystem/watch_handle.py index 33759d88ce..8f02e645e0 100644 --- a/packages/python-sdk/e2b/sandbox_async/filesystem/watch_handle.py +++ b/packages/python-sdk/e2b/sandbox_async/filesystem/watch_handle.py @@ -1,7 +1,7 @@ import asyncio import inspect -from typing import Any, AsyncGenerator, Optional +from typing import AsyncIterator, Optional from e2b.envd.rpc import handle_rpc_exception from e2b.envd.filesystem.filesystem_pb2 import WatchDirResponse @@ -18,7 +18,7 @@ class AsyncWatchHandle: def __init__( self, - events: AsyncGenerator[WatchDirResponse, Any], + events: AsyncIterator[WatchDirResponse], on_event: OutputHandler[FilesystemEvent], on_exit: Optional[OutputHandler[Exception]] = None, ): diff --git a/packages/python-sdk/e2b/sandbox_async/main.py b/packages/python-sdk/e2b/sandbox_async/main.py index 3dd7044ea7..bd9716dcf8 100644 --- a/packages/python-sdk/e2b/sandbox_async/main.py +++ b/packages/python-sdk/e2b/sandbox_async/main.py @@ -13,6 +13,7 @@ from e2b.api.client_async import get_transport from e2b.connection_config import ApiParams, ConnectionConfig from e2b.envd.api import ENVD_API_HEALTH_ROUTE, ahandle_envd_api_exception +from e2b.envd.pyqwest_httpx_adapter import AsyncPyqwestHTTPXAdapter from e2b.envd.versions import ENVD_DEBUG_FALLBACK from e2b.exceptions import ( SandboxException, @@ -102,6 +103,7 @@ def __init__( super().__init__(**opts) self._transport = get_transport(self.connection_config) + self._rpc_client = AsyncPyqwestHTTPXAdapter(self._transport) self._envd_api = httpx.AsyncClient( base_url=self.connection_config.get_sandbox_url( self.sandbox_id, self.sandbox_domain @@ -113,19 +115,19 @@ def __init__( self.envd_api_url, self._envd_version, self.connection_config, - self._transport.pool, + self._rpc_client, self._envd_api, ) self._commands = Commands( self.envd_api_url, self.connection_config, - self._transport.pool, + self._rpc_client, self._envd_version, ) self._pty = Pty( self.envd_api_url, self.connection_config, - self._transport.pool, + self._rpc_client, self._envd_version, ) self._git = Git(self._commands) diff --git a/packages/python-sdk/e2b/sandbox_sync/commands/command.py b/packages/python-sdk/e2b/sandbox_sync/commands/command.py index 512b7d9923..e5756534ef 100644 --- a/packages/python-sdk/e2b/sandbox_sync/commands/command.py +++ b/packages/python-sdk/e2b/sandbox_sync/commands/command.py @@ -1,7 +1,7 @@ from typing import Callable, Dict, List, Literal, Optional, Union, overload -import e2b_connect -import httpcore +from connectrpc.code import Code +from connectrpc.errors import ConnectError from packaging.version import Version from e2b.connection_config import ( ConnectionConfig, @@ -9,8 +9,16 @@ KEEPALIVE_PING_HEADER, KEEPALIVE_PING_INTERVAL_SEC, ) +from e2b.envd.pyqwest_httpx_adapter import PyqwestHTTPXAdapter from e2b.envd.process import process_connect, process_pb2 -from e2b.envd.rpc import authentication_header, handle_rpc_exception +from e2b.envd.rpc import ( + authentication_header, + connect_client_kwargs, + handle_rpc_exception, + request_timeout_ms, + stream_request_headers, + stream_timeout_ms, +) from e2b.envd.versions import ENVD_COMMANDS_STDIN from e2b.exceptions import SandboxException from e2b.sandbox.commands.main import ProcessInfo @@ -27,18 +35,14 @@ def __init__( self, envd_api_url: str, connection_config: ConnectionConfig, - pool: httpcore.ConnectionPool, + rpc_client: PyqwestHTTPXAdapter, envd_version: Version, ) -> None: self._connection_config = connection_config self._envd_version = envd_version - self._rpc = process_connect.ProcessClient( + self._rpc = process_connect.ProcessClientSync( envd_api_url, - # TODO: Fix and enable compression again — the headers compression is not solved for streaming. - # compressor=e2b_connect.GzipCompressor, - pool=pool, - json=True, - headers=connection_config.sandbox_headers, + **connect_client_kwargs(connection_config.sandbox_headers, rpc_client), ) def list( @@ -55,8 +59,8 @@ def list( try: res = self._rpc.list( process_pb2.ListRequest(), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) return [ @@ -93,14 +97,14 @@ def kill( process=process_pb2.ProcessSelector(pid=pid), signal=process_pb2.Signal.SIGNAL_SIGKILL, ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) return True except Exception as e: - if isinstance(e, e2b_connect.ConnectException): - if e.status == e2b_connect.Code.not_found: + if isinstance(e, ConnectError): + if e.code == Code.NOT_FOUND: return False raise handle_rpc_exception(e) @@ -125,8 +129,8 @@ def send_stdin( stdin=data.encode(), ), ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) except Exception as e: @@ -256,14 +260,14 @@ def _start( ), stdin=stdin, ), - headers={ - **authentication_header(self._envd_version, user), - KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), - }, - timeout=timeout, - request_timeout=self._connection_config.get_request_timeout( - request_timeout + headers=stream_request_headers( + { + **authentication_header(self._envd_version, user), + KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), + }, + self._connection_config.get_request_timeout(request_timeout), ), + timeout_ms=stream_timeout_ms(timeout), ) try: @@ -302,13 +306,13 @@ def connect( process_pb2.ConnectRequest( process=process_pb2.ProcessSelector(pid=pid), ), - headers={ - KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), - }, - timeout=timeout, - request_timeout=self._connection_config.get_request_timeout( - request_timeout + headers=stream_request_headers( + { + KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), + }, + self._connection_config.get_request_timeout(request_timeout), ), + timeout_ms=stream_timeout_ms(timeout), ) try: diff --git a/packages/python-sdk/e2b/sandbox_sync/commands/command_handle.py b/packages/python-sdk/e2b/sandbox_sync/commands/command_handle.py index a58a613e02..c8351d7725 100644 --- a/packages/python-sdk/e2b/sandbox_sync/commands/command_handle.py +++ b/packages/python-sdk/e2b/sandbox_sync/commands/command_handle.py @@ -1,4 +1,4 @@ -from typing import Optional, Callable, Any, Generator, Union, Tuple +from typing import Optional, Callable, Generator, Iterator, Union, Tuple from e2b.envd.rpc import handle_rpc_exception from e2b.envd.process import process_pb2 @@ -29,9 +29,7 @@ def __init__( self, pid: int, handle_kill: Callable[[], bool], - events: Generator[ - Union[process_pb2.StartResponse, process_pb2.ConnectResponse], Any, None - ], + events: Iterator[Union[process_pb2.StartResponse, process_pb2.ConnectResponse]], ): self._pid = pid self._handle_kill = handle_kill @@ -92,7 +90,9 @@ def disconnect(self) -> None: The command is not killed, but SDK stops receiving events from the command. You can reconnect to the command using `sandbox.commands.connect` method. """ - self._events.close() + close = getattr(self._events, "close", None) + if close: + close() def wait( self, diff --git a/packages/python-sdk/e2b/sandbox_sync/commands/pty.py b/packages/python-sdk/e2b/sandbox_sync/commands/pty.py index fd936ef404..eb682ca5ae 100644 --- a/packages/python-sdk/e2b/sandbox_sync/commands/pty.py +++ b/packages/python-sdk/e2b/sandbox_sync/commands/pty.py @@ -1,8 +1,7 @@ -import e2b_connect -import httpcore - from typing import Dict, Optional +from connectrpc.code import Code +from connectrpc.errors import ConnectError from packaging.version import Version from e2b.envd.process import process_connect, process_pb2 from e2b.connection_config import ( @@ -12,7 +11,15 @@ KEEPALIVE_PING_INTERVAL_SEC, ) from e2b.exceptions import SandboxException -from e2b.envd.rpc import authentication_header, handle_rpc_exception +from e2b.envd.pyqwest_httpx_adapter import PyqwestHTTPXAdapter +from e2b.envd.rpc import ( + authentication_header, + connect_client_kwargs, + handle_rpc_exception, + request_timeout_ms, + stream_request_headers, + stream_timeout_ms, +) from e2b.sandbox.commands.command_handle import PtySize from e2b.sandbox_sync.commands.command_handle import CommandHandle @@ -26,18 +33,14 @@ def __init__( self, envd_api_url: str, connection_config: ConnectionConfig, - pool: httpcore.ConnectionPool, + rpc_client: PyqwestHTTPXAdapter, envd_version: Version, ) -> None: self._connection_config = connection_config self._envd_version = envd_version - self._rpc = process_connect.ProcessClient( + self._rpc = process_connect.ProcessClientSync( envd_api_url, - # TODO: Fix and enable compression again — the headers compression is not solved for streaming. - # compressor=e2b_connect.GzipCompressor, - pool=pool, - json=True, - headers=connection_config.sandbox_headers, + **connect_client_kwargs(connection_config.sandbox_headers, rpc_client), ) def kill( @@ -59,14 +62,14 @@ def kill( process=process_pb2.ProcessSelector(pid=pid), signal=process_pb2.Signal.SIGNAL_SIGKILL, ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) return True except Exception as e: - if isinstance(e, e2b_connect.ConnectException): - if e.status == e2b_connect.Code.not_found: + if isinstance(e, ConnectError): + if e.code == Code.NOT_FOUND: return False raise handle_rpc_exception(e) @@ -91,8 +94,8 @@ def send_stdin( pty=data, ), ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) except Exception as e: @@ -135,14 +138,14 @@ def create( size=process_pb2.PTY.Size(rows=size.rows, cols=size.cols) ), ), - headers={ - **authentication_header(self._envd_version, user), - KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), - }, - timeout=timeout, - request_timeout=self._connection_config.get_request_timeout( - request_timeout + headers=stream_request_headers( + { + **authentication_header(self._envd_version, user), + KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), + }, + self._connection_config.get_request_timeout(request_timeout), ), + timeout_ms=stream_timeout_ms(timeout), ) try: @@ -180,13 +183,13 @@ def connect( process_pb2.ConnectRequest( process=process_pb2.ProcessSelector(pid=pid), ), - headers={ - KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), - }, - timeout=timeout, - request_timeout=self._connection_config.get_request_timeout( - request_timeout + headers=stream_request_headers( + { + KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC), + }, + self._connection_config.get_request_timeout(request_timeout), ), + timeout_ms=stream_timeout_ms(timeout), ) try: @@ -226,7 +229,7 @@ def resize( size=process_pb2.PTY.Size(rows=size.rows, cols=size.cols), ), ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), ) diff --git a/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py b/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py index b145200e41..8991ced031 100644 --- a/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py +++ b/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py @@ -1,11 +1,11 @@ from io import IOBase, TextIOBase from typing import IO, Iterator, List, Literal, Optional, Union, overload -import httpcore import httpx +from connectrpc.code import Code +from connectrpc.errors import ConnectError from packaging.version import Version -import e2b_connect from e2b.connection_config import ( KEEPALIVE_PING_HEADER, KEEPALIVE_PING_INTERVAL_SEC, @@ -13,11 +13,16 @@ Username, default_username, ) -from e2b_connect.client import Code from e2b.envd.api import ENVD_API_FILES_ROUTE, handle_envd_api_exception from e2b.envd.filesystem import filesystem_connect, filesystem_pb2 -from e2b.envd.rpc import authentication_header, handle_rpc_exception +from e2b.envd.pyqwest_httpx_adapter import PyqwestHTTPXAdapter +from e2b.envd.rpc import ( + authentication_header, + connect_client_kwargs, + handle_rpc_exception, + request_timeout_ms, +) from e2b.envd.versions import ( ENVD_DEFAULT_USER, ENVD_OCTET_STREAM_UPLOAD, @@ -40,15 +45,28 @@ _FILESYSTEM_RPC_ERROR_MAP = { - Code.not_found: FileNotFoundException, + Code.NOT_FOUND: FileNotFoundException, } _FILESYSTEM_HTTP_ERROR_MAP = { 404: FileNotFoundException, } +_ENOENT_MESSAGE = "no such file or directory" + def _handle_filesystem_rpc_exception(e: Exception) -> Exception: + if isinstance(e, ConnectError): + if e.code == Code.NOT_FOUND: + return FileNotFoundException(e.message) + + if ( + e.code == Code.UNKNOWN + # Older envd versions returned ENOENT as UNKNOWN instead of NOT_FOUND. + and _ENOENT_MESSAGE in e.message.lower() + ): + return FileNotFoundException(e.message) + return handle_rpc_exception(e, _FILESYSTEM_RPC_ERROR_MAP) @@ -66,22 +84,17 @@ def __init__( envd_api_url: str, envd_version: Version, connection_config: ConnectionConfig, - pool: httpcore.ConnectionPool, + rpc_client: PyqwestHTTPXAdapter, envd_api: httpx.Client, ) -> None: self._envd_api_url = envd_api_url self._envd_version = envd_version self._connection_config = connection_config - self._pool = pool self._envd_api = envd_api - self._rpc = filesystem_connect.FilesystemClient( + self._rpc = filesystem_connect.FilesystemClientSync( envd_api_url, - # TODO: Fix and enable compression again — the headers compression is not solved for streaming. - # compressor=e2b_connect.GzipCompressor, - pool=pool, - json=True, - headers=connection_config.sandbox_headers, + **connect_client_kwargs(connection_config.sandbox_headers, rpc_client), ) @overload @@ -363,8 +376,8 @@ def list( try: res = self._rpc.list_dir( filesystem_pb2.ListDirRequest(path=path, depth=depth), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) @@ -416,18 +429,18 @@ def exists( try: self._rpc.stat( filesystem_pb2.StatRequest(path=path), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) return True except Exception as e: - if isinstance(e, e2b_connect.ConnectException): - if e.status == e2b_connect.Code.not_found: - return False - raise _handle_filesystem_rpc_exception(e) + err = _handle_filesystem_rpc_exception(e) + if isinstance(err, FileNotFoundException): + return False + raise err def get_info( self, @@ -447,8 +460,8 @@ def get_info( try: r = self._rpc.stat( filesystem_pb2.StatRequest(path=path), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) @@ -489,8 +502,8 @@ def remove( try: self._rpc.remove( filesystem_pb2.RemoveRequest(path=path), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) @@ -520,8 +533,8 @@ def rename( source=old_path, destination=new_path, ), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) @@ -564,16 +577,16 @@ def make_dir( try: self._rpc.make_dir( filesystem_pb2.MakeDirRequest(path=path), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers=authentication_header(self._envd_version, user), ) return True except Exception as e: - if isinstance(e, e2b_connect.ConnectException): - if e.status == e2b_connect.Code.already_exists: + if isinstance(e, ConnectError): + if e.code == Code.ALREADY_EXISTS: return False raise _handle_filesystem_rpc_exception(e) @@ -603,8 +616,8 @@ def watch_dir( try: r = self._rpc.create_watcher( filesystem_pb2.CreateWatcherRequest(path=path, recursive=recursive), - request_timeout=self._connection_config.get_request_timeout( - request_timeout + timeout_ms=request_timeout_ms( + self._connection_config.get_request_timeout(request_timeout) ), headers={ **authentication_header(self._envd_version, user), diff --git a/packages/python-sdk/e2b/sandbox_sync/filesystem/watch_handle.py b/packages/python-sdk/e2b/sandbox_sync/filesystem/watch_handle.py index bbf531c57d..c07d3c34c5 100644 --- a/packages/python-sdk/e2b/sandbox_sync/filesystem/watch_handle.py +++ b/packages/python-sdk/e2b/sandbox_sync/filesystem/watch_handle.py @@ -20,7 +20,7 @@ class WatchHandle: def __init__( self, - rpc: filesystem_connect.FilesystemClient, + rpc: filesystem_connect.FilesystemClientSync, watcher_id: str, ): self._rpc = rpc diff --git a/packages/python-sdk/e2b/sandbox_sync/main.py b/packages/python-sdk/e2b/sandbox_sync/main.py index 43f3a858c5..d99e69917c 100644 --- a/packages/python-sdk/e2b/sandbox_sync/main.py +++ b/packages/python-sdk/e2b/sandbox_sync/main.py @@ -13,6 +13,7 @@ from e2b.api.client_sync import get_transport from e2b.connection_config import ApiParams, ConnectionConfig from e2b.envd.api import ENVD_API_HEALTH_ROUTE, handle_envd_api_exception +from e2b.envd.pyqwest_httpx_adapter import PyqwestHTTPXAdapter from e2b.envd.versions import ENVD_DEBUG_FALLBACK from e2b.exceptions import ( SandboxException, @@ -101,6 +102,7 @@ def __init__(self, **opts: Unpack[SandboxOpts]): super().__init__(**opts) self._transport = get_transport(self.connection_config) + self._rpc_client = PyqwestHTTPXAdapter(self._transport) self._envd_api = httpx.Client( base_url=self.envd_api_url, @@ -111,19 +113,19 @@ def __init__(self, **opts: Unpack[SandboxOpts]): self.envd_api_url, self._envd_version, self.connection_config, - self._transport.pool, + self._rpc_client, self._envd_api, ) self._commands = Commands( self.envd_api_url, self.connection_config, - self._transport.pool, + self._rpc_client, self._envd_version, ) self._pty = Pty( self.envd_api_url, self.connection_config, - self._transport.pool, + self._rpc_client, self._envd_version, ) self._git = Git(self._commands) diff --git a/packages/python-sdk/e2b_connect/__init__.py b/packages/python-sdk/e2b_connect/__init__.py deleted file mode 100644 index 6f31c8cee1..0000000000 --- a/packages/python-sdk/e2b_connect/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .client import Client, GzipCompressor, ConnectException, Code # noqa: F401 diff --git a/packages/python-sdk/e2b_connect/client.py b/packages/python-sdk/e2b_connect/client.py deleted file mode 100644 index b41fedb03c..0000000000 --- a/packages/python-sdk/e2b_connect/client.py +++ /dev/null @@ -1,499 +0,0 @@ -import gzip -import inspect -import json -import struct -import typing - -from httpcore import ( - ConnectionPool, - AsyncConnectionPool, - RemoteProtocolError, - Response, -) -from enum import Flag, Enum -from typing import Callable, Optional, Dict, Any, Generator, Tuple -from google.protobuf import json_format - - -class EnvelopeFlags(Flag): - compressed = 0b00000001 - end_stream = 0b00000010 - - -class Code(Enum): - canceled = "canceled" - unknown = "unknown" - invalid_argument = "invalid_argument" - deadline_exceeded = "deadline_exceeded" - not_found = "not_found" - already_exists = "already_exists" - permission_denied = "permission_denied" - resource_exhausted = "resource_exhausted" - failed_precondition = "failed_precondition" - aborted = "aborted" - out_of_range = "out_of_range" - unimplemented = "unimplemented" - internal = "internal" - unavailable = "unavailable" - data_loss = "data_loss" - unauthenticated = "unauthenticated" - - -def make_error_from_http_code(http_code: int): - error_code_map = { - 400: Code.invalid_argument, - 401: Code.unauthenticated, - 403: Code.permission_denied, - 404: Code.not_found, - 409: Code.already_exists, - 413: Code.resource_exhausted, - 429: Code.resource_exhausted, - 499: Code.canceled, - 500: Code.internal, - 501: Code.unimplemented, - 502: Code.unavailable, - 503: Code.unavailable, - 504: Code.deadline_exceeded, - 505: Code.unimplemented, - } - - return error_code_map.get(http_code, Code.unknown) - - -class ConnectException(Exception): - def __init__(self, status: Code, message: str): - self.status = status - self.message = message - - -envelope_header_length = 5 -envelope_header_pack = ">BI" - - -def encode_envelope(*, flags: EnvelopeFlags, data): - return encode_envelope_header(flags=flags.value, data=data) + data - - -def encode_envelope_header(*, flags, data): - return struct.pack(envelope_header_pack, flags, len(data)) - - -def decode_envelope_header(header): - flags, data_len = struct.unpack(envelope_header_pack, header) - return EnvelopeFlags(flags), data_len - - -def error_for_response(http_resp: Response): - try: - error = json.loads(http_resp.content) - return make_error(error) - except (json.decoder.JSONDecodeError, KeyError): - error = {"code": http_resp.status, "message": http_resp.content.decode("utf-8")} - return make_error(error) - - -def make_error(error): - status = None - try: - code_value = error.get("code") - # return error code from http status code - if isinstance(code_value, int): - status = make_error_from_http_code(code_value) - else: - status = Code(code_value) - except (KeyError, ValueError): - status = Code.unknown - - return ConnectException(status, error.get("message", "")) - - -def _sync_retry(func, exc, retries): - def retry(*args, **kwargs): - for _ in range(retries): - try: - return func(*args, **kwargs) - except exc: - continue - - return func(*args, **kwargs) - - return retry - - -def _async_retry(func, exc, retries): - async def retry(*args, **kwargs): - for _ in range(retries): - try: - return await func(*args, **kwargs) - except exc: - continue - - return await func(*args, **kwargs) - - return retry - - -def _retry(exc: typing.Type[Exception], retries: int): - def decorator(func): - if inspect.iscoroutinefunction(func): - return _async_retry(func, exc, retries) - - return _sync_retry(func, exc, retries) - - return decorator - - -class GzipCompressor: - name = "gzip" - decompress = gzip.decompress - compress = gzip.compress - - -class JSONCodec: - content_type = "json" - - @staticmethod - def encode(msg): - return json_format.MessageToJson(msg).encode("utf8") - - @staticmethod - def decode(data, *, msg_type): - msg = msg_type() - json_format.Parse(data.decode("utf8"), msg, ignore_unknown_fields=True) - return msg - - -class ProtobufCodec: - content_type = "proto" - - @staticmethod - def encode(msg): - return msg.SerializeToString() - - @staticmethod - def decode(data, *, msg_type): - msg = msg_type() - msg.ParseFromString(data) - return msg - - -class Client: - def __init__( - self, - *, - pool: Optional[ConnectionPool] = None, - async_pool: Optional[AsyncConnectionPool] = None, - url: str, - response_type, - compressor=None, - json: Optional[bool] = False, - headers: Optional[Dict[str, str]] = None, - ): - if headers is None: - headers = {} - - self.pool = pool - self.async_pool = async_pool - self.url = url - self._codec = JSONCodec if json else ProtobufCodec - self._response_type = response_type - self._compressor = compressor - self._headers = headers - self._connection_retries = 3 - - def _prepare_unary_request( - self, - req, - request_timeout=None, - headers: Optional[dict] = None, - **opts, - ) -> dict: - data = self._codec.encode(req) - - if self._compressor is not None: - data = self._compressor.compress(data) - - if headers is None: - headers = {} - - extensions = ( - None - if request_timeout is None - else { - "timeout": { - "connect": request_timeout, - "pool": request_timeout, - "read": request_timeout, - "write": request_timeout, - } - } - ) - - return { - "method": "POST", - "url": self.url, - "content": data, - "extensions": extensions, - "headers": { - **self._headers, - **headers, - **opts.get("headers", {}), - "connect-protocol-version": "1", - "content-encoding": ( - "identity" if self._compressor is None else self._compressor.name - ), - "content-type": f"application/{self._codec.content_type}", - }, - } - - def _process_unary_response( - self, - http_resp: Response, - ): - if http_resp.status != 200: - raise error_for_response(http_resp) - - content = http_resp.content - - if self._compressor is not None: - content = self._compressor.decompress(content) - - return self._codec.decode( - content, - msg_type=self._response_type, - ) - - @_retry(RemoteProtocolError, 3) - async def acall_unary( - self, - req, - request_timeout=None, - headers: Optional[dict] = None, - **opts, - ): - if self.async_pool is None: - raise ValueError("async_pool is required") - - req_data = self._prepare_unary_request( - req, - request_timeout, - headers, - **opts, - ) - - res = await self.async_pool.request(**req_data) - return self._process_unary_response(res) - - @_retry(RemoteProtocolError, 3) - def call_unary( - self, - req, - request_timeout=None, - headers: Optional[dict] = None, - **opts, - ): - if self.pool is None: - raise ValueError("pool is required") - - req_data = self._prepare_unary_request( - req, - request_timeout, - headers, - **opts, - ) - - res = self.pool.request(**req_data) - return self._process_unary_response(res) - - def _create_stream_timeout(self, timeout: Optional[float]): - if timeout: - return {"connect-timeout-ms": str(int(timeout * 1000))} - return {} - - def _prepare_server_stream_request( - self, - req, - request_timeout=None, - timeout=None, - headers: Optional[dict] = None, - **opts, - ) -> dict: - headers = headers or {} - data = self._codec.encode(req) - flags = EnvelopeFlags(0) - - timeout_ext = {} - if request_timeout is not None: - timeout_ext["connect"] = request_timeout - timeout_ext["pool"] = request_timeout - timeout_ext["write"] = request_timeout - if timeout: - # This is not actually timeout for the whole stream read, but timeout from the last read chunk. - # At worst then, the timeout of a hanging stream could be 2 * timeout (reading body until timeout-ϵ, then waiting for the read timeout). - # However, this is still better than no timeout at all and the full timeout in sync python might be way more complicated. - timeout_ext["read"] = timeout - extensions = {"timeout": timeout_ext} if timeout_ext else None - - if self._compressor is not None: - data = self._compressor.compress(data) - flags |= EnvelopeFlags.compressed - - stream_timeout = self._create_stream_timeout(timeout) - - return { - "method": "POST", - "url": self.url, - "content": encode_envelope( - flags=flags, - data=data, - ), - "extensions": extensions, - "headers": { - **self._headers, - **headers, - **opts.get("headers", {}), - **stream_timeout, - "connect-protocol-version": "1", - "connect-content-encoding": ( - "identity" if self._compressor is None else self._compressor.name - ), - "content-type": f"application/connect+{self._codec.content_type}", - }, - } - - @_retry(RemoteProtocolError, 3) - async def acall_server_stream( - self, - req, - request_timeout=None, - timeout=None, - headers: Optional[dict] = None, - **opts, - ): - if self.async_pool is None: - raise ValueError("async_pool is required") - - req_data = self._prepare_server_stream_request( - req, - request_timeout, - timeout, - headers, - **opts, - ) - - parser = ServerStreamParser( - decode=self._codec.decode, - response_type=self._response_type, - ) - - async with self.async_pool.stream(**req_data) as http_resp: - if http_resp.status != 200: - await http_resp.aread() - raise error_for_response(http_resp) - - async for chunk in http_resp.aiter_stream(): - for parsed in parser.parse(chunk): - yield parsed - - @_retry(RemoteProtocolError, 3) - def call_server_stream( - self, - req, - request_timeout=None, - timeout=None, - headers: Optional[dict] = None, - **opts, - ): - if self.pool is None: - raise ValueError("pool is required") - - req_data = self._prepare_server_stream_request( - req, - request_timeout, - timeout, - headers, - **opts, - ) - - parser = ServerStreamParser( - decode=self._codec.decode, - response_type=self._response_type, - ) - - with self.pool.stream(**req_data) as http_resp: - if http_resp.status != 200: - http_resp.read() - raise error_for_response(http_resp) - - for chunk in http_resp.iter_stream(): - for parsed in parser.parse(chunk): - yield parsed - - def call_client_stream(self, req, **opts): - raise NotImplementedError("client stream not supported") - - def acall_client_stream(self, req, **opts): - raise NotImplementedError("client stream not supported") - - def call_bidi_stream(self, req, **opts): - raise NotImplementedError("bidi stream not supported") - - def acall_bidi_stream(self, req, **opts): - raise NotImplementedError("bidi stream not supported") - - -DataLen = int - - -class ServerStreamParser: - def __init__( - self, - decode: Callable, - response_type: Any, - ): - self.decode = decode - self.response_type = response_type - - self.buffer: bytes = b"" - self._header: Optional[tuple[EnvelopeFlags, DataLen]] = None - - def shift_buffer(self, size: int): - buffer = self.buffer[:size] - self.buffer = self.buffer[size:] - return buffer - - @property - def header(self) -> Tuple[EnvelopeFlags, DataLen]: - if self._header: - return self._header - - header_data = self.shift_buffer(envelope_header_length) - self._header = decode_envelope_header(header_data) - - return self._header - - @header.deleter - def header(self): - self._header = None - - def parse(self, chunk: bytes) -> Generator[Any, None, None]: - self.buffer += chunk - - while len(self.buffer) >= envelope_header_length: - flags, data_len = self.header - - if data_len > len(self.buffer): - break - - data = self.shift_buffer(data_len) - - if EnvelopeFlags.end_stream in flags: - data = json.loads(data) - - if "error" in data: - raise make_error(data["error"]) - - return - - yield self.decode(data, msg_type=self.response_type) - del self.header diff --git a/packages/python-sdk/poetry.lock b/packages/python-sdk/poetry.lock index 8a5186b132..fa309ac955 100644 --- a/packages/python-sdk/poetry.lock +++ b/packages/python-sdk/poetry.lock @@ -301,6 +301,22 @@ files = [ {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] +[[package]] +name = "connectrpc" +version = "0.10.0" +description = "Server and client runtime library for Connect RPC" +optional = false +python-versions = ">=3.10" +groups = ["main"] +files = [ + {file = "connectrpc-0.10.0-py3-none-any.whl", hash = "sha256:fedaf5abb8c90ab69662958b88f23f4cc18b7a0b0359401d8e5776ccfceefc38"}, + {file = "connectrpc-0.10.0.tar.gz", hash = "sha256:44f89b70abae9f8192883e79e8a3317bdca555036ed3bd4a20f398c0bac7162f"}, +] + +[package.dependencies] +protobuf = ">=5.28" +pyqwest = ">=0.5.1" + [[package]] name = "databind" version = "4.5.2" @@ -620,6 +636,30 @@ files = [ [package.extras] all = ["flake8 (>=7.1.1)", "mypy (>=1.11.2)", "pytest (>=8.3.2)", "ruff (>=0.6.2)"] +[[package]] +name = "importlib-metadata" +version = "8.7.1" +description = "Read metadata from Python packages" +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "importlib_metadata-8.7.1-py3-none-any.whl", hash = "sha256:5a1f80bf1daa489495071efbb095d75a634cf28a8bc299581244063b53176151"}, + {file = "importlib_metadata-8.7.1.tar.gz", hash = "sha256:49fef1ae6440c182052f407c8d34a68f72efc36db9ca90dc0113398f2fdde8bb"}, +] + +[package.dependencies] +zipp = ">=3.20" + +[package.extras] +check = ["pytest-checkdocs (>=2.4)", "pytest-ruff (>=0.2.1) ; sys_platform != \"cygwin\""] +cover = ["pytest-cov"] +doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] +enabler = ["pytest-enabler (>=3.4)"] +perf = ["ipython"] +test = ["flufl.flake8", "jaraco.test (>=5.4)", "packaging", "pyfakefs", "pytest (>=6,!=8.1.*)", "pytest-perf (>=0.9.2)"] +type = ["mypy (<1.19) ; platform_python_implementation == \"PyPy\"", "pytest-mypy (>=1.0.1)"] + [[package]] name = "inflect" version = "7.5.0" @@ -890,6 +930,22 @@ files = [ deprecated = ">=1.2.0,<2.0.0" typing-extensions = ">=3.0.0" +[[package]] +name = "opentelemetry-api" +version = "1.41.1" +description = "OpenTelemetry Python API" +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "opentelemetry_api-1.41.1-py3-none-any.whl", hash = "sha256:a22df900e75c76dc08440710e51f52f1aa6b451b429298896023e60db5b3139f"}, + {file = "opentelemetry_api-1.41.1.tar.gz", hash = "sha256:0ad1814d73b875f84494387dae86ce0b12c68556331ce6ce8fe789197c949621"}, +] + +[package.dependencies] +importlib-metadata = ">=6.0,<8.8.0" +typing-extensions = ">=4.5.0" + [[package]] name = "packaging" version = "25.0" @@ -1167,6 +1223,56 @@ files = [ [package.extras] windows-terminal = ["colorama (>=0.4.6)"] +[[package]] +name = "pyqwest" +version = "0.5.1" +description = "" +optional = false +python-versions = ">=3.10" +groups = ["main"] +files = [ + {file = "pyqwest-0.5.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:047078c4c3f7c93d8a1df07138c471cfdb234577f60563afe090b160ada1a132"}, + {file = "pyqwest-0.5.1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3f732aed3498abf5742f0c952f24fc21625d62fe0a701ec8e472257709540d61"}, + {file = "pyqwest-0.5.1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:77565ad8e551a101eddbf08328b0e3d54a179efec97a3da195ce831f0474ef92"}, + {file = "pyqwest-0.5.1-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:57ea14094f841d735fcf401a6428e94a39aabf22d15762a979cf8d6096fea90b"}, + {file = "pyqwest-0.5.1-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e6716309fc179b0714978d1f8c32a70724b3b9f43acc9aea3cf9507872e631d6"}, + {file = "pyqwest-0.5.1-cp310-abi3-win_amd64.whl", hash = "sha256:736b560fd2256a41264f554243edf3ff872ead4da347392531576203cf97a4ce"}, + {file = "pyqwest-0.5.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c30bb16da6580d085d86a46e28c3b59813c634c5dd7363faa980c6e53d42ea21"}, + {file = "pyqwest-0.5.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4540be138c7e4ee498c084fa9f19f34dd135f263653619ec91ccebecefd32a3a"}, + {file = "pyqwest-0.5.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ae56109419b37be3e90a134f827fda26c6338c432bae5ecef561e93c3df275ab"}, + {file = "pyqwest-0.5.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:82fbf17cf45b5a67c95734934f584045f114a1aa59d179fb79d80e56e653e052"}, + {file = "pyqwest-0.5.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:717b65ee898bc3e4f29a74e068bb753bbb8a338456f1fb65edc2edabc3f3eaa5"}, + {file = "pyqwest-0.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:26498ead2aaf491430980c06cbee3bd72a26a8f33a025e9a19ddcd7e781a069a"}, + {file = "pyqwest-0.5.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:81bf0feb6d07fc25ffbe96be09ef64ff732d444e26c80db1b1de999d74a7a7ed"}, + {file = "pyqwest-0.5.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3a7d2c4a162e72b57dfd175f53882558702c3ddd5704f536599f463e46081bd9"}, + {file = "pyqwest-0.5.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b33bcba509e07ee79e61f9311880a22e8c11a562382648166f624f0b480266c1"}, + {file = "pyqwest-0.5.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:3a7a7b56e0f9099b50dae2e5907f05fc5e61252cd48cc1561f156765b233dd7b"}, + {file = "pyqwest-0.5.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:a3cd54af62f01fd0daf694218074325ccbbea6527f27738a85fc1e38880119f0"}, + {file = "pyqwest-0.5.1-cp313-cp313-win_amd64.whl", hash = "sha256:938ac244f96730a8c1addb260cbec93e44ef63d73358285bafb00f18526c106b"}, + {file = "pyqwest-0.5.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:7755dd658c612016c38f0e26938b80dafd37d80ee074c5cc4301f54f385b1e52"}, + {file = "pyqwest-0.5.1-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:29015b12242f2feeb24cbad7d1025a78b3239872efcbfa975ce23527e8c05ec7"}, + {file = "pyqwest-0.5.1-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1271fb28a78484e93d09b7c78df5213900538af8cf5a53ce377ed2cbb82e72b2"}, + {file = "pyqwest-0.5.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:eaf970412356f9d6482045a8e7b7c9d36c07b6a7352c96c90a811e730eba906f"}, + {file = "pyqwest-0.5.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:986b787138893459e02485e89a7d4a383ec2864194a5ab7e0650dce9635f5215"}, + {file = "pyqwest-0.5.1-cp314-cp314-win_amd64.whl", hash = "sha256:4d9fb2bcbea6adc365ea3c69727b34dd21dee29edb1551b64531b25babb68f91"}, + {file = "pyqwest-0.5.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:dc7a7c65e839057bcc3a5d4f48c6c2202f89fc13ddb12bd8e2f20d95378a95d8"}, + {file = "pyqwest-0.5.1-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e30c2268e9f7ec5509874047ebf83ecde1cab29d0bd7dc352fa2a782da2b2abf"}, + {file = "pyqwest-0.5.1-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f632070abff4d0bfd0f6deada3beaaaa42a042809e50c7f2e2b25014687e070"}, + {file = "pyqwest-0.5.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:78f8b6b41c0206e9605d75b9fc3f20db97c266829f5a0bee8eff7d524de2bf20"}, + {file = "pyqwest-0.5.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:a1d95a8a7e4246ba37dc3115b10b72c5dcdf4b226447d1a7a51973d686460fa4"}, + {file = "pyqwest-0.5.1-cp314-cp314t-win_amd64.whl", hash = "sha256:fbd3232cca327bc591d47b0368ba712d929f66658d4822ac5129f5ede1a581df"}, + {file = "pyqwest-0.5.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:45f9d448fedd6d85c4c9a5cb8e743b83fdd798c335936588f5d90628c6163f73"}, + {file = "pyqwest-0.5.1-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:144c22aa18de37f5b6f877939c50fb2aafb94a6433ac47ddab3850793402e46b"}, + {file = "pyqwest-0.5.1-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:850c58da7df5c0a4c6351ffe66f99768cb5d92e882b55daf9984a96a5c46e11e"}, + {file = "pyqwest-0.5.1-pp311-pypy311_pp73-musllinux_1_2_aarch64.whl", hash = "sha256:03e57c1d58d7f236a08dfe5b27493ee56ae2d5e4b40e5c4bd9e6309000b4b86f"}, + {file = "pyqwest-0.5.1-pp311-pypy311_pp73-musllinux_1_2_x86_64.whl", hash = "sha256:958540d4861b9e7a228a775ea8b937c019bdc769d52b6e2e872b5bdc1ebf6870"}, + {file = "pyqwest-0.5.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:469dfd7baaac6fd6565d396c8d05f19cb5ba279856af38b5f69ff8278e991c49"}, + {file = "pyqwest-0.5.1.tar.gz", hash = "sha256:49535565a55a23830d376c6c0b8ca1276f19bb871e39330e1fbb3df69d9f02df"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.39.1" + [[package]] name = "pytest" version = "9.0.3" @@ -1856,7 +1962,27 @@ files = [ platformdirs = ">=3.5.1" tomli = {version = ">=2.0.1", markers = "python_version < \"3.11\""} +[[package]] +name = "zipp" +version = "3.23.1" +description = "Backport of pathlib-compatible object wrapper for zip files" +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "zipp-3.23.1-py3-none-any.whl", hash = "sha256:0b3596c50a5c700c9cb40ba8d86d9f2cc4807e9bedb06bcdf7fac85633e444dc"}, + {file = "zipp-3.23.1.tar.gz", hash = "sha256:32120e378d32cd9714ad503c1d024619063ec28aad2248dc6672ad13edfa5110"}, +] + +[package.extras] +check = ["pytest-checkdocs (>=2.4)", "pytest-ruff (>=0.2.1) ; sys_platform != \"cygwin\""] +cover = ["pytest-cov"] +doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] +enabler = ["pytest-enabler (>=2.2)"] +test = ["big-O", "jaraco.functools", "jaraco.itertools", "jaraco.test", "more_itertools", "pytest (>=6,!=8.1.*)", "pytest-ignore-flaky"] +type = ["pytest-mypy"] + [metadata] lock-version = "2.1" python-versions = "^3.10" -content-hash = "10b34c02d9b97fb1fd229107629913601c669a7d81fc558d11dce61a92812743" +content-hash = "7e4cb41b37f445098bcdf477491dd49c032a77021e3970a7e9b8f84dca6133f9" diff --git a/packages/python-sdk/pyproject.toml b/packages/python-sdk/pyproject.toml index 5b2167bc2c..71b4add539 100644 --- a/packages/python-sdk/pyproject.toml +++ b/packages/python-sdk/pyproject.toml @@ -7,13 +7,15 @@ license = "MIT" readme = "README.md" homepage = "https://e2b.dev/" repository = "https://github.com/e2b-dev/e2b/tree/main/packages/python-sdk" -packages = [{ include = "e2b" }, { include = "e2b_connect" }] +packages = [{ include = "e2b" }] [tool.poetry.dependencies] python = "^3.10" python-dateutil = ">=2.8.2" wcmatch = "^10.1" -protobuf = ">=4.21.0" +protobuf = ">=5.28" +connectrpc = "^0.10.0" +pyqwest = ">=0.5.1" httpcore = "^1.0.5" httpx = ">=0.27.0, <1.0.0" h2 = ">=4,<5" diff --git a/packages/python-sdk/scripts/fix-python-pb.sh b/packages/python-sdk/scripts/fix-python-pb.sh index 7bdea385a3..a0be898614 100755 --- a/packages/python-sdk/scripts/fix-python-pb.sh +++ b/packages/python-sdk/scripts/fix-python-pb.sh @@ -6,6 +6,8 @@ rm -rf e2b/envd/process/__pycache__ sed -i.bak 's/from\ process\ import/from e2b.envd.process import/g' e2b/envd/process/* e2b/envd/filesystem/* sed -i.bak 's/from\ filesystem\ import/from e2b.envd.filesystem import/g' e2b/envd/process/* e2b/envd/filesystem/* +sed -i.bak 's/import\ process\.process_pb2\ as/from e2b.envd.process import process_pb2 as/g' e2b/envd/process/* e2b/envd/filesystem/* +sed -i.bak 's/import\ filesystem\.filesystem_pb2\ as/from e2b.envd.filesystem import filesystem_pb2 as/g' e2b/envd/process/* e2b/envd/filesystem/* rm -f e2b/envd/process/*.bak rm -f e2b/envd/filesystem/*.bak diff --git a/packages/python-sdk/tests/async/sandbox_async/files/test_exists.py b/packages/python-sdk/tests/async/sandbox_async/files/test_exists.py index 203dd19a34..4a5a4b0c59 100644 --- a/packages/python-sdk/tests/async/sandbox_async/files/test_exists.py +++ b/packages/python-sdk/tests/async/sandbox_async/files/test_exists.py @@ -1,4 +1,8 @@ -from e2b import AsyncSandbox +from connectrpc.code import Code +from connectrpc.errors import ConnectError + +from e2b import AsyncSandbox, FileNotFoundException +from e2b.sandbox_async.filesystem.filesystem import _handle_filesystem_rpc_exception async def test_exists(async_sandbox: AsyncSandbox): @@ -6,3 +10,23 @@ async def test_exists(async_sandbox: AsyncSandbox): await async_sandbox.files.write(filename, "test") assert await async_sandbox.files.exists(filename) + + +async def test_exists_non_existing_file(async_sandbox: AsyncSandbox): + assert not await async_sandbox.files.exists("non_existing_file.txt") + + +def test_unknown_enoent_maps_to_file_not_found(): + err = ConnectError(Code.UNKNOWN, "open /tmp/file: No such file or directory") + + mapped = _handle_filesystem_rpc_exception(err) + + assert isinstance(mapped, FileNotFoundException) + + +def test_not_found_maps_to_file_not_found(): + err = ConnectError(Code.NOT_FOUND, "file not found") + + mapped = _handle_filesystem_rpc_exception(err) + + assert isinstance(mapped, FileNotFoundException) diff --git a/packages/python-sdk/tests/e2b_connect/__init__.py b/packages/python-sdk/tests/e2b_connect/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/packages/python-sdk/tests/e2b_connect/test_client.py b/packages/python-sdk/tests/e2b_connect/test_client.py deleted file mode 100644 index ea1f272603..0000000000 --- a/packages/python-sdk/tests/e2b_connect/test_client.py +++ /dev/null @@ -1,134 +0,0 @@ -import asyncio - -import pytest - -from e2b_connect.client import _retry - - -class GoodError(Exception): - pass - - -class BadError(Exception): - pass - - -def test_sync_retry_after_expected_exception(): - total = 0 - - @_retry(GoodError, 1) - def f(): - nonlocal total - total += 1 - raise GoodError() - - with pytest.raises(GoodError): - f() - - assert total == 2 - - -def test_sync_do_not_retry_on_unexpected_exception(): - total = 0 - - @_retry(GoodError, 1) - def f(): - nonlocal total - total += 1 - raise BadError() - - with pytest.raises(BadError): - f() - - assert total == 1 - - -def test_sync_do_not_throw_when_retry_works(): - total = 0 - - @_retry(GoodError, 1) - def f(): - nonlocal total - total += 1 - - if total < 2: - raise GoodError() - - return True - - result = f() - assert result is True - assert total == 2 - - -async def test_async_retry_after_expected_exception(): - total = 0 - - @_retry(GoodError, 1) - async def f(): - nonlocal total - total += 1 - raise GoodError() - - with pytest.raises(GoodError): - await f() - - assert total == 2 - - -async def test_async_do_not_retry_on_unexpected_exception(): - total = 0 - - @_retry(GoodError, 1) - async def f(): - nonlocal total - total += 1 - raise BadError() - - with pytest.raises(BadError): - await f() - - assert total == 1 - - -async def test_async_do_not_throw_when_retry_works(): - total = 0 - - @_retry(GoodError, 1) - async def f(): - nonlocal total - total += 1 - - if total < 2: - raise GoodError() - - return True - - result = await f() - assert result is True - assert total == 2 - - -async def test_async_with_multiple_await_calls(): - total = 0 - - async def a(): - await asyncio.sleep(0.001) - - @_retry(GoodError, 1) - async def f(): - nonlocal total - total += 1 - - await a() - - if total < 2: - raise GoodError() - - await a() - - return True - - result = await f() - assert result is True - assert total == 2 diff --git a/packages/python-sdk/tests/sync/sandbox_sync/files/test_exists.py b/packages/python-sdk/tests/sync/sandbox_sync/files/test_exists.py index 2a04aec5d7..c25e2a67ba 100644 --- a/packages/python-sdk/tests/sync/sandbox_sync/files/test_exists.py +++ b/packages/python-sdk/tests/sync/sandbox_sync/files/test_exists.py @@ -1,4 +1,8 @@ -from e2b import Sandbox +from connectrpc.code import Code +from connectrpc.errors import ConnectError + +from e2b import FileNotFoundException, Sandbox +from e2b.sandbox_sync.filesystem.filesystem import _handle_filesystem_rpc_exception def test_exists(sandbox: Sandbox): @@ -6,3 +10,23 @@ def test_exists(sandbox: Sandbox): sandbox.files.write(filename, "test") assert sandbox.files.exists(filename) + + +def test_exists_non_existing_file(sandbox: Sandbox): + assert not sandbox.files.exists("non_existing_file.txt") + + +def test_unknown_enoent_maps_to_file_not_found(): + err = ConnectError(Code.UNKNOWN, "open /tmp/file: No such file or directory") + + mapped = _handle_filesystem_rpc_exception(err) + + assert isinstance(mapped, FileNotFoundException) + + +def test_not_found_maps_to_file_not_found(): + err = ConnectError(Code.NOT_FOUND, "file not found") + + mapped = _handle_filesystem_rpc_exception(err) + + assert isinstance(mapped, FileNotFoundException) diff --git a/packages/python-sdk/tests/test_envd_pyqwest_httpx_adapter.py b/packages/python-sdk/tests/test_envd_pyqwest_httpx_adapter.py new file mode 100644 index 0000000000..29ac273276 --- /dev/null +++ b/packages/python-sdk/tests/test_envd_pyqwest_httpx_adapter.py @@ -0,0 +1,729 @@ +import socket +import threading +from dataclasses import dataclass + +import httpx +import pytest +from connectrpc.code import Code +from connectrpc.errors import ConnectError +from pyqwest import Headers + +from e2b.envd.pyqwest_httpx_adapter import AsyncPyqwestHTTPXAdapter, PyqwestHTTPXAdapter +from e2b.envd.process import process_connect, process_pb2 +from e2b.envd.rpc import ( + STREAM_REQUEST_TIMEOUT_HEADER, + connect_client_kwargs, + stream_request_headers, + stream_timeout_ms, +) + + +class AsyncBytes(httpx.AsyncByteStream): + def __init__(self, chunks: list[bytes]) -> None: + self._chunks = chunks + + async def __aiter__(self): + for chunk in self._chunks: + yield chunk + + +class SyncTimeoutBytes(httpx.SyncByteStream): + def __iter__(self): + raise httpx.ReadTimeout("read timed out") + + +class AsyncTimeoutBytes(httpx.AsyncByteStream): + def __aiter__(self): + async def iterator(): + raise httpx.ReadTimeout("read timed out") + yield b"" + + return iterator() + + +@dataclass +class ProxyRequest: + method: str + target: str + body: bytes + + +class RecordingProxy: + def __init__(self) -> None: + self.requests: list[ProxyRequest] = [] + self._ready = threading.Event() + self._closed = threading.Event() + self._thread = threading.Thread(target=self._serve, daemon=True) + + @property + def url(self) -> str: + self._ready.wait(timeout=5) + return f"http://{self.host}:{self.port}" + + def __enter__(self): + self._thread.start() + self._ready.wait(timeout=5) + return self + + def __exit__(self, exc_type, exc, tb): + self._closed.set() + try: + with socket.create_connection((self.host, self.port), timeout=1): + pass + except OSError: + pass + self._thread.join(timeout=5) + + def _serve(self) -> None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server: + server.bind(("127.0.0.1", 0)) + server.listen() + self.host, self.port = server.getsockname() + self._ready.set() + + while not self._closed.is_set(): + conn, _ = server.accept() + with conn: + data = self._read_request(conn) + if not data: + continue + + request_head, body = data.split(b"\r\n\r\n", 1) + lines = request_head.decode().splitlines() + method, target, _ = lines[0].split(" ", 2) + content_length = self._content_length(lines[1:]) + + while len(body) < content_length: + body += conn.recv(65536) + + self.requests.append( + ProxyRequest(method=method, target=target, body=body) + ) + response = b'{"ok":true}' + conn.sendall( + b"HTTP/1.1 200 OK\r\n" + b"Content-Type: application/json\r\n" + + f"Content-Length: {len(response)}\r\n".encode() + + b"Connection: close\r\n\r\n" + + response + ) + + def _read_request(self, conn: socket.socket) -> bytes: + data = b"" + while b"\r\n\r\n" not in data: + chunk = conn.recv(65536) + if not chunk: + return data + data += chunk + return data + + def _content_length(self, header_lines: list[str]) -> int: + for line in header_lines: + name, _, value = line.partition(":") + if name.lower() == "content-length": + return int(value.strip()) + return 0 + + +def test_sync_pyqwest_httpx_adapter_uses_httpx_transport(): + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=b'{"ok":true}', + ) + + client = PyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + response = client.post( + "https://sandbox.test/process.Process/List", + headers=Headers({"x-test": "1"}), + content=b"payload", + timeout=1, + ) + + assert response.status == 200 + assert response.content == b'{"ok":true}' + assert requests[0].headers["x-test"] == "1" + assert requests[0].content == b"payload" + + +def test_sync_pyqwest_httpx_adapter_disables_httpx_default_timeout(): + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=b'{"ok":true}', + ) + + client = PyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + client.post("https://sandbox.test/process.Process/List", content=b"payload") + + assert requests[0].extensions["timeout"] == { + "connect": None, + "read": None, + "write": None, + "pool": None, + } + + +def test_sync_pyqwest_httpx_adapter_retries_remote_protocol_errors(): + attempts: list[bytes] = [] + + def handler(request: httpx.Request) -> httpx.Response: + attempts.append(request.content) + if len(attempts) <= 3: + raise httpx.RemoteProtocolError("connection reset") + + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=b'{"ok":true}', + ) + + client = PyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + response = client.post( + "https://sandbox.test/process.Process/List", + content=b"payload", + ) + + assert response.status == 200 + assert attempts == [b"payload", b"payload", b"payload", b"payload"] + + +def test_sync_pyqwest_httpx_adapter_maps_request_timeout_errors(): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectTimeout("connect timed out") + + client = PyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + with pytest.raises(ConnectError) as exc_info: + client.post("https://sandbox.test/process.Process/List", content=b"payload") + + assert exc_info.value.code == Code.CANCELED + + +def test_sync_pyqwest_httpx_adapter_streams_response(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "application/connect+json"}, + content=[b"one", b"two"], + ) + + client = PyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + with client.stream( + "POST", + "https://sandbox.test/process.Process/Start", + headers=Headers({"x-test": "1"}), + content=[b"payload"], + timeout=1, + ) as response: + assert response.status == 200 + assert list(response.content) == [b"one", b"two"] + + +def test_sync_pyqwest_httpx_adapter_does_not_retry_stream_open_errors(): + attempts: list[bytes] = [] + + def content(): + yield b"payload" + + def handler(request: httpx.Request) -> httpx.Response: + attempts.append(request.content) + raise httpx.RemoteProtocolError("connection reset") + + client = PyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + with pytest.raises(httpx.RemoteProtocolError): + with client.stream( + "POST", + "https://sandbox.test/process.Process/Start", + content=content(), + ) as response: + list(response.content) + + assert attempts == [b"payload"] + + +def test_sync_pyqwest_httpx_adapter_maps_stream_read_timeout_errors(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "application/connect+json"}, + stream=SyncTimeoutBytes(), + ) + + client = PyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + with client.stream( + "POST", + "https://sandbox.test/process.Process/Start", + content=[b"payload"], + timeout=60, + ) as response: + with pytest.raises(ConnectError) as exc_info: + list(response.content) + + assert exc_info.value.code == Code.DEADLINE_EXCEEDED + + +def test_sync_pyqwest_httpx_adapter_applies_stream_request_timeout(): + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={"content-type": "application/connect+json"}, + content=[b"one"], + ) + + client = PyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + with client.stream( + "POST", + "https://sandbox.test/process.Process/Start", + headers=Headers( + { + "x-test": "1", + STREAM_REQUEST_TIMEOUT_HEADER: "5", + } + ), + content=[b"payload"], + timeout=60, + ) as response: + assert list(response.content) == [b"one"] + + assert STREAM_REQUEST_TIMEOUT_HEADER not in requests[0].headers + assert requests[0].extensions["timeout"] == { + "connect": 5.0, + "read": 60, + "write": 5.0, + "pool": 5.0, + } + + +def test_sync_pyqwest_httpx_adapter_ignores_unlimited_stream_request_timeout(): + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={"content-type": "application/connect+json"}, + content=[b"one"], + ) + + client = PyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + with client.stream( + "POST", + "https://sandbox.test/process.Process/Start", + headers=Headers({STREAM_REQUEST_TIMEOUT_HEADER: "0"}), + content=[b"payload"], + timeout=60, + ) as response: + assert list(response.content) == [b"one"] + + assert STREAM_REQUEST_TIMEOUT_HEADER not in requests[0].headers + assert requests[0].extensions["timeout"] == { + "connect": None, + "read": 60, + "write": None, + "pool": None, + } + + +def test_sync_pyqwest_httpx_adapter_uses_configured_proxy(): + with RecordingProxy() as proxy: + transport = httpx.HTTPTransport(proxy=proxy.url) + client = PyqwestHTTPXAdapter(transport) + + response = client.post( + "http://sandbox.test/process.Process/List", + headers=Headers({"x-test": "1"}), + content=b"payload", + timeout=1, + ) + + assert response.status == 200 + assert proxy.requests == [ + ProxyRequest( + method="POST", + target="http://sandbox.test/process.Process/List", + body=b"payload", + ) + ] + + +def test_sync_generated_stream_request_shape(): + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={"content-type": "application/connect+json"}, + content=b"{}", + ) + + client = process_connect.ProcessClientSync( + "https://sandbox.test", + **connect_client_kwargs( + {"x-sandbox": "1"}, + PyqwestHTTPXAdapter(httpx.MockTransport(handler)), + ), + ) + events = client.start( + process_pb2.StartRequest( + process=process_pb2.ProcessConfig(cmd="/bin/bash"), + ), + headers=stream_request_headers({"E2B-Keepalive-Ping": "15"}, 5), + timeout_ms=stream_timeout_ms(60), + ) + + with pytest.raises(StopIteration): + next(events) + + request = requests[0] + assert request.method == "POST" + assert str(request.url) == "https://sandbox.test/process.Process/Start" + assert request.headers["x-sandbox"] == "1" + assert request.headers["e2b-keepalive-ping"] == "15" + assert request.headers["connect-timeout-ms"] == "60000" + assert STREAM_REQUEST_TIMEOUT_HEADER not in request.headers + assert request.extensions["timeout"]["connect"] == 5.0 + assert 0 < request.extensions["timeout"]["read"] <= 60 + assert request.extensions["timeout"]["write"] == 5.0 + assert request.extensions["timeout"]["pool"] == 5.0 + assert b'"cmd": "/bin/bash"' in request.content + + +def test_sync_generated_stream_with_unlimited_request_timeout_shape(): + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={"content-type": "application/connect+json"}, + content=b"{}", + ) + + client = process_connect.ProcessClientSync( + "https://sandbox.test", + **connect_client_kwargs( + {"x-sandbox": "1"}, + PyqwestHTTPXAdapter(httpx.MockTransport(handler)), + ), + ) + events = client.start( + process_pb2.StartRequest( + process=process_pb2.ProcessConfig(cmd="/bin/bash"), + ), + headers=stream_request_headers({"E2B-Keepalive-Ping": "15"}, None), + timeout_ms=stream_timeout_ms(60), + ) + + with pytest.raises(StopIteration): + next(events) + + request = requests[0] + assert request.headers["connect-timeout-ms"] == "60000" + assert STREAM_REQUEST_TIMEOUT_HEADER not in request.headers + assert request.extensions["timeout"]["connect"] is None + assert 0 < request.extensions["timeout"]["read"] <= 60 + assert request.extensions["timeout"]["write"] is None + assert request.extensions["timeout"]["pool"] is None + + +@pytest.mark.asyncio +async def test_async_pyqwest_httpx_adapter_uses_httpx_transport(): + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + stream=AsyncBytes([b'{"ok":true}']), + ) + + client = AsyncPyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + response = await client.post( + "https://sandbox.test/process.Process/List", + headers=Headers({"x-test": "1"}), + content=b"payload", + ) + + assert response.status == 200 + assert response.content == b'{"ok":true}' + assert requests[0].headers["x-test"] == "1" + assert await requests[0].aread() == b"payload" + + +@pytest.mark.asyncio +async def test_async_pyqwest_httpx_adapter_disables_httpx_default_timeout(): + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + stream=AsyncBytes([b'{"ok":true}']), + ) + + client = AsyncPyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + await client.post("https://sandbox.test/process.Process/List", content=b"payload") + + assert requests[0].extensions["timeout"] == { + "connect": None, + "read": None, + "write": None, + "pool": None, + } + + +@pytest.mark.asyncio +async def test_async_pyqwest_httpx_adapter_retries_remote_protocol_errors(): + attempts: list[bytes] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + attempts.append(await request.aread()) + if len(attempts) <= 3: + raise httpx.RemoteProtocolError("connection reset") + + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + stream=AsyncBytes([b'{"ok":true}']), + ) + + client = AsyncPyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + response = await client.post( + "https://sandbox.test/process.Process/List", + content=b"payload", + ) + + assert response.status == 200 + assert attempts == [b"payload", b"payload", b"payload", b"payload"] + + +@pytest.mark.asyncio +async def test_async_pyqwest_httpx_adapter_maps_request_timeout_errors(): + async def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectTimeout("connect timed out") + + client = AsyncPyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + with pytest.raises(ConnectError) as exc_info: + await client.post( + "https://sandbox.test/process.Process/List", content=b"payload" + ) + + assert exc_info.value.code == Code.CANCELED + + +@pytest.mark.asyncio +async def test_async_pyqwest_httpx_adapter_streams_response(): + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "application/connect+json"}, + stream=AsyncBytes([b"one", b"two"]), + ) + + async def content(): + yield b"payload" + + client = AsyncPyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + async with client.stream( + "POST", + "https://sandbox.test/process.Process/Start", + headers=Headers({"x-test": "1"}), + content=content(), + ) as response: + chunks = [] + async for chunk in response.content: + chunks.append(chunk) + + assert response.status == 200 + assert chunks == [b"one", b"two"] + + +@pytest.mark.asyncio +async def test_async_pyqwest_httpx_adapter_does_not_retry_stream_open_errors(): + attempts: list[bytes] = [] + + async def content(): + yield b"payload" + + async def handler(request: httpx.Request) -> httpx.Response: + attempts.append(await request.aread()) + raise httpx.RemoteProtocolError("connection reset") + + client = AsyncPyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + with pytest.raises(httpx.RemoteProtocolError): + async with client.stream( + "POST", + "https://sandbox.test/process.Process/Start", + content=content(), + ) as response: + async for _ in response.content: + pass + + assert attempts == [b"payload"] + + +@pytest.mark.asyncio +async def test_async_pyqwest_httpx_adapter_maps_stream_read_timeout_errors(): + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "application/connect+json"}, + stream=AsyncTimeoutBytes(), + ) + + async def content(): + yield b"payload" + + client = AsyncPyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + async with client.stream( + "POST", + "https://sandbox.test/process.Process/Start", + content=content(), + timeout=60, + ) as response: + with pytest.raises(ConnectError) as exc_info: + async for _ in response.content: + pass + + assert exc_info.value.code == Code.DEADLINE_EXCEEDED + + +@pytest.mark.asyncio +async def test_async_pyqwest_httpx_adapter_applies_stream_request_timeout(): + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={"content-type": "application/connect+json"}, + stream=AsyncBytes([b"one"]), + ) + + async def content(): + yield b"payload" + + client = AsyncPyqwestHTTPXAdapter(httpx.MockTransport(handler)) + + async with client.stream( + "POST", + "https://sandbox.test/process.Process/Start", + headers=Headers( + { + "x-test": "1", + STREAM_REQUEST_TIMEOUT_HEADER: "5", + } + ), + content=content(), + timeout=60, + ) as response: + chunks = [] + async for chunk in response.content: + chunks.append(chunk) + + assert chunks == [b"one"] + assert STREAM_REQUEST_TIMEOUT_HEADER not in requests[0].headers + assert requests[0].extensions["timeout"] == { + "connect": 5.0, + "read": 60, + "write": 5.0, + "pool": 5.0, + } + + +@pytest.mark.asyncio +async def test_async_pyqwest_httpx_adapter_uses_configured_proxy(): + with RecordingProxy() as proxy: + transport = httpx.AsyncHTTPTransport(proxy=proxy.url) + client = AsyncPyqwestHTTPXAdapter(transport) + + response = await client.post( + "http://sandbox.test/process.Process/List", + headers=Headers({"x-test": "1"}), + content=b"payload", + ) + + assert response.status == 200 + assert proxy.requests == [ + ProxyRequest( + method="POST", + target="http://sandbox.test/process.Process/List", + body=b"payload", + ) + ] + + +@pytest.mark.asyncio +async def test_async_generated_unlimited_stream_request_shape(): + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={"content-type": "application/connect+json"}, + content=b"{}", + ) + + client = process_connect.ProcessClient( + "https://sandbox.test", + **connect_client_kwargs( + {"x-sandbox": "1"}, + AsyncPyqwestHTTPXAdapter(httpx.MockTransport(handler)), + ), + ) + events = client.start( + process_pb2.StartRequest( + process=process_pb2.ProcessConfig(cmd="/bin/bash"), + ), + headers=stream_request_headers({"E2B-Keepalive-Ping": "15"}, 5), + timeout_ms=stream_timeout_ms(0), + ) + + with pytest.raises(StopAsyncIteration): + await events.__anext__() + + request = requests[0] + assert request.method == "POST" + assert str(request.url) == "https://sandbox.test/process.Process/Start" + assert request.headers["x-sandbox"] == "1" + assert request.headers["e2b-keepalive-ping"] == "15" + assert "connect-timeout-ms" not in request.headers + assert STREAM_REQUEST_TIMEOUT_HEADER not in request.headers + assert request.extensions["timeout"] == { + "connect": 5.0, + "read": None, + "write": 5.0, + "pool": 5.0, + } + assert b'"cmd": "/bin/bash"' in await request.aread() diff --git a/packages/python-sdk/tests/test_envd_rpc.py b/packages/python-sdk/tests/test_envd_rpc.py new file mode 100644 index 0000000000..840d6b8c8b --- /dev/null +++ b/packages/python-sdk/tests/test_envd_rpc.py @@ -0,0 +1,47 @@ +from e2b.envd.process import process_pb2 +from e2b.envd.rpc import ( + STREAM_REQUEST_TIMEOUT_HEADER, + ProtoJSONCodec, + request_timeout_ms, + stream_request_headers, + stream_timeout_ms, +) + + +def test_request_timeout_ms_preserves_zero(): + assert request_timeout_ms(None) is None + assert request_timeout_ms(0) == 0 + assert request_timeout_ms(0.0) == 0 + assert request_timeout_ms(1.25) == 1250 + + +def test_stream_timeout_ms_uses_stream_timeout(): + assert stream_timeout_ms(60) == 60000 + assert stream_timeout_ms(300) == 300000 + assert stream_timeout_ms(None) is None + assert stream_timeout_ms(0) is None + + +def test_stream_request_headers_adds_transport_timeout(): + headers = stream_request_headers({"x-test": "1"}, 5) + + assert headers == {"x-test": "1", STREAM_REQUEST_TIMEOUT_HEADER: "5"} + + +def test_stream_request_headers_ignores_unlimited_timeout(): + headers = {"x-test": "1"} + + assert stream_request_headers(headers, None) is headers + assert stream_request_headers(headers, 0) is headers + + +def test_proto_json_codec_ignores_unknown_fields(): + response = process_pb2.ListResponse() + + decoded = ProtoJSONCodec().decode( + b'{"processes":[],"serverAddedField":"ignored"}', + response, + ) + + assert decoded is response + assert list(response.processes) == [] diff --git a/spec/envd/buf-python.gen.yaml b/spec/envd/buf-python.gen.yaml index e84a2a852f..e9046546d6 100644 --- a/spec/envd/buf-python.gen.yaml +++ b/spec/envd/buf-python.gen.yaml @@ -1,15 +1,16 @@ # buf.gen.yaml defines a local generation template. # For details, see https://buf.build/docs/configuration/v1/buf-gen-yaml -version: v1 +version: v2 plugins: - - plugin: python + - remote: buf.build/protocolbuffers/python:v26.1 out: ../../packages/python-sdk/e2b/envd - opt: - - pyi_out=../../packages/python-sdk/e2b/envd - - name: connect-python + - remote: buf.build/protocolbuffers/pyi:v26.1 + out: ../../packages/python-sdk/e2b/envd + - remote: buf.build/connectrpc/python:v0.10.0 out: ../../packages/python-sdk/e2b/envd - path: protoc-gen-connect-python managed: enabled: true - optimize_for: SPEED + override: + - file_option: optimize_for + value: SPEED