hyperpocket 0.0.1__py3-none-any.whl

hyperpocket/pocket_auth.py

@@ -0,0 +1,344 @@
+import asyncio
+import enum
+import uuid
+from typing import Optional, Type
+
+from hyperpocket.auth import AuthProvider, PREBUILT_AUTH_HANDLERS
+from hyperpocket.auth.context import AuthContext
+from hyperpocket.auth.handler import AuthHandlerInterface, AuthenticateRequest
+from hyperpocket.config import config, pocket_logger
+from hyperpocket.futures import FutureStore
+from hyperpocket.session import SESSION_STORAGE_LIST
+from hyperpocket.session.interface import SessionStorageInterface
+
+
+class AuthState(enum.Enum):
+    SKIP_AUTH = "skip_auth"
+    DO_AUTH = "do_auth"
+    DO_REFRESH = "do_refresh"
+    NO_SESSION = "no_session"
+    PENDING_RESOLVE = "pending_resolve"
+    RESOLVED = "resolved"
+
+
+class PocketAuth(object):
+    handlers: dict[str, AuthHandlerInterface]
+    session_storage: SessionStorageInterface
+
+    def __init__(
+            self,
+            handlers: Optional[list[Type[AuthHandlerInterface]]] = None,
+            session_storage: Optional[SessionStorageInterface] = None,
+            use_prebuilt_handlers: bool = None):
+        if config.auth.use_prebuilt_auth or use_prebuilt_handlers:
+            handlers = PREBUILT_AUTH_HANDLERS + (handlers or [])
+        handler_impls = [C() for C in handlers] if handlers else []
+        self.handlers = {handler.name: handler for handler in handler_impls}
+        if session_storage:
+            self.session_storage = session_storage
+        else:
+            for session_type in SESSION_STORAGE_LIST:
+                if session_type.session_storage_type() == config.session.session_type:
+                    session_config = getattr(config.session, config.session.session_type.value)
+
+                    pocket_logger.info(f"init {session_type.session_storage_type()} session storage..")
+                    self.session_storage = session_type(session_config)
+
+            if self.session_storage is None:
+                pocket_logger.error(f"not supported session type({config.session.session_type})")
+                raise RuntimeError(f"Not Supported Session Type({config.session.session_type})")
+
+    def make_request(self,
+                     auth_scopes: list[str] = None,
+                     auth_handler_name: Optional[str] = None,
+                     auth_provider: Optional[AuthProvider] = None,
+                     **kwargs) -> AuthenticateRequest:
+        """
+        Make AuthenticationRequest based on authentication handler.
+
+        Args:
+            auth_scopes (list[str]): list of auth scopes
+            auth_handler_name (Optional[str]): auth handler name
+            auth_provider (Optional[AuthProvider]): auth provider
+
+        Returns:
+            AuthenticateRequest: authenticate request of the handler
+        """
+        handler = self.find_handler_instance(auth_handler_name, auth_provider)
+        return handler.make_request(auth_scopes, **kwargs)
+
+    def check(self,
+              auth_req: AuthenticateRequest,
+              auth_handler_name: Optional[str] = None,
+              auth_provider: Optional[AuthProvider] = None,
+              thread_id: str = "default",
+              profile: str = 'default',
+              *args, **kwargs
+              ) -> AuthState:
+        """
+        Check current authentication state.
+
+        The `AuthState` includes 5 states:
+
+        - `NO_SESSION` : the session does not exist. it needs to create new session.
+        - `PENDING_RESOLVE` : the session already exists. but it needs user interaction.
+        - `RESOLVED` : the session already exists, and user interaction has already been completed. waiting for authentication to be continue.
+        - `SKIP_AUTH` : the session already exists, and the request can be processed using this session. so authentication is not required.
+        - `DO_AUTH` : the session already exists, but the request can't be processed using this session. so authentication is required.
+        - `DO_REFRESH` : the session already exists, but the session is expired. so refreshing session is required.
+
+        Args:
+            auth_req (AuthenticateRequest): authenticate request
+            auth_handler_name (Optional[str]): auth handler name
+            auth_provider (Optional[AuthProvider]): auth provider
+            thread_id (Optional[str]): thread id
+            profile (Optional[str]): profile name
+
+        Returns:
+            AuthState: current authentication state
+        """
+        handler = self.find_handler_instance(auth_handler_name, auth_provider)
+        session = self.session_storage.get(handler.provider(), thread_id, profile)
+
+        if not session:
+            return AuthState.NO_SESSION
+
+        if session.auth_resolve_uid:
+            future_data = FutureStore.get_future(session.auth_resolve_uid)
+            if future_data is not None and future_data.future.done():
+                return AuthState.RESOLVED
+
+            return AuthState.PENDING_RESOLVE
+
+        if not session.is_auth_applicable(auth_provider_name=handler.provider().name, auth_req=auth_req):
+            return AuthState.DO_AUTH
+
+        if session.is_near_expires():
+            return AuthState.DO_REFRESH
+
+        return AuthState.SKIP_AUTH
+
+    def prepare(self,
+                auth_req: AuthenticateRequest,
+                auth_handler_name: Optional[str] = None,
+                auth_provider: Optional[AuthProvider] = None,
+                thread_id: str = "default",
+                profile: str = 'default',
+                **kwargs) -> Optional[str]:
+        """
+        Prepare authentication.
+
+        - If the session is pending(e.g., PENDING_RESOLVE), return the existing URL.
+        - If the session is not created or not applicable(e.g., NO_SESSION, DO_AUTH), create a new session and future, then return the authentication URL.
+        - Other cases (e.g., DO_REFRESH, SKIP_AUTH, RESOLVED) is not handled in this method , it just returns None.
+
+        Args:
+            auth_req (AuthenticateRequest): authenticate request
+            auth_handler_name (Optional[str]): auth handler name
+            auth_provider (Optional[AuthProvider]): auth provider
+            thread_id (Optional[str]): thread id
+            profile (Optional[str]): profile name
+
+        Returns:
+            Optional[str]: authentication URL
+        """
+        auth_state = self.check(
+            auth_req=auth_req,
+            auth_handler_name=auth_handler_name,
+            auth_provider=auth_provider,
+            thread_id=thread_id,
+            profile=profile,
+            **kwargs,
+        )
+
+        pocket_logger.debug(
+            f"[thread_id({thread_id}):profile({profile})] {auth_provider.name} provider current auth state in prepare : {auth_state}")
+        if auth_state in [AuthState.SKIP_AUTH, AuthState.DO_REFRESH, AuthState.RESOLVED]:
+            return None
+
+        handler = self.find_handler_instance(auth_handler_name, auth_provider)
+        scope = handler.recommended_scopes().union(auth_req.auth_scopes)
+        session = self.session_storage.get(handler.provider(), thread_id, profile)
+        if session:
+            scope = scope.union(session.auth_scopes)
+
+        modified_req = auth_req.model_copy(update={'auth_scopes': scope})
+
+        # session in pending
+        if session and session.auth_resolve_uid:
+            pocket_logger.debug(
+                f"[thread_id({thread_id}):profile({profile})] already exists pending session(auth_resolve_uid:{session.auth_resolve_uid}).")
+            future_uid = session.auth_resolve_uid
+
+            # update session, in case of requesting new scopes before session pending resolved.
+            self._upsert_pending_session(
+                auth_handler=handler,
+                future_uid=session.auth_resolve_uid,
+                profile=profile,
+                thread_id=thread_id,
+                scope=scope
+            )
+        else:  # create new pending session
+            future_uid = str(uuid.uuid4())
+            self._upsert_pending_session(
+                auth_handler=handler,
+                future_uid=future_uid,
+                profile=profile,
+                thread_id=thread_id,
+                scope=scope)
+
+            pocket_logger.debug(
+                f"[thread_id({thread_id}):profile({profile})] create new pending session(auth_resolve_uid:{future_uid}).")
+            asyncio.create_task(self._check_session_pending_resolved(handler, thread_id, profile))
+
+        prepare_url = handler.prepare(modified_req, thread_id, profile, future_uid, **kwargs)
+        pocket_logger.debug(
+            f"[thread_id({thread_id}):profile({profile})] auth_handler({auth_handler_name})'s prepare_url : {prepare_url}.")
+
+        return prepare_url
+
+    async def authenticate_async(self,
+                                 auth_req: AuthenticateRequest,
+                                 auth_handler_name: Optional[str] = None,
+                                 auth_provider: Optional[AuthProvider] = None,
+                                 thread_id: str = "default",
+                                 profile: str = 'default',
+                                 **kwargs) -> AuthContext:
+        """
+        Performing authentication.
+        It is performing authentication. and save the session in session storage.
+        And return `AuthContext`. `AuthContext` has only necessary fields of session to invoke tool.
+
+        - If auth state is SKIP_AUTH, return the existing AuthContext.
+        - If auth state is DO_REFRESH, it refreshes session, and update session.
+        - If auth state is PENDING_RESOLVE or RESOLVED, it performs authentication. in pending_resolve state, it waits for user interaction.
+        - Other cases(NO_SESSION or DO_AUTH) is not handled in this method , raise RuntimeError
+
+        Args:
+            auth_req (AuthenticateRequest): authenticate request
+            auth_handler_name (Optional[str]): auth handler name
+            auth_provider (Optional[AuthProvider]): auth provider
+            thread_id (Optional[str]): thread id
+            profile (Optional[str]): profile name
+
+        Returns:
+            AuthContext: authentication context
+        """
+        auth_state = self.check(
+            auth_req=auth_req,
+            auth_handler_name=auth_handler_name,
+            auth_provider=auth_provider,
+            thread_id=thread_id,
+            profile=profile,
+            **kwargs,
+        )
+        handler = self.find_handler_instance(auth_handler_name, auth_provider)
+
+        session = self.session_storage.get(handler.provider(), thread_id, profile)
+        if session is None:
+            pocket_logger.warning(
+                f"[thread_id({thread_id}):profile({profile})] Session can't find. session should exist in 'authenticate'.")
+            raise RuntimeError(
+                f"[thread_id({thread_id}):profile({profile})] Session can't find. session should exist in 'authenticate'.")
+
+        pocket_logger.debug(
+            f"[thread_id({thread_id}):profile({profile})] auth_handler({auth_handler_name})'s auth state : {auth_state}")
+        try:
+            if auth_state == AuthState.SKIP_AUTH:
+                context = session.auth_context
+            elif auth_state == AuthState.DO_REFRESH:
+                try:
+                    context = await asyncio.wait_for(
+                        handler.refresh(auth_req=auth_req, context=session.auth_context, **kwargs), timeout=300)
+                except Exception as e:
+                    self.session_storage.delete(handler.provider(), thread_id, profile)
+                    FutureStore.delete_future(session.auth_resolve_uid)
+
+                    pocket_logger.warning(
+                        f"[thread_id({thread_id}):profile({profile})] auth_handler({auth_handler_name}) failed to refresh the token.")
+                    raise RuntimeError("Failed to refresh the token. Please re-authenticate.") from e
+            elif auth_state == AuthState.PENDING_RESOLVE or auth_state == AuthState.RESOLVED:
+                future_uid = session.auth_resolve_uid
+                context = await asyncio.wait_for(
+                    handler.authenticate(auth_req=auth_req, future_uid=future_uid, **kwargs), timeout=300)
+            else:
+                # maybe auth_state is either AuthState.NO_SESSION or AuthState.DO_AUTH
+                pocket_logger.warning(
+                    f"[thread_id({thread_id}):profile({profile})] Invalid State. 'authenticate' cannot be reached while in state {auth_state}")
+                raise RuntimeError(
+                    f"[thread_id({thread_id}):profile({profile})] Invalid State. 'authenticate' cannot be reached while in state {auth_state}")
+
+            session = await self._set_session_active(
+                context=context,
+                provider=handler.provider(),
+                profile=profile,
+                thread_id=thread_id
+            )
+
+            return session.auth_context
+        except asyncio.TimeoutError as e:
+            pocket_logger.warning(f"Authentication Timeout. {session.auth_resolve_uid}")
+            self.delete_session(handler.provider(), thread_id, profile)
+            FutureStore.delete_future(session.auth_resolve_uid)
+            raise e
+
+    def delete_session(self, auth_provider: AuthProvider, thread_id: str = "default", profile: str = "default") -> bool:
+        return self.session_storage.delete(auth_provider, thread_id, profile)
+
+    def get_auth_context(self, auth_provider: AuthProvider, thread_id: str = "default", profile: str = "default",
+                         **kwargs) -> Optional[AuthContext]:
+        session = self.session_storage.get(auth_provider, thread_id, profile, **kwargs)
+        if session is None:
+            return None
+
+        return session.auth_context
+
+    def find_handler_instance(self, name: Optional[str] = None,
+                              auth_provider: Optional[AuthProvider] = None) -> AuthHandlerInterface:
+        if name:
+            return self.handlers[name]
+        if auth_provider:
+            for handler in self.handlers.values():
+                if handler.provider() == auth_provider and handler.provider_default():
+                    return handler
+        raise ValueError("No handler found")
+
+    async def _check_session_pending_resolved(self, auth_handler: AuthHandlerInterface, thread_id: str = "default",
+                                              profile: str = "default", timeout_seconds=300, **kwargs):
+        await asyncio.sleep(timeout_seconds)
+        session = self.session_storage.get(auth_handler.provider(), thread_id, profile, **kwargs)
+        if session.auth_resolve_uid is not None:
+            pocket_logger.info(f"session({session.auth_resolve_uid}) is not resolved yet and timeout. remove session")
+            self.delete_session(auth_handler.provider(), thread_id, profile)
+            FutureStore.delete_future(session.auth_resolve_uid)
+
+        return
+
+    def _upsert_pending_session(
+            self, auth_handler: AuthHandlerInterface, future_uid: str, profile: str, thread_id: str, scope: set[str]):
+        return self.session_storage.set(
+            auth_provider=auth_handler.provider(),
+            thread_id=thread_id,
+            profile=profile,
+            auth_scopes=list(scope),
+            auth_context=None,
+            is_auth_scope_universal=auth_handler.scoped,
+            auth_resolve_uid=future_uid,
+        )
+
+    async def _set_session_active(self, context: AuthContext, provider: AuthProvider, profile: str, thread_id: str):
+        session = self.session_storage.get(provider, thread_id, profile)
+        if session is None:
+            pocket_logger.error("the session to be active doesn't exist.")
+            return None
+
+        active_session = self.session_storage.set(
+            auth_provider=provider,
+            thread_id=thread_id,
+            profile=profile,
+            auth_scopes=list(session.auth_scopes),
+            is_auth_scope_universal=session.scoped,
+            auth_resolve_uid=None,
+            auth_context=context,
+        )
+        return active_session

