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