agentscope-runtime 0.1.0__py3-none-any.whl

agentscope_runtime/sandbox/box/shared/routers/mcp_utils.py

@@ -0,0 +1,153 @@
+# -*- coding: utf-8 -*-
+import asyncio
+import logging
+import os
+import shutil
+import traceback
+from contextlib import AsyncExitStack
+from typing import Any
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.sse import sse_client
+from mcp.client.stdio import stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class MCPSessionHandler:
+    """Manages MCP server connections and tool execution."""
+
+    def __init__(self, name: str, config: dict[str, Any]) -> None:
+        self.name: str = name
+        self.config: dict[str, Any] = config
+        self.stdio_context: Any | None = None
+        self.session: ClientSession | None = None
+        self._cleanup_lock: asyncio.Lock = asyncio.Lock()
+        self._exit_stack: AsyncExitStack = AsyncExitStack()
+
+    async def initialize(self) -> None:
+        """Initialize the server connection."""
+        command = (
+            shutil.which("npx")
+            if self.config.get("command") == "npx"
+            else self.config.get("command")
+        )
+
+        try:
+            if command:
+                server_params = StdioServerParameters(
+                    command=command,
+                    args=self.config.get("args", []),
+                    env={**os.environ, **self.config.get("env", {})},
+                )
+
+                streams = await self._exit_stack.enter_async_context(
+                    stdio_client(server_params),
+                )
+            else:
+                if self.config.get("type") in [
+                    "streamable_http",
+                    "streamableHttp",
+                ]:
+                    streams = await self._exit_stack.enter_async_context(
+                        streamablehttp_client(url=self.config["url"]),
+                    )
+                    streams = (streams[0], streams[1])
+                else:
+                    streams = await self._exit_stack.enter_async_context(
+                        sse_client(url=self.config["url"]),
+                    )
+            session = await self._exit_stack.enter_async_context(
+                ClientSession(*streams),
+            )
+            await session.initialize()
+            self.session = session
+        except Exception as e:
+            logging.error(f"Error initializing server {self.name}: {e}")
+            await self.cleanup()
+            raise
+
+    async def list_tools(self) -> list[Any]:
+        """List available tools from the server.
+
+        Returns:
+            A list of available tools.
+
+        Raises:
+            RuntimeError: If the server is not initialized.
+        """
+        if not self.session:
+            raise RuntimeError(f"Server {self.name} not initialized")
+
+        tools_response = await self.session.list_tools()
+        tools = [
+            tool
+            for item in tools_response
+            if isinstance(item, tuple) and item[0] == "tools"
+            for tool in item[1]
+        ]
+
+        return tools
+
+    async def call_tool(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        retries: int = 2,
+        delay: float = 1.0,
+    ) -> Any:
+        """Execute a tool with retry mechanism.
+
+        Args:
+            tool_name: Name of the tool to execute.
+            arguments: tool arguments.
+            retries: Number of retry attempts.
+            delay: Delay between retries in seconds.
+
+        Returns:
+            Tool execution result.
+
+        Raises:
+            RuntimeError: If server is not initialized.
+            Exception: If tool execution fails after all retries.
+        """
+        if not self.session:
+            raise RuntimeError(f"Server {self.name} not initialized")
+
+        attempt = 0
+
+        while attempt < retries:
+            try:
+                logging.info(f"Executing {tool_name}...")
+                result = await self.session.call_tool(tool_name, arguments)
+                return result
+
+            except Exception as e:
+                attempt += 1
+                logging.warning(
+                    f"Error executing tool: {e} {traceback.format_exc()}."
+                    f" Attempt {attempt} of {retries}.",
+                )
+                if attempt >= retries:
+                    logging.error("Max retries reached. Failing.")
+                    raise
+                logging.info(f"Retrying in {delay} seconds...")
+                await asyncio.sleep(delay)
+        return None
+
+    async def cleanup(self) -> None:
+        """Clean up server resources."""
+        async with self._cleanup_lock:
+            try:
+                await self._exit_stack.aclose()
+            except Exception as e:
+                if (
+                    "Attempted to exit cancel scope in a different task"
+                    in str(e)
+                ):
+                    pass
+            finally:
+                self.session = None
+                self.stdio_context = None

agentscope_runtime/sandbox/box/shared/routers/runtime_watcher.py

