airbyte_ops_mcp.mcp.organization_admin

   1# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
   2"""MCP tools for organization administration: is_agentic flag, payment config, and customer-tier lookup/cache.
   3
   4## MCP reference
   5
   6.. include:: ../../../docs/mcp-generated/organization_admin.md
   7    :start-line: 2
   8"""
   9
  10# NOTE: We intentionally do NOT use `from __future__ import annotations` here.
  11# FastMCP has issues resolving forward references when PEP 563 deferred annotations
  12# are used. See: https://github.com/jlowin/fastmcp/issues/905
  13# Python 3.12+ supports modern type hint syntax natively, so this is not needed.
  14
  15__all__: list[str] = []
  16
  17import logging
  18import re
  19from datetime import date, datetime, timedelta, timezone
  20from typing import Annotated, Any, Literal
  21
  22import requests
  23from airbyte import constants
  24from fastmcp import Context, FastMCP
  25from fastmcp_extensions import get_mcp_config, mcp_tool, register_mcp_tools
  26from pydantic import BaseModel, Field
  27from zoneinfo import ZoneInfo
  28
  29from airbyte_ops_mcp.approval_resolution import (
  30    ApprovalResolutionError,
  31    resolve_admin_email_from_approval,
  32)
  33from airbyte_ops_mcp.cloud_admin.auth import (
  34    CloudAuthError,
  35    require_internal_admin_flag_only,
  36)
  37from airbyte_ops_mcp.cloud_admin.entitlements import (
  38    WAIVER_TYPE_TO_ENTITLEMENT_PLAN,
  39    EntitlementAPIError,
  40    update_entitlement_plan,
  41)
  42from airbyte_ops_mcp.cloud_admin.models import (
  43    OrbSubscriptionInfo,
  44    OrganizationAgenticFlagBatchInfo,
  45    OrganizationAgenticFlagBatchUpdateResult,
  46    OrganizationAgenticFlagInfo,
  47    OrganizationAgenticFlagUpdateResult,
  48    OrganizationInfo,
  49    OrganizationPaymentConfigInfo,
  50    OrganizationPaymentConfigUpdateResult,
  51)
  52from airbyte_ops_mcp.cloud_admin.orb_billing import (
  53    OrbAPIError,
  54    _get_orb_api_key,
  55    _resolve_plan_id,
  56    extract_subscription_summary,
  57    get_active_subscription,
  58    schedule_plan_change,
  59)
  60from airbyte_ops_mcp.cloud_admin.organization_agentic_flag import (
  61    OrganizationAgenticFlagAPIError,
  62)
  63from airbyte_ops_mcp.cloud_admin.organization_agentic_flag import (
  64    get_organization_info as _get_organization_info,
  65)
  66from airbyte_ops_mcp.cloud_admin.organization_agentic_flag import (
  67    set_organization_agentic_status as _set_organization_agentic_status,
  68)
  69from airbyte_ops_mcp.cloud_admin.payment_config import (
  70    PaymentConfigAPIError,
  71    get_organization_info,
  72)
  73from airbyte_ops_mcp.cloud_admin.payment_config import (
  74    get_organization_payment_config as _get_organization_payment_config,
  75)
  76from airbyte_ops_mcp.cloud_admin.payment_config import (
  77    update_organization_payment_config as _update_organization_payment_config,
  78)
  79from airbyte_ops_mcp.cloud_admin.version_overrides import validate_tier_filter
  80from airbyte_ops_mcp.constants import ServerConfigKey
  81from airbyte_ops_mcp.prod_db_access.queries import (
  82    query_connection_workspace_details,
  83    query_organization_agentic_flags,
  84)
  85from airbyte_ops_mcp.tier_cache import (
  86    CustomerTier,
  87    TierCacheStats,
  88    TierFilter,
  89    TierSummary,
  90    build_tier_summary,
  91    get_cache_stats,
  92    get_org_tier,
  93    get_org_tiers,
  94    refresh_tier_cache,
  95    resolve_workspaces,
  96)
  97
  98
  99def _normalize_organization_ids(organization_ids: str | list[str]) -> list[str]:
 100    """Normalize one organization ID or a list of organization IDs."""
 101    ids = [organization_ids] if isinstance(organization_ids, str) else organization_ids
 102
 103    normalized = [org_id.strip() for org_id in ids if org_id and org_id.strip()]
 104    return list(dict.fromkeys(normalized))
 105
 106
 107def _build_update_failure_results(
 108    organization_ids: list[str],
 109    message: str,
 110) -> list[OrganizationAgenticFlagUpdateResult]:
 111    """Build per-org failure results for request-level validation failures."""
 112    return [
 113        OrganizationAgenticFlagUpdateResult(
 114            success=False,
 115            message=message,
 116            organization_id=organization_id,
 117        )
 118        for organization_id in organization_ids
 119    ]
 120
 121
 122def _resolve_cloud_auth(ctx: Context) -> tuple[str | None, str | None, str | None]:
 123    """Resolve auth credentials, returning `(bearer_token, client_id, client_secret)`."""
 124    bearer_token = get_mcp_config(ctx, ServerConfigKey.BEARER_TOKEN)
 125    if bearer_token:
 126        return bearer_token, None, None
 127
 128    client_id = get_mcp_config(ctx, ServerConfigKey.CLIENT_ID)
 129    client_secret = get_mcp_config(ctx, ServerConfigKey.CLIENT_SECRET)
 130    return None, client_id, client_secret
 131
 132
 133def _org_info_to_flag_info(
 134    org_info: OrganizationInfo,
 135    customer_tier: str | None = None,
 136    tier_warning: str | None = None,
 137) -> OrganizationAgenticFlagInfo:
 138    """Convert Config API organization info into agentic flag info."""
 139    return OrganizationAgenticFlagInfo(
 140        organization_id=org_info.organization_id,
 141        organization_name=org_info.organization_name,
 142        email=org_info.email,
 143        tombstone=False,
 144        is_agentic=org_info.is_agentic or False,
 145        customer_tier=customer_tier,
 146        tier_warning=tier_warning,
 147    )
 148
 149
 150def _row_to_flag_info(
 151    row: dict[str, object],
 152    customer_tier: str | None = None,
 153    tier_warning: str | None = None,
 154) -> OrganizationAgenticFlagInfo:
 155    """Convert a DB row into agentic flag info."""
 156    return OrganizationAgenticFlagInfo(
 157        organization_id=str(row["organization_id"]),
 158        organization_name=str(row["organization_name"])
 159        if row["organization_name"] is not None
 160        else None,
 161        email=str(row["email"]) if row["email"] is not None else None,
 162        tombstone=bool(row["tombstone"]),
 163        is_agentic=bool(row["is_agentic"]),
 164        customer_tier=customer_tier,
 165        tier_warning=tier_warning,
 166    )
 167
 168
 169@mcp_tool(
 170    destructive=False,
 171    idempotent=True,
 172    open_world=True,
 173)
 174def get_organization_agentic_flag(
 175    organization_ids: Annotated[
 176        str | list[str],
 177        Field(
 178            description="One organization UUID, or a list of organization UUIDs.",
 179        ),
 180    ],
 181    config_api_root: Annotated[
 182        str | None,
 183        Field(
 184            description="Optional Config API root URL override. When omitted, "
 185            "the tool reads from the Prod DB replica.",
 186            default=None,
 187        ),
 188    ] = None,
 189    *,
 190    ctx: Context,
 191) -> OrganizationAgenticFlagBatchInfo:
 192    """Get the current `is_agentic` flag for one or more organizations."""
 193    normalized_ids = _normalize_organization_ids(organization_ids)
 194    if not normalized_ids:
 195        return OrganizationAgenticFlagBatchInfo(
 196            organizations=[],
 197            missing_organization_ids=[],
 198        )
 199
 200    if config_api_root is None:
 201        rows = query_organization_agentic_flags(normalized_ids)
 202        rows_by_org_id = {str(row["organization_id"]): row for row in rows}
 203        organizations: list[OrganizationAgenticFlagInfo] = []
 204        for organization_id in normalized_ids:
 205            row = rows_by_org_id.get(organization_id)
 206            if row is None:
 207                continue
 208            tier_result = get_org_tier(organization_id)
 209            organizations.append(
 210                _row_to_flag_info(
 211                    row,
 212                    customer_tier=tier_result.customer_tier,
 213                    tier_warning=_build_tier_warning(tier_result.customer_tier),
 214                )
 215            )
 216        return OrganizationAgenticFlagBatchInfo(
 217            organizations=organizations,
 218            missing_organization_ids=[
 219                organization_id
 220                for organization_id in normalized_ids
 221                if organization_id not in rows_by_org_id
 222            ],
 223        )
 224
 225    resolved_api_root = config_api_root or constants.CLOUD_CONFIG_API_ROOT
 226    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
 227
 228    organizations: list[OrganizationAgenticFlagInfo] = []
 229    missing_organization_ids: list[str] = []
 230    for organization_id in normalized_ids:
 231        try:
 232            org_info = _get_organization_info(
 233                organization_id=organization_id,
 234                config_api_root=resolved_api_root,
 235                client_id=client_id,
 236                client_secret=client_secret,
 237                bearer_token=bearer_token,
 238            )
 239        except OrganizationAgenticFlagAPIError:
 240            missing_organization_ids.append(organization_id)
 241            continue
 242        if org_info is None:
 243            missing_organization_ids.append(organization_id)
 244            continue
 245
 246        tier_result = get_org_tier(organization_id)
 247        organizations.append(
 248            _org_info_to_flag_info(
 249                org_info,
 250                customer_tier=tier_result.customer_tier,
 251                tier_warning=_build_tier_warning(tier_result.customer_tier),
 252            )
 253        )
 254
 255    return OrganizationAgenticFlagBatchInfo(
 256        organizations=organizations,
 257        missing_organization_ids=missing_organization_ids,
 258    )
 259
 260
 261@mcp_tool(
 262    destructive=True,
 263    idempotent=True,
 264    open_world=True,
 265)
 266def update_organization_agentic_flag(
 267    organization_ids: Annotated[
 268        str | list[str],
 269        Field(
 270            description="One organization UUID, or a list of organization UUIDs.",
 271        ),
 272    ],
 273    is_agentic: Annotated[
 274        bool,
 275        Field(description="The desired `is_agentic` value."),
 276    ],
 277    approval_comment_url: Annotated[
 278        str,
 279        Field(
 280            description="URL to the Slack approval record. Obtain this by calling "
 281            "`escalate_to_human` with `approval_requested=True`; the backend "
 282            "delivers the approval record URL when a human clicks Approve.",
 283        ),
 284    ],
 285    organization_name: Annotated[
 286        str | None,
 287        Field(
 288            description="Confirmation of the target organization. Accepts the "
 289            "organization name, email address, or email domain. Required when "
 290            "updating a single organization. Ignored for multi-org updates; use "
 291            "`organization_names` instead.",
 292            default=None,
 293        ),
 294    ] = None,
 295    organization_names: Annotated[
 296        dict[str, str] | None,
 297        Field(
 298            description="Per-organization confirmation for multi-org updates. "
 299            "Keys are organization UUIDs; values may be organization names, "
 300            "email addresses, or email domains.",
 301            default=None,
 302        ),
 303    ] = None,
 304    customer_tier_filter: Annotated[
 305        TierFilter,
 306        Field(
 307            description=(
 308                "Required tier filter: 'TIER_0', 'TIER_1', 'TIER_2', or 'ALL'. "
 309                "The operation is rejected if the actual customer tier does not match. "
 310                "Use 'ALL' to proceed regardless of tier after human approval."
 311            ),
 312        ),
 313    ] = "TIER_2",
 314    *,
 315    ctx: Context,
 316) -> OrganizationAgenticFlagBatchUpdateResult:
 317    """Update the `is_agentic` flag for one or more organizations.
 318
 319    All updates require human-in-the-loop approval via `escalate_to_human`.
 320    The tool validates each organization identity before writing the flag.
 321    """
 322    normalized_ids = _normalize_organization_ids(organization_ids)
 323    if not normalized_ids:
 324        return OrganizationAgenticFlagBatchUpdateResult(
 325            success=False,
 326            message="No organization IDs provided.",
 327            results=[],
 328        )
 329
 330    try:
 331        require_internal_admin_flag_only()
 332    except CloudAuthError as e:
 333        return OrganizationAgenticFlagBatchUpdateResult(
 334            success=False,
 335            message=f"Admin authentication failed: {e}",
 336            results=_build_update_failure_results(
 337                normalized_ids, f"Admin authentication failed: {e}"
 338            ),
 339        )
 340
 341    try:
 342        resolve_admin_email_from_approval(
 343            approval_comment_url=approval_comment_url,
 344        )
 345    except ApprovalResolutionError as e:
 346        return OrganizationAgenticFlagBatchUpdateResult(
 347            success=False,
 348            message=str(e),
 349            results=_build_update_failure_results(normalized_ids, str(e)),
 350        )
 351
 352    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
 353
 354    results: list[OrganizationAgenticFlagUpdateResult] = []
 355    for organization_id in normalized_ids:
 356        rows = query_organization_agentic_flags([organization_id])
 357        if not rows:
 358            results.append(
 359                OrganizationAgenticFlagUpdateResult(
 360                    success=False,
 361                    message=f"Organization {organization_id} not found.",
 362                    organization_id=organization_id,
 363                )
 364            )
 365            continue
 366
 367        current_info = _row_to_flag_info(rows[0])
 368        if current_info.tombstone:
 369            results.append(
 370                OrganizationAgenticFlagUpdateResult(
 371                    success=False,
 372                    message=f"Organization {organization_id} is tombstoned.",
 373                    organization_id=organization_id,
 374                    organization_name=current_info.organization_name,
 375                    email=current_info.email,
 376                    previous_is_agentic=current_info.is_agentic,
 377                )
 378            )
 379            continue
 380
 381        org_info = OrganizationInfo.model_validate(
 382            {
 383                "organizationId": current_info.organization_id,
 384                "organizationName": current_info.organization_name or "",
 385                "email": current_info.email,
 386                "isAgentic": current_info.is_agentic,
 387            }
 388        )
 389        expected_name = (
 390            organization_names.get(organization_id)
 391            if organization_names is not None
 392            else organization_name
 393        )
 394        name_ok, name_error = _validate_organization_name(
 395            organization_id, expected_name, org_info
 396        )
 397        if not name_ok:
 398            results.append(
 399                OrganizationAgenticFlagUpdateResult(
 400                    success=False,
 401                    message=name_error or "Organization name validation failed.",
 402                    organization_id=organization_id,
 403                    organization_name=current_info.organization_name,
 404                    email=current_info.email,
 405                    previous_is_agentic=current_info.is_agentic,
 406                )
 407            )
 408            continue
 409
 410        tier_result = get_org_tier(organization_id)
 411        customer_tier = tier_result.customer_tier
 412        tier_warning = _build_tier_warning(customer_tier)
 413        tier_ok, tier_error = validate_tier_filter(customer_tier, customer_tier_filter)
 414        if not tier_ok:
 415            results.append(
 416                OrganizationAgenticFlagUpdateResult(
 417                    success=False,
 418                    message=tier_error or "Customer tier validation failed.",
 419                    organization_id=organization_id,
 420                    organization_name=current_info.organization_name,
 421                    email=current_info.email,
 422                    previous_is_agentic=current_info.is_agentic,
 423                    customer_tier=customer_tier,
 424                    tier_warning=tier_warning,
 425                )
 426            )
 427            continue
 428        previous_is_agentic = current_info.is_agentic
 429
 430        if previous_is_agentic == is_agentic:
 431            results.append(
 432                OrganizationAgenticFlagUpdateResult(
 433                    success=True,
 434                    message=(
 435                        f"Organization {organization_id} already has managed agentic "
 436                        f"org status {is_agentic}."
 437                    ),
 438                    organization_id=organization_id,
 439                    organization_name=current_info.organization_name,
 440                    email=current_info.email,
 441                    previous_is_agentic=previous_is_agentic,
 442                    new_is_agentic=is_agentic,
 443                    customer_tier=customer_tier,
 444                    tier_warning=tier_warning,
 445                )
 446            )
 447            continue
 448
 449        try:
 450            updated = _set_organization_agentic_status(
 451                organization_id=organization_id,
 452                is_agentic=is_agentic,
 453                config_api_root=constants.CLOUD_CONFIG_API_ROOT,
 454                client_id=client_id,
 455                client_secret=client_secret,
 456                bearer_token=bearer_token,
 457            )
 458        except Exception as e:
 459            results.append(
 460                OrganizationAgenticFlagUpdateResult(
 461                    success=False,
 462                    message=str(e),
 463                    organization_id=organization_id,
 464                    organization_name=current_info.organization_name,
 465                    email=current_info.email,
 466                    previous_is_agentic=previous_is_agentic,
 467                    customer_tier=customer_tier,
 468                    tier_warning=tier_warning,
 469                )
 470            )
 471            continue
 472
 473        updated_info = _org_info_to_flag_info(
 474            updated,
 475            customer_tier=customer_tier,
 476            tier_warning=tier_warning,
 477        )
 478        results.append(
 479            OrganizationAgenticFlagUpdateResult(
 480                success=True,
 481                message=(
 482                    f"Organization {organization_id} updated managed agentic org "
 483                    f"status to {updated_info.is_agentic}."
 484                ),
 485                organization_id=organization_id,
 486                organization_name=updated_info.organization_name,
 487                email=updated_info.email,
 488                previous_is_agentic=previous_is_agentic,
 489                new_is_agentic=updated_info.is_agentic,
 490                customer_tier=customer_tier,
 491                tier_warning=tier_warning,
 492            )
 493        )
 494
 495    success = all(result.success for result in results)
 496    return OrganizationAgenticFlagBatchUpdateResult(
 497        success=success,
 498        message=(
 499            f"Updated {sum(result.success for result in results)} of "
 500            f"{len(results)} organization agentic flag request(s)."
 501        ),
 502        results=results,
 503    )
 504
 505
 506logger = logging.getLogger(__name__)
 507
 508_DATE_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}$")
 509
 510_DAYS_PATTERN = re.compile(r"^\d+$")
 511
 512_PACIFIC = ZoneInfo("America/Los_Angeles")
 513
 514
 515def _build_tier_warning(customer_tier: str) -> str | None:
 516    """Build a warning message for sensitive customer tiers."""
 517    if customer_tier == "TIER_0":
 518        return (
 519            "WARNING: This is a TIER_0 (highest-value) customer. "
 520            "Proceed with extreme caution."
 521        )
 522    if customer_tier == "TIER_1":
 523        return "WARNING: This is a TIER_1 (high-value) customer. Proceed with caution."
 524    return None
 525
 526
 527def _fetch_orb_subscription_info(
 528    organization_id: str,
 529) -> OrbSubscriptionInfo | None:
 530    """Try to fetch the active Orb subscription for an organization.
 531
 532    Returns `None` silently if the Orb API key is not configured or if no
 533    active subscription is found. Logs warnings on API errors but does not
 534    raise — Orb data is supplemental, not required.
 535    """
 536    orb_api_key = _get_orb_api_key()
 537    if not orb_api_key:
 538        return None
 539
 540    try:
 541        active_sub = get_active_subscription(organization_id, orb_api_key)
 542    except (OrbAPIError, requests.RequestException):
 543        logger.warning(
 544            "Failed to fetch Orb subscription for org %s",
 545            organization_id,
 546            exc_info=True,
 547        )
 548        return None
 549
 550    if active_sub is None:
 551        return None
 552
 553    summary = extract_subscription_summary(active_sub)
 554    return OrbSubscriptionInfo(
 555        subscription_id=summary["subscription_id"],
 556        status=summary["status"],
 557        plan_name=summary.get("plan_name"),
 558        plan_id=summary.get("plan_id"),
 559        external_plan_id=summary.get("external_plan_id"),
 560        start_date=summary.get("start_date"),
 561        end_date=summary.get("end_date"),
 562        orb_customer_id=summary.get("orb_customer_id"),
 563    )
 564
 565
 566def _validate_organization_name(
 567    organization_id: str,
 568    organization_name: str | None,
 569    org_info: OrganizationInfo,
 570) -> tuple[bool, str | None]:
 571    """Validate `organization_name` against the org record from the Config API.
 572
 573    Accepts the org's literal name, email, or email domain as valid inputs.
 574
 575    Returns `(ok, error_message)`. On success `error_message` is `None`.
 576    """
 577    org_db_name: str = org_info.organization_name or ""
 578    org_db_email: str = org_info.email or ""
 579    org_db_domain: str = org_db_email.split("@", 1)[-1] if "@" in org_db_email else ""
 580
 581    valid_identifiers = [
 582        v
 583        for v in [org_db_name.lower(), org_db_email.lower(), org_db_domain.lower()]
 584        if v
 585    ]
 586
 587    identifier_parts = []
 588    if org_db_name:
 589        identifier_parts.append(f"'{org_db_name}' (org name)")
 590    if org_db_email:
 591        identifier_parts.append(f"'{org_db_email}' (email)")
 592    if org_db_domain:
 593        identifier_parts.append(f"'{org_db_domain}' (email domain)")
 594
 595    if not identifier_parts:
 596        return (
 597            False,
 598            f"Organization {organization_id} has no name or email on record.",
 599        )
 600
 601    hint_message = (
 602        f"To confirm, resend with `organization_name` set to one of: "
 603        f"{', '.join(identifier_parts)}. "
 604        f"Double-check that this is the correct organization before retrying. "
 605        f"If there is any doubt, confirm with your user."
 606    )
 607
 608    if organization_name is None:
 609        return (
 610            False,
 611            (
 612                f"`organization_name` is required to confirm the target organization. "
 613                f"{hint_message}"
 614            ),
 615        )
 616
 617    if organization_name.strip().lower().lstrip("@") not in valid_identifiers:
 618        return (
 619            False,
 620            (
 621                f"Organization name mismatch: '{organization_name}' does not match "
 622                f"organization {organization_id}. {hint_message}"
 623            ),
 624        )
 625
 626    return True, None
 627
 628
 629def _format_grace_period_end(target_date: date) -> str:
 630    """Build a backend-compatible datetime string for 11:59 PM Pacific on `target_date`.
 631
 632    The result is converted to UTC and formatted as
 633    `yyyy-MM-dd'T'HH:mm:ss.SSS+0000` to match the Java `@JsonFormat` annotation.
 634    """
 635    end_pacific = datetime(
 636        target_date.year,
 637        target_date.month,
 638        target_date.day,
 639        23,
 640        59,
 641        59,
 642        tzinfo=_PACIFIC,
 643    )
 644    end_utc = end_pacific.astimezone(timezone.utc)
 645    return end_utc.strftime("%Y-%m-%dT%H:%M:%S.000+0000")
 646
 647
 648def _parse_grace_period_value(
 649    value: str,
 650) -> tuple[str | None, str | None]:
 651    """Parse `set_grace_period` into an ISO 8601 datetime string or an action keyword.
 652
 653    Dates and day offsets resolve to 11:59 PM Pacific Time on the target date.
 654    Day offsets use the current date in Pacific Time as the starting point.
 655
 656    Returns `(iso_datetime_or_action, error_message)`.
 657    On success `error_message` is `None`.
 658    """
 659    value = value.strip()
 660
 661    if value.lower() == "cancel":
 662        return "cancel", None
 663
 664    today_pacific = datetime.now(tz=_PACIFIC).date()
 665
 666    if _DATE_PATTERN.match(value):
 667        try:
 668            parsed = date.fromisoformat(value)
 669        except ValueError:
 670            return None, f"Invalid date: {value}. Use YYYY-MM-DD format."
 671        if parsed < today_pacific:
 672            return None, f"Grace period end date {value} is in the past."
 673        if (parsed - today_pacific).days > 90:
 674            return None, (
 675                f"Grace period end date {value} is more than 90 days in the future. "
 676                "Maximum grace period is 90 days."
 677            )
 678        return _format_grace_period_end(parsed), None
 679
 680    if _DAYS_PATTERN.match(value):
 681        days = int(value)
 682        if days < 1 or days > 90:
 683            return None, f"Days must be between 1 and 90, got {days}."
 684        target_date = today_pacific + timedelta(days=days)
 685        return _format_grace_period_end(target_date), None
 686
 687    return None, (
 688        f"Invalid `set_grace_period` value: '{value}'. "
 689        "Expected a date (YYYY-MM-DD), an integer number of days, or 'cancel'."
 690    )
 691
 692
 693@mcp_tool(
 694    read_only=True,
 695    idempotent=True,
 696    open_world=True,
 697)
 698def get_organization_payment_config(
 699    organization_id: Annotated[
 700        str,
 701        "The organization UUID.",
 702    ],
 703    config_api_root: Annotated[
 704        str | None,
 705        Field(
 706            description="Optional Config API root URL override. "
 707            "Defaults to Airbyte Cloud (`https://cloud.airbyte.com/api/v1`).",
 708            default=None,
 709        ),
 710    ] = None,
 711    *,
 712    ctx: Context,
 713) -> OrganizationPaymentConfigInfo:
 714    """Get the current payment configuration for an organization.
 715
 716    Returns payment status, subscription status, grace period info,
 717    usage category override, and current Orb billing plan (when
 718    `ORB_API_KEY` is configured). No PII or sensitive payment details
 719    are included in the response.
 720
 721    Authentication credentials are resolved in priority order:
 722    1. Bearer token (Authorization header or AIRBYTE_CLOUD_BEARER_TOKEN env var)
 723    2. HTTP headers: X-Airbyte-Cloud-Client-Id, X-Airbyte-Cloud-Client-Secret
 724    3. Environment variables: AIRBYTE_CLOUD_CLIENT_ID, AIRBYTE_CLOUD_CLIENT_SECRET
 725    """
 726    resolved_api_root = config_api_root or constants.CLOUD_CONFIG_API_ROOT
 727    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
 728
 729    data = _get_organization_payment_config(
 730        organization_id=organization_id,
 731        config_api_root=resolved_api_root,
 732        client_id=client_id,
 733        client_secret=client_secret,
 734        bearer_token=bearer_token,
 735    )
 736
 737    # Enrich with tier info
 738    tier_result = get_org_tier(organization_id)
 739    tier_warning = _build_tier_warning(tier_result.customer_tier)
 740
 741    # Enrich with Orb subscription info (best-effort)
 742    orb_subscription = _fetch_orb_subscription_info(organization_id)
 743
 744    return OrganizationPaymentConfigInfo(
 745        organization_id=data["organizationId"],
 746        payment_status=data["paymentStatus"],
 747        subscription_status=data["subscriptionStatus"],
 748        payment_provider_id=data.get("paymentProviderId"),
 749        grace_period_end_at=data.get("gracePeriodEndAt"),
 750        usage_category_overwrite=data.get("usageCategoryOverwrite"),
 751        customer_tier=tier_result.customer_tier,
 752        tier_warning=tier_warning,
 753        orb_subscription=orb_subscription,
 754    )
 755
 756
 757@mcp_tool(
 758    destructive=True,
 759    idempotent=False,
 760    open_world=True,
 761)
 762def update_organization_payment_config(
 763    organization_id: Annotated[
 764        str,
 765        "The organization UUID.",
 766    ],
 767    approval_comment_url: Annotated[
 768        str,
 769        Field(
 770            description="URL to the Slack approval record. Obtain this by calling "
 771            "`escalate_to_human` with `approval_requested=True`; the backend "
 772            "delivers the approval record URL when a human clicks Approve.",
 773        ),
 774    ],
 775    organization_name: Annotated[
 776        str | None,
 777        Field(
 778            description="Confirmation of the target organization. Accepts the "
 779            "organization name, email address, or email domain. Required to "
 780            "prevent accidental modifications to the wrong organization. "
 781            "If omitted or mismatched, the tool returns an error with the valid "
 782            "identifiers so you can verify and retry.",
 783            default=None,
 784        ),
 785    ] = None,
 786    set_grace_period: Annotated[
 787        str | None,
 788        Field(
 789            description="Set or modify the grace period. Accepts three forms: "
 790            "(1) A date in `YYYY-MM-DD` format — grace period ends at 11:59 PM Pacific on that date. "
 791            "(2) An integer number of days (1-90) from today (Pacific Time) — "
 792            "grace period ends at 11:59 PM Pacific on the resulting date. "
 793            "(3) `'cancel'` to terminate the current grace period "
 794            "(sets status to `manual`). "
 795            "Requires `set_grace_period_reason` when setting or extending.",
 796            default=None,
 797        ),
 798    ] = None,
 799    set_grace_period_reason: Annotated[
 800        str | None,
 801        Field(
 802            description="Reason for starting, extending, or canceling the grace "
 803            "period. Required when `set_grace_period` is a date or number of days.",
 804            default=None,
 805        ),
 806    ] = None,
 807    set_permanent_waiver_type: Annotated[
 808        Literal["free", "internal", "none"] | None,
 809        Field(
 810            description="Set a permanent billing waiver for the organization. "
 811            "Use `'free'` for partner accounts that should not be billed, "
 812            "`'internal'` for Airbyte-internal organizations, "
 813            "or `'none'` to remove an existing waiver. "
 814            "Mutually exclusive with `set_grace_period`.",
 815            default=None,
 816        ),
 817    ] = None,
 818    set_permanent_waiver_reason: Annotated[
 819        str | None,
 820        Field(
 821            description="Reason for setting the permanent billing waiver. "
 822            "Required when `set_permanent_waiver_type` is provided.",
 823            default=None,
 824        ),
 825    ] = None,
 826    config_api_root: Annotated[
 827        str | None,
 828        Field(
 829            description="Optional Config API root URL override. "
 830            "Defaults to Airbyte Cloud (`https://cloud.airbyte.com/api/v1`).",
 831            default=None,
 832        ),
 833    ] = None,
 834    *,
 835    ctx: Context,
 836) -> OrganizationPaymentConfigUpdateResult:
 837    """Update the payment configuration for an organization.
 838
 839    All updates require human-in-the-loop approval via `escalate_to_human`.
 840
 841    Use `set_grace_period` to start, extend, or cancel a grace period.
 842    If the org is not already in `manual` status, the tool automatically
 843    transitions to `manual` first before setting the grace period.
 844
 845    Use `set_permanent_waiver_type` to mark an organization as a partner (`free`)
 846    or internal (`internal`) account. This is mutually exclusive with
 847    `set_grace_period` — only one may be provided per call. Setting the waiver
 848    type to `free` or `internal` also changes the Orb billing plan (`free` →
 849    Airbyte Partner, `internal` → Airbyte Internal). The `ORB_API_KEY`
 850    environment variable must be configured for waiver type changes.
 851
 852    The `organization_name` parameter is a safety check: the tool looks up the
 853    organization via the Config API and verifies that the provided name, email, or
 854    email domain matches. If omitted or mismatched, the tool returns the valid
 855    identifiers so the caller can verify and retry.
 856    """
 857    # --- Validate that an action was specified ---
 858    if set_grace_period is None and set_permanent_waiver_type is None:
 859        return OrganizationPaymentConfigUpdateResult(
 860            success=False,
 861            message="No action specified. Provide `set_grace_period` with a date "
 862            "(YYYY-MM-DD), number of days (1-90), or 'cancel'; or "
 863            "`set_permanent_waiver_type` with 'free' (partner) or 'internal'.",
 864            organization_id=organization_id,
 865        )
 866
 867    # --- Validate mutual exclusivity ---
 868    if set_grace_period is not None and set_permanent_waiver_type is not None:
 869        return OrganizationPaymentConfigUpdateResult(
 870            success=False,
 871            message="`set_grace_period` and `set_permanent_waiver_type` are "
 872            "mutually exclusive. Provide only one per call.",
 873            organization_id=organization_id,
 874        )
 875
 876    # --- Validate waiver reason is provided when waiver type is set ---
 877    if set_permanent_waiver_type is not None and not set_permanent_waiver_reason:
 878        return OrganizationPaymentConfigUpdateResult(
 879            success=False,
 880            message="`set_permanent_waiver_reason` is required when "
 881            "`set_permanent_waiver_type` is provided.",
 882            organization_id=organization_id,
 883        )
 884
 885    # --- Validate ORB_API_KEY is configured for waiver type changes ---
 886    if set_permanent_waiver_type in ("free", "internal") and not _get_orb_api_key():
 887        return OrganizationPaymentConfigUpdateResult(
 888            success=False,
 889            message="`ORB_API_KEY` environment variable is not configured. "
 890            "It is required when setting `set_permanent_waiver_type` to "
 891            "'free' or 'internal' because the Orb billing plan must also "
 892            "be changed.",
 893            organization_id=organization_id,
 894        )
 895
 896    # --- Validate admin access ---
 897    try:
 898        require_internal_admin_flag_only()
 899    except CloudAuthError as e:
 900        return OrganizationPaymentConfigUpdateResult(
 901            success=False,
 902            message=f"Admin authentication failed: {e}",
 903            organization_id=organization_id,
 904        )
 905
 906    # --- Resolve approval ---
 907    try:
 908        resolve_admin_email_from_approval(
 909            approval_comment_url=approval_comment_url,
 910        )
 911    except ApprovalResolutionError as e:
 912        return OrganizationPaymentConfigUpdateResult(
 913            success=False,
 914            message=str(e),
 915            organization_id=organization_id,
 916        )
 917
 918    # --- Resolve auth ---
 919    resolved_api_root = config_api_root or constants.CLOUD_CONFIG_API_ROOT
 920    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
 921
 922    # --- Look up organization info via Config API ---
 923    try:
 924        org_info = get_organization_info(
 925            organization_id=organization_id,
 926            config_api_root=resolved_api_root,
 927            client_id=client_id,
 928            client_secret=client_secret,
 929            bearer_token=bearer_token,
 930        )
 931    except PaymentConfigAPIError as e:
 932        return OrganizationPaymentConfigUpdateResult(
 933            success=False,
 934            message=f"Failed to fetch organization info: {e}",
 935            organization_id=organization_id,
 936        )
 937    if org_info is None:
 938        return OrganizationPaymentConfigUpdateResult(
 939            success=False,
 940            message=f"Organization {organization_id} not found.",
 941            organization_id=organization_id,
 942        )
 943
 944    # --- Validate organization name (safety check) ---
 945    name_ok, name_error = _validate_organization_name(
 946        organization_id, organization_name, org_info
 947    )
 948    if not name_ok:
 949        return OrganizationPaymentConfigUpdateResult(
 950            success=False,
 951            message=name_error or "Organization name validation failed.",
 952            organization_id=organization_id,
 953        )
 954
 955    # --- Enrich with tier info ---
 956    tier_result = get_org_tier(organization_id)
 957    customer_tier = tier_result.customer_tier
 958    tier_warning = _build_tier_warning(customer_tier)
 959
 960    # --- Permanent-waiver-only path (no grace period change) ---
 961    #
 962    # Statuses that the API cannot set back: uninitialized, okay, disabled.
 963    # If the org is in one of these, we transition to 'manual' first (same
 964    # pattern the grace-period path uses).
 965    _api_nonsettable_statuses = ("uninitialized", "okay", "disabled")
 966
 967    if set_grace_period is None:
 968        # Only set_permanent_waiver_type was requested
 969        assert set_permanent_waiver_type is not None
 970        try:
 971            current_config = _get_organization_payment_config(
 972                organization_id=organization_id,
 973                config_api_root=resolved_api_root,
 974                client_id=client_id,
 975                client_secret=client_secret,
 976                bearer_token=bearer_token,
 977            )
 978        except PaymentConfigAPIError as e:
 979            return OrganizationPaymentConfigUpdateResult(
 980                success=False,
 981                message=f"Failed to fetch current config: {e}",
 982                organization_id=organization_id,
 983                customer_tier=customer_tier,
 984                tier_warning=tier_warning,
 985            )
 986        current_status = current_config["paymentStatus"]
 987
 988        # Transition to 'manual' if current status is not API-settable
 989        target_status = current_status
 990        if current_status in _api_nonsettable_statuses:
 991            try:
 992                _update_organization_payment_config(
 993                    organization_id=organization_id,
 994                    payment_status="manual",
 995                    config_api_root=resolved_api_root,
 996                    client_id=client_id,
 997                    client_secret=client_secret,
 998                    bearer_token=bearer_token,
 999                    new_grace_period_reason=(
1000                        f"Transitioned to manual for permanent waiver: "
1001                        f"{set_permanent_waiver_reason}"
1002                    ),
1003                )
1004                target_status = "manual"
1005            except PaymentConfigAPIError as e:
1006                return OrganizationPaymentConfigUpdateResult(
1007                    success=False,
1008                    message=f"Failed to transition from '{current_status}' "
1009                    f"to 'manual': {e}",
1010                    organization_id=organization_id,
1011                    payment_status=current_status,
1012                    customer_tier=customer_tier,
1013                    tier_warning=tier_warning,
1014                )
1015
1016        try:
1017            data = _update_organization_payment_config(
1018                organization_id=organization_id,
1019                payment_status=target_status,
1020                config_api_root=resolved_api_root,
1021                client_id=client_id,
1022                client_secret=client_secret,
1023                bearer_token=bearer_token,
1024                usage_category_overwrite=(
1025                    set_permanent_waiver_type
1026                    if set_permanent_waiver_type != "none"
1027                    else ""
1028                ),
1029            )
1030        except PaymentConfigAPIError as e:
1031            return OrganizationPaymentConfigUpdateResult(
1032                success=False,
1033                message=f"Failed to set permanent waiver type: {e}",
1034                organization_id=organization_id,
1035                payment_status=target_status,
1036                customer_tier=customer_tier,
1037                tier_warning=tier_warning,
1038            )
1039        parts = [
1040            f"Permanent waiver type set to '{set_permanent_waiver_type}' "
1041            f"for org {organization_id}.",
1042        ]
1043        if current_status in _api_nonsettable_statuses:
1044            parts.append(
1045                f"Payment status transitioned from '{current_status}' to 'manual'."
1046            )
1047
1048        # --- Orb plan change (required for "free" / "internal") ---
1049        orb_plan_change_result: str | None = None
1050        if set_permanent_waiver_type in ("free", "internal"):
1051            # ORB_API_KEY is validated at the top of the function, so this
1052            # is guaranteed to be non-None here.
1053            orb_api_key = _get_orb_api_key()
1054            assert orb_api_key, "ORB_API_KEY should have been validated earlier"
1055            try:
1056                active_sub = get_active_subscription(organization_id, orb_api_key)
1057                if active_sub is None:
1058                    parts.append(
1059                        "Orb plan change skipped: no active subscription "
1060                        "found for this organization in Orb."
1061                    )
1062                    orb_plan_change_result = "Skipped: no active Orb subscription"
1063                else:
1064                    target_plan_id = _resolve_plan_id(set_permanent_waiver_type)
1065                    current_plan_id = (active_sub.get("plan") or {}).get("id")
1066                    current_plan_name = (active_sub.get("plan") or {}).get(
1067                        "name", current_plan_id or "unknown"
1068                    )
1069                    if current_plan_id == target_plan_id:
1070                        orb_plan_change_result = (
1071                            f"Already on plan '{current_plan_name}'"
1072                        )
1073                        parts.append(f"Orb plan already set to '{current_plan_name}'.")
1074                    else:
1075                        schedule_plan_change(
1076                            subscription_id=active_sub["id"],
1077                            plan_id=target_plan_id,
1078                            api_key=orb_api_key,
1079                        )
1080                        orb_plan_change_result = (
1081                            f"Changed from '{current_plan_name}' to '{target_plan_id}'"
1082                        )
1083                        parts.append(f"Orb plan changed to '{target_plan_id}'.")
1084            except (OrbAPIError, requests.RequestException) as e:
1085                parts.append(f"Orb plan change failed: {e}")
1086                orb_plan_change_result = f"Failed: {e}"
1087
1088        # --- Entitlement plan update (Stigg) ---
1089        entitlement_plan_change_result: str | None = None
1090        target_entitlement_plan = WAIVER_TYPE_TO_ENTITLEMENT_PLAN.get(
1091            set_permanent_waiver_type
1092        )
1093        if target_entitlement_plan:
1094            try:
1095                update_entitlement_plan(
1096                    organization_id=organization_id,
1097                    plan_name=target_entitlement_plan,
1098                    config_api_root=resolved_api_root,
1099                    client_id=client_id,
1100                    client_secret=client_secret,
1101                    bearer_token=bearer_token,
1102                )
1103                entitlement_plan_change_result = f"Changed to {target_entitlement_plan}"
1104                parts.append(
1105                    f"Entitlement plan updated to '{target_entitlement_plan}'."
1106                )
1107            except EntitlementAPIError as e:
1108                parts.append(f"Entitlement plan update failed: {e}")
1109                entitlement_plan_change_result = f"Failed: {e}"
1110
1111        return OrganizationPaymentConfigUpdateResult(
1112            success=True,
1113            message=" ".join(parts),
1114            organization_id=organization_id,
1115            payment_status=data["paymentStatus"],
1116            grace_period_end_at=data.get("gracePeriodEndAt"),
1117            permanent_waiver_type=data.get("usageCategoryOverwrite"),
1118            customer_tier=customer_tier,
1119            tier_warning=tier_warning,
1120            orb_plan_change=orb_plan_change_result,
1121            entitlement_plan_change=entitlement_plan_change_result,
1122        )
1123
1124    # --- Parse grace period value ---
1125    parsed_value, parse_error = _parse_grace_period_value(set_grace_period)
1126    if parse_error is not None:
1127        return OrganizationPaymentConfigUpdateResult(
1128            success=False,
1129            message=parse_error,
1130            organization_id=organization_id,
1131            customer_tier=customer_tier,
1132            tier_warning=tier_warning,
1133        )
1134
1135    assert parsed_value is not None
1136
1137    # --- Fetch current payment config (needed for both cancel and set paths) ---
1138    try:
1139        current_config = _get_organization_payment_config(
1140            organization_id=organization_id,
1141            config_api_root=resolved_api_root,
1142            client_id=client_id,
1143            client_secret=client_secret,
1144            bearer_token=bearer_token,
1145        )
1146    except PaymentConfigAPIError as e:
1147        return OrganizationPaymentConfigUpdateResult(
1148            success=False,
1149            message=f"Failed to fetch current config: {e}",
1150            organization_id=organization_id,
1151            customer_tier=customer_tier,
1152            tier_warning=tier_warning,
1153        )
1154
1155    current_status = current_config["paymentStatus"]
1156
1157    if parsed_value == "cancel":
1158        # Only allow cancel when org is actually in grace_period (or manual with
1159        # an active grace period end date).
1160        if current_status not in ("grace_period", "manual"):
1161            return OrganizationPaymentConfigUpdateResult(
1162                success=False,
1163                message=f"Cannot cancel grace period: organization is in "
1164                f"'{current_status}' status, not 'grace_period'.",
1165                organization_id=organization_id,
1166                payment_status=current_status,
1167                customer_tier=customer_tier,
1168                tier_warning=tier_warning,
1169            )
1170
1171        # Cancel grace period by setting status to manual
1172        try:
1173            data = _update_organization_payment_config(
1174                organization_id=organization_id,
1175                payment_status="manual",
1176                config_api_root=resolved_api_root,
1177                client_id=client_id,
1178                client_secret=client_secret,
1179                bearer_token=bearer_token,
1180                new_grace_period_reason=(
1181                    set_grace_period_reason or "Grace period canceled via MCP tool"
1182                ),
1183            )
1184        except PaymentConfigAPIError as e:
1185            return OrganizationPaymentConfigUpdateResult(
1186                success=False,
1187                message=str(e),
1188                organization_id=organization_id,
1189                customer_tier=customer_tier,
1190                tier_warning=tier_warning,
1191            )
1192
1193        return OrganizationPaymentConfigUpdateResult(
1194            success=True,
1195            message=f"Grace period canceled for org {organization_id}. "
1196            f"New status: {data['paymentStatus']}.",
1197            organization_id=organization_id,
1198            payment_status=data["paymentStatus"],
1199            grace_period_end_at=data.get("gracePeriodEndAt"),
1200            permanent_waiver_type=data.get("usageCategoryOverwrite"),
1201            customer_tier=customer_tier,
1202            tier_warning=tier_warning,
1203        )
1204
1205    # --- Setting/extending grace period ---
1206    if not set_grace_period_reason:
1207        return OrganizationPaymentConfigUpdateResult(
1208            success=False,
1209            message="`set_grace_period_reason` is required when setting or "
1210            "extending a grace period.",
1211            organization_id=organization_id,
1212            customer_tier=customer_tier,
1213            tier_warning=tier_warning,
1214        )
1215
1216    # The API only allows setting grace_period from manual status.
1217    # If not already in manual, transition to manual first.
1218    transitioned_to_manual = False
1219    if current_status != "manual":
1220        logger.info(
1221            "Org %s is in '%s' status; transitioning to 'manual' first.",
1222            organization_id,
1223            current_status,
1224        )
1225        try:
1226            _update_organization_payment_config(
1227                organization_id=organization_id,
1228                payment_status="manual",
1229                config_api_root=resolved_api_root,
1230                client_id=client_id,
1231                client_secret=client_secret,
1232                bearer_token=bearer_token,
1233            )
1234            transitioned_to_manual = True
1235        except PaymentConfigAPIError as e:
1236            return OrganizationPaymentConfigUpdateResult(
1237                success=False,
1238                message=f"Failed to transition to 'manual' from '{current_status}': {e}",
1239                organization_id=organization_id,
1240                customer_tier=customer_tier,
1241                tier_warning=tier_warning,
1242            )
1243
1244    # Now set the grace period
1245    try:
1246        data = _update_organization_payment_config(
1247            organization_id=organization_id,
1248            payment_status="grace_period",
1249            config_api_root=resolved_api_root,
1250            client_id=client_id,
1251            client_secret=client_secret,
1252            bearer_token=bearer_token,
1253            grace_period_end_at=parsed_value,
1254            new_grace_period_reason=set_grace_period_reason,
1255        )
1256    except PaymentConfigAPIError as e:
1257        if transitioned_to_manual:
1258            msg = (
1259                f"Failed to set grace period after transitioning to 'manual' "
1260                f"from '{current_status}'. The organization is now in 'manual' "
1261                f"status. Original error: {e}"
1262            )
1263        else:
1264            msg = str(e)
1265        return OrganizationPaymentConfigUpdateResult(
1266            success=False,
1267            message=msg,
1268            organization_id=organization_id,
1269            payment_status="manual" if transitioned_to_manual else None,
1270            customer_tier=customer_tier,
1271            tier_warning=tier_warning,
1272        )
1273
1274    return OrganizationPaymentConfigUpdateResult(
1275        success=True,
1276        message=f"Grace period set for org {organization_id}. "
1277        f"New status: {data['paymentStatus']}.",
1278        organization_id=organization_id,
1279        payment_status=data["paymentStatus"],
1280        grace_period_end_at=data.get("gracePeriodEndAt"),
1281        permanent_waiver_type=data.get("usageCategoryOverwrite"),
1282        customer_tier=customer_tier,
1283        tier_warning=tier_warning,
1284    )
1285
1286
1287class TierLookupEntry(BaseModel):
1288    """A single resolved tier entry for an org, workspace, or connection."""
1289
1290    input_id: str = Field(description="The original ID that was looked up")
1291    input_type: str = Field(
1292        description="Type of the input ID: 'organization', 'workspace', or 'connection'"
1293    )
1294    organization_id: str | None = Field(
1295        default=None, description="Resolved organization UUID"
1296    )
1297    workspace_id: str | None = Field(
1298        default=None,
1299        description="Workspace UUID (if input was workspace or connection)",
1300    )
1301    connection_id: str | None = Field(
1302        default=None, description="Connection UUID (if input was connection)"
1303    )
1304    customer_tier: CustomerTier = Field(
1305        description="Resolved tier: TIER_0, TIER_1, or TIER_2"
1306    )
1307    dataplane_name: str | None = Field(
1308        default=None, description="Dataplane region name (e.g., 'US', 'EU')"
1309    )
1310    is_eu: bool = Field(
1311        default=False, description="Whether the entity is in the EU region"
1312    )
1313    resolved: bool = Field(
1314        default=True, description="Whether the ID was successfully resolved"
1315    )
1316
1317
1318class TierLookupResult(BaseModel):
1319    """Result of a customer tier lookup across multiple IDs."""
1320
1321    entries: list[TierLookupEntry] = Field(
1322        description="Individual tier lookup results for each input ID"
1323    )
1324    summary: TierSummary = Field(
1325        description="Tier distribution summary across all resolved entries"
1326    )
1327    summary_text: str = Field(description="Human-readable tier distribution summary")
1328
1329
1330class TierCacheRefreshResult(BaseModel):
1331    """Result of refreshing the tier cache."""
1332
1333    stats: TierCacheStats = Field(description="Cache statistics after refresh")
1334    message: str = Field(description="Human-readable result message")
1335
1336
1337@mcp_tool(
1338    read_only=True,
1339    idempotent=True,
1340)
1341def lookup_customer_tiers(
1342    organization_ids: Annotated[
1343        list[str] | None,
1344        Field(
1345            description=(
1346                "List of organization UUIDs to look up tiers for. "
1347                "Example: ['664c690e-5263-49ba-b01f-4a6759b3330a']"
1348            ),
1349            default=None,
1350        ),
1351    ] = None,
1352    workspace_ids: Annotated[
1353        list[str] | None,
1354        Field(
1355            description=(
1356                "List of workspace UUIDs to look up tiers for. "
1357                "Each workspace will be resolved to its organization and tier. "
1358                "Example: ['266ebdfe-0d7b-4540-9817-de7e4505ba61']"
1359            ),
1360            default=None,
1361        ),
1362    ] = None,
1363    connection_ids: Annotated[
1364        list[str] | None,
1365        Field(
1366            description=(
1367                "List of connection UUIDs to look up tiers for. "
1368                "Each connection will be resolved to its workspace, organization, and tier."
1369            ),
1370            default=None,
1371        ),
1372    ] = None,
1373) -> TierLookupResult:
1374    """Look up customer tier classification for organizations, workspaces, and/or connections.
1375
1376    Accepts mixed lists of organization IDs, workspace IDs, and connection IDs.
1377    Resolves each to its organization and maps to a customer tier (TIER_0, TIER_1, or TIER_2).
1378
1379    Tier 0 and Tier 1 orgs are explicitly tracked in Salesforce and cached from the GCS export.
1380    Any org not in the cache defaults to TIER_2.
1381
1382    Returns enriched entries with tier, region (EU/US), and a summary of the distribution.
1383    """
1384    entries: list[TierLookupEntry] = []
1385
1386    # Resolve organization IDs directly
1387    org_ids = organization_ids or []
1388    if org_ids:
1389        org_results = get_org_tiers(org_ids)
1390        for result in org_results:
1391            entries.append(
1392                TierLookupEntry(
1393                    input_id=result.organization_id,
1394                    input_type="organization",
1395                    organization_id=result.organization_id,
1396                    customer_tier=result.customer_tier,
1397                    resolved=True,
1398                )
1399            )
1400
1401    # Resolve workspace IDs -> org -> tier
1402    ws_ids = workspace_ids or []
1403    if ws_ids:
1404        ws_results = resolve_workspaces(ws_ids)
1405        for result in ws_results:
1406            entries.append(
1407                TierLookupEntry(
1408                    input_id=result.workspace_id,
1409                    input_type="workspace",
1410                    organization_id=result.organization_id,
1411                    workspace_id=result.workspace_id,
1412                    customer_tier=result.customer_tier,
1413                    dataplane_name=result.dataplane_name,
1414                    is_eu=result.is_eu,
1415                    resolved=result.resolved,
1416                )
1417            )
1418
1419    # Resolve connection IDs -> workspace -> org -> tier
1420    conn_ids = connection_ids or []
1421    if conn_ids:
1422        conn_details = query_connection_workspace_details(conn_ids)
1423        conn_map: dict[str, dict[str, Any]] = {
1424            str(row["connection_id"]): row for row in conn_details
1425        }
1426
1427        for conn_id in conn_ids:
1428            row = conn_map.get(conn_id)
1429            if row is None:
1430                entries.append(
1431                    TierLookupEntry(
1432                        input_id=conn_id,
1433                        input_type="connection",
1434                        connection_id=conn_id,
1435                        customer_tier="TIER_2",
1436                        resolved=False,
1437                    )
1438                )
1439                continue
1440
1441            org_id = str(row["organization_id"])
1442            dataplane_name = row.get("dataplane_name") or "US"
1443            org_result = get_org_tier(org_id)
1444
1445            entries.append(
1446                TierLookupEntry(
1447                    input_id=conn_id,
1448                    input_type="connection",
1449                    organization_id=org_id,
1450                    workspace_id=str(row["workspace_id"]),
1451                    connection_id=conn_id,
1452                    customer_tier=org_result.customer_tier,
1453                    dataplane_name=dataplane_name,
1454                    is_eu=dataplane_name == "EU",
1455                    resolved=True,
1456                )
1457            )
1458
1459    # Build summary from enriched entries
1460    summary_rows = [{"customer_tier": e.customer_tier} for e in entries if e.resolved]
1461    summary = build_tier_summary(summary_rows)
1462
1463    return TierLookupResult(
1464        entries=entries,
1465        summary=summary,
1466        summary_text=str(summary),
1467    )
1468
1469
1470@mcp_tool(
1471    read_only=False,
1472    idempotent=True,
1473)
1474def refresh_customer_tier_cache() -> TierCacheRefreshResult:
1475    """Force-refresh the customer tier cache from the GCS export.
1476
1477    The tier cache is automatically refreshed every 24 hours. Use this tool to
1478    manually trigger a refresh if you need the latest tier data immediately
1479    (e.g., after a Salesforce update).
1480
1481    Returns cache statistics after the refresh.
1482    """
1483    stats = refresh_tier_cache()
1484    return TierCacheRefreshResult(
1485        stats=stats,
1486        message=(
1487            f"Tier cache refreshed: {stats.tier_cache_size} orgs cached, "
1488            f"{stats.workspace_cache_size} workspaces cached."
1489        ),
1490    )
1491
1492
1493@mcp_tool(
1494    read_only=True,
1495    idempotent=True,
1496)
1497def get_customer_tier_cache_stats() -> TierCacheStats:
1498    """Get current statistics about the customer tier cache.
1499
1500    Returns cache size, age, and file paths for both the tier cache
1501    (org -> tier) and workspace cache (workspace -> org + region).
1502
1503    Useful for checking cache freshness and diagnosing issues.
1504    """
1505    return get_cache_stats()
1506
1507
1508def register_organization_admin_tools(app: FastMCP) -> None:
1509    """Register organization_admin tools with the FastMCP app."""
1510    register_mcp_tools(app, mcp_module=__name__)