@josephyan/qingflow-cli 0.2.0-beta.1000

package/src/qingflow_mcp/tools/base.py

@@ -0,0 +1,281 @@
+from __future__ import annotations
+
+import inspect
+from contextvars import ContextVar
+from typing import Any, Callable, TypeVar
+
+from ..backend_client import BackendRequestContext, BackendClient
+from ..config import (
+    get_credit_meter_enabled,
+    get_credit_usage_base_url,
+    get_credit_usage_path,
+)
+from ..errors import QingflowApiError, raise_tool_error
+from ..json_types import JSONObject
+from ..session_store import BackendSession, SessionProfile, SessionStore
+
+
+T = TypeVar("T")
+_RUN_DEPTH: ContextVar[int] = ContextVar("qingflow_mcp_tool_run_depth", default=0)
+
+
+def tool_cn_name(name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """为 tools 方法声明中文展示名（用于计费 extraInfo 记录）。"""
+    normalized = str(name).strip()
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        setattr(func, "__tool_cn_name__", normalized)
+        return func
+
+    return decorator
+
+
+class ToolBase:
+    def __init__(self, sessions: SessionStore, backend: BackendClient) -> None:
+        self.sessions = sessions
+        self.backend = backend
+
+    def _require_context(self, profile: str, *, require_workspace: bool) -> tuple[SessionProfile, BackendSession, BackendRequestContext]:
+        session_profile = self.sessions.get_profile(profile)
+        if session_profile is None:
+            raise QingflowApiError.auth_required(profile)
+        backend_session = self.sessions.get_backend_session(profile)
+        if backend_session is None:
+            raise QingflowApiError.auth_required(profile)
+        if require_workspace and session_profile.selected_ws_id is None:
+            raise QingflowApiError.workspace_not_selected(profile)
+        context = BackendRequestContext(
+            base_url=backend_session.base_url,
+            token=backend_session.token,
+            # Keep the selected workspace in context whenever we know it.
+            # Some top-level tools (for example workspace browsing) do not
+            # require a workspace to execute, but credit metering still needs
+            # a stable wsId when the profile already has one selected.
+            ws_id=session_profile.selected_ws_id,
+            qf_version=backend_session.qf_version,
+            qf_version_source=backend_session.qf_version_source,
+        )
+        return session_profile, backend_session, context
+
+    def _run(self, profile: str, func: Callable[[SessionProfile, BackendRequestContext], T], *, require_workspace: bool = True) -> T:
+        depth = _RUN_DEPTH.get()
+        depth_token = _RUN_DEPTH.set(depth + 1)
+        is_top_level_call = depth == 0
+        session_profile: SessionProfile | None = None
+        backend_session: BackendSession | None = None
+        tool_name = self._resolve_calling_tool_name()
+        try:
+            session_profile, backend_session, context = self._require_context(profile, require_workspace=require_workspace)
+            if is_top_level_call:
+                self._charge_tool_call(
+                    profile=profile,
+                    tool_name=tool_name,
+                    session_profile=session_profile,
+                    context=context,
+                )
+            return func(session_profile, context)
+        except QingflowApiError as error:
+            if (
+                error.looks_like_invalid_token()
+                and session_profile is not None
+                and backend_session is not None
+                and self._refresh_session_from_credential(
+                    profile,
+                    session_profile=session_profile,
+                    backend_session=backend_session,
+                )
+            ):
+                try:
+                    refreshed_profile, _, refreshed_context = self._require_context(profile, require_workspace=require_workspace)
+                    return func(refreshed_profile, refreshed_context)
+                except QingflowApiError as refreshed_error:
+                    self._handle_error(profile, refreshed_error)
+            self._handle_error(profile, error)
+        finally:
+            _RUN_DEPTH.reset(depth_token)
+        raise AssertionError("unreachable")
+
+    def _handle_error(self, profile: str, error: QingflowApiError) -> None:
+        if error.looks_like_invalid_token():
+            self.sessions.invalidate(profile)
+            error = QingflowApiError(
+                category="auth",
+                message=f"Qingflow session for profile '{profile}' has expired. Run auth login or auth_use_credential again.",
+                backend_code=error.backend_code,
+                request_id=error.request_id,
+                http_status=error.http_status,
+            )
+        raise_tool_error(error)
+
+    def _refresh_session_from_credential(
+        self,
+        profile: str,
+        *,
+        session_profile: SessionProfile,
+        backend_session: BackendSession,
+    ) -> bool:
+        credential = (backend_session.credential or "").strip()
+        if not credential:
+            return False
+        try:
+            response = self.backend.public_request_with_meta(
+                "POST",
+                session_profile.base_url,
+                "/mcp/auth/context",
+                json_body={"credential": credential},
+                qf_version=backend_session.qf_version,
+            )
+            payload = self._unwrap_context_payload(response.data)
+            token = self._normalize_text(payload.get("token"))
+            ws_id = self._coerce_positive_int(payload.get("wsId"))
+            uid = self._coerce_positive_int(payload.get("uid")) or session_profile.uid
+            if not token or ws_id is None:
+                return False
+            backend_qf_version = self._normalize_text(payload.get("qfVersion")) or response.qf_response_version
+            qf_version = backend_qf_version if backend_qf_version is not None else backend_session.qf_version
+            qf_version_source = (
+                "backend_response"
+                if backend_qf_version is not None
+                else backend_session.qf_version_source
+            )
+            base_url = self._normalize_text(payload.get("baseUrl")) or session_profile.base_url
+            ws_name = self._normalize_text(payload.get("wsName")) or session_profile.selected_ws_name
+            self.sessions.save_session(
+                profile=profile,
+                base_url=base_url,
+                qf_version=qf_version,
+                qf_version_source=qf_version_source,
+                token=token,
+                login_token=None,
+                credential=credential,
+                uid=uid,
+                email=self._normalize_text(payload.get("email")) or session_profile.email,
+                nick_name=(
+                    self._normalize_text(payload.get("nickName"))
+                    or self._normalize_text(payload.get("displayName"))
+                    or self._normalize_text(payload.get("name"))
+                    or session_profile.nick_name
+                ),
+                persist=session_profile.persisted,
+            )
+            self.sessions.select_workspace(profile, ws_id=ws_id, ws_name=ws_name)
+            return True
+        except QingflowApiError:
+            return False
+
+    def _unwrap_context_payload(self, payload: Any) -> dict[str, Any]:
+        if not isinstance(payload, dict):
+            return {}
+        nested = payload.get("data")
+        if isinstance(nested, dict):
+            return nested
+        nested = payload.get("result")
+        if isinstance(nested, dict):
+            return nested
+        return payload
+
+    def _normalize_text(self, value: Any) -> str | None:
+        if value is None:
+            return None
+        text = str(value).strip()
+        return text or None
+
+    def _coerce_positive_int(self, value: Any) -> int | None:
+        if isinstance(value, bool) or value is None:
+            return None
+        try:
+            parsed = int(value)
+        except (TypeError, ValueError):
+            return None
+        return parsed if parsed > 0 else None
+
+    def _resolve_calling_tool_name(self) -> str:
+        frame = inspect.currentframe()
+        if frame is None:
+            return "unknown_tool"
+        caller = frame.f_back.f_back if frame.f_back is not None else None
+        try:
+            if caller is None:
+                return "unknown_tool"
+            name = caller.f_code.co_name
+            if not name:
+                return "unknown_tool"
+            method = getattr(self, name, None)
+            cn_name = getattr(method, "__tool_cn_name__", None)
+            if isinstance(cn_name, str) and cn_name.strip():
+                return cn_name.strip()
+            return name
+        finally:
+            del frame
+
+    def _charge_tool_call(
+        self,
+        *,
+        profile: str,
+        tool_name: str,
+        session_profile: SessionProfile,
+        context: BackendRequestContext,
+    ) -> None:
+        if not get_credit_meter_enabled():
+            return
+        if context.ws_id is None:
+            raise_tool_error(QingflowApiError(category="payment", message="credit meter requires wsId in current session context"))
+
+        self._record_credit_usage(
+            tool_name=tool_name,
+            context=context,
+        )
+
+    def _record_credit_usage(
+        self,
+        *,
+        tool_name: str,
+        context: BackendRequestContext,
+    ) -> None:
+        usage_context = BackendRequestContext(
+            base_url=get_credit_usage_base_url() or context.base_url,
+            token=context.token,
+            ws_id=context.ws_id,
+            qf_request_id=context.qf_request_id,
+            qf_version=context.qf_version,
+            qf_version_source=context.qf_version_source,
+        )
+        self.backend.request(
+            "POST",
+            usage_context,
+            get_credit_usage_path(),
+            json_body={
+                "skuType": "MCP",
+                "skuName": "MCP",
+                "modelName": "MCP",
+                "scene": "MCP",
+                "aiBiz": "MCP",
+                "extraInfo": tool_name,
+            },
+        )
+
+    def _require_dict(self, payload: JSONObject | None, field_name: str = "payload") -> JSONObject:
+        if not isinstance(payload, dict) or not payload:
+            raise_tool_error(QingflowApiError.config_error(f"{field_name} must be a non-empty object"))
+        return payload
+
+    def _high_risk_tool_description(self, *, operation: str, target: str) -> str:
+        return (
+            f"High-risk {operation} operation for {target}. Read the current state first, "
+            "confirm the exact target IDs and intended diff with a human, and avoid running "
+            "against production without explicit approval."
+        )
+
+    def _attach_human_review_notice(self, response: JSONObject, *, operation: str, target: str) -> JSONObject:
+        payload = dict(response)
+        payload["requires_human_review"] = True
+        payload["risk_notice"] = {
+            "operation": operation,
+            "target": target,
+            "severity": "high",
+            "guidance": (
+                "Read the current state first, confirm the exact target IDs and intended diff with a human, "
+                "and require explicit approval before running against production."
+            ),
+        }
+        return payload