@@ -0,0 +1,187 @@
+# -*- coding: utf-8 -*-
+import difflib
+import logging
+import traceback
+
+import git
+from fastapi import APIRouter, Body, HTTPException
+
+watcher_router = APIRouter()
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def initialize_git_user(repo):
+    repo.config_writer().set_value("user", "name", "User").release()
+    repo.config_writer().set_value(
+        "user",
+        "email",
+        "user@example.com",
+    ).release()
+    return repo
+
+
+@watcher_router.post(
+    "/watcher/commit_changes",
+    summary="...",
+)
+async def commit_changes(
+    commit_message: str = Body(
+        "Automated commit",
+        example="Your commit message",
+        embed=True,
+    ),
+):
+    """
+    Commit the uncommitted changes.
+    """
+    try:
+        repo_path = "."
+
+        repo = git.Repo(repo_path)
+        repo = initialize_git_user(repo)
+
+        # Add all changes to the staging area
+        repo.git.add(A=True)
+
+        # Commit the changes
+        commit = repo.index.commit(commit_message)
+        return {"commit": commit.hexsha, "message": commit_message}
+
+    except Exception as e:
+        logger.error(f"{str(e)}:\n{traceback.format_exc()}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"{str(e)}: {traceback.format_exc()}",
+        ) from e
+
+
+@watcher_router.post(
+    "/watcher/generate_diff",
+    summary="...",
+)
+async def generate_diff(
+    commit_a: str = Body(..., embed=True),
+    commit_b: str = Body(..., embed=True),
+):
+    """
+    Generate the diff of the uncommitted changes or two commits.
+    """
+    try:
+        repo_path = "."
+        repo = git.Repo(repo_path)
+        repo = initialize_git_user(repo)
+
+        if not commit_a and not commit_b:
+            # Default to uncommitted changes compared to the last commit
+            repo.git.add(A=True)
+            diff_index = repo.index.diff("HEAD")
+            print(diff_index, repo.git.status())
+        elif commit_a and commit_b:
+            # Get diff between two commits
+            diff_index = repo.commit(commit_a).diff(commit_b)
+        else:
+            return HTTPException(
+                detail="Invalid commit range",
+                status_code=400,
+            )
+        diffs = {}
+        for diff in diff_index:
+            if diff.a_blob and diff.b_blob:
+                # Both files are present in commits; perform a diff
+                a_content = (
+                    diff.a_blob.data_stream.read()
+                    .decode(
+                        "utf-8",
+                    )
+                    .splitlines()
+                )
+                b_content = (
+                    diff.b_blob.data_stream.read()
+                    .decode(
+                        "utf-8",
+                    )
+                    .splitlines()
+                )
+            elif diff.a_blob:  # File was deleted
+                # Only 'a' file is present; 'b' file is empty
+                a_content = (
+                    diff.a_blob.data_stream.read()
+                    .decode(
+                        "utf-8",
+                    )
+                    .splitlines()
+                )
+                b_content = []
+            elif diff.b_blob:  # File was added
+                # Only 'b' file is present; 'a' file is empty
+                a_content = []
+                b_content = (
+                    diff.b_blob.data_stream.read()
+                    .decode(
+                        "utf-8",
+                    )
+                    .splitlines()
+                )
+            else:
+                continue
+
+            # Generate the diff content
+            diff_text = "\n".join(
+                difflib.unified_diff(
+                    a_content,
+                    b_content,
+                    fromfile=f"a/{diff.a_path}",
+                    tofile=f"b/{diff.b_path}",
+                    lineterm="",
+                ),
+            )
+            diffs[diff.b_path or diff.a_path] = diff_text
+        return {"diffs": diffs}
+
+    except Exception as e:
+        logger.error(f"{str(e)}:\n{traceback.format_exc()}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"{str(e)}: {traceback.format_exc()}",
+        ) from e
+
+
+@watcher_router.get(
+    "/watcher/git_logs",
+    summary="...",
+)
+async def git_logs():
+    """
+    Return the git logs.
+    """
+    try:
+        repo = git.Repo(".")
+        repo = initialize_git_user(repo)
+        logs = []
+        for commit in repo.iter_commits():
+            diff_result = {"diffs": {}}
+            if commit.parents:
+                parent_commit = commit.parents[0]
+                diff_result = await generate_diff(
+                    commit.hexsha,
+                    parent_commit.hexsha,
+                )
+
+            log_entry = {
+                "commit": commit.hexsha,
+                "author": commit.author.name,
+                "date": commit.committed_datetime.isoformat(),
+                "message": commit.message.strip(),
+                "diff": diff_result["diffs"],
+            }
+            logs.append(log_entry)
+        return {"logs": logs}
+    except Exception as e:
+        logger.error(f"{str(e)}:\n{traceback.format_exc()}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"{str(e)}: {traceback.format_exc()}",
+        ) from e