hyperpocket/pocket_main.py

@@ -0,0 +1,351 @@
+import asyncio
+import pathlib
+from typing import Union, Any, Optional, Callable
+
+from hyperpocket.config import pocket_logger
+from hyperpocket.pocket_auth import PocketAuth
+from hyperpocket.repository import Lockfile
+from hyperpocket.repository.lock import LocalLock, GitLock
+from hyperpocket.server.server import PocketServer, PocketServerOperations
+from hyperpocket.tool import Tool, ToolRequest
+from hyperpocket.tool.function.tool import FunctionTool
+from hyperpocket.tool.wasm import WasmTool
+from hyperpocket.tool.wasm.tool import WasmToolRequest
+
+ToolLike = Union[Tool, str, Callable, ToolRequest]
+
+
+class Pocket(object):
+    auth: PocketAuth
+    server: PocketServer
+    tools: dict[str, Tool]
+
+    def __init__(self,
+                 tools: list[ToolLike],
+                 auth: PocketAuth = None,
+                 lockfile_path: Optional[str] = None,
+                 force_update: bool = False):
+        if auth is None:
+            auth = PocketAuth()
+
+        if lockfile_path is None:
+            lockfile_path = "./pocket.lock"
+        lockfile_pathlib_path = pathlib.Path(lockfile_path)
+        lockfile = Lockfile(lockfile_pathlib_path)
+        tool_likes = []
+        for tool_like in tools:
+            if isinstance(tool_like, str):
+                if pathlib.Path(tool_like).exists():
+                    lock = LocalLock(tool_like)
+                    req = WasmToolRequest(lock, "")
+                else:
+                    lock = GitLock(repository_url=tool_like, git_ref='HEAD')
+                    req = WasmToolRequest(lock, "")
+                lockfile.add_lock(lock)
+                tool_likes.append(req)
+            elif isinstance(tool_like, WasmToolRequest):
+                lockfile.add_lock(tool_like.lock)
+                tool_likes.append(tool_like)
+            elif isinstance(tool_like, ToolRequest):
+                raise ValueError(f"unreachable. tool_like:{tool_like}")
+            elif isinstance(tool_like, WasmTool):
+                raise ValueError("WasmTool should pass ToolRequest instance instead.")
+            else:
+                tool_likes.append(tool_like)
+        lockfile.sync(force_update=force_update, referenced_only=True)
+
+        self.tools = dict()
+        for tool_like in tool_likes:
+            tool = self._load_tool(tool_like, lockfile)
+            if tool.name in self.tools:
+                pocket_logger.error(f"Duplicate tool name: {tool.name}.")
+                raise ValueError(f"Duplicate tool name: {tool.name}")
+            self.tools[tool.name] = tool
+
+        pocket_logger.info(f"All Tools Loaded successfully. total registered tools : {len(self.tools)}")
+
+        self.auth = auth
+        self.server = PocketServer()
+        self.server.run(self)
+
+    def _load_tool(self, tool_like: ToolLike, lockfile: Lockfile) -> Tool:
+        pocket_logger.info(f"Loading Tool {tool_like}")
+        if isinstance(tool_like, Tool):
+            tool = tool_like
+        elif isinstance(tool_like, ToolRequest):
+            tool = Tool.from_tool_request(tool_like, lockfile=lockfile)
+        elif isinstance(tool_like, Callable):
+            tool = FunctionTool.from_func(tool_like)
+        else:
+            raise ValueError(f"Invalid tool type: {type(tool_like)}")
+
+        pocket_logger.info(f"Complete Loading Tool {tool.name}")
+        return tool
+
+    def invoke(self,
+               tool_name: str,
+               body: Any,
+               thread_id: str = 'default',
+               profile: str = 'default',
+               *args, **kwargs) -> str:
+        """
+        Invoke Tool synchronously
+
+        Args:
+            tool_name(str): tool name to invoke
+            body(Any): tool arguments. should be json format
+            thread_id(str): thread id
+            profile(str): profile name
+
+        Returns:
+            str: tool result
+        """
+        result = asyncio.run(self.ainvoke(tool_name, body, thread_id, profile, *args, **kwargs))
+        return result
+
+    async def ainvoke(self,
+                      tool_name: str,
+                      body: Any,
+                      thread_id: str = 'default',
+                      profile: str = 'default',
+                      *args, **kwargs) -> str:
+        """
+        Invoke Tool asynchronously
+
+        Args:
+            tool_name(str): tool name to invoke
+            body(Any): tool arguments. should be json format
+            thread_id(str): thread id
+            profile(str): profile name
+
+        Returns:
+            str: tool result
+        """
+        result, _ = await self.ainvoke_with_state(
+            tool_name=tool_name,
+            body=body,
+            thread_id=thread_id,
+            profile=profile,
+            *args,
+            **kwargs,
+        )
+        return result
+
+    def invoke_with_state(self,
+                          tool_name: str,
+                          body: Any,
+                          thread_id: str = 'default',
+                          profile: str = 'default',
+                          *args, **kwargs) -> tuple[str, bool]:
+        """
+        Invoke Tool with state synchronously
+        State indicates whether this tool is paused or not.
+        If the tool needs user's interaction or waiting for some process, this tool is paused.
+
+        Args:
+            tool_name(str): tool name to invoke
+            body(Any): tool arguments. should be json format
+            thread_id(str): thread id
+            profile(str): profile name
+
+        Returns:
+            tuple[str, bool]: tool result and state.
+        """
+        try:
+            loop = asyncio.new_event_loop()
+            result = loop.run_until_complete(
+                self.ainvoke_with_state(tool_name, body, thread_id, profile, *args, **kwargs))
+
+        except RuntimeError as e:
+            pocket_logger.warning("Can't execute sync def in event loop. use nest-asyncio")
+
+            import nest_asyncio
+            loop = asyncio.new_event_loop()
+            nest_asyncio.apply(loop=loop)
+
+            result = loop.run_until_complete(
+                self.ainvoke_with_state(tool_name, body, thread_id, profile, *args, **kwargs))
+
+        return result
+
+    async def ainvoke_with_state(self,
+                                 tool_name: str,
+                                 body: Any,
+                                 thread_id: str = 'default',
+                                 profile: str = 'default',
+                                 *args, **kwargs) -> tuple[str, bool]:
+        """
+        Invoke Tool with state synchronously
+        State indicates whether this tool is paused or not.
+        If the tool needs user's interaction or waiting for some process, this tool is paused.
+
+        Args:
+            tool_name(str): tool name to invoke
+            body(Any): tool arguments. should be json format
+            thread_id(str): thread id
+            profile(str): profile name
+
+        Returns:
+            tuple[str, bool]: tool result and state.
+        """
+        result, paused = await self.server.call_in_subprocess(
+            PocketServerOperations.CALL,
+            args,
+            {
+                'tool_name': tool_name,
+                'body': body,
+                'thread_id': thread_id,
+                'profile': profile,
+                **kwargs,
+            },
+        )
+        if not isinstance(result, str):
+            result = str(result)
+
+        return result, paused
+
+    # DO NOT EVER THINK ABOUT USING SERVER PROPERTY BELOW HERE
+    async def acall(self,
+                    tool_name: str,
+                    body: Any,
+                    thread_id: str = 'default',
+                    profile: str = 'default',
+                    *args, **kwargs) -> tuple[str, bool]:
+        """
+        Invoke tool asynchronously, not that different from `Pocket.invoke`
+        But this method is called only in subprocess.
+
+        This function performs the following steps:
+            1. `prepare_auth` : preparing the authentication process for the tool if necessary.
+            2. `authenticate` : performing authentication that needs to invoke tool.
+            3. `tool_call` : Executing tool actually with authentication information.
+
+        Args:
+            tool_name(str): tool name to invoke
+            body(Any): tool arguments. should be json format
+            thread_id(str): thread id
+            profile(str): profile name
+
+        Returns:
+            tuple[str, bool]: tool result and state.
+        """
+        tool = self._tool_instance(tool_name)
+        if tool.auth is not None:
+            callback_info = self.prepare_auth(tool_name, thread_id, profile, **kwargs)
+            if callback_info:
+                return callback_info, True
+        # 02. authenticate
+        credentials = await self.authenticate(tool_name, thread_id, profile, **kwargs)
+        # 03. call tool
+        result = await self.tool_call(tool_name, body=body, envs=credentials, **kwargs)
+        return result, False
+
+    def prepare_auth(self,
+                     tool_name: str,
+                     thread_id: str = 'default',
+                     profile: str = 'default',
+                     **kwargs) -> Optional[str]:
+        """
+        Prepares the authentication process for the tool if necessary.
+        Returns callback URL and whether the tool requires authentication.
+
+        Args:
+            tool_name(str): tool name to invoke
+            thread_id(str): thread id
+            profile(str): profile name
+
+        Returns:
+            Optional[str]: callback URI if necessary
+        """
+        tool = self._tool_instance(tool_name)
+        if tool.auth is None:
+            return None
+
+        auth_req = self.auth.make_request(
+            auth_handler_name=tool.auth.auth_handler,
+            auth_provider=tool.auth.auth_provider,
+            auth_scopes=tool.auth.scopes)
+
+        return self.auth.prepare(
+            auth_req=auth_req,
+            auth_handler_name=tool.auth.auth_handler,
+            auth_provider=tool.auth.auth_provider,
+            thread_id=thread_id,
+            profile=profile,
+            **kwargs
+        )
+
+    async def authenticate(
+            self,
+            tool_name: str,
+            thread_id: str = 'default',
+            profile: str = 'default',
+            **kwargs) -> dict[str, str]:
+        """
+        Authenticates the handler included in the tool and returns credentials.
+
+        Args:
+            tool_name(str): tool name to invoke
+            thread_id(str): thread id
+            profile(str): profile name
+
+        Returns:
+            dict[str, str]: credentials
+        """
+        tool = self._tool_instance(tool_name)
+        if tool.auth is None:
+            return {}
+        auth_req = self.auth.make_request(
+            auth_handler_name=tool.auth.auth_handler,
+            auth_provider=tool.auth.auth_provider,
+            auth_scopes=tool.auth.scopes)
+        auth_ctx = await self.auth.authenticate_async(
+            auth_req=auth_req,
+            auth_handler_name=tool.auth.auth_handler,
+            auth_provider=tool.auth.auth_provider,
+            thread_id=thread_id,
+            profile=profile,
+            **kwargs,
+        )
+        return auth_ctx.to_dict()
+
+    async def tool_call(self, tool_name: str, **kwargs) -> str:
+        """
+        Executing tool actually
+
+        Args:
+            tool_name(str): tool name to invoke
+            kwargs(dict): keyword arguments. authentication information is passed through this.
+
+        Returns:
+            str: tool result
+        """
+        tool = self._tool_instance(tool_name)
+        try:
+            return await asyncio.wait_for(tool.ainvoke(**kwargs), timeout=180)
+        except asyncio.TimeoutError:
+            pocket_logger.warning("Timeout tool call.")
+            return "timeout tool call"
+
+    def _tool_instance(self, tool_name: str) -> Tool:
+        return self.tools[tool_name]
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.__dict__.get('server'):
+            self.server.teardown()
+
+    def __del__(self):
+        if self.__dict__.get('server'):
+            self.server.teardown()
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        if 'server' in state:
+            del state['server']
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)