airbyte_ops_mcp.mcp.organization_payment_config

  1# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
  2"""MCP tools for organization payment config management.
  3
  4Provides tools to read and update organization payment configurations
  5in Airbyte Cloud, including grace period management and payment status changes.
  6
  7## MCP reference
  8
  9.. include:: ../../../docs/mcp-generated/organization_payment_config.md
 10    :start-line: 2
 11"""
 12
 13# NOTE: We intentionally do NOT use `from __future__ import annotations` here.
 14# FastMCP has issues resolving forward references when PEP 563 deferred annotations
 15# are used. See: https://github.com/jlowin/fastmcp/issues/905
 16
 17__all__: list[str] = []
 18
 19import logging
 20import re
 21from datetime import date, datetime, timedelta, timezone
 22from typing import Annotated, Literal
 23
 24import requests
 25from airbyte import constants
 26from fastmcp import Context, FastMCP
 27from fastmcp_extensions import get_mcp_config, mcp_tool, register_mcp_tools
 28from pydantic import Field
 29from zoneinfo import ZoneInfo
 30
 31from airbyte_ops_mcp.approval_resolution import (
 32    ApprovalResolutionError,
 33    resolve_admin_email_from_approval,
 34)
 35from airbyte_ops_mcp.cloud_admin.auth import (
 36    CloudAuthError,
 37    require_internal_admin_flag_only,
 38)
 39from airbyte_ops_mcp.cloud_admin.entitlements import (
 40    WAIVER_TYPE_TO_ENTITLEMENT_PLAN,
 41    EntitlementAPIError,
 42    update_entitlement_plan,
 43)
 44from airbyte_ops_mcp.cloud_admin.models import (
 45    OrbSubscriptionInfo,
 46    OrganizationInfo,
 47    OrganizationPaymentConfigInfo,
 48    OrganizationPaymentConfigUpdateResult,
 49)
 50from airbyte_ops_mcp.cloud_admin.orb_billing import (
 51    OrbAPIError,
 52    _get_orb_api_key,
 53    _resolve_plan_id,
 54    extract_subscription_summary,
 55    get_active_subscription,
 56    schedule_plan_change,
 57)
 58from airbyte_ops_mcp.cloud_admin.payment_config import (
 59    PaymentConfigAPIError,
 60    get_organization_info,
 61)
 62from airbyte_ops_mcp.cloud_admin.payment_config import (
 63    get_organization_payment_config as _get_organization_payment_config,
 64)
 65from airbyte_ops_mcp.cloud_admin.payment_config import (
 66    update_organization_payment_config as _update_organization_payment_config,
 67)
 68from airbyte_ops_mcp.constants import ServerConfigKey
 69from airbyte_ops_mcp.tier_cache import get_org_tier
 70
 71logger = logging.getLogger(__name__)
 72
 73_DATE_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}$")
 74_DAYS_PATTERN = re.compile(r"^\d+$")
 75_PACIFIC = ZoneInfo("America/Los_Angeles")
 76
 77
 78def _build_tier_warning(customer_tier: str) -> str | None:
 79    """Build a warning message for sensitive customer tiers."""
 80    if customer_tier == "TIER_0":
 81        return (
 82            "WARNING: This is a TIER_0 (highest-value) customer. "
 83            "Proceed with extreme caution."
 84        )
 85    if customer_tier == "TIER_1":
 86        return "WARNING: This is a TIER_1 (high-value) customer. Proceed with caution."
 87    return None
 88
 89
 90def _resolve_cloud_auth(ctx: Context) -> tuple[str | None, str | None, str | None]:
 91    """Resolve auth credentials, returning `(bearer_token, client_id, client_secret)`."""
 92    bearer_token = get_mcp_config(ctx, ServerConfigKey.BEARER_TOKEN)
 93    if bearer_token:
 94        return bearer_token, None, None
 95
 96    client_id = get_mcp_config(ctx, ServerConfigKey.CLIENT_ID)
 97    client_secret = get_mcp_config(ctx, ServerConfigKey.CLIENT_SECRET)
 98    return None, client_id, client_secret
 99
