lybic-guiagents 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

gui_agents/grpc_app.py

@@ -29,6 +29,7 @@ import platform
 from concurrent import futures
 import grpc
 import uuid
+import datetime
 
 from lybic import LybicClient, LybicAuth, Sandbox
 import gui_agents.cli_app as app
@@ -44,20 +45,19 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
     Implements the Agent gRPC service.
     """
 
-    def __init__(self, max_concurrent_task_num = 1):
+    def __init__(self, max_concurrent_task_num: int = 1, log_dir: str = "runtime"):
         """
         Initialize the AgentServicer with concurrency and runtime state.
         
         Parameters:
             max_concurrent_task_num (int): Maximum number of agent tasks allowed to run concurrently; defaults to 1.
+            log_dir (str): Directory for logging and task-related files.
         """
-        self.lybic_auth: LybicAuth | None = None
         self.max_concurrent_task_num = max_concurrent_task_num
         self.tasks = {}
         self.global_common_config = agent_pb2.CommonConfig(id="global")
         self.task_lock = asyncio.Lock()
-        self.lybic_client: LybicClient | None = None
-        self.sandbox: Sandbox | None = None
+        self.log_dir = log_dir
 
     async def GetAgentTaskStream(self, request, context):
         """
@@ -110,18 +110,55 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
             domain=platform.node(),
         )
 
+    def _setup_task_state(self, task_id: str) -> Registry:
+        """Setup global state and registry for task execution with task isolation"""
+        # Create timestamp-based directory structure like cli_app.py
+        datetime_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        timestamp_dir = Path(self.log_dir) / f"{datetime_str}_{task_id[:8]}"  # Include task_id prefix
+        cache_dir = timestamp_dir / "cache" / "screens"
+        state_dir = timestamp_dir / "state"
+
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        state_dir.mkdir(parents=True, exist_ok=True)
+
+        # Create task-specific registry
+        task_registry = Registry()
+
+        # Register global state for this task in task-specific registry
+        global_state = GlobalState(
+            screenshot_dir=str(cache_dir),
+            tu_path=str(state_dir / "tu.json"),
+            search_query_path=str(state_dir / "search_query.json"),
+            completed_subtasks_path=str(state_dir / "completed_subtasks.json"),
+            failed_subtasks_path=str(state_dir / "failed_subtasks.json"),
+            remaining_subtasks_path=str(state_dir / "remaining_subtasks.json"),
+            termination_flag_path=str(state_dir / "termination_flag.json"),
+            running_state_path=str(state_dir / "running_state.json"),
+            display_info_path=str(timestamp_dir / "display.json"),
+            agent_log_path=str(timestamp_dir / "agent_log.json")
+        )
+
+        # Register in task-specific registry using instance method
+        registry_key = "GlobalStateStore"
+        task_registry.register_instance(registry_key, global_state)
+
+        logger.info(f"Created task-specific registry for task {task_id}")
+
+        return task_registry
+
     async def _run_task(self, task_id: str, backend_kwargs):
         """
         Run the lifecycle of a single agent task: mark it running, execute the agent, record final state, emit stream messages, and unregister the task.
-        
+
         Parameters:
         	task_id (str): Identifier of the task to run.
         	backend_kwargs (dict): Backend configuration passed to the HardwareInterface (e.g., platform, org/api fields, sandbox id).
-        
+
         Notes:
         	- Updates the task entry in self.tasks (status and final_state).
         	- Emits task lifecycle messages via stream_manager and unregisters the task when finished.
         	- Exceptions are caught, the task status is set to "error", and an error message is emitted.
+        	- Supports task cancellation via asyncio.CancelledError.
         """
         async with self.task_lock:
             self.tasks[task_id]["status"] = "running"
@@ -136,23 +173,29 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
             # Send message through stream manager
             await stream_manager.add_message(task_id, "starting", "Task starting")
 
-            # Set task_id for the agent to enable streaming from within the agent
+            # Create task-specific registry
+            task_registry = self._setup_task_state(task_id)
+
+            # Set task_id for the agent. This is needed so that agent.reset() can find the right components.
             if hasattr(agent, 'set_task_id'):
                 agent.set_task_id(task_id)
 
             hwi = HardwareInterface(backend='lybic', **backend_kwargs)
 
+            # We need to set the registry for the main thread context before reset
+            Registry.set_task_registry(task_id, task_registry)
             agent.reset()
+            Registry.remove_task_registry(task_id) # Clean up main thread's local
 