agentscope_runtime/sandbox/box/shared/routers/workspace.py

@@ -0,0 +1,325 @@
+# -*- coding: utf-8 -*-
+import shutil
+import os
+import logging
+import traceback
+
+import aiofiles
+
+from fastapi import APIRouter, HTTPException, Query, Body
+from fastapi.responses import FileResponse
+
+workspace_router = APIRouter()
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def ensure_within_workspace(
+    path: str,
+    base_directory: str = "/workspace",
+) -> str:
+    """
+    Ensure the provided path is within the /workspace directory.
+    """
+    base_directory = os.path.abspath(base_directory)
+
+    # Determine if the input path is absolute or relative
+    if os.path.isabs(path):
+        full_path = os.path.abspath(path)
+    else:
+        full_path = os.path.abspath(os.path.join(base_directory, path))
+
+    # Check for path traversal attacks and ensure path is within base_directory
+    if not full_path.startswith(base_directory):
+        raise HTTPException(
+            status_code=403,
+            detail="Permission error. Access restricted to /workspace "
+            "directory.",
+        )
+
+    return full_path
+
+
+@workspace_router.get(
+    "/workspace/files",
+    summary="Retrieve a file within the /workspace directory",
+)
+async def get_workspace_file(
+    file_path: str = Query(
+        ...,
+        description="Path to the file within /workspace relative to its root",
+    ),
+):
+    """
+    Get a file within the /workspace directory.
+    """
+    try:
+        # Ensure the file path is within the /workspace directory
+        full_path = ensure_within_workspace(file_path)
+
+        # Check if the file exists
+        if not os.path.isfile(full_path):
+            raise HTTPException(status_code=404, detail="File not found.")
+
+        # Return the file using FileResponse
+        return FileResponse(
+            full_path,
+            media_type="application/octet-stream",
+            filename=os.path.basename(full_path),
+        )
+
+    except Exception as e:
+        logger.error(f"{str(e)}:\n{traceback.format_exc()}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"{str(e)}: {traceback.format_exc()}",
+        ) from e
+
+
+@workspace_router.post(
+    "/workspace/files",
+    summary="Create or edit a file within the /workspace directory",
+)
+async def create_or_edit_file(
+    file_path: str = Query(
+        ...,
+        description="Path to the file within /workspace",
+    ),
+    content: str = Body(..., description="Content to write to the file"),
+):
+    try:
+        full_path = ensure_within_workspace(file_path)
+        async with aiofiles.open(full_path, "w", encoding="utf-8") as f:
+            await f.write(content)
+        return {"message": "File created or edited successfully."}
+    except Exception as e:
+        logger.error(
+            f"Error creating or editing file: {str(e)}:\
+            n{traceback.format_exc()}",
+        )
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error creating or editing file: {str(e)}",
+        ) from e
+
+
+@workspace_router.get(
+    "/workspace/list-directories",
+    summary="List file items in the /workspace directory, including nested "
+    "files and directories",
+)
+async def list_workspace_files(
+    directory: str = Query(
+        "/workspace",
+        description="Directory to list files and directories from, default "
+        "is /workspace.",
+    ),
+):
+    """
+    List all files and directories in the specified directory, including
+    nested items, with type indication and statistics.
+    """
+    try:
+        target_directory = ensure_within_workspace(directory)
+
+        # Verify if the specified directory exists
+        if not os.path.isdir(target_directory):
+            raise HTTPException(status_code=404, detail="Directory not found.")
+
+        nested_items = []
+        file_count = 0
+        directory_count = 0
+
+        for root, dirs, files in os.walk(target_directory):
+            for d in dirs:
+                dir_path = os.path.join(root, d)
+                nested_items.append(
+                    {
+                        "type": "directory",
+                        "path": os.path.relpath(dir_path, target_directory),
+                    },
+                )
+                directory_count += 1
+
+            for f in files:
+                file_path = os.path.join(root, f)
+                nested_items.append(
+                    {
+                        "type": "file",
+                        "path": os.path.relpath(file_path, target_directory),
+                    },
+                )
+                file_count += 1
+
+        return {
+            "items": nested_items,
+            "statistics": {
+                "total_directories": directory_count,
+                "total_files": file_count,
+            },
+        }
+
+    except Exception as e:
+        logger.error(
+            f"Error listing files: {str(e)}:\n{traceback.format_exc()}",
+        )
+        raise HTTPException(
+            status_code=500,
+            detail=f"An error occurred while listing files: {str(e)}",
+        ) from e
+
+
+@workspace_router.post(
+    "/workspace/directories",
+    summary="Create a directory within the /workspace directory",
+)
+async def create_directory(
+    directory_path: str = Query(
+        ...,
+        description="Path to the directory within /workspace",
+    ),
+):
+    try:
+        full_path = ensure_within_workspace(directory_path)
+        os.makedirs(full_path, exist_ok=True)
+        return {"message": "Directory created successfully."}
+    except Exception as e:
+        logger.error(
+            f"Error creating directory: {str(e)}:\n{traceback.format_exc()}",
+        )
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error creating directory: {str(e)}",
+        ) from e
+
+
+@workspace_router.delete(
+    "/workspace/files",
+    summary="Delete a file within the /workspace directory",
+)
+async def delete_file(
+    file_path: str = Query(
+        ...,
+        description="Path to the file within /workspace",
+    ),
+):
+    try:
+        full_path = ensure_within_workspace(file_path)
+        if os.path.isfile(full_path):
+            os.remove(full_path)
+            return {"message": "File deleted successfully."}
+        else:
+            raise HTTPException(status_code=404, detail="File not found.")
+    except Exception as e:
+        logger.error(
+            f"Error deleting file: {str(e)}:\n{traceback.format_exc()}",
+        )
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error deleting file: {str(e)}",
+        ) from e
+
+
+@workspace_router.delete(
+    "/workspace/directories",
+    summary="Delete a directory within the /workspace directory",
+)
+async def delete_directory(
+    directory_path: str = Query(
+        ...,
+        description="Path to the directory within /workspace",
+    ),
+    recursive: bool = Query(
+        False,
+        description="Recursively delete directory contents",
+    ),
+):
+    try:
+        full_path = ensure_within_workspace(directory_path)
+        if recursive:
+            shutil.rmtree(full_path)
+        else:
+            os.rmdir(full_path)
+        return {"message": "Directory deleted successfully."}
+    except Exception as e:
+        logger.error(
+            f"Error deleting directory: {str(e)}:\n{traceback.format_exc()}",
+        )
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error deleting directory: {str(e)}",
+        ) from e
+
+
+@workspace_router.put(
+    "/workspace/move",
+    summary="Move or rename a file or directory within the /workspace "
+    "directory",
+)
+async def move_or_rename(
+    source_path: str = Query(
+        ...,
+        description="Source path within /workspace",
+    ),
+    destination_path: str = Query(
+        ...,
+        description="Destination path within /workspace",
+    ),
+):
+    try:
+        full_source_path = ensure_within_workspace(source_path)
+        full_destination_path = ensure_within_workspace(destination_path)
+        if not os.path.exists(full_source_path):
+            raise HTTPException(
+                status_code=404,
+                detail="Source file or directory not found.",
+            )
+        os.rename(full_source_path, full_destination_path)
+        return {"message": "Move or rename operation successful."}
+    except Exception as e:
+        logger.error(
+            f"Error moving or renaming: {str(e)}:\n{traceback.format_exc()}",
+        )
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error moving or renaming: {str(e)}",
+        ) from e
+
+
+@workspace_router.post(
+    "/workspace/copy",
+    summary="Copy a file or directory within the /workspace directory",
+)
+async def copy(
+    source_path: str = Query(
+        ...,
+        description="Source path within /workspace",
+    ),
+    destination_path: str = Query(
+        ...,
+        description="Destination path within /workspace",
+    ),
+):
+    try:
+        full_source_path = ensure_within_workspace(source_path)
+        full_destination_path = ensure_within_workspace(destination_path)
+        if not os.path.exists(full_source_path):
+            raise HTTPException(
+                status_code=404,
+                detail="Source file or directory not found.",
+            )
+
+        if os.path.isdir(full_source_path):
+            shutil.copytree(full_source_path, full_destination_path)
+        else:
+            shutil.copy2(full_source_path, full_destination_path)
+
+        return {"message": "Copy operation successful."}
+    except Exception as e:
+        logger.error(f"Error copying: {str(e)}:\n{traceback.format_exc()}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error copying: " f"{str(e)}",
+        ) from e