100
101def _fetch_orb_subscription_info(
102    organization_id: str,
103) -> OrbSubscriptionInfo | None:
104    """Try to fetch the active Orb subscription for an organization.
105
106    Returns `None` silently if the Orb API key is not configured or if no
107    active subscription is found. Logs warnings on API errors but does not
108    raise — Orb data is supplemental, not required.
109    """
110    orb_api_key = _get_orb_api_key()
111    if not orb_api_key:
112        return None
113
114    try:
115        active_sub = get_active_subscription(organization_id, orb_api_key)
116    except (OrbAPIError, requests.RequestException):
117        logger.warning(
118            "Failed to fetch Orb subscription for org %s",
119            organization_id,
120            exc_info=True,
121        )
122        return None
123
124    if active_sub is None:
125        return None
126
127    summary = extract_subscription_summary(active_sub)
128    return OrbSubscriptionInfo(
129        subscription_id=summary["subscription_id"],
130        status=summary["status"],
131        plan_name=summary.get("plan_name"),
132        plan_id=summary.get("plan_id"),
133        external_plan_id=summary.get("external_plan_id"),
134        start_date=summary.get("start_date"),
135        end_date=summary.get("end_date"),
136        orb_customer_id=summary.get("orb_customer_id"),
137    )
138
139
140def _validate_organization_name(
141    organization_id: str,
142    organization_name: str | None,
143    org_info: OrganizationInfo,
144) -> tuple[bool, str | None]:
145    """Validate `organization_name` against the org record from the Config API.
146
147    Accepts the org's literal name, email, or email domain as valid inputs.
148
149    Returns `(ok, error_message)`. On success `error_message` is `None`.
150    """
151    org_db_name: str = org_info.organization_name or ""
152    org_db_email: str = org_info.email or ""
153    org_db_domain: str = org_db_email.split("@", 1)[-1] if "@" in org_db_email else ""
154
155    valid_identifiers = [
156        v
157        for v in [org_db_name.lower(), org_db_email.lower(), org_db_domain.lower()]
158        if v
159    ]
160
161    identifier_parts = []
162    if org_db_name:
163        identifier_parts.append(f"'{org_db_name}' (org name)")
164    if org_db_email:
165        identifier_parts.append(f"'{org_db_email}' (email)")
166    if org_db_domain:
167        identifier_parts.append(f"'{org_db_domain}' (email domain)")
168
169    if not identifier_parts:
170        return (
171            False,
172            f"Organization {organization_id} has no name or email on record.",
173        )
174
175    hint_message = (
176        f"To confirm, resend with `organization_name` set to one of: "
177        f"{', '.join(identifier_parts)}. "
178        f"Double-check that this is the correct organization before retrying. "
179        f"If there is any doubt, confirm with your user."
180    )
181
182    if organization_name is None:
183        return (
184            False,
185            (
186                f"`organization_name` is required to confirm the target organization. "
187                f"{hint_message}"
188            ),
189        )
190
191    if organization_name.strip().lower().lstrip("@") not in valid_identifiers:
192        return (
193            False,
194            (
195                f"Organization name mismatch: '{organization_name}' does not match "
196                f"organization {organization_id}. {hint_message}"
197            ),
198        )
199
200    return True, None
201
202
203def _format_grace_period_end(target_date: date) -> str:
204    """Build a backend-compatible datetime string for 11:59 PM Pacific on `target_date`.
205
206    The result is converted to UTC and formatted as
207    `yyyy-MM-dd'T'HH:mm:ss.SSS+0000` to match the Java `@JsonFormat` annotation.
208    """
209    end_pacific = datetime(
210        target_date.year,
211        target_date.month,
212        target_date.day,
213        23,
214        59,
215        59,
216        tzinfo=_PACIFIC,
217    )
218    end_utc = end_pacific.astimezone(timezone.utc)
219    return end_utc.strftime("%Y-%m-%dT%H:%M:%S.000+0000")
220
221
222def _parse_grace_period_value(
223    value: str,
224) -> tuple[str | None, str | None]:
225    """Parse `set_grace_period` into an ISO 8601 datetime string or an action keyword.
226
227    Dates and day offsets resolve to 11:59 PM Pacific Time on the target date.
228    Day offsets use the current date in Pacific Time as the starting point.
229
230    Returns `(iso_datetime_or_action, error_message)`.
231    On success `error_message` is `None`.
232    """
233    value = value.strip()
234
235    if value.lower() == "cancel":
236        return "cancel", None
237
238    today_pacific = datetime.now(tz=_PACIFIC).date()
239
240    if _DATE_PATTERN.match(value):
241        try:
242            parsed = date.fromisoformat(value)
243        except ValueError:
244            return None, f"Invalid date: {value}. Use YYYY-MM-DD format."
245        if parsed < today_pacific:
246            return None, f"Grace period end date {value} is in the past."
247        if (parsed - today_pacific).days > 90:
248            return None, (
249                f"Grace period end date {value} is more than 90 days in the future. "
250                "Maximum grace period is 90 days."
251            )
252        return _format_grace_period_end(parsed), None
253
254    if _DAYS_PATTERN.match(value):
255        days = int(value)
256        if days < 1 or days > 90:
257            return None, f"Days must be between 1 and 90, got {days}."
258        target_date = today_pacific + timedelta(days=days)
259        return _format_grace_period_end(target_date), None
260
261    return None, (
262        f"Invalid `set_grace_period` value: '{value}'. "
263        "Expected a date (YYYY-MM-DD), an integer number of days, or 'cancel'."
264    )
265
266
267# =============================================================================
268# MCP Tools
269# =============================================================================
270
271
272@mcp_tool(
273    read_only=True,
274    idempotent=True,
275    open_world=True,
276)
277def get_organization_payment_config(
278    organization_id: Annotated[
279        str,
280        "The organization UUID.",
281    ],
282    config_api_root: Annotated[
283        str | None,
284        Field(
285            description="Optional Config API root URL override. "
286            "Defaults to Airbyte Cloud (`https://cloud.airbyte.com/api/v1`).",
287            default=None,
288        ),
289    ] = None,
290    *,
291    ctx: Context,
292) -> OrganizationPaymentConfigInfo:
293    """Get the current payment configuration for an organization.
294
295    Returns payment status, subscription status, grace period info,
296    usage category override, and current Orb billing plan (when
297    `ORB_API_KEY` is configured). No PII or sensitive payment details
298    are included in the response.
299
300    Authentication credentials are resolved in priority order:
301    1. Bearer token (Authorization header or AIRBYTE_CLOUD_BEARER_TOKEN env var)
302    2. HTTP headers: X-Airbyte-Cloud-Client-Id, X-Airbyte-Cloud-Client-Secret
303    3. Environment variables: AIRBYTE_CLOUD_CLIENT_ID, AIRBYTE_CLOUD_CLIENT_SECRET
304    """
305    resolved_api_root = config_api_root or constants.CLOUD_CONFIG_API_ROOT
306    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
307
308    data = _get_organization_payment_config(
309        organization_id=organization_id,
310        config_api_root=resolved_api_root,
311        client_id=client_id,
312        client_secret=client_secret,
313        bearer_token=bearer_token,
314    )
315
316    # Enrich with tier info
317    tier_result = get_org_tier(organization_id)
318    tier_warning = _build_tier_warning(tier_result.customer_tier)
319
320    # Enrich with Orb subscription info (best-effort)
321    orb_subscription = _fetch_orb_subscription_info(organization_id)
322
323    return OrganizationPaymentConfigInfo(
324        organization_id=data["organizationId"],
325        payment_status=data["paymentStatus"],
326        subscription_status=data["subscriptionStatus"],
327        payment_provider_id=data.get("paymentProviderId"),
328        grace_period_end_at=data.get("gracePeriodEndAt"),
329        usage_category_overwrite=data.get("usageCategoryOverwrite"),
330        customer_tier=tier_result.customer_tier,
331        tier_warning=tier_warning,
332        orb_subscription=orb_subscription,
333    )
334
335
336@mcp_tool(
337    destructive=True,
338    idempotent=False,
339    open_world=True,
340)
341def update_organization_payment_config(
342    organization_id: Annotated[
343        str,
344        "The organization UUID.",
345    ],
346    approval_comment_url: Annotated[
347        str,
348        Field(
349            description="URL to the Slack approval record. Obtain this by calling "
350            "`escalate_to_human` with `approval_requested=True`; the backend "
351            "delivers the approval record URL when a human clicks Approve.",
352        ),
353    ],
354    organization_name: Annotated[
355        str | None,
356        Field(
357            description="Confirmation of the target organization. Accepts the "
358            "organization name, email address, or email domain. Required to "
359            "prevent accidental modifications to the wrong organization. "
360            "If omitted or mismatched, the tool returns an error with the valid "
361            "identifiers so you can verify and retry.",
362            default=None,
363        ),
364    ] = None,
365    set_grace_period: Annotated[
366        str | None,
367        Field(
368            description="Set or modify the grace period. Accepts three forms: "
369            "(1) A date in `YYYY-MM-DD` format — grace period ends at 11:59 PM Pacific on that date. "
370            "(2) An integer number of days (1-90) from today (Pacific Time) — "
371            "grace period ends at 11:59 PM Pacific on the resulting date. "
372            "(3) `'cancel'` to terminate the current grace period "
373            "(sets status to `manual`). "
374            "Requires `set_grace_period_reason` when setting or extending.",
375            default=None,
376        ),
377    ] = None,
378    set_grace_period_reason: Annotated[
379        str | None,
380        Field(
381            description="Reason for starting, extending, or canceling the grace "
382            "period. Required when `set_grace_period` is a date or number of days.",
383            default=None,
384        ),
385    ] = None,
386    set_permanent_waiver_type: Annotated[
387        Literal["free", "internal", "none"] | None,
388        Field(
389            description="Set a permanent billing waiver for the organization. "
390            "Use `'free'` for partner accounts that should not be billed, "
391            "`'internal'` for Airbyte-internal organizations, "
392            "or `'none'` to remove an existing waiver. "
393            "Mutually exclusive with `set_grace_period`.",
394            default=None,
395        ),
396    ] = None,
397    set_permanent_waiver_reason: Annotated[
398        str | None,
399        Field(
400            description="Reason for setting the permanent billing waiver. "
401            "Required when `set_permanent_waiver_type` is provided.",
402            default=None,
403        ),
404    ] = None,
405    config_api_root: Annotated[
406        str | None,
407        Field(
408            description="Optional Config API root URL override. "
409            "Defaults to Airbyte Cloud (`https://cloud.airbyte.com/api/v1`).",
410            default=None,
411        ),
412    ] = None,
413    *,
414    ctx: Context,
415) -> OrganizationPaymentConfigUpdateResult:
416    """Update the payment configuration for an organization.
417
418    All updates require human-in-the-loop approval via `escalate_to_human`.
419
420    Use `set_grace_period` to start, extend, or cancel a grace period.
421    If the org is not already in `manual` status, the tool automatically
422    transitions to `manual` first before setting the grace period.
423
424    Use `set_permanent_waiver_type` to mark an organization as a partner (`free`)
425    or internal (`internal`) account. This is mutually exclusive with
426    `set_grace_period` — only one may be provided per call. Setting the waiver
427    type to `free` or `internal` also changes the Orb billing plan (`free` →
428    Airbyte Partner, `internal` → Airbyte Internal). The `ORB_API_KEY`
429    environment variable must be configured for waiver type changes.
430
431    The `organization_name` parameter is a safety check: the tool looks up the
432    organization via the Config API and verifies that the provided name, email, or
433    email domain matches. If omitted or mismatched, the tool returns the valid
434    identifiers so the caller can verify and retry.
435    """
436    # --- Validate that an action was specified ---
437    if set_grace_period is None and set_permanent_waiver_type is None:
438        return OrganizationPaymentConfigUpdateResult(
439            success=False,
440            message="No action specified. Provide `set_grace_period` with a date "
441            "(YYYY-MM-DD), number of days (1-90), or 'cancel'; or "
442            "`set_permanent_waiver_type` with 'free' (partner) or 'internal'.",
443            organization_id=organization_id,
444        )
445
446    # --- Validate mutual exclusivity ---
447    if set_grace_period is not None and set_permanent_waiver_type is not None:
448        return OrganizationPaymentConfigUpdateResult(
449            success=False,
450            message="`set_grace_period` and `set_permanent_waiver_type` are "
451            "mutually exclusive. Provide only one per call.",
452            organization_id=organization_id,
453        )
454
455    # --- Validate waiver reason is provided when waiver type is set ---
456    if set_permanent_waiver_type is not None and not set_permanent_waiver_reason:
457        return OrganizationPaymentConfigUpdateResult(
458            success=False,
459            message="`set_permanent_waiver_reason` is required when "
460            "`set_permanent_waiver_type` is provided.",
461            organization_id=organization_id,
462        )
463
464    # --- Validate ORB_API_KEY is configured for waiver type changes ---
465    if set_permanent_waiver_type in ("free", "internal") and not _get_orb_api_key():
466        return OrganizationPaymentConfigUpdateResult(
467            success=False,
468            message="`ORB_API_KEY` environment variable is not configured. "
469            "It is required when setting `set_permanent_waiver_type` to "
470            "'free' or 'internal' because the Orb billing plan must also "
471            "be changed.",
472            organization_id=organization_id,
473        )
474
475    # --- Validate admin access ---
476    try:
477        require_internal_admin_flag_only()
478    except CloudAuthError as e:
479        return OrganizationPaymentConfigUpdateResult(
480            success=False,
481            message=f"Admin authentication failed: {e}",
482            organization_id=organization_id,
483        )
484
485    # --- Resolve approval ---
486    try:
487        resolve_admin_email_from_approval(
488            approval_comment_url=approval_comment_url,
489        )
490    except ApprovalResolutionError as e:
491        return OrganizationPaymentConfigUpdateResult(
492            success=False,
493            message=str(e),
494            organization_id=organization_id,
495        )
496
497    # --- Resolve auth ---
498    resolved_api_root = config_api_root or constants.CLOUD_CONFIG_API_ROOT
499    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
500
501    # --- Look up organization info via Config API ---
502    try:
503        org_info = get_organization_info(
504            organization_id=organization_id,
505            config_api_root=resolved_api_root,
506            client_id=client_id,
507            client_secret=client_secret,
508            bearer_token=bearer_token,
509        )
510    except PaymentConfigAPIError as e:
511        return OrganizationPaymentConfigUpdateResult(
512            success=False,
513            message=f"Failed to fetch organization info: {e}",
514            organization_id=organization_id,
515        )
516    if org_info is None:
517        return OrganizationPaymentConfigUpdateResult(
518            success=False,
519            message=f"Organization {organization_id} not found.",
520            organization_id=organization_id,
521        )
522
523    # --- Validate organization name (safety check) ---
524    name_ok, name_error = _validate_organization_name(
525        organization_id, organization_name, org_info
526    )
527    if not name_ok:
528        return OrganizationPaymentConfigUpdateResult(
529            success=False,
530            message=name_error or "Organization name validation failed.",
531            organization_id=organization_id,
532        )
533
534    # --- Enrich with tier info ---
535    tier_result = get_org_tier(organization_id)
536    customer_tier = tier_result.customer_tier
537    tier_warning = _build_tier_warning(customer_tier)
538
539    # --- Permanent-waiver-only path (no grace period change) ---
540    #
541    # Statuses that the API cannot set back: uninitialized, okay, disabled.
542    # If the org is in one of these, we transition to 'manual' first (same
543    # pattern the grace-period path uses).
544    _api_nonsettable_statuses = ("uninitialized", "okay", "disabled")
545
546    if set_grace_period is None:
547        # Only set_permanent_waiver_type was requested
548        assert set_permanent_waiver_type is not None
549        try:
550            current_config = _get_organization_payment_config(
551                organization_id=organization_id,
552                config_api_root=resolved_api_root,
553                client_id=client_id,
554                client_secret=client_secret,
555                bearer_token=bearer_token,
556            )
557        except PaymentConfigAPIError as e:
558            return OrganizationPaymentConfigUpdateResult(
559                success=False,
560                message=f"Failed to fetch current config: {e}",
561                organization_id=organization_id,
562                customer_tier=customer_tier,
563                tier_warning=tier_warning,
564            )
565        current_status = current_config["paymentStatus"]
566
567        # Transition to 'manual' if current status is not API-settable
568        target_status = current_status
569        if current_status in _api_nonsettable_statuses:
570            try:
571                _update_organization_payment_config(
572                    organization_id=organization_id,
573                    payment_status="manual",
574                    config_api_root=resolved_api_root,
575                    client_id=client_id,
576                    client_secret=client_secret,
577                    bearer_token=bearer_token,
578                    new_grace_period_reason=(
579                        f"Transitioned to manual for permanent waiver: "
580                        f"{set_permanent_waiver_reason}"
581                    ),
582                )
583                target_status = "manual"
584            except PaymentConfigAPIError as e:
585                return OrganizationPaymentConfigUpdateResult(
586                    success=False,
587                    message=f"Failed to transition from '{current_status}' "
588                    f"to 'manual': {e}",
589                    organization_id=organization_id,
590                    payment_status=current_status,
591                    customer_tier=customer_tier,
592                    tier_warning=tier_warning,
593                )
594
595        try:
596            data = _update_organization_payment_config(
597                organization_id=organization_id,
598                payment_status=target_status,
599                config_api_root=resolved_api_root,
600                client_id=client_id,
601                client_secret=client_secret,
602                bearer_token=bearer_token,
603                usage_category_overwrite=(
604                    set_permanent_waiver_type
605                    if set_permanent_waiver_type != "none"
606                    else ""
607                ),
608            )
609        except PaymentConfigAPIError as e:
610            return OrganizationPaymentConfigUpdateResult(
611                success=False,
612                message=f"Failed to set permanent waiver type: {e}",
613                organization_id=organization_id,
614                payment_status=target_status,
615                customer_tier=customer_tier,
616                tier_warning=tier_warning,
617            )
618        parts = [
619            f"Permanent waiver type set to '{set_permanent_waiver_type}' "
620            f"for org {organization_id}.",
621        ]
622        if current_status in _api_nonsettable_statuses:
623            parts.append(
624                f"Payment status transitioned from '{current_status}' to 'manual'."
625            )
626
627        # --- Orb plan change (required for "free" / "internal") ---
628        orb_plan_change_result: str | None = None
629        if set_permanent_waiver_type in ("free", "internal"):
630            # ORB_API_KEY is validated at the top of the function, so this
631            # is guaranteed to be non-None here.
632            orb_api_key = _get_orb_api_key()
633            assert orb_api_key, "ORB_API_KEY should have been validated earlier"
634            try:
635                active_sub = get_active_subscription(organization_id, orb_api_key)
636                if active_sub is None:
637                    parts.append(
638                        "Orb plan change skipped: no active subscription "
639                        "found for this organization in Orb."
640                    )
641                    orb_plan_change_result = "Skipped: no active Orb subscription"
642                else:
643                    target_plan_id = _resolve_plan_id(set_permanent_waiver_type)
644                    current_plan_id = (active_sub.get("plan") or {}).get("id")
645                    current_plan_name = (active_sub.get("plan") or {}).get(
646                        "name", current_plan_id or "unknown"
647                    )
648                    if current_plan_id == target_plan_id:
649                        orb_plan_change_result = (
650                            f"Already on plan '{current_plan_name}'"
651                        )
652                        parts.append(f"Orb plan already set to '{current_plan_name}'.")
653                    else:
654                        schedule_plan_change(
655                            subscription_id=active_sub["id"],
656                            plan_id=target_plan_id,
657                            api_key=orb_api_key,
658                        )
659                        orb_plan_change_result = (
660                            f"Changed from '{current_plan_name}' to '{target_plan_id}'"
661                        )
662                        parts.append(f"Orb plan changed to '{target_plan_id}'.")
663            except (OrbAPIError, requests.RequestException) as e:
664                parts.append(f"Orb plan change failed: {e}")
665                orb_plan_change_result = f"Failed: {e}"
666
667        # --- Entitlement plan update (Stigg) ---
668        entitlement_plan_change_result: str | None = None
669        target_entitlement_plan = WAIVER_TYPE_TO_ENTITLEMENT_PLAN.get(
670            set_permanent_waiver_type
671        )
672        if target_entitlement_plan:
673            try:
674                update_entitlement_plan(
675                    organization_id=organization_id,
676                    plan_name=target_entitlement_plan,
677                    config_api_root=resolved_api_root,
678                    client_id=client_id,
679                    client_secret=client_secret,
680                    bearer_token=bearer_token,
681                )
682                entitlement_plan_change_result = f"Changed to {target_entitlement_plan}"
683                parts.append(
684                    f"Entitlement plan updated to '{target_entitlement_plan}'."
685                )
686            except EntitlementAPIError as e:
687                parts.append(f"Entitlement plan update failed: {e}")
688                entitlement_plan_change_result = f"Failed: {e}"
689
690        return OrganizationPaymentConfigUpdateResult(
691            success=True,
692            message=" ".join(parts),
693            organization_id=organization_id,
694            payment_status=data["paymentStatus"],
695            grace_period_end_at=data.get("gracePeriodEndAt"),
696            permanent_waiver_type=data.get("usageCategoryOverwrite"),
697            customer_tier=customer_tier,
698            tier_warning=tier_warning,
699            orb_plan_change=orb_plan_change_result,
700            entitlement_plan_change=entitlement_plan_change_result,
701        )
702
703    # --- Parse grace period value ---
704    parsed_value, parse_error = _parse_grace_period_value(set_grace_period)
705    if parse_error is not None:
706        return OrganizationPaymentConfigUpdateResult(
707            success=False,
708            message=parse_error,
709            organization_id=organization_id,
710            customer_tier=customer_tier,
711            tier_warning=tier_warning,
712        )
713
714    assert parsed_value is not None
715
716    # --- Fetch current payment config (needed for both cancel and set paths) ---
717    try:
718        current_config = _get_organization_payment_config(
719            organization_id=organization_id,
720            config_api_root=resolved_api_root,
721            client_id=client_id,
722            client_secret=client_secret,
723            bearer_token=bearer_token,
724        )
725    except PaymentConfigAPIError as e:
726        return OrganizationPaymentConfigUpdateResult(
727            success=False,
728            message=f"Failed to fetch current config: {e}",
729            organization_id=organization_id,
730            customer_tier=customer_tier,
731            tier_warning=tier_warning,
732        )
733
734    current_status = current_config["paymentStatus"]
735
736    if parsed_value == "cancel":
737        # Only allow cancel when org is actually in grace_period (or manual with
738        # an active grace period end date).
739        if current_status not in ("grace_period", "manual"):
740            return OrganizationPaymentConfigUpdateResult(
741                success=False,
742                message=f"Cannot cancel grace period: organization is in "
743                f"'{current_status}' status, not 'grace_period'.",
744                organization_id=organization_id,
745                payment_status=current_status,
746                customer_tier=customer_tier,
747                tier_warning=tier_warning,
748            )
749
750        # Cancel grace period by setting status to manual
751        try:
752            data = _update_organization_payment_config(
753                organization_id=organization_id,
754                payment_status="manual",
755                config_api_root=resolved_api_root,
756                client_id=client_id,
757                client_secret=client_secret,
758                bearer_token=bearer_token,
759                new_grace_period_reason=(
760                    set_grace_period_reason or "Grace period canceled via MCP tool"
761                ),
762            )
763        except PaymentConfigAPIError as e:
764            return OrganizationPaymentConfigUpdateResult(
765                success=False,
766                message=str(e),
767                organization_id=organization_id,
768                customer_tier=customer_tier,
769                tier_warning=tier_warning,
770            )
771
772        return OrganizationPaymentConfigUpdateResult(
773            success=True,
774            message=f"Grace period canceled for org {organization_id}. "
775            f"New status: {data['paymentStatus']}.",
776            organization_id=organization_id,
777            payment_status=data["paymentStatus"],
778            grace_period_end_at=data.get("gracePeriodEndAt"),
779            permanent_waiver_type=data.get("usageCategoryOverwrite"),
780            customer_tier=customer_tier,
781            tier_warning=tier_warning,
782        )
783
784    # --- Setting/extending grace period ---
785    if not set_grace_period_reason:
786        return OrganizationPaymentConfigUpdateResult(
787            success=False,
788            message="`set_grace_period_reason` is required when setting or "
789            "extending a grace period.",
790            organization_id=organization_id,
791            customer_tier=customer_tier,
792            tier_warning=tier_warning,
793        )
794
795    # The API only allows setting grace_period from manual status.
796    # If not already in manual, transition to manual first.
797    transitioned_to_manual = False
798    if current_status != "manual":
799        logger.info(
800            "Org %s is in '%s' status; transitioning to 'manual' first.",
801            organization_id,
802            current_status,
803        )
804        try:
805            _update_organization_payment_config(
806                organization_id=organization_id,
807                payment_status="manual",
808                config_api_root=resolved_api_root,
809                client_id=client_id,
810                client_secret=client_secret,
811                bearer_token=bearer_token,
812            )
813            transitioned_to_manual = True
814        except PaymentConfigAPIError as e:
815            return OrganizationPaymentConfigUpdateResult(
816                success=False,
817                message=f"Failed to transition to 'manual' from '{current_status}': {e}",
818                organization_id=organization_id,
819                customer_tier=customer_tier,
820                tier_warning=tier_warning,
821            )
822
823    # Now set the grace period
824    try:
825        data = _update_organization_payment_config(
826            organization_id=organization_id,
827            payment_status="grace_period",
828            config_api_root=resolved_api_root,
829            client_id=client_id,
830            client_secret=client_secret,
831            bearer_token=bearer_token,
832            grace_period_end_at=parsed_value,
833            new_grace_period_reason=set_grace_period_reason,
834        )
835    except PaymentConfigAPIError as e:
836        if transitioned_to_manual:
837            msg = (
838                f"Failed to set grace period after transitioning to 'manual' "
839                f"from '{current_status}'. The organization is now in 'manual' "
840                f"status. Original error: {e}"
841            )
842        else:
843            msg = str(e)
844        return OrganizationPaymentConfigUpdateResult(
845            success=False,
846            message=msg,
847            organization_id=organization_id,
848            payment_status="manual" if transitioned_to_manual else None,
849            customer_tier=customer_tier,
850            tier_warning=tier_warning,
851        )
852
853    return OrganizationPaymentConfigUpdateResult(
854        success=True,
855        message=f"Grace period set for org {organization_id}. "
856        f"New status: {data['paymentStatus']}.",
857        organization_id=organization_id,
858        payment_status=data["paymentStatus"],
859        grace_period_end_at=data.get("gracePeriodEndAt"),
860        permanent_waiver_type=data.get("usageCategoryOverwrite"),
861        customer_tier=customer_tier,
862        tier_warning=tier_warning,
863    )
864
865
866def register_organization_payment_config_tools(app: FastMCP) -> None:
867    """Register organization payment config tools with the FastMCP app."""
868    register_mcp_tools(app, mcp_module=__name__)