-            # Run the blocking function in a separate thread to avoid blocking the event loop
-            mode: InstanceMode | None = backend_kwargs["mode"]
+            # Run the blocking function in a separate thread, passing the context
+            mode: InstanceMode | None = backend_kwargs.get("mode")
             if mode and mode == InstanceMode.NORMAL:
-                await asyncio.to_thread(app.run_agent_normal,agent, query, hwi, steps, False)
+                await asyncio.to_thread(app.run_agent_normal, agent, query, hwi, steps, False, task_id=task_id, task_registry=task_registry)
             else:
-                await asyncio.to_thread(app.run_agent_fast, agent, query, hwi, steps, False)
+                await asyncio.to_thread(app.run_agent_fast, agent, query, hwi, steps, False, task_id=task_id, task_registry=task_registry)
 
-            global_state: GlobalState = Registry.get("GlobalStateStore")  # type: ignore
-            final_state = global_state.get_running_state()
+            # The final state is now determined inside the thread. We'll assume success if no exception.
+            final_state = "completed"
 
             async with self.task_lock:
                 self.tasks[task_id]["final_state"] = final_state
@@ -164,6 +207,11 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
                 status = final_state if final_state else 'unknown'
                 await stream_manager.add_message(task_id, "finished", f"Task finished with status: {status}")
 
+        except asyncio.CancelledError:
+            logger.info(f"Task {task_id} was cancelled")
+            async with self.task_lock:
+                self.tasks[task_id]["status"] = "cancelled"
+            await stream_manager.add_message(task_id, "cancelled", "Task was cancelled by user request")
         except Exception as e:
             logger.exception(f"Error during task execution for {task_id}: {e}")
             async with self.task_lock:
@@ -171,6 +219,7 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
             await stream_manager.add_message(task_id, "error", f"An error occurred: {e}")
         finally:
             logger.info(f"Task {task_id} processing finished.")
+            # Registry cleanup is now handled within the worker thread
             await stream_manager.unregister_task(task_id)
 
     async def _make_backend_kwargs(self, request):
@@ -190,7 +239,6 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
                 - "endpoint": Lybic API endpoint.
         
         Side effects:
-            - May initialize or replace self.lybic_auth from request.runningConfig.authorizationInfo.
             - May call self._create_sandbox(...) to create or retrieve a sandbox and determine the platform.
         """
         backend_kwargs = {}
@@ -203,33 +251,46 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
         if request.HasField("runningConfig"):
             if request.runningConfig.backend:
                 backend = request.runningConfig.backend
-            if request.runningConfig.HasField("authorizationInfo"):
-                self.lybic_auth = LybicAuth(
-                    org_id=request.runningConfig.authorizationInfo.orgID,
-                    api_key=request.runningConfig.authorizationInfo.apiKey,
-                    endpoint=request.runningConfig.authorizationInfo.apiEndpoint or "https://api.lybic.cn/"
-                )
             backend_kwargs["mode"] = request.runningConfig.mode
 
         platform_str = platform.system()
         sid = ''
+        sandbox_pb = None
+
         if backend == 'lybic':
+            auth_info = (request.runningConfig.authorizationInfo
+                         if request.HasField("runningConfig") and request.runningConfig.HasField("authorizationInfo")
+                         else self.global_common_config.authorizationInfo)
+            if not auth_info or not auth_info.orgID or not auth_info.apiKey:
+                raise ValueError("Lybic backend requires valid authorization (orgID and apiKey)")
+
+            lybic_auth = LybicAuth(
+                org_id=auth_info.orgID,
+                api_key=auth_info.apiKey,
+                endpoint=auth_info.apiEndpoint or "https://api.lybic.cn/"
+            )
+
             if request.HasField("sandbox"):
                 shape = request.sandbox.shapeName
                 sid = request.sandbox.id
                 if sid:
                     logger.info(f"Using existing sandbox with id: {sid}")
+                    sandbox_pb = await self._get_sandbox_pb(sid, lybic_auth)  # if not exist raise NotFound
+                    platform_str = platform_map.get(sandbox_pb.os, platform.system())
                 else:
-                    sid, platform_str = await self._create_sandbox(shape)
+                    sandbox_pb = await self._create_sandbox(shape, lybic_auth)
+                    sid, platform_str = sandbox_pb.id, platform_map.get(sandbox_pb.os, platform.system())
 
                 if request.sandbox.os != agent_pb2.SandboxOS.OSUNDEFINED:
                     platform_str = platform_map.get(request.sandbox.os, platform.system())
             else:
-                sid, platform_str = await self._create_sandbox(shape)
-
+                sandbox_pb = await self._create_sandbox(shape, lybic_auth)
+                sid, platform_str = sandbox_pb.id, platform_map.get(sandbox_pb.os, platform.system())
         else:
-            platform_str = platform_map.get(request.sandbox.os, platform.system())
+            if request.HasField("sandbox") and request.sandbox.os != agent_pb2.SandboxOS.OSUNDEFINED:
+                platform_str = platform_map.get(request.sandbox.os, platform.system())
 
+        backend_kwargs["sandbox"] = sandbox_pb
         backend_kwargs["platform"] = platform_str
         backend_kwargs["precreate_sid"] = sid
 
@@ -250,12 +311,6 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
         return backend_kwargs
 
     async def _make_agent(self,request):
-        # todo: add max_steps support
-        # max_steps = 50
-        # if request.HasField("runningConfig") and request.runningConfig.steps:
-        #     max_steps = request.runningConfig.steps
-
-        # Dynamically build tools_dict based on global config
         """
         Builds and returns an AgentS2 configured for the incoming request by applying model and provider overrides to the tool configurations.
         
