plato-sdk-v2 2.7.9__py3-none-any.whl → 2.8.0__py3-none-any.whl

plato/v2/async_/session.py

@@ -23,6 +23,8 @@ if TYPE_CHECKING:
 
 from plato._generated.api.v2.jobs import get_flows as jobs_get_flows
 from plato._generated.api.v2.jobs import public_url as jobs_public_url
+from plato._generated.api.v2.jobs import wait_for_ready as jobs_wait_for_ready
+from plato._generated.api.v2.sessions import add_job as sessions_add_job
 from plato._generated.api.v2.sessions import close as sessions_close
 from plato._generated.api.v2.sessions import connect_network as sessions_connect_network
 from plato._generated.api.v2.sessions import disk_snapshot as sessions_disk_snapshot
@@ -31,6 +33,7 @@ from plato._generated.api.v2.sessions import execute as sessions_execute
 from plato._generated.api.v2.sessions import get_public_url as sessions_get_public_url
 from plato._generated.api.v2.sessions import heartbeat as sessions_heartbeat
 from plato._generated.api.v2.sessions import make as sessions_make
+from plato._generated.api.v2.sessions import remove_job as sessions_remove_job
 from plato._generated.api.v2.sessions import reset as sessions_reset
 from plato._generated.api.v2.sessions import set_date as sessions_set_date
 from plato._generated.api.v2.sessions import setup_sandbox as sessions_setup_sandbox
