stackless-mcp 1.0.0b2__py3-none-any.whl

stackless_mcp/server.py

@@ -0,0 +1,1158 @@
+from __future__ import annotations
+
+import argparse
+import atexit
+import json
+import os
+import sys
+from collections.abc import Callable
+from datetime import datetime, timezone
+from typing import Any
+
+from dotenv import load_dotenv
+from fastmcp import FastMCP
+
+from stackless_mcp import __version__
+from stackless_mcp import tool_handlers as handlers
+from stackless_mcp.auth import CredentialStore, StoredCredential
+from stackless_mcp.auth import login as mcp_login
+from stackless_mcp.auth import logout as mcp_logout
+from stackless_mcp.client import DEFAULT_TIMEOUT_SECONDS, StacklessClient
+
+mcp = FastMCP("stackless-mcp")
+
+_client: StacklessClient | None = None
+_client_config: dict[str, str | float] = {}
+
+TOOLSET_ALL = "all"
+TOOLSET_CORE = "core"
+OPTIONAL_TOOLSETS = (
+    "catalog",
+    "runs",
+    "transformations",
+    "semantic",
+    "dashboards",
+    "lifecycle",
+)
+SELECTABLE_TOOLSETS = (TOOLSET_CORE, *OPTIONAL_TOOLSETS)
+VALID_TOOLSETS = frozenset((TOOLSET_ALL, TOOLSET_CORE, *OPTIONAL_TOOLSETS))
+DEFAULT_TOOLSETS = TOOLSET_ALL
+
+_ToolFn = Callable[..., Any]
+_TOOL_REGISTRY: dict[str, tuple[_ToolFn, frozenset[str]]] = {}
+_active_toolsets: frozenset[str] = frozenset()
+
+
+def _mcp_tool(*toolsets: str) -> Callable[[_ToolFn], _ToolFn]:
+    normalized = frozenset(toolsets)
+    if not normalized:
+        raise ValueError("MCP tools must declare at least one toolset")
+    unknown = normalized - VALID_TOOLSETS
+    if unknown:
+        raise ValueError(f"Unknown MCP toolsets: {', '.join(sorted(unknown))}")
+
+    def decorator(fn: _ToolFn) -> _ToolFn:
+        if fn.__name__ in _TOOL_REGISTRY:
+            raise ValueError(f"Duplicate MCP tool registration: {fn.__name__}")
+        _TOOL_REGISTRY[fn.__name__] = (fn, normalized)
+        return fn
+
+    return decorator
+
+
+def parse_toolsets(value: str | None) -> frozenset[str]:
+    raw = value or DEFAULT_TOOLSETS
+    requested = {item.strip() for item in raw.split(",") if item.strip()}
+    if not requested or TOOLSET_ALL in requested:
+        if requested - {TOOLSET_ALL}:
+            raise ValueError("MCP toolset 'all' cannot be combined with other toolsets")
+        return frozenset((TOOLSET_CORE, *OPTIONAL_TOOLSETS))
+    unknown = requested - VALID_TOOLSETS
+    if unknown:
+        raise ValueError(
+            "Unknown MCP toolsets: "
+            + ", ".join(sorted(unknown))
+            + ". Valid values are: "
+            + ", ".join((TOOLSET_ALL, *SELECTABLE_TOOLSETS))
+        )
+    return frozenset((TOOLSET_CORE, *requested))
+
+
+def tool_names_for_toolsets(value: str | None) -> set[str]:
+    selected = parse_toolsets(value)
+    return {
+        name for name, (_, toolsets) in _TOOL_REGISTRY.items() if toolsets & selected
+    }
+
+
+def _workflow_tool_names_for_toolsets(value: str | None) -> set[str]:
+    return tool_names_for_toolsets(value) - _core_tool_names()
+
+
+def _core_tool_names() -> set[str]:
+    return {
+        name
+        for name, (_, toolsets) in _TOOL_REGISTRY.items()
+        if TOOLSET_CORE in toolsets
+    }
+
+
+def _validate_capabilities(payload: Any) -> tuple[str, frozenset[str]]:
+    if not isinstance(payload, dict):
+        raise ValueError("capabilities response is not an object")
+    server_version = payload.get("server_version")
+    tools = payload.get("tools")
+    if not isinstance(server_version, str) or not server_version:
+        raise ValueError("capabilities response is missing server_version")
+    if not isinstance(tools, list) or not all(isinstance(tool, str) for tool in tools):
+        raise ValueError("capabilities response tools must be a list of strings")
+    return server_version, frozenset(tools)
+
+
+def registered_tool_names() -> set[str]:
+    return {tool.name for tool in mcp.list_tools()}
+
+
+def configure_toolsets(
+    value: str | None = DEFAULT_TOOLSETS,
+    *,
+    server_supported_tools: frozenset[str] | set[str] | None = None,
+) -> frozenset[str]:
+    global mcp, _active_toolsets
+    selected = parse_toolsets(value)
+    next_mcp = FastMCP("stackless-mcp")
+    for name, (fn, toolsets) in _TOOL_REGISTRY.items():
+        is_core_tool = TOOLSET_CORE in toolsets
+        is_supported = server_supported_tools is None or name in server_supported_tools
+        if is_core_tool or (toolsets & selected and is_supported):
+            next_mcp.add_tool(fn, tags=set(toolsets))
+    mcp = next_mcp
+    _active_toolsets = selected
+    return selected
+
+
+def configure_toolsets_from_server(
+    value: str | None = DEFAULT_TOOLSETS,
+) -> frozenset[str]:
+    try:
+        _, server_supported_tools = _validate_capabilities(get_client().capabilities())
+    except Exception as exc:
+        print(
+            "Warning: Stackless server did not report MCP capabilities; "
+            f"registering core tools only. Discovery error: {exc}",
+            file=sys.stderr,
+        )
+        server_supported_tools = frozenset()
+    return configure_toolsets(value, server_supported_tools=server_supported_tools)
+
+
+def configure_client(
+    *,
+    base_url: str,
+    token: str,
+    cookie: str,
+    timeout: float,
+    auth_storage: str = "keyring",
+) -> None:
+    global _client, _client_config
+    if _client is not None:
+        _client.close()
+    _client = None
+    _client_config = {
+        "base_url": base_url,
+        "token": token,
+        "cookie": cookie,
+        "timeout": timeout,
+        "auth_storage": auth_storage,
+    }
+
+
+def _configured_client() -> StacklessClient:
+    base_url = str(_client_config.get("base_url") or "")
+    token = str(_client_config.get("token") or "")
+    cookie = str(_client_config.get("cookie") or "")
+    timeout = float(_client_config.get("timeout") or DEFAULT_TIMEOUT_SECONDS)
+    auth_storage = str(_client_config.get("auth_storage") or "keyring")
+
+    if not base_url:
+        raise ValueError(
+            "Stackless base URL is required. Pass --base-url or " "STACKLESS_BASE_URL."
+        )
+    if not token and not cookie:
+        credential = CredentialStore(auth_storage).load(base_url)
+        if credential:
+            token = credential.token
+    if not token and not cookie:
+        raise ValueError(
+            "Stackless auth is required. Run `stackless-mcp login --base-url ...`, "
+            "or pass --token/STACKLESS_TOKEN or --cookie/STACKLESS_COOKIE."
+        )
+
+    return StacklessClient(
+        base_url=base_url,
+        token=token,
+        cookie=cookie,
+        timeout=timeout,
+    )
+
+
+def get_client() -> StacklessClient:
+    global _client
+    if _client is None:
+        if not _client_config:
+            configure_client(
+                base_url=os.environ.get("STACKLESS_BASE_URL", ""),
+                token=os.environ.get("STACKLESS_TOKEN", ""),
+                cookie=os.environ.get("STACKLESS_COOKIE", ""),
+                timeout=float(
+                    os.environ.get(
+                        "STACKLESS_TIMEOUT_SECONDS",
+                        DEFAULT_TIMEOUT_SECONDS,
+                    )
+                ),
+                auth_storage=os.environ.get("STACKLESS_AUTH_STORAGE", "keyring"),
+            )
+        _client = _configured_client()
+    assert _client is not None
+    return _client
+
+
+@atexit.register
+def _close_client() -> None:
+    if _client is not None:
+        _client.close()
+
+
+@_mcp_tool(TOOLSET_CORE)
+def check_stackless_connection() -> dict:
+    """Check Stackless API auth with the configured Bearer token or ALB cookies."""
+    return handlers.check_stackless_connection(get_client())
+
+
+def _run_workflow(
+    tool_name: str,
+    arguments: dict[str, Any],
+    idempotency_key: str | None = None,
+    confirmation: dict[str, Any] | None = None,
+) -> dict:
+    return handlers.run_mcp_workflow(
+        get_client(),
+        tool_name,
+        arguments,
+        idempotency_key,
+        confirmation,
+    )
+
+
+@_mcp_tool("catalog")
+def find_relevant_data(
+    question: str,
+    filters: dict[str, Any] | None = None,
+    limit: int = 10,
+) -> dict:
+    """Find catalog assets relevant to a business question."""
+    return _run_workflow(
+        "find_relevant_data",
+        {"question": question, "filters": filters or {}, "limit": limit},
+    )
+
+
+@_mcp_tool("catalog")
+def explain_stackless_asset(asset_ref: str, expand: bool = False) -> dict:
+    """Explain a Stackless catalog asset, columns, and lineage handles."""
+    return _run_workflow(
+        "explain_stackless_asset",
+        {"asset_ref": asset_ref, "expand": expand},
+    )
+
+
+@_mcp_tool("catalog")
+def list_stackless_assets(
+    asset_type: str | None = None,
+    domain: str | None = None,
+    tag: str | None = None,
+    limit: int = 25,
+) -> dict:
+    """List visible catalog assets filtered by type, domain, or tag."""
+    return _run_workflow(
+        "list_stackless_assets",
+        {
+            "asset_type": asset_type,
+            "domain": domain,
+            "tag": tag,
+            "limit": limit,
+        },
+    )
+
+
+@_mcp_tool("catalog")
+def trace_stackless_lineage(
+    asset_ref: str,
+    direction: str = "upstream",
+    max_depth: int | None = None,
+) -> dict:
+    """Trace upstream or downstream lineage for a visible catalog asset."""
+    return _run_workflow(
+        "trace_stackless_lineage",
+        {
+            "asset_ref": asset_ref,
+            "direction": direction,
+            "max_depth": max_depth,
+        },
+    )
+
+
+@_mcp_tool("catalog")
+def get_stackless_column_details(asset_ref: str) -> dict:
+    """Get column metadata for a visible catalog asset."""
+    return _run_workflow(
+        "get_stackless_column_details",
+        {"asset_ref": asset_ref},
+    )
+
+
+@_mcp_tool("catalog")
+def refresh_stackless_catalog(mode: str = "light") -> dict:
+    """Refresh the Stackless catalog snapshot in light or full mode."""
+    return _run_workflow("refresh_stackless_catalog", {"mode": mode})
+
+
+@_mcp_tool("catalog")
+def get_stackless_connector_status(
+    connector_ref: str | None = None,
+    connector_name: str | None = None,
+) -> dict:
+    """Get Fivetran connector sync status from the catalog snapshot."""
+    return _run_workflow(
+        "get_stackless_connector_status",
+        {
+            "connector_ref": connector_ref,
+            "connector_name": connector_name,
+        },
+    )
+
+
+@_mcp_tool("catalog")
+def get_stackless_asset_freshness(
+    asset_ref: str,
+    max_depth: int | None = None,
+) -> dict:
+    """Summarize upstream freshness for a visible catalog asset."""
+    return _run_workflow(
+        "get_stackless_asset_freshness",
+        {"asset_ref": asset_ref, "max_depth": max_depth},
+    )
+
+
+@_mcp_tool("catalog")
+def list_stackless_snowflake_schemas() -> dict:
+    """List Snowflake schemas visible in the catalog snapshot."""
+    return _run_workflow("list_stackless_snowflake_schemas", {})
+
+
+@_mcp_tool("catalog")
+def query_snowflake(sql: str, schema_name: str = "") -> dict:
+    """Run direct read-only Snowflake SQL through Stackless guardrails."""
+    return _run_workflow(
+        "query_snowflake",
+        {"sql": sql, "schema_name": schema_name},
+    )
+
+
+@_mcp_tool("runs")
+def list_stackless_runs(
+    status: str | None = None,
+    limit: int = 20,
+) -> dict:
+    """List recent Stackless scheduler runs with stable run refs."""
+    return _run_workflow("list_stackless_runs", {"status": status, "limit": limit})
+
+
+@_mcp_tool("runs")
+def diagnose_stackless_run(
+    run_ref: str | None = None,
+    job_name: str | None = None,
+) -> dict:
+    """Diagnose a Stackless run from run ref or job name with redacted logs."""
+    return _run_workflow(
+        "diagnose_stackless_run",
+        {"run_ref": run_ref, "job_name": job_name},
+    )
+
+
+@_mcp_tool("transformations")
+def draft_transformation_model(
+    goal: str,
+    parent_refs: list[dict[str, Any]] | None = None,
+    spec: dict[str, Any] | None = None,
+    name: str | None = None,
+    title: str | None = None,
+    description: str | None = None,
+    idempotency_key: str | None = None,
+) -> dict:
+    """Draft or scaffold a Stackless Transformation Model."""
+    return _run_workflow(
+        "draft_transformation_model",
+        {
+            "goal": goal,
+            "parent_refs": parent_refs or [],
+            "spec": spec,
+            "name": name,
+            "title": title,
+            "description": description,
+        },
+        idempotency_key,
+    )
+
+
+@_mcp_tool("transformations")
+def validate_transformation_model_draft(draft_ref: str) -> dict:
+    """Validate a Transformation Model draft without previewing or publishing."""
+    return _run_workflow(
+        "validate_transformation_model_draft",
+        {"draft_ref": draft_ref},
+    )
+
+
+@_mcp_tool("transformations")
+def preview_transformation_model_draft(
+    draft_ref: str,
+    idempotency_key: str,
+    mode: str = "view_preview",
+) -> dict:
+    """Queue a guarded preview for a Transformation Model draft."""
+    return _run_workflow(
+        "preview_transformation_model_draft",
+        {"draft_ref": draft_ref, "mode": mode},
+        idempotency_key,
+    )
+
+
+@_mcp_tool("semantic")
+def draft_semantic_model(
+    source_ref: str | None = None,
+    intent: str | None = None,
+    model: dict[str, Any] | None = None,
+    idempotency_key: str | None = None,
+) -> dict:
+    """Draft or scaffold a Stackless Semantic Model."""
+    return _run_workflow(
+        "draft_semantic_model",
+        {"source_ref": source_ref, "intent": intent, "model": model},
+        idempotency_key,
+    )
+
+
+@_mcp_tool("semantic")
+def validate_semantic_model_draft(
+    model_ref: str | None = None,
+    model: dict[str, Any] | None = None,
+) -> dict:
+    """Validate a Semantic Model draft or payload without publishing."""
+    return _run_workflow(
+        "validate_semantic_model_draft",
+        {"model_ref": model_ref, "model": model},
+    )
+
+
+@_mcp_tool("semantic")
+def preview_semantic_model_query(
+    model_ref: str,
+    query: dict[str, Any] | None = None,
+) -> dict:
+    """Run a Semantic Model preview query through Stackless guardrails."""
+    return _run_workflow(
+        "preview_semantic_model_query",
+        {"model_ref": model_ref, "query": query or {}},
+    )
+
+
+@_mcp_tool("semantic")
+def list_stackless_semantic_models(
+    status: str | None = None,
+    source_table_asset: str | None = None,
+    limit: int = 25,
+) -> dict:
+    """List Semantic Models visible in the Stackless webapp."""
+    return _run_workflow(
+        "list_stackless_semantic_models",
+        {
+            "status": status,
+            "source_table_asset": source_table_asset,
+            "limit": limit,
+        },
+    )
+
+
+@_mcp_tool("semantic")
+def explain_stackless_semantic_model(model_ref: str) -> dict:
+    """Explain a Stackless Semantic Model and its queryable members."""
+    return _run_workflow(
+        "explain_stackless_semantic_model",
+        {"model_ref": model_ref},
+    )
+
+
+@_mcp_tool("dashboards")
+def list_stackless_dashboards(
+    status: str | None = None,
+    workspace_type: str | None = None,
+    provider: str | None = None,
+    conversation_id: str | None = None,
+    limit: int = 25,
+) -> dict:
+    """List dashboards visible in the Stackless webapp."""
+    return _run_workflow(
+        "list_stackless_dashboards",
+        {
+            "status": status,
+            "workspace_type": workspace_type,
+            "provider": provider,
+            "conversation_id": conversation_id,
+            "limit": limit,
+        },
+    )
+
+
+@_mcp_tool("dashboards")
+def get_stackless_dashboard(dashboard_ref: str) -> dict:
+    """Get a dashboard, its custom spec, and the recommended MCP workflow."""
+    return _run_workflow(
+        "get_stackless_dashboard",
+        {"dashboard_ref": dashboard_ref},
+    )
+
+
+@_mcp_tool("dashboards")
+def diagnose_stackless_dashboard(
+    dashboard_ref: str,
+    request: dict[str, Any] | None = None,
+    max_widgets: int = 25,
+) -> dict:
+    """Diagnose custom dashboard widget hydration failures through Stackless/Cube."""
+    return _run_workflow(
+        "diagnose_stackless_dashboard",
+        {
+            "dashboard_ref": dashboard_ref,
+            "request": request or {},
+            "max_widgets": max_widgets,
+        },
+    )
+
+
+@_mcp_tool("dashboards")
+def draft_dashboard_from_goal(
+    goal: str,
+    model_refs: list[str] | None = None,
+    spec: dict[str, Any] | None = None,
+    title: str | None = None,
+    description: str | None = None,
+    idempotency_key: str | None = None,
+) -> dict:
+    """Draft or scaffold a Stackless custom dashboard."""
+    return _run_workflow(
+        "draft_dashboard_from_goal",
+        {
+            "goal": goal,
+            "model_refs": model_refs or [],
+            "spec": spec,
+            "title": title,
+            "description": description,
+        },
+        idempotency_key,
+    )
+
+
+@_mcp_tool("dashboards")
+def fork_stackless_dashboard(
+    dashboard_ref: str,
+    idempotency_key: str,
+    title: str | None = None,
+    description: str | None = None,
+    conversation_id: str | None = None,
+) -> dict:
+    """Fork a published Stackless custom dashboard into a private editable draft."""
+    return _run_workflow(
+        "fork_stackless_dashboard",
+        {
+            "dashboard_ref": dashboard_ref,
+            "title": title,
+            "description": description,
+            "conversation_id": conversation_id,
+        },
+        idempotency_key,
+    )
+
+
+@_mcp_tool("dashboards")
+def update_stackless_dashboard_draft(
+    dashboard_ref: str,
+    idempotency_key: str,
+    title: str | None = None,
+    description: str | None = None,
+    spec: dict[str, Any] | None = None,
+    conversation_id: str | None = None,
+) -> dict:
+    """Update a private Stackless custom dashboard draft after validation."""
+    return _run_workflow(
+        "update_stackless_dashboard_draft",
+        {
+            "dashboard_ref": dashboard_ref,
+            "title": title,
+            "description": description,
+            "spec": spec,
+            "conversation_id": conversation_id,
+        },
+        idempotency_key,
+    )
+
+
+@_mcp_tool("dashboards")
+def hydrate_dashboard(
+    dashboard_ref: str,
+    request: dict[str, Any] | None = None,
+) -> dict:
+    """Hydrate custom dashboard widgets through Stackless/Cube guardrails."""
+    return _run_workflow(
+        "hydrate_dashboard",
+        {"dashboard_ref": dashboard_ref, "request": request or {}},
+    )
+
+
+@_mcp_tool("dashboards")
+def validate_dashboard_draft(
+    dashboard_ref: str | None = None,
+    spec: dict[str, Any] | None = None,
+) -> dict:
+    """Validate a custom dashboard draft or spec without publishing."""
+    return _run_workflow(
+        "validate_dashboard_draft",
+        {"dashboard_ref": dashboard_ref, "spec": spec},
+    )
+
+
+@_mcp_tool("dashboards")
+def preview_dashboard(
+    dashboard_ref: str,
+    request: dict[str, Any] | None = None,
+) -> dict:
+    """Hydrate a custom dashboard preview through Stackless/Cube."""
+    return _run_workflow(
+        "preview_dashboard",
+        {"dashboard_ref": dashboard_ref, "request": request or {}},
+    )
+
+
+@_mcp_tool("lifecycle")
+def diff_stackless_draft(draft_ref: str, mode: str = "auto") -> dict:
+    """Inspect draft diff/review output before risky publication."""
+    return _run_workflow("diff_stackless_draft", {"draft_ref": draft_ref, "mode": mode})
+
+
+@_mcp_tool("lifecycle")
+def publish_stackless_draft(
+    draft_ref: str,
+    confirmation: dict[str, Any],
+    idempotency_key: str,
+    publish_context: dict[str, Any] | None = None,
+) -> dict:
+    """Publish a Stackless draft only after explicit confirmation."""
+    return _run_workflow(
+        "publish_stackless_draft",
+        {"draft_ref": draft_ref, "publish_context": publish_context or {}},
+        idempotency_key,
+        confirmation,
+    )
+
+
+@_mcp_tool("runs", "lifecycle")
+def cancel_stackless_run(
+    run_ref: str,
+    confirmation: dict[str, Any],
+    idempotency_key: str,
+) -> dict:
+    """Cancel a pending or running Stackless run after explicit confirmation."""
+    return _run_workflow(
+        "cancel_stackless_run",
+        {"run_ref": run_ref},
+        idempotency_key,
+        confirmation,
+    )
+
+
+@_mcp_tool("lifecycle")
+def unpublish_stackless_resource(
+    resource_ref: str,
+    confirmation: dict[str, Any],
+    idempotency_key: str,
+) -> dict:
+    """Unpublish a Stackless resource after explicit confirmation."""
+    return _run_workflow(
+        "unpublish_stackless_resource",
+        {"resource_ref": resource_ref},
+        idempotency_key,
+        confirmation,
+    )
+
+
+@_mcp_tool("lifecycle")
+def delete_stackless_draft(
+    draft_ref: str,
+    confirmation: dict[str, Any],
+    idempotency_key: str,
+) -> dict:
+    """Delete a draft resource after explicit confirmation."""
+    return _run_workflow(
+        "delete_stackless_draft",
+        {"draft_ref": draft_ref},
+        idempotency_key,
+        confirmation,
+    )
+
+
+@_mcp_tool(TOOLSET_CORE)
+def get_stackless_operation(operation_id: str) -> dict:
+    """Get MCP operation status from a returned operation_id."""
+    return handlers.get_mcp_operation(get_client(), operation_id)
+
+
+configure_toolsets(os.environ.get("STACKLESS_TOOLSETS", DEFAULT_TOOLSETS))
+
+
+def _parse_expires_at(value: str | None) -> datetime | None:
+    if not value:
+        return None
+    normalized = value.replace("Z", "+00:00")
+    parsed = datetime.fromisoformat(normalized)
+    if parsed.tzinfo is None:
+        return parsed.replace(tzinfo=timezone.utc)
+    return parsed.astimezone(timezone.utc)
+
+
+def _stored_credential_diagnostic(
+    *,
+    base_url: str,
+    auth_storage: str,
+) -> tuple[StoredCredential | None, dict[str, Any], list[str]]:
+    try:
+        credential = CredentialStore(auth_storage).load(base_url)
+    except Exception as exc:
+        return (
+            None,
+            {
+                "mode": "stored_token",
+                "storage_backend": auth_storage,
+                "credential_present": False,
+                "token_present": False,
+                "expires_at": None,
+                "expiry_status": "unknown",
+                "error": str(exc),
+            },
+            [f"auth storage check failed: {exc}"],
+        )
+
+    if credential is None:
+        return (
+            None,
+            {
+                "mode": "stored_token",
+                "storage_backend": auth_storage,
+                "credential_present": False,
+                "token_present": False,
+                "expires_at": None,
+                "expiry_status": "missing",
+            },
+            ["no token or cookie is configured"],
+        )
+
+    expires_at = credential.expires_at
+    expiry_status = "unknown"
+    errors: list[str] = []
+    if expires_at:
+        try:
+            expires_at_dt = _parse_expires_at(expires_at)
+        except ValueError:
+            expiry_status = "invalid"
+            errors.append("stored token expiry is invalid")
+        else:
+            if expires_at_dt is not None and expires_at_dt <= datetime.now(
+                timezone.utc
+            ):
+                expiry_status = "expired"
+                errors.append("stored token is expired")
+            else:
+                expiry_status = "unexpired"
+
+    return (
+        credential,
+        {
+            "mode": "stored_token",
+            "storage_backend": auth_storage,
+            "credential_present": True,
+            "token_present": bool(credential.token),
+            "expires_at": expires_at,
+            "expiry_status": expiry_status,
+            "user_id": credential.user_id,
+        },
+        errors,
+    )
+
+
+def _auth_diagnostic(
+    *,
+    base_url: str,
+    auth_storage: str,
+    token: str | None = None,
+    cookie: str | None = None,
+) -> tuple[str, str, dict[str, Any], list[str]]:
+    cli_token = (token or "").strip()
+    cli_cookie = (cookie or "").strip()
+    if cli_token:
+        return (
+            cli_token,
+            "",
+            {
+                "mode": "cli_token",
+                "storage_backend": auth_storage,
+                "credential_present": True,
+                "token_present": True,
+                "expires_at": None,
+                "expiry_status": "unknown",
+            },
+            [],
+        )
+    if cli_cookie:
+        return (
+            "",
+            cli_cookie,
+            {
+                "mode": "cli_cookie",
+                "storage_backend": auth_storage,
+                "credential_present": True,
+                "token_present": False,
+                "expires_at": None,
+                "expiry_status": "unknown",
+            },
+            [],
+        )
+
+    env_token = (os.environ.get("STACKLESS_TOKEN") or "").strip()
+    env_cookie = (os.environ.get("STACKLESS_COOKIE") or "").strip()
+    if env_token:
+        return (
+            env_token,
+            "",
+            {
+                "mode": "env_token",
+                "storage_backend": auth_storage,
+                "credential_present": True,
+                "token_present": True,
+                "expires_at": None,
+                "expiry_status": "unknown",
+            },
+            [],
+        )
+    if env_cookie:
+        return (
+            "",
+            env_cookie,
+            {
+                "mode": "env_cookie",
+                "storage_backend": auth_storage,
+                "credential_present": True,
+                "token_present": False,
+                "expires_at": None,
+                "expiry_status": "unknown",
+            },
+            [],
+        )
+
+    credential, auth, errors = _stored_credential_diagnostic(
+        base_url=base_url,
+        auth_storage=auth_storage,
+    )
+    return (credential.token if credential else "", "", auth, errors)
+
+
+def _doctor(args: argparse.Namespace) -> tuple[int, dict[str, Any]]:
+    base_url = (args.base_url or "").rstrip("/")
+    errors: list[str] = []
+    try:
+        selected_toolsets = sorted(parse_toolsets(args.toolsets))
+    except ValueError as exc:
+        selected_toolsets = []
+        errors.append(str(exc))
+
+    result: dict[str, Any] = {
+        "ok": False,
+        "base_url": base_url or None,
+        "client_version": __version__,
+        "auth_storage": args.auth_storage,
+        "toolsets": selected_toolsets,
+        "auth": {
+            "mode": "none",
+            "storage_backend": args.auth_storage,
+            "credential_present": False,
+            "token_present": False,
+            "expires_at": None,
+            "expiry_status": "unknown",
+        },
+        "connectivity": {
+            "ok": False,
+            "checked": False,
+            "error": None,
+        },
+        "capabilities": {
+            "ok": False,
+            "checked": False,
+            "server_version": None,
+            "server_supported_unknown_to_client": [],
+            "client_known_unsupported_by_server": [],
+            "error": None,
+        },
+        "errors": errors,
+    }
+
+    if not base_url:
+        result["errors"].append("--base-url or STACKLESS_BASE_URL is required")
+        return 1, result
+
+    token, cookie, auth, auth_errors = _auth_diagnostic(
+        base_url=base_url,
+        auth_storage=args.auth_storage,
+        token=getattr(args, "token", None),
+        cookie=getattr(args, "cookie", None),
+    )
+    result["auth"] = auth
+    result["errors"].extend(auth_errors)
+    if auth_errors:
+        return 1, result
+
+    try:
+        client = StacklessClient(
+            base_url=base_url,
+            token=token,
+            cookie=cookie,
+            timeout=args.timeout,
+        )
+    except Exception as exc:
+        result["connectivity"] = {
+            "ok": False,
+            "checked": False,
+            "error": str(exc),
+        }
+        result["errors"].append(str(exc))
+        return 1, result
+
+    try:
+        connectivity = handlers.check_stackless_connection(client)
+    except Exception as exc:
+        client.close()
+        result["connectivity"] = {
+            "ok": False,
+            "checked": True,
+            "error": str(exc),
+        }
+        result["errors"].append(f"connectivity check failed: {exc}")
+        return 1, result
+
+    if not isinstance(connectivity, dict):
+        connectivity = {"ok": bool(connectivity), "result": connectivity}
+    result["connectivity"] = {"checked": True, **connectivity}
+    if not connectivity.get("ok"):
+        client.close()
+        result["errors"].append("connectivity check failed")
+        return 1, result
+
+    try:
+        server_version, server_supported_tools = _validate_capabilities(
+            client.capabilities()
+        )
+    except Exception as exc:
+        result["capabilities"] = {
+            "ok": False,
+            "checked": True,
+            "server_version": None,
+            "server_supported_unknown_to_client": [],
+            "client_known_unsupported_by_server": [],
+            "error": str(exc),
+        }
+    else:
+        client_known_tools = _workflow_tool_names_for_toolsets(args.toolsets)
+        all_client_workflow_tools = _workflow_tool_names_for_toolsets(TOOLSET_ALL)
+        result["capabilities"] = {
+            "ok": True,
+            "checked": True,
+            "server_version": server_version,
+            "server_supported_unknown_to_client": sorted(
+                server_supported_tools - all_client_workflow_tools
+            ),
+            "client_known_unsupported_by_server": sorted(
+                client_known_tools - server_supported_tools
+            ),
+            "error": None,
+        }
+    finally:
+        client.close()
+
+    result["ok"] = not result["errors"]
+    return (0 if result["ok"] else 1), result
+
+
+def main() -> None:
+    load_dotenv()
+    if len(sys.argv) > 1 and sys.argv[1] in {
+        "login",
+        "logout",
+        "reset-auth",
+        "doctor",
+    }:
+        command = sys.argv[1]
+        parser = argparse.ArgumentParser(description=f"stackless-mcp {command}")
+        parser.add_argument("--base-url", default=os.environ.get("STACKLESS_BASE_URL"))
+        parser.add_argument(
+            "--auth-storage",
+            choices=["keyring", "file"],
+            default=os.environ.get("STACKLESS_AUTH_STORAGE", "keyring"),
+        )
+        parser.add_argument(
+            "--timeout",
+            type=float,
+            default=float(
+                os.environ.get("STACKLESS_TIMEOUT_SECONDS", DEFAULT_TIMEOUT_SECONDS)
+            ),
+        )
+        if command == "login":
+            parser.add_argument(
+                "--no-browser",
+                action="store_true",
+                help="Print the login URL instead of opening a browser.",
+            )
+            parser.add_argument("--client-name", default="stackless-mcp")
+        if command == "doctor":
+            parser.add_argument(
+                "--token",
+                default=None,
+                help="Stackless Bearer token for this diagnostic run.",
+            )
+            parser.add_argument(
+                "--cookie",
+                default=None,
+                help="Stackless Cookie header for this diagnostic run.",
+            )
+            parser.add_argument(
+                "--toolsets",
+                default=os.environ.get("STACKLESS_TOOLSETS", DEFAULT_TOOLSETS),
+                help=(
+                    "Comma-separated MCP toolsets to register. Use 'all' or any of: "
+                    + ", ".join(SELECTABLE_TOOLSETS)
+                    + ". Core auth/operation tools are always registered."
+                ),
+            )
+        args = parser.parse_args(sys.argv[2:])
+        if command == "doctor":
+            exit_code, result = _doctor(args)
+            print(json.dumps(result, indent=2, sort_keys=True))
+            if exit_code:
+                raise SystemExit(exit_code)
+            return
+        if not args.base_url:
+            print("--base-url or STACKLESS_BASE_URL is required", file=sys.stderr)
+            raise SystemExit(2)
+        try:
+            if command == "login":
+                credential = mcp_login(
+                    base_url=args.base_url,
+                    storage=args.auth_storage,
+                    timeout=args.timeout,
+                    open_browser=not args.no_browser,
+                    client_name=args.client_name,
+                )
+                print(
+                    json.dumps(
+                        {
+                            "ok": True,
+                            "base_url": credential.base_url,
+                            "expires_at": credential.expires_at,
+                            "user_id": credential.user_id,
+                        },
+                        indent=2,
+                        sort_keys=True,
+                    )
+                )
+            elif command == "logout":
+                revoked = mcp_logout(
+                    base_url=args.base_url,
+                    storage=args.auth_storage,
+                    timeout=args.timeout,
+                )
+                print(json.dumps({"ok": True, "revoked": revoked}, indent=2))
+            else:
+                CredentialStore(args.auth_storage).delete(args.base_url)
+                print(json.dumps({"ok": True, "reset": True}, indent=2))
+        except Exception as exc:
+            print(str(exc), file=sys.stderr)
+            raise SystemExit(1) from exc
+        return
+
+    parser = argparse.ArgumentParser(description="Run the Stackless MCP server")
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=f"stackless-mcp {__version__}",
+    )
+    parser.add_argument("--base-url", default=os.environ.get("STACKLESS_BASE_URL"))
+    parser.add_argument("--token", default=os.environ.get("STACKLESS_TOKEN"))
+    parser.add_argument("--cookie", default=os.environ.get("STACKLESS_COOKIE"))
+    parser.add_argument(
+        "--auth-storage",
+        choices=["keyring", "file"],
+        default=os.environ.get("STACKLESS_AUTH_STORAGE", "keyring"),
+    )
+    parser.add_argument(
+        "--toolsets",
+        default=os.environ.get("STACKLESS_TOOLSETS", DEFAULT_TOOLSETS),
+        help=(
+            "Comma-separated MCP toolsets to register. Use 'all' or any of: "
+            + ", ".join(SELECTABLE_TOOLSETS)
+            + ". Core auth/operation tools are always registered."
+        ),
+    )
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        default=float(
+            os.environ.get("STACKLESS_TIMEOUT_SECONDS", DEFAULT_TIMEOUT_SECONDS)
+        ),
+    )
+    parser.add_argument(
+        "--check",
+        action="store_true",
+        help="Check Stackless auth/API connectivity and exit without starting MCP.",
+    )
+    args = parser.parse_args()
+
+    configure_client(
+        base_url=args.base_url or "",
+        token=args.token or "",
+        cookie=args.cookie or "",
+        timeout=args.timeout,
+        auth_storage=args.auth_storage,
+    )
+    try:
+        configure_toolsets_from_server(args.toolsets)
+    except ValueError as exc:
+        parser.error(str(exc))
+    if args.check:
+        try:
+            result = handlers.check_stackless_connection(get_client())
+        except Exception as exc:
+            print(str(exc), file=sys.stderr)
+            raise SystemExit(1) from exc
+        print(json.dumps(result, indent=2, sort_keys=True))
+        return
+    mcp.run(transport="stdio")
+
+
+if __name__ == "__main__":
+    main()