@@ -345,7 +400,6 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
             tools_config=tools_config,
         )
 
-
     async def RunAgentInstruction(self, request, context):
         """
         Stream task progress for a newly created instruction-run agent while managing task lifecycle and concurrency.
@@ -390,8 +444,13 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
                 "query": request.instruction,
                 "agent": agent,
                 "max_steps": max_steps,
+                "sandbox": backend_kwargs["sandbox"],
             }
 
+            # This property is used to pass sandbox information.
+            # It has now completed its mission and needs to be deleted, otherwise other backends may crash.
+            del backend_kwargs["sandbox"]
+
             task_future = asyncio.create_task(self._run_task(task_id, backend_kwargs))
             self.tasks[task_id]["future"] = task_future
         try:
@@ -406,6 +465,16 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
             logger.info(f"RunAgentInstruction stream for {task_id} cancelled by client.")
             if task_future:
                 task_future.cancel()
+                # Set cancellation flag in global state for agents to check
+                try:
+                    global_state: GlobalState = Registry.get_from_context("GlobalStateStore", task_id)
+                    if global_state:
+                        global_state.set_running_state("cancelled")
+                        logger.info(f"Set running state to 'cancelled' for task {task_id} due to client disconnect.")
+                    else:
+                        logger.warning(f"Could not find GlobalState for task {task_id} to set cancellation flag on client disconnect.")
+                except Exception as e:
+                    logger.error(f"Error setting cancellation flag for task {task_id} on client disconnect: {e}")
         except Exception as e:
             logger.exception(f"Error in RunAgentInstruction stream for task {task_id}")
             context.set_code(grpc.StatusCode.INTERNAL)
@@ -449,7 +518,11 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
                 "query": request.instruction,
                 "agent": agent,
                 "max_steps": max_steps,
+                "sandbox": backend_kwargs["sandbox"],
             }
+            # This property is used to pass sandbox information.
+            # It has now completed its mission and needs to be deleted, otherwise other backends may crash.
+            del backend_kwargs["sandbox"]
 
             # Start the task in background
             task_future = asyncio.create_task(self._run_task(task_id, backend_kwargs))
@@ -490,6 +563,7 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
             "running": agent_pb2.TaskStatus.RUNNING,
             "fulfilled": agent_pb2.TaskStatus.SUCCESS,
             "rejected": agent_pb2.TaskStatus.FAILURE,
+            "cancelled": agent_pb2.TaskStatus.CANCELLED,
         }
 
         if status == "finished":
@@ -500,6 +574,10 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
             task_status = agent_pb2.TaskStatus.FAILURE
             message = "Task failed with an exception."
             result = ""
+        elif status == "cancelled":
+            task_status = agent_pb2.TaskStatus.CANCELLED
+            message = "Task was cancelled by user request."
+            result = ""
         else:  # pending or running
             task_status = status_map.get(status, agent_pb2.TaskStatus.TASKSTATUSUNDEFINED)
             message = "Task is running."
@@ -510,9 +588,85 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
             status=task_status,
             message=message,
             result=result,
-            sandbox=task_info["request"].sandbox
+            sandbox=task_info["sandbox"]
         )
 
+    async def CancelTask(self, request, context):
+        """
+        Cancel a running task by its taskId.
+
+        If the task exists and is running, it will be cancelled and a success response is returned.
+        If the task is not found or already completed, an appropriate response is returned.
+
+        Parameters:
+            request: CancelTaskRequest containing the taskId to cancel
+            context: gRPC context for setting status codes and details
+
+        Returns:
+            CancelTaskResponse: Response containing taskId, success status, and message
+        """
+        task_id = request.taskId
+        logger.info(f"Received CancelTask request for taskId: {task_id}")
+
+        async with self.task_lock:
+            task_info = self.tasks.get(task_id)
+
+        if not task_info:
+            return agent_pb2.CancelTaskResponse(
+                taskId=task_id,
+                success=False,
+                message=f"Task with ID {task_id} not found."
+            )
+
+        status = task_info["status"]
+        task_future = task_info.get("future")
+
+        # Check if task can be cancelled
+        if status in ["finished", "error"]:
+            return agent_pb2.CancelTaskResponse(
+                taskId=task_id,
+                success=False,
+                message=f"Task {task_id} is already {status} and cannot be cancelled."
+            )
+        elif status == "cancelled":
+            return agent_pb2.CancelTaskResponse(
+                taskId=task_id,
+                success=True,
+                message=f"Task {task_id} was already cancelled."
+            )
+        elif status in ["pending", "running"] and task_future:
+            try:
+                # Cancel the task future
+                task_future.cancel()
+                task_info["status"] = "cancelled"
+
+                # Set cancellation flag in global state for agents to check
+                global_state: GlobalState = Registry.get_from_context("GlobalStateStore", task_id)  # type: ignore
+                global_state.set_running_state("cancelled")
+
+                # Send cancellation message through stream manager
+                await stream_manager.add_message(task_id, "cancelled", "Task was cancelled by user request")
+
+                logger.info(f"Task {task_id} successfully cancelled")
+                return agent_pb2.CancelTaskResponse(
+                    taskId=task_id,
+                    success=True,
+                    message=f"Task {task_id} has been successfully cancelled."
+                )
+            except Exception as e:
+                logger.error(f"Failed to cancel task {task_id}: {e}")
+                return agent_pb2.CancelTaskResponse(
+                    taskId=task_id,
+                    success=False,
+                    message=f"Failed to cancel task {task_id}: {e}"
+                )
+        else:
+            return agent_pb2.CancelTaskResponse(
+                taskId=task_id,
+                success=False,
+                message=f"Task {task_id} is in state '{status}' and cannot be cancelled."
+            )
+
     def _mask_config_secrets(self, config: CommonConfig) -> CommonConfig:
         """
         Return a deep copy of a CommonConfig with sensitive API keys replaced by "********".
