```diff
--- test_files/356-original.txt	2025-03-07 19:06:49
+++ test_files/356-modified.txt	2025-03-07 19:06:49
@@ -5,20 +5,28 @@
 import typing as t
 import uuid
 
-import gql
-import gql.transport
 import requests
-from gql.transport.requests import RequestsHTTPTransport
+import typing_extensions as te
 
-from composio.constants import ENV_COMPOSIO_API_KEY, ENV_COMPOSIO_BASE_URL
+from composio.tools.env.constants import DEFAULT_IMAGE
 from composio.utils.logging import WithLogger
 
 
+try:
+    import gql
+    import gql.transport
+    from gql.transport.requests import RequestsHTTPTransport
+
+    FLYIO_DEPENDENCIES_INSTALLED = True
+except ImportError:
+    RequestsHTTPTransport = t.Any
+    FLYIO_DEPENDENCIES_INSTALLED = False
+
+
 FLY_API = "https://api.machines.dev"
 FLY_GRAPHQL_API = "https://api.fly.io/graphql"
 API_VERSION = "/v1"
 BASE_URL = FLY_API + API_VERSION
-TOOLSERVER_IMAGE = "angrybayblade/composio"
 
 
 ALLOCATE_IP_QUERY = """mutation {
@@ -51,11 +59,44 @@
 """
 
 ENV_FLY_API_TOKEN = "FLY_API_TOKEN"
-ENV_GITHUB_ACCESS_TOKEN = "GITHUB_ACCESS_TOKEN"
-ENV_ACCESS_TOKEN = "ACCESS_TOKEN"
-ENV_E2B_TEMPLATE = "E2B_TEMPLATE"
 
 
+TOOLING_ERVICE = {
+    "ports": [
+        {"port": 8000, "handlers": ["tls", "http"]},
+    ],
+    "protocol": "tcp",
+    "internal_port": 8000,
+}
+
+
+class ExternalPortConfig(te.TypedDict):
+    """External port config."""
+
+    port: int
+    """External port to map (Port 8000 is reserved for the tooling server)."""
+
+    handlers: t.List[str]
+    """List of protocol handlers (`http` will be used if not provided)."""
+
+
+class PortRequest(te.TypedDict):
+    """
+    Port request
+
+    Read more at: https://fly.io/docs/machines/api/machines-resource/#create-a-machine-with-services
+    """
+
+    ports: t.List[ExternalPortConfig]
+    """List of public port configurations (Port 8000 is reserved for the tooling server)."""
+
+    internal_port: int
+    """Internal port number (Port 8000 is reserved for the tooling server)."""
+
+    protocol: t.Literal["tcp", "udp"]
+    """List of protocol handlers."""
+
+
 class FlyIO(WithLogger):
     """FlyIO client."""
 
@@ -66,39 +107,25 @@
         access_token: str,
         image: t.Optional[str] = None,
         flyio_token: t.Optional[str] = None,
-        api_key: t.Optional[str] = None,
-        base_url: t.Optional[str] = None,
-        env: t.Optional[t.Dict[str, str]] = None,
+        environment: t.Optional[t.Dict] = None,
+        ports: t.Optional[t.List[PortRequest]] = None,
     ) -> None:
         """Initialize FlyIO client."""
         super().__init__()
-        api_key = api_key or os.environ.get(ENV_COMPOSIO_API_KEY)
-        if api_key is None:
+        flyio_token = flyio_token or os.environ.get(ENV_FLY_API_TOKEN)
+        if flyio_token is None:
             raise ValueError(
-                "`api_key` cannot be `None` when initializing E2BWorkspace"
+                "FlyIO API Key is required for using FlyIO workspace, "
+                f"You can export it as `{ENV_FLY_API_TOKEN}`"
             )
 
-        base_url = base_url or os.environ.get(ENV_COMPOSIO_BASE_URL)
-        if base_url is None:
-            raise ValueError(
-                "`base_url` cannot be `None` when initializing E2BWorkspace"
-            )
-
-        github_access_token = os.environ.get(ENV_GITHUB_ACCESS_TOKEN)
-        if github_access_token is None:
-            raise ValueError(
-                f"Please export your github access token as `{ENV_GITHUB_ACCESS_TOKEN}`"
-            )
-
-        flyio_token = flyio_token or os.environ.get(ENV_FLY_API_TOKEN)
-        if flyio_token is None:
-            raise ValueError("FlyIO API Key is required for using FlyIO workspace")
-
+        self.ports = ports or []
+        self.environment = environment or {}
         self.flyio_token = flyio_token
         self.access_token = access_token
-        self.image = image or TOOLSERVER_IMAGE
+        self.image = image or DEFAULT_IMAGE
         self.app_name = f"composio-{uuid.uuid4().hex.replace('-', '')}"
-        self.url = f"https://{self.app_name}.fly.dev/api"
+        self.url = f"https://{self.app_name}.fly.dev:8000/api"
         self.gql = gql.Client(
             transport=RequestsHTTPTransport(
                 url=FLY_GRAPHQL_API,
@@ -107,13 +134,6 @@
                 },
             )
         )
-        self.env = {
-            **(env or {}),
-            ENV_COMPOSIO_API_KEY: api_key,
-            ENV_COMPOSIO_BASE_URL: base_url,
-            ENV_GITHUB_ACCESS_TOKEN: github_access_token,
-            ENV_ACCESS_TOKEN: self.access_token,
-        }
 
     def _request(
         self,
@@ -188,13 +208,12 @@
 
     def _wait_for_machine(self) -> None:
         """Wait for machine to get started."""
-        while True:
+        deadline = time.time() + float(os.environ.get("WORKSPACE_WAIT_TIMEOUT", 60.0))
+        while time.time() < deadline:
             try:
                 requests.get(
                     self.url,
-                    headers={
-                        "x-api-key": self.access_token,
-                    },
+                    headers={"x-api-key": self.access_token},
                     timeout=30.0,
                 ).content.decode()
                 return
@@ -209,17 +228,8 @@
                 json={
                     "config": {
                         "image": self.image,
-                        "env": self.env,
-                        "services": [
-                            {
-                                "ports": [
-                                    {"port": 443, "handlers": ["tls", "http"]},
-                                    {"port": 80, "handlers": ["http"]},
-                                ],
-                                "protocol": "tcp",
-                                "internal_port": 8000,
-                            }
-                        ],
+                        "env": self.environment,
+                        "services": [TOOLING_ERVICE, *self.ports],
                         "guest": {
                             "cpu_kind": "shared",
                             "cpus": 1,
```
