@josephyan/qingflow-cli 0.2.0-beta.55

package/src/qingflow_mcp/tools/workspace_tools.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+from typing import Any
+
+from mcp.server.fastmcp import FastMCP
+
+from ..backend_client import BackendRequestContext
+from ..config import DEFAULT_PROFILE
+from ..errors import QingflowApiError, raise_tool_error
+from .base import ToolBase
+
+
+class WorkspaceTools(ToolBase):
+    def register(self, mcp: FastMCP) -> None:
+        @mcp.tool()
+        def workspace_list(
+            profile: str = DEFAULT_PROFILE,
+            page_num: int = 1,
+            page_size: int = 20,
+            include_external: bool = False,
+        ) -> dict[str, Any]:
+            return self.workspace_list(
+                profile=profile,
+                page_num=page_num,
+                page_size=page_size,
+                include_external=include_external,
+            )
+
+        @mcp.tool()
+        def workspace_select(profile: str = DEFAULT_PROFILE, ws_id: int = 0) -> dict[str, Any]:
+            return self.workspace_select(profile=profile, ws_id=ws_id)
+
+        @mcp.tool()
+        def workspace_set_plugin_status(
+            profile: str = DEFAULT_PROFILE,
+            plugin_id: int = 0,
+            being_installed: bool = True,
+        ) -> dict[str, Any]:
+            return self.workspace_set_plugin_status(
+                profile=profile,
+                plugin_id=plugin_id,
+                being_installed=being_installed,
+            )
+
+    def workspace_list(
+        self,
+        *,
+        profile: str = DEFAULT_PROFILE,
+        page_num: int = 1,
+        page_size: int = 20,
+        include_external: bool = False,
+    ) -> dict[str, Any]:
+        if page_num <= 0 or page_size <= 0:
+            raise_tool_error(QingflowApiError.config_error("page_num and page_size must be positive"))
+
+        def runner(_, context):
+            path = "/user/allWorkspaceList/pageQuery" if include_external else "/user/workspaceList/pageQuery"
+            result = self.backend.request(
+                "POST",
+                context,
+                path,
+                json_body={"pageNum": page_num, "pageSize": page_size, "authList": [0, 1, 2]},
+            )
+            normalized = result if isinstance(result, dict) else {"list": result if isinstance(result, list) else []}
+            workspace_list = normalized.get("list") if isinstance(normalized.get("list"), list) else []
+            selected_ws_id = _.selected_ws_id if _ is not None else None
+            if selected_ws_id and not any(isinstance(item, dict) and item.get("wsId") == selected_ws_id for item in workspace_list):
+                try:
+                    selected_workspace = self.backend.request("GET", context, f"/user/workspace/{selected_ws_id}")
+                except QingflowApiError:
+                    selected_workspace = {
+                        "wsId": selected_ws_id,
+                        "workspaceName": _.selected_ws_name if _ is not None else None,
+                    }
+                if isinstance(selected_workspace, dict):
+                    workspace_list.append(selected_workspace)
+                    normalized["list"] = workspace_list
+                    total = normalized.get("total")
+                    if isinstance(total, int):
+                        normalized["total"] = total + 1
+            return {
+                "profile": profile,
+                "include_external": include_external,
+                "page": normalized,
+            }
+
+        return self._run(profile, runner, require_workspace=False)
+
+    def workspace_select(self, *, profile: str = DEFAULT_PROFILE, ws_id: int) -> dict[str, Any]:
+        if ws_id <= 0:
+            raise_tool_error(QingflowApiError.config_error("ws_id must be positive"))
+
+        def runner(_, context):
+            # Create a context with the target ws_id for the API call
+            # This is necessary because the API requires wsId header even when getting workspace details
+            call_context = BackendRequestContext(
+                base_url=context.base_url,
+                token=context.token,
+                ws_id=ws_id,
+                qf_version=context.qf_version,
+                qf_version_source=context.qf_version_source,
+            )
+            
+            try:
+                result = self.backend.request("GET", call_context, f"/user/workspace/{ws_id}")
+            except QingflowApiError as e:
+                if e.http_status == 404:
+                    result = self._workspace_from_list(call_context, ws_id) or {"wsId": ws_id, "workspaceName": f"Workspace {ws_id}"}
+                else:
+                    raise
+
+            if isinstance(result, dict):
+                ws_name = (
+                    str(result.get("workspaceName") or result.get("wsName") or "").strip()
+                    or None
+                )
+                if ws_name is None:
+                    fallback_workspace = self._workspace_from_list(call_context, ws_id)
+                    if isinstance(fallback_workspace, dict):
+                        result = fallback_workspace
+                        ws_name = (
+                            str(result.get("workspaceName") or result.get("wsName") or "").strip()
+                            or None
+                        )
+            else:
+                ws_name = None
+            if ws_name is None:
+                ws_name = f"Workspace {ws_id}"
+            session_profile = self.sessions.select_workspace(profile, ws_id=ws_id, ws_name=ws_name)
+            workspace_qf_version = self._workspace_system_version(result)
+            should_refresh_route = (
+                workspace_qf_version is not None
+                and (context.qf_version_source or "unset") != "explicit"
+                and workspace_qf_version != session_profile.qf_version
+            )
+            if should_refresh_route:
+                session_profile = self.sessions.update_route(
+                    profile,
+                    qf_version=workspace_qf_version,
+                    qf_version_source="workspace_system_version",
+                )
+            return {
+                "profile": profile,
+                "selected_ws_id": session_profile.selected_ws_id,
+                "selected_ws_name": session_profile.selected_ws_name,
+                "workspace": result,
+                "qf_version": session_profile.qf_version,
+                "qf_version_source": session_profile.qf_version_source,
+                "request_route": self.backend.describe_route(
+                    BackendRequestContext(
+                        base_url=session_profile.base_url,
+                        token=context.token,
+                        ws_id=session_profile.selected_ws_id,
+                        qf_version=session_profile.qf_version,
+                        qf_version_source=session_profile.qf_version_source,
+                    )
+                ),
+            }
+
+        return self._run(profile, runner, require_workspace=False)
+
+    def _workspace_from_list(self, context: BackendRequestContext, ws_id: int) -> dict[str, Any] | None:
+        try:
+            payload = self.backend.request(
+                "POST",
+                context,
+                "/user/workspaceList/pageQuery",
+                json_body={"pageNum": 1, "pageSize": 100, "authList": [0, 1, 2]},
+            )
+        except Exception:
+            return None
+        workspaces = payload.get("list", []) if isinstance(payload, dict) else []
+        found = next(
+            (
+                item
+                for item in workspaces
+                if isinstance(item, dict) and item.get("wsId") == ws_id
+            ),
+            None,
+        )
+        return found if isinstance(found, dict) else None
+
+    def _workspace_system_version(self, workspace: Any) -> str | None:
+        if not isinstance(workspace, dict):
+            return None
+        value = workspace.get("systemVersion")
+        if value is None:
+            return None
+        normalized = str(value).strip()
+        return normalized or None
+
+    def workspace_set_plugin_status(
+        self,
+        *,
+        profile: str = DEFAULT_PROFILE,
+        plugin_id: int,
+        being_installed: bool = True,
+    ) -> dict[str, Any]:
+        if plugin_id <= 0:
+            raise_tool_error(QingflowApiError.config_error("plugin_id must be positive"))
+
+        def runner(_, context):
+            result = self.backend.request(
+                "POST",
+                context,
+                "/ws/plugin",
+                json_body={
+                    "pluginId": plugin_id,
+                    "beingInstalled": being_installed,
+                },
+            )
+            return {
+                "profile": profile,
+                "plugin_id": plugin_id,
+                "being_installed": being_installed,
+                "result": result,
+            }
+
+        return self._run(profile, runner)