@@ -615,20 +769,17 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
         context.set_details(f"Config for task {request.id} not found.")
         return agent_pb2.CommonConfig()
 
-    async def _new_lybic_client(self):
+    def _new_lybic_client(self, lybic_auth: LybicAuth) -> LybicClient:
         """
-        Create and store a new LybicClient using the servicer's current LybicAuth.
-        
-        This replaces the servicer's `lybic_client` attribute with a newly constructed
-        LybicClient initialized from `self.lybic_auth`.
+        Create and return a new LybicClient.
         """
-        self.lybic_client = LybicClient(self.lybic_auth)
+        return LybicClient(lybic_auth)
 
     async def SetGlobalCommonConfig(self, request, context):
         """
-        Set the server's global common configuration and initialize Lybic authorization if provided.
+        Set the server's global common configuration.
         
-        Sets request.commonConfig.id to "global" and stores it as the servicer's global_common_config. If the provided config contains authorizationInfo, initializes or updates self.lybic_auth with the org ID, API key, and API endpoint (defaulting to "https://api.lybic.cn/" when endpoint is empty).
+        Sets request.commonConfig.id to "global" and stores it as the servicer's global_common_config.
         
         Parameters:
             request: gRPC request containing `commonConfig` to apply.
@@ -636,17 +787,14 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
         Returns:
             agent_pb2.SetCommonConfigResponse: Response with `success=True` and the configuration `id`.
         """
+        if os.environ.get("ALLOW_SET_GLOBAL_CONFIG", "0")=="0":
+            context.set_code(grpc.StatusCode.PERMISSION_DENIED)
+            context.set_details("Permission denied.")
+            return agent_pb2.SetCommonConfigResponse()
         logger.info("Setting new global common config.")
         request.commonConfig.id = "global"
         self.global_common_config = request.commonConfig
 