@@ -39,6 +42,7 @@ from plato._generated.api.v2.sessions import snapshot_store as sessions_snapshot
 from plato._generated.api.v2.sessions import state as sessions_state
 from plato._generated.api.v2.sessions import wait_for_ready as sessions_wait_for_ready
 from plato._generated.models import (
+    AddJobRequest,
     AppApiV2SchemasSessionCreateSnapshotRequest,
     AppApiV2SchemasSessionCreateSnapshotResponse,
     AppApiV2SchemasSessionEvaluateResponse,
@@ -49,10 +53,12 @@ from plato._generated.models import (
     CreateDiskSnapshotResponse,
     CreateSessionFromEnvs,
     CreateSessionFromTask,
+    EnvironmentContext,
     Envs,
     ExecuteCommandRequest,
     ExecuteCommandResponse,
     Flow,
+    RemoveJobRequest,
     ResetSessionRequest,
     ResetSessionResponse,
     RunSessionSource,
@@ -780,6 +786,158 @@ class Session:
 
         return result
 
+    async def add_env(
+        self,
+        env: EnvFromSimulator | EnvFromArtifact | EnvFromResource,
+        *,
+        timeout: int = 1800,
+        heartbeat_timeout: int | None = None,
+        wait_for_ready: bool = True,
+    ) -> Environment:
+        """Add a new environment to this session.
+
+        The new environment will:
+        1. Become part of the session's job group
+        2. Be matched to an available VM via the resource matcher
+        3. Automatically join the session's WireGuard network if one exists
+
+        Args:
+            env: Environment configuration (from Env.simulator(), Env.artifact(), or Env.resource()).
+            timeout: VM timeout in seconds (default: 1800).
+            heartbeat_timeout: Per-VM heartbeat timeout. None=use default (300s), 0=disabled.
+            wait_for_ready: If True, wait for the job to be ready before returning (default: True).
+
+        Returns:
+            Environment object for the new job.
+
+        Raises:
+            RuntimeError: If session is closed or job creation fails.
+            TimeoutError: If wait_for_ready=True and the job doesn't become ready within timeout.
+        """
+        self._check_closed()
+
+        # Auto-generate alias if not set
+        if env.alias is None:
+            existing_aliases = {e.alias for e in self.envs}
+            unique_alias = f"env-{uuid.uuid4().hex[:8]}"
+            while unique_alias in existing_aliases:
+                unique_alias = f"env-{uuid.uuid4().hex[:8]}"
+            env.alias = unique_alias
+
+        # Build request
+        request = AddJobRequest(
+            env=env,
+            timeout=timeout,
+            heartbeat_timeout=heartbeat_timeout,
+        )
+
+        # Call the add_job API
+        response = await sessions_add_job.asyncio(
+            client=self._http,
+            session_id=self.session_id,
+            body=request,
+            x_api_key=self._api_key,
+        )
+
+        # Check for failures
+        if not response.env.success:
+            raise RuntimeError(f"Failed to add job: {response.env.error}")
+
+        if not response.env.job_id:
+            raise RuntimeError("Backend did not return job_id for new environment")
+
+        job_id = response.env.job_id
+
+        # Wait for the job to be ready if requested
+        if wait_for_ready:
+            ready_response = await jobs_wait_for_ready.asyncio(
+                client=self._http,
+                job_id=job_id,
+                timeout=timeout,
+                x_api_key=self._api_key,
+            )
+
+            if not ready_response.ready:
+                error = ready_response.error or "Unknown error"
+                raise TimeoutError(f"Job {job_id} did not become ready: {error}")
+
+        # Update internal context with the new environment
+        new_env_context = EnvironmentContext(
+            job_id=job_id,
+            alias=env.alias,
+            artifact_id=response.env.artifact_id,
+            simulator=getattr(env, "simulator", None),
+        )
+
+        # Add to context's envs list
+        if self._context.envs is None:
+            self._context.envs = []
+        self._context.envs.append(new_env_context)
+
+        # Reset cached envs to force rebuild
+        self._envs = None
+
+        # Create and return the Environment object
+        new_environment = Environment(
+            session=self,
+            job_id=job_id,
+            alias=env.alias,
+            artifact_id=response.env.artifact_id,
+        )
+
+        logger.info(f"Added job {job_id} (alias={env.alias}) to session {self.session_id}")
+        return new_environment
+
+    async def remove_env(self, env: Environment | str) -> None:
+        """Remove an environment from this session.
+
+        This will:
+        1. Remove the job from the session's network (if connected)
+        2. Shut down the VM associated with the job
+        3. Cancel the job in the system
+
+        Args:
+            env: Environment object or alias string to remove.
+
+        Raises:
+            RuntimeError: If session is closed or removal fails.
+            ValueError: If environment not found in session.
+        """
+        self._check_closed()
+
+        # Resolve to job_id
+        if isinstance(env, str):
+            # Find by alias
+            found_env = self.get_env(env)
+            if not found_env:
+                raise ValueError(f"Environment with alias '{env}' not found in session")
+            job_id = found_env.job_id
+            alias = env
+        else:
+            job_id = env.job_id
+            alias = env.alias
+
+        # Call the remove_job API
+        request = RemoveJobRequest(job_id=job_id)
+        response = await sessions_remove_job.asyncio(
+            client=self._http,
+            session_id=self.session_id,
+            body=request,
+            x_api_key=self._api_key,
+        )
+
+        if not response.success:
+            raise RuntimeError(f"Failed to remove job {job_id}")
+
+        # Update internal context - remove the environment
+        if self._context.envs:
+            self._context.envs = [e for e in self._context.envs if e.job_id != job_id]
+
+        # Reset cached envs to force rebuild
+        self._envs = None
+
+        logger.info(f"Removed job {job_id} (alias={alias}) from session {self.session_id}")
+
     async def cleanup_databases(self) -> SessionCleanupResult:
         """Clean up database audit logs for all environments.
 

plato/v2/sync/session.py

@@ -22,6 +22,8 @@ if TYPE_CHECKING:
 
 from plato._generated.api.v2.jobs import get_flows as jobs_get_flows
 from plato._generated.api.v2.jobs import public_url as jobs_public_url
+from plato._generated.api.v2.jobs import wait_for_ready as jobs_wait_for_ready
+from plato._generated.api.v2.sessions import add_job as sessions_add_job
 from plato._generated.api.v2.sessions import close as sessions_close
 from plato._generated.api.v2.sessions import connect_network as sessions_connect_network
 from plato._generated.api.v2.sessions import disk_snapshot as sessions_disk_snapshot
@@ -29,6 +31,7 @@ from plato._generated.api.v2.sessions import evaluate as sessions_evaluate
 from plato._generated.api.v2.sessions import execute as sessions_execute
 from plato._generated.api.v2.sessions import heartbeat as sessions_heartbeat
 from plato._generated.api.v2.sessions import make as sessions_make
+from plato._generated.api.v2.sessions import remove_job as sessions_remove_job
 from plato._generated.api.v2.sessions import reset as sessions_reset
 from plato._generated.api.v2.sessions import set_date as sessions_set_date
 from plato._generated.api.v2.sessions import setup_sandbox as sessions_setup_sandbox
@@ -37,6 +40,7 @@ from plato._generated.api.v2.sessions import snapshot_store as sessions_snapshot
 from plato._generated.api.v2.sessions import state as sessions_state
 from plato._generated.api.v2.sessions import wait_for_ready as sessions_wait_for_ready
 from plato._generated.models import (
+    AddJobRequest,
     AppApiV2SchemasSessionCreateSnapshotRequest,
     AppApiV2SchemasSessionCreateSnapshotResponse,
     AppApiV2SchemasSessionEvaluateResponse,
@@ -47,10 +51,12 @@ from plato._generated.models import (
     CreateDiskSnapshotResponse,
     CreateSessionFromEnvs,
     CreateSessionFromTask,
+    EnvironmentContext,
     Envs,
     ExecuteCommandRequest,
     ExecuteCommandResponse,
     Flow,
+    RemoveJobRequest,
     ResetSessionRequest,
     ResetSessionResponse,
     RunSessionSource,
@@ -690,6 +696,158 @@ class Session:
 
         return result
 
+    def add_env(
+        self,
+        env: EnvFromSimulator | EnvFromArtifact | EnvFromResource,
+        *,
+        timeout: int = 1800,
+        heartbeat_timeout: int | None = None,
+        wait_for_ready: bool = True,
+    ) -> Environment:
+        """Add a new environment to this session.
+
+        The new environment will:
+        1. Become part of the session's job group
+        2. Be matched to an available VM via the resource matcher
+        3. Automatically join the session's WireGuard network if one exists
+
+        Args:
+            env: Environment configuration (from Env.simulator(), Env.artifact(), or Env.resource()).
+            timeout: VM timeout in seconds (default: 1800).
+            heartbeat_timeout: Per-VM heartbeat timeout. None=use default (300s), 0=disabled.
+            wait_for_ready: If True, wait for the job to be ready before returning (default: True).
+
+        Returns:
+            Environment object for the new job.
+
+        Raises:
+            RuntimeError: If session is closed or job creation fails.
+            TimeoutError: If wait_for_ready=True and the job doesn't become ready within timeout.
+        """
+        self._check_closed()
+
+        # Auto-generate alias if not set
+        if env.alias is None:
+            existing_aliases = {e.alias for e in self.envs}
+            unique_alias = f"env-{uuid.uuid4().hex[:8]}"
+            while unique_alias in existing_aliases:
+                unique_alias = f"env-{uuid.uuid4().hex[:8]}"
+            env.alias = unique_alias
+
+        # Build request
+        request = AddJobRequest(
+            env=env,
+            timeout=timeout,
+            heartbeat_timeout=heartbeat_timeout,
+        )
+
+        # Call the add_job API
+        response = sessions_add_job.sync(
+            client=self._http,
+            session_id=self.session_id,
+            body=request,
+            x_api_key=self._api_key,
+        )
+
+        # Check for failures
+        if not response.env.success:
+            raise RuntimeError(f"Failed to add job: {response.env.error}")
+
+        if not response.env.job_id:
+            raise RuntimeError("Backend did not return job_id for new environment")
+
+        job_id = response.env.job_id
+
+        # Wait for the job to be ready if requested
+        if wait_for_ready:
+            ready_response = jobs_wait_for_ready.sync(
+                client=self._http,
+                job_id=job_id,
+                timeout=timeout,
+                x_api_key=self._api_key,
+            )
+
+            if not ready_response.ready:
+                error = ready_response.error or "Unknown error"
+                raise TimeoutError(f"Job {job_id} did not become ready: {error}")
+
+        # Update internal context with the new environment
+        new_env_context = EnvironmentContext(
+            job_id=job_id,
+            alias=env.alias,
+            artifact_id=response.env.artifact_id,
+            simulator=getattr(env, "simulator", None),
+        )
+
+        # Add to context's envs list
+        if self._context.envs is None:
+            self._context.envs = []
+        self._context.envs.append(new_env_context)
+
+        # Reset cached envs to force rebuild
+        self._envs = None
+
+        # Create and return the Environment object
+        new_environment = Environment(
+            session=self,
+            job_id=job_id,
+            alias=env.alias,
+            artifact_id=response.env.artifact_id,
+        )
+
+        logger.info(f"Added job {job_id} (alias={env.alias}) to session {self.session_id}")
+        return new_environment
+
+    def remove_env(self, env: Environment | str) -> None:
+        """Remove an environment from this session.
+
+        This will:
+        1. Remove the job from the session's network (if connected)
+        2. Shut down the VM associated with the job
+        3. Cancel the job in the system
+
+        Args:
+            env: Environment object or alias string to remove.
+
+        Raises:
+            RuntimeError: If session is closed or removal fails.
+            ValueError: If environment not found in session.
+        """
+        self._check_closed()
+
+        # Resolve to job_id
+        if isinstance(env, str):
+            # Find by alias
+            found_env = self.get_env(env)
+            if not found_env:
+                raise ValueError(f"Environment with alias '{env}' not found in session")
+            job_id = found_env.job_id
+            alias = env
+        else:
+            job_id = env.job_id
+            alias = env.alias
+
+        # Call the remove_job API
+        request = RemoveJobRequest(job_id=job_id)
+        response = sessions_remove_job.sync(
+            client=self._http,
+            session_id=self.session_id,
+            body=request,
+            x_api_key=self._api_key,
+        )
+
+        if not response.success:
+            raise RuntimeError(f"Failed to remove job {job_id}")
+
+        # Update internal context - remove the environment
+        if self._context.envs:
+            self._context.envs = [e for e in self._context.envs if e.job_id != job_id]
+
+        # Reset cached envs to force rebuild
+        self._envs = None
+
+        logger.info(f"Removed job {job_id} (alias={alias}) from session {self.session_id}")
+
     def login(
         self,
         browser: Browser,

{plato_sdk_v2-2.7.9.dist-info → plato_sdk_v2-2.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plato-sdk-v2
-Version: 2.7.9
+Version: 2.8.0
 Summary: Python SDK for the Plato API
 Author-email: Plato <support@plato.so>
 License-Expression: MIT

{plato_sdk_v2-2.7.9.dist-info → plato_sdk_v2-2.8.0.dist-info}/RECORD

@@ -506,7 +506,7 @@ plato/v2/async_/chronos.py,sha256=WeqYF3HIKs7hV9LNZb2GlDS1yP6b422DZKtNuPxdL34,12
 plato/v2/async_/client.py,sha256=IhiEiwbLNPBr9JJilw4uz7MLKXY_rUZpYGYC1dX-UfA,5186
 plato/v2/async_/environment.py,sha256=M5IeWYLwREOIyuS2zqgBSqHE_x66_OZXrevA9Rkc8Is,5825
 plato/v2/async_/flow_executor.py,sha256=Tl4nRu1ZPWJFNNxyTGy-PxvebZEUD18ZDaz8T2chtzU,14188
-plato/v2/async_/session.py,sha256=0zErrTEm9-xxIz6pzp_y2stzn498KXVle7n51HVhyhc,37910
+plato/v2/async_/session.py,sha256=31sJeQNKlUL2g0GRsJAG1OmoQp5LJxwfCECuUBh87CQ,43459
 plato/v2/sync/__init__.py,sha256=3WXLqem7GbicVevLD1lCarr7YI1m6j7-AmJf9OVKKgc,521
 plato/v2/sync/artifact.py,sha256=wTLC-tugG128wLvh-JqNPb0zsw5FXEJlZNahurSWink,1169
 plato/v2/sync/chronos.py,sha256=ChXpasjRzAZjoYTimpPqYydnwEk-IgdxR0SDXDOZbUM,12078
@@ -514,7 +514,7 @@ plato/v2/sync/client.py,sha256=rsrU7_RhE-syf3FMNw5LaxmF7rYw2GBzC_TPpd-6thk,4986
 plato/v2/sync/environment.py,sha256=WnDzbyEHpwCSEP8XnfNSjIYS7rt7lYR4HGJjzprZmTQ,5066
 plato/v2/sync/flow_executor.py,sha256=N41-WCWIJVcCR2UmPUEiK7roNacYoeONkRXpR7lUgT8,13941
 plato/v2/sync/sandbox.py,sha256=yRwkms1_qsHjhbLZQ7FiaLgFme2tMw8YM2xg2-ZZS04,53734
-plato/v2/sync/session.py,sha256=NL_qEXnXjZfAsuFgpq72pVmHa9ouo8W3HkJ5vYaNSQU,29555
+plato/v2/sync/session.py,sha256=FYVxgGZ3eL_YpoJiUgJamrt-jVaBqJITzFzepOKMRjA,35065
 plato/v2/utils/__init__.py,sha256=XLeFFsjXkm9g2raMmo7Wt4QN4hhCrNZDJKnpffJ4LtM,38
 plato/v2/utils/db_cleanup.py,sha256=JMzAAJz0ZnoUXtd8F4jpQmBpJpos2__RkgN_cuEearg,8692
 plato/v2/utils/gateway_tunnel.py,sha256=eWgwf4VV8-jx6iCuHFgCISsAOVmNOOjCB56EuZLsnOA,7171
@@ -526,7 +526,7 @@ plato/worlds/base.py,sha256=-RR71bSxEFI5yydtrtq-AAbuw98CIjvmrbztqzB9oIc,31041
 plato/worlds/build_hook.py,sha256=KSoW0kqa5b7NyZ7MYOw2qsZ_2FkWuz0M3Ru7AKOP7Qw,3486
 plato/worlds/config.py,sha256=O1lUXzxp-Z_M7izslT8naXgE6XujjzwYFFrDDzUOueI,12736
 plato/worlds/runner.py,sha256=r9B2BxBae8_dM7y5cJf9xhThp_I1Qvf_tlPq2rs8qC8,4013
-plato_sdk_v2-2.7.9.dist-info/METADATA,sha256=kEHstjuAW3N6cxIjGfaXdAHMfwfVzugWzZFAHkoaVsg,8652
-plato_sdk_v2-2.7.9.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-plato_sdk_v2-2.7.9.dist-info/entry_points.txt,sha256=iynJvTkU7E4MZNtSozVF0Wh083yPm6cuKV362Ol_ez8,133
-plato_sdk_v2-2.7.9.dist-info/RECORD,,
+plato_sdk_v2-2.8.0.dist-info/METADATA,sha256=ouPYAxE-Lp1XW9dhodGVz5E-cPdN-jEFRPjkYnsooNM,8652
+plato_sdk_v2-2.8.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+plato_sdk_v2-2.8.0.dist-info/entry_points.txt,sha256=iynJvTkU7E4MZNtSozVF0Wh083yPm6cuKV362Ol_ez8,133
+plato_sdk_v2-2.8.0.dist-info/RECORD,,