-        if self.global_common_config.HasField("authorizationInfo"):  # lybic
-            self.lybic_auth = LybicAuth(
-                org_id=self.global_common_config.authorizationInfo.orgID,
-                api_key=self.global_common_config.authorizationInfo.apiKey,
-                endpoint=self.global_common_config.authorizationInfo.apiEndpoint or "https://api.lybic.cn/"
-            )
-
         return agent_pb2.SetCommonConfigResponse(success=True, id=self.global_common_config.id)
 
     async def SetGlobalCommonLLMConfig(self, request, context):
@@ -658,6 +806,10 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
         Returns:
             llmConfig: The `LLMConfig` message that was stored in the global configuration.
         """
+        if os.environ.get("ALLOW_SET_GLOBAL_CONFIG", "0")=="0":
+            context.set_code(grpc.StatusCode.PERMISSION_DENIED)
+            context.set_details("Permission denied.")
+            return agent_pb2.LLMConfig()
         if not self.global_common_config.HasField("stageModelConfig"):
             self.global_common_config.stageModelConfig.SetInParent()
         self.global_common_config.stageModelConfig.actionGeneratorModel.CopyFrom(request.llmConfig)
@@ -678,6 +830,10 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
         Returns:
         	LLMConfig: The `llmConfig` that was applied.
         """
+        if os.environ.get("ALLOW_SET_GLOBAL_CONFIG", "0")=="0":
+            context.set_code(grpc.StatusCode.PERMISSION_DENIED)
+            context.set_details("Permission denied.")
+            return agent_pb2.LLMConfig()
         if not self.global_common_config.HasField("stageModelConfig"):
             self.global_common_config.stageModelConfig.SetInParent()
         self.global_common_config.stageModelConfig.groundingModel.CopyFrom(request.llmConfig)
@@ -694,34 +850,79 @@ class AgentServicer(agent_pb2_grpc.AgentServicer):
         Returns:
             The `llmConfig` that was set as the global embedding model.
         """
+        if os.environ.get("ALLOW_SET_GLOBAL_CONFIG", "0")=="0":
+            context.set_code(grpc.StatusCode.PERMISSION_DENIED)
+            context.set_details("Permission denied.")
+            return agent_pb2.LLMConfig()
         if not self.global_common_config.HasField("stageModelConfig"):
             self.global_common_config.stageModelConfig.SetInParent()
         self.global_common_config.stageModelConfig.embeddingModel.CopyFrom(request.llmConfig)
         logger.info(f"Global embedding LLM config updated to: {request.llmConfig.modelName}")
         return request.llmConfig
 
-    async def _create_sandbox(self,shape:str):
+    async def _create_sandbox(self, shape: str, lybic_auth: LybicAuth) -> agent_pb2.Sandbox:
         """
         Create a sandbox with the given shape via the Lybic service and return its identifier and operating system.
-        
+
         Parameters:
             shape (str): The sandbox shape to create (provider-specific size/OS configuration).
-        
+            lybic_auth (LybicAuth): The authentication object for Lybic.
+
         Returns:
-            tuple: A pair (sandbox_id, platform_os) where `sandbox_id` is the created sandbox identifier and `platform_os` is the sandbox operating system string.
-        
-        Raises:
-            Exception: If Lybic authorization is not initialized (call SetGlobalCommonConfig first).
+            agent_pb2.Sandbox: A protobuf message containing sandbox details.
+        """
+        lybic_client = self._new_lybic_client(lybic_auth)
+        sandbox_service = Sandbox(lybic_client)
+        result = await sandbox_service.create(shape=shape)
+        sandbox = await sandbox_service.get(result.id)
+        await lybic_client.close()
+
+        return agent_pb2.Sandbox(
+            id=sandbox.sandbox.id,
+            os=self._lybic_sandbox_os_to_pb_enum(sandbox.sandbox.shape),
+            shapeName=sandbox.sandbox.shapeName,
+            hardwareAcceleratedEncoding=sandbox.sandbox.shape.hardwareAcceleratedEncoding,
+            virtualization=sandbox.sandbox.shape.virtualization,
+            architecture=sandbox.sandbox.shape.architecture,
+        )
+
+    @staticmethod
+    def _lybic_sandbox_os_to_pb_enum(os) -> agent_pb2.SandboxOS:
+        """
+        Converts a sandbox OS string to an enum value.
+        """
+        os_raw = getattr(os, "os", "") or ""
+        os_upper = str(os_raw).upper()
+        if "WIN" in os_upper:
+            os_enum = agent_pb2.SandboxOS.WINDOWS
+        elif "LINUX" in os_upper or "UBUNTU" in os_upper:
+            os_enum = agent_pb2.SandboxOS.LINUX
+        elif "ANDROID" in os_upper:
+            os_enum = agent_pb2.SandboxOS.ANDROID
+        else:
+            os_enum = agent_pb2.SandboxOS.OSUNDEFINED
+        return os_enum
+
+    async def _get_sandbox_pb(self, sid: str, lybic_auth: LybicAuth) -> agent_pb2.Sandbox:
+        """
+        Retrieves sandbox details for a given sandbox ID and returns them as a protobuf message.
         """
-        if not self.lybic_auth:
-            raise Exception("Lybic client not initialized. Please call SetGlobalCommonConfig before")
+        if not lybic_auth:
+            raise ValueError("Lybic client not initialized. Please call SetGlobalCommonConfig before")
 
-        await self._new_lybic_client()
-        if not self.sandbox:
-            self.sandbox = Sandbox(self.lybic_client)
-        result = await self.sandbox.create(shape=shape)
-        sandbox = await self.sandbox.get(result.id)
-        return sandbox.sandbox.id, sandbox.sandbox.shape.os
+        lybic_client = self._new_lybic_client(lybic_auth)
+        sandbox_service = Sandbox(lybic_client)
+        sandbox_details = await sandbox_service.get(sid)
+        await lybic_client.close()
+
+        return agent_pb2.Sandbox(
+            id=sandbox_details.sandbox.id,
+            os=self._lybic_sandbox_os_to_pb_enum(sandbox_details.sandbox.shape),
+            shapeName=sandbox_details.sandbox.shapeName,
+            hardwareAcceleratedEncoding=sandbox_details.sandbox.shape.hardwareAcceleratedEncoding,
+            virtualization=sandbox_details.sandbox.shape.virtualization,
+            architecture=sandbox_details.sandbox.shape.architecture,
+        )
 
 async def serve():
     """
@@ -735,8 +936,8 @@ async def serve():
     """
     port = os.environ.get("GRPC_PORT", 50051)
     max_workers = int(os.environ.get("GRPC_MAX_WORKER_THREADS", 100))
-    # task_num = int(os.environ.get("TASK_MAX_TASKS", 5))
-    servicer = AgentServicer(max_concurrent_task_num=1)
+    task_num = int(os.environ.get("TASK_MAX_TASKS", 5))
+    servicer = AgentServicer(max_concurrent_task_num=task_num, log_dir=app.log_dir)
     server = grpc.aio.server(futures.ThreadPoolExecutor(max_workers))
     agent_pb2_grpc.add_AgentServicer_to_server(servicer, server)
     server.add_insecure_port(f'[::]:{port}')
@@ -752,32 +953,6 @@ def main():
     has_display, pyautogui_available, _ = app.check_display_environment()
     compatible_backends, incompatible_backends = app.get_compatible_backends(has_display, pyautogui_available)
     app.validate_backend_compatibility('lybic', compatible_backends, incompatible_backends)
-    timestamp_dir = os.path.join(app.log_dir, app.datetime_str)
-    cache_dir = os.path.join(timestamp_dir, "cache", "screens")
-    state_dir = os.path.join(timestamp_dir, "state")
-
-    os.makedirs(cache_dir, exist_ok=True)
-    os.makedirs(state_dir, exist_ok=True)
-
-    Registry.register(
-        "GlobalStateStore",
-        GlobalState(
-            screenshot_dir=cache_dir,
-            tu_path=os.path.join(state_dir, "tu.json"),
-            search_query_path=os.path.join(state_dir, "search_query.json"),
-            completed_subtasks_path=os.path.join(state_dir,
-                                                 "completed_subtasks.json"),
-            failed_subtasks_path=os.path.join(state_dir,
-                                              "failed_subtasks.json"),
-            remaining_subtasks_path=os.path.join(state_dir,
-                                                 "remaining_subtasks.json"),
-            termination_flag_path=os.path.join(state_dir,
-                                               "termination_flag.json"),
-            running_state_path=os.path.join(state_dir, "running_state.json"),
-            display_info_path=os.path.join(timestamp_dir, "display.json"),
-            agent_log_path=os.path.join(timestamp_dir, "agent_log.json")
-        )
-    )
     asyncio.run(serve())
 
 if __name__ == '__main__':