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
 23
 24from airbyte import constants
 25from fastmcp import Context, FastMCP
 26from fastmcp_extensions import get_mcp_config, mcp_tool, register_mcp_tools
 27from pydantic import Field
 28from zoneinfo import ZoneInfo
 29
 30from airbyte_ops_mcp.approval_resolution import (
 31    ApprovalResolutionError,
 32    resolve_admin_email_from_approval,
 33)
 34from airbyte_ops_mcp.cloud_admin.auth import (
 35    CloudAuthError,
 36    require_internal_admin_flag_only,
 37)
 38from airbyte_ops_mcp.cloud_admin.models import (
 39    OrganizationInfo,
 40    OrganizationPaymentConfigInfo,
 41    OrganizationPaymentConfigUpdateResult,
 42)
 43from airbyte_ops_mcp.cloud_admin.payment_config import (
 44    PaymentConfigAPIError,
 45    get_organization_info,
 46)
 47from airbyte_ops_mcp.cloud_admin.payment_config import (
 48    get_organization_payment_config as _get_organization_payment_config,
 49)
 50from airbyte_ops_mcp.cloud_admin.payment_config import (
 51    update_organization_payment_config as _update_organization_payment_config,
 52)
 53from airbyte_ops_mcp.constants import ServerConfigKey
 54from airbyte_ops_mcp.tier_cache import get_org_tier
 55
 56logger = logging.getLogger(__name__)
 57
 58_DATE_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}$")
 59_DAYS_PATTERN = re.compile(r"^\d+$")
 60_PACIFIC = ZoneInfo("America/Los_Angeles")
 61
 62
 63def _build_tier_warning(customer_tier: str) -> str | None:
 64    """Build a warning message for sensitive customer tiers."""
 65    if customer_tier == "TIER_0":
 66        return (
 67            "WARNING: This is a TIER_0 (highest-value) customer. "
 68            "Proceed with extreme caution."
 69        )
 70    if customer_tier == "TIER_1":
 71        return "WARNING: This is a TIER_1 (high-value) customer. Proceed with caution."
 72    return None
 73
 74
 75def _resolve_cloud_auth(ctx: Context) -> tuple[str | None, str | None, str | None]:
 76    """Resolve auth credentials, returning `(bearer_token, client_id, client_secret)`."""
 77    bearer_token = get_mcp_config(ctx, ServerConfigKey.BEARER_TOKEN)
 78    if bearer_token:
 79        return bearer_token, None, None
 80
 81    client_id = get_mcp_config(ctx, ServerConfigKey.CLIENT_ID)
 82    client_secret = get_mcp_config(ctx, ServerConfigKey.CLIENT_SECRET)
 83    return None, client_id, client_secret
 84
 85
 86def _validate_organization_name(
 87    organization_id: str,
 88    organization_name: str | None,
 89    org_info: OrganizationInfo,
 90) -> tuple[bool, str | None]:
 91    """Validate `organization_name` against the org record from the Config API.
 92
 93    Accepts the org's literal name, email, or email domain as valid inputs.
 94
 95    Returns `(ok, error_message)`. On success `error_message` is `None`.
 96    """
 97    org_db_name: str = org_info.organization_name or ""
 98    org_db_email: str = org_info.email or ""
 99    org_db_domain: str = org_db_email.split("@", 1)[-1] if "@" in org_db_email else ""
100
101    valid_identifiers = [
102        v
103        for v in [org_db_name.lower(), org_db_email.lower(), org_db_domain.lower()]
104        if v
105    ]
106
107    identifier_parts = []
108    if org_db_name:
109        identifier_parts.append(f"'{org_db_name}' (org name)")
110    if org_db_email:
111        identifier_parts.append(f"'{org_db_email}' (email)")
112    if org_db_domain:
113        identifier_parts.append(f"'{org_db_domain}' (email domain)")
114
115    if not identifier_parts:
116        return (
117            False,
118            f"Organization {organization_id} has no name or email on record.",
119        )
120
121    hint_message = (
122        f"To confirm, resend with `organization_name` set to one of: "
123        f"{', '.join(identifier_parts)}. "
124        f"Double-check that this is the correct organization before retrying. "
125        f"If there is any doubt, confirm with your user."
126    )
127
128    if organization_name is None:
129        return (
130            False,
131            (
132                f"`organization_name` is required to confirm the target organization. "
133                f"{hint_message}"
134            ),
135        )
136
137    if organization_name.strip().lower().lstrip("@") not in valid_identifiers:
138        return (
139            False,
140            (
141                f"Organization name mismatch: '{organization_name}' does not match "
142                f"organization {organization_id}. {hint_message}"
143            ),
144        )
145
146    return True, None
147
148
149def _format_grace_period_end(target_date: date) -> str:
150    """Build a backend-compatible datetime string for 11:59 PM Pacific on `target_date`.
151
152    The result is converted to UTC and formatted as
153    `yyyy-MM-dd'T'HH:mm:ss.SSS+0000` to match the Java `@JsonFormat` annotation.
154    """
155    end_pacific = datetime(
156        target_date.year,
157        target_date.month,
158        target_date.day,
159        23,
160        59,
161        59,
162        tzinfo=_PACIFIC,
163    )
164    end_utc = end_pacific.astimezone(timezone.utc)
165    return end_utc.strftime("%Y-%m-%dT%H:%M:%S.000+0000")
166
167
168def _parse_grace_period_value(
169    value: str,
170) -> tuple[str | None, str | None]:
171    """Parse `set_grace_period` into an ISO 8601 datetime string or an action keyword.
172
173    Dates and day offsets resolve to 11:59 PM Pacific Time on the target date.
174    Day offsets use the current date in Pacific Time as the starting point.
175
176    Returns `(iso_datetime_or_action, error_message)`.
177    On success `error_message` is `None`.
178    """
179    value = value.strip()
180
181    if value.lower() == "cancel":
182        return "cancel", None
183
184    today_pacific = datetime.now(tz=_PACIFIC).date()
185
186    if _DATE_PATTERN.match(value):
187        try:
188            parsed = date.fromisoformat(value)
189        except ValueError:
190            return None, f"Invalid date: {value}. Use YYYY-MM-DD format."
191        if parsed < today_pacific:
192            return None, f"Grace period end date {value} is in the past."
193        if (parsed - today_pacific).days > 90:
194            return None, (
195                f"Grace period end date {value} is more than 90 days in the future. "
196                "Maximum grace period is 90 days."
197            )
198        return _format_grace_period_end(parsed), None
199
200    if _DAYS_PATTERN.match(value):
201        days = int(value)
202        if days < 1 or days > 90:
203            return None, f"Days must be between 1 and 90, got {days}."
204        target_date = today_pacific + timedelta(days=days)
205        return _format_grace_period_end(target_date), None
206
207    return None, (
208        f"Invalid `set_grace_period` value: '{value}'. "
209        "Expected a date (YYYY-MM-DD), an integer number of days, or 'cancel'."
210    )
211
212
213# =============================================================================
214# MCP Tools
215# =============================================================================
216
217
218@mcp_tool(
219    read_only=True,
220    idempotent=True,
221    open_world=True,
222)
223def get_organization_payment_config(
224    organization_id: Annotated[
225        str,
226        "The organization UUID.",
227    ],
228    config_api_root: Annotated[
229        str | None,
230        Field(
231            description="Optional Config API root URL override. "
232            "Defaults to Airbyte Cloud (`https://cloud.airbyte.com/api/v1`).",
233            default=None,
234        ),
235    ] = None,
236    *,
237    ctx: Context,
238) -> OrganizationPaymentConfigInfo:
239    """Get the current payment configuration for an organization.
240
241    Returns payment status, subscription status, grace period info,
242    and usage category override. No PII or sensitive payment details
243    are included in the response.
244
245    Authentication credentials are resolved in priority order:
246    1. Bearer token (Authorization header or AIRBYTE_CLOUD_BEARER_TOKEN env var)
247    2. HTTP headers: X-Airbyte-Cloud-Client-Id, X-Airbyte-Cloud-Client-Secret
248    3. Environment variables: AIRBYTE_CLOUD_CLIENT_ID, AIRBYTE_CLOUD_CLIENT_SECRET
249    """
250    resolved_api_root = config_api_root or constants.CLOUD_CONFIG_API_ROOT
251    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
252
253    data = _get_organization_payment_config(
254        organization_id=organization_id,
255        config_api_root=resolved_api_root,
256        client_id=client_id,
257        client_secret=client_secret,
258        bearer_token=bearer_token,
259    )
260
261    # Enrich with tier info
262    tier_result = get_org_tier(organization_id)
263    tier_warning = _build_tier_warning(tier_result.customer_tier)
264
265    return OrganizationPaymentConfigInfo(
266        organization_id=data["organizationId"],
267        payment_status=data["paymentStatus"],
268        subscription_status=data["subscriptionStatus"],
269        payment_provider_id=data.get("paymentProviderId"),
270        grace_period_end_at=data.get("gracePeriodEndAt"),
271        usage_category_overwrite=data.get("usageCategoryOverwrite"),
272        customer_tier=tier_result.customer_tier,
273        tier_warning=tier_warning,
274    )
275
276
277@mcp_tool(
278    destructive=True,
279    idempotent=False,
280    open_world=True,
281)
282def update_organization_payment_config(
283    organization_id: Annotated[
284        str,
285        "The organization UUID.",
286    ],
287    approval_comment_url: Annotated[
288        str,
289        Field(
290            description="URL to the Slack approval record. Obtain this by calling "
291            "`escalate_to_human` with `approval_requested=True`; the backend "
292            "delivers the approval record URL when a human clicks Approve.",
293        ),
294    ],
295    organization_name: Annotated[
296        str | None,
297        Field(
298            description="Confirmation of the target organization. Accepts the "
299            "organization name, email address, or email domain. Required to "
300            "prevent accidental modifications to the wrong organization. "
301            "If omitted or mismatched, the tool returns an error with the valid "
302            "identifiers so you can verify and retry.",
303            default=None,
304        ),
305    ] = None,
306    set_grace_period: Annotated[
307        str | None,
308        Field(
309            description="Set or modify the grace period. Accepts three forms: "
310            "(1) A date in `YYYY-MM-DD` format — grace period ends at 11:59 PM Pacific on that date. "
311            "(2) An integer number of days (1-90) from today (Pacific Time) — "
312            "grace period ends at 11:59 PM Pacific on the resulting date. "
313            "(3) `'cancel'` to terminate the current grace period "
314            "(sets status to `manual`). "
315            "Requires `set_grace_period_reason` when setting or extending.",
316            default=None,
317        ),
318    ] = None,
319    set_grace_period_reason: Annotated[
320        str | None,
321        Field(
322            description="Reason for starting, extending, or canceling the grace "
323            "period. Required when `set_grace_period` is a date or number of days.",
324            default=None,
325        ),
326    ] = None,
327    config_api_root: Annotated[
328        str | None,
329        Field(
330            description="Optional Config API root URL override. "
331            "Defaults to Airbyte Cloud (`https://cloud.airbyte.com/api/v1`).",
332            default=None,
333        ),
334    ] = None,
335    *,
336    ctx: Context,
337) -> OrganizationPaymentConfigUpdateResult:
338    """Update the payment configuration for an organization.
339
340    All updates require human-in-the-loop approval via `escalate_to_human`.
341
342    Use `set_grace_period` to start, extend, or cancel a grace period.
343    If the org is not already in `manual` status, the tool automatically
344    transitions to `manual` first before setting the grace period.
345
346    The `organization_name` parameter is a safety check: the tool looks up the
347    organization via the Config API and verifies that the provided name, email, or
348    email domain matches. If omitted or mismatched, the tool returns the valid
349    identifiers so the caller can verify and retry.
350    """
351    # --- Validate that an action was specified ---
352    if set_grace_period is None:
353        return OrganizationPaymentConfigUpdateResult(
354            success=False,
355            message="No action specified. Provide `set_grace_period` with a date "
356            "(YYYY-MM-DD), number of days (1-90), or 'cancel'.",
357            organization_id=organization_id,
358        )
359
360    # --- Validate admin access ---
361    try:
362        require_internal_admin_flag_only()
363    except CloudAuthError as e:
364        return OrganizationPaymentConfigUpdateResult(
365            success=False,
366            message=f"Admin authentication failed: {e}",
367            organization_id=organization_id,
368        )
369
370    # --- Resolve approval ---
371    try:
372        resolve_admin_email_from_approval(
373            approval_comment_url=approval_comment_url,
374        )
375    except ApprovalResolutionError as e:
376        return OrganizationPaymentConfigUpdateResult(
377            success=False,
378            message=str(e),
379            organization_id=organization_id,
380        )
381
382    # --- Resolve auth ---
383    resolved_api_root = config_api_root or constants.CLOUD_CONFIG_API_ROOT
384    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
385
386    # --- Look up organization info via Config API ---
387    try:
388        org_info = get_organization_info(
389            organization_id=organization_id,
390            config_api_root=resolved_api_root,
391            client_id=client_id,
392            client_secret=client_secret,
393            bearer_token=bearer_token,
394        )
395    except PaymentConfigAPIError as e:
396        return OrganizationPaymentConfigUpdateResult(
397            success=False,
398            message=f"Failed to fetch organization info: {e}",
399            organization_id=organization_id,
400        )
401    if org_info is None:
402        return OrganizationPaymentConfigUpdateResult(
403            success=False,
404            message=f"Organization {organization_id} not found.",
405            organization_id=organization_id,
406        )
407
408    # --- Validate organization name (safety check) ---
409    name_ok, name_error = _validate_organization_name(
410        organization_id, organization_name, org_info
411    )
412    if not name_ok:
413        return OrganizationPaymentConfigUpdateResult(
414            success=False,
415            message=name_error or "Organization name validation failed.",
416            organization_id=organization_id,
417        )
418
419    # --- Enrich with tier info ---
420    tier_result = get_org_tier(organization_id)
421    customer_tier = tier_result.customer_tier
422    tier_warning = _build_tier_warning(customer_tier)
423
424    # --- Parse grace period value ---
425    parsed_value, parse_error = _parse_grace_period_value(set_grace_period)
426    if parse_error is not None:
427        return OrganizationPaymentConfigUpdateResult(
428            success=False,
429            message=parse_error,
430            organization_id=organization_id,
431            customer_tier=customer_tier,
432            tier_warning=tier_warning,
433        )
434
435    assert parsed_value is not None
436
437    # --- Fetch current payment config (needed for both cancel and set paths) ---
438    try:
439        current_config = _get_organization_payment_config(
440            organization_id=organization_id,
441            config_api_root=resolved_api_root,
442            client_id=client_id,
443            client_secret=client_secret,
444            bearer_token=bearer_token,
445        )
446    except PaymentConfigAPIError as e:
447        return OrganizationPaymentConfigUpdateResult(
448            success=False,
449            message=f"Failed to fetch current config: {e}",
450            organization_id=organization_id,
451            customer_tier=customer_tier,
452            tier_warning=tier_warning,
453        )
454
455    current_status = current_config["paymentStatus"]
456
457    if parsed_value == "cancel":
458        # Only allow cancel when org is actually in grace_period (or manual with
459        # an active grace period end date).
460        if current_status not in ("grace_period", "manual"):
461            return OrganizationPaymentConfigUpdateResult(
462                success=False,
463                message=f"Cannot cancel grace period: organization is in "
464                f"'{current_status}' status, not 'grace_period'.",
465                organization_id=organization_id,
466                payment_status=current_status,
467                customer_tier=customer_tier,
468                tier_warning=tier_warning,
469            )
470
471        # Cancel grace period by setting status to manual
472        try:
473            data = _update_organization_payment_config(
474                organization_id=organization_id,
475                payment_status="manual",
476                config_api_root=resolved_api_root,
477                client_id=client_id,
478                client_secret=client_secret,
479                bearer_token=bearer_token,
480                new_grace_period_reason=(
481                    set_grace_period_reason or "Grace period canceled via MCP tool"
482                ),
483            )
484        except PaymentConfigAPIError as e:
485            return OrganizationPaymentConfigUpdateResult(
486                success=False,
487                message=str(e),
488                organization_id=organization_id,
489                customer_tier=customer_tier,
490                tier_warning=tier_warning,
491            )
492
493        return OrganizationPaymentConfigUpdateResult(
494            success=True,
495            message=f"Grace period canceled for org {organization_id}. "
496            f"New status: {data['paymentStatus']}.",
497            organization_id=organization_id,
498            payment_status=data["paymentStatus"],
499            grace_period_end_at=data.get("gracePeriodEndAt"),
500            customer_tier=customer_tier,
501            tier_warning=tier_warning,
502        )
503
504    # --- Setting/extending grace period ---
505    if not set_grace_period_reason:
506        return OrganizationPaymentConfigUpdateResult(
507            success=False,
508            message="`set_grace_period_reason` is required when setting or "
509            "extending a grace period.",
510            organization_id=organization_id,
511            customer_tier=customer_tier,
512            tier_warning=tier_warning,
513        )
514
515    # The API only allows setting grace_period from manual status.
516    # If not already in manual, transition to manual first.
517    transitioned_to_manual = False
518    if current_status != "manual":
519        logger.info(
520            "Org %s is in '%s' status; transitioning to 'manual' first.",
521            organization_id,
522            current_status,
523        )
524        try:
525            _update_organization_payment_config(
526                organization_id=organization_id,
527                payment_status="manual",
528                config_api_root=resolved_api_root,
529                client_id=client_id,
530                client_secret=client_secret,
531                bearer_token=bearer_token,
532            )
533            transitioned_to_manual = True
534        except PaymentConfigAPIError as e:
535            return OrganizationPaymentConfigUpdateResult(
536                success=False,
537                message=f"Failed to transition to 'manual' from '{current_status}': {e}",
538                organization_id=organization_id,
539                customer_tier=customer_tier,
540                tier_warning=tier_warning,
541            )
542
543    # Now set the grace period
544    try:
545        data = _update_organization_payment_config(
546            organization_id=organization_id,
547            payment_status="grace_period",
548            config_api_root=resolved_api_root,
549            client_id=client_id,
550            client_secret=client_secret,
551            bearer_token=bearer_token,
552            grace_period_end_at=parsed_value,
553            new_grace_period_reason=set_grace_period_reason,
554        )
555    except PaymentConfigAPIError as e:
556        if transitioned_to_manual:
557            msg = (
558                f"Failed to set grace period after transitioning to 'manual' "
559                f"from '{current_status}'. The organization is now in 'manual' "
560                f"status. Original error: {e}"
561            )
562        else:
563            msg = str(e)
564        return OrganizationPaymentConfigUpdateResult(
565            success=False,
566            message=msg,
567            organization_id=organization_id,
568            payment_status="manual" if transitioned_to_manual else None,
569            customer_tier=customer_tier,
570            tier_warning=tier_warning,
571        )
572
573    return OrganizationPaymentConfigUpdateResult(
574        success=True,
575        message=f"Grace period set for org {organization_id}. "
576        f"New status: {data['paymentStatus']}.",
577        organization_id=organization_id,
578        payment_status=data["paymentStatus"],
579        grace_period_end_at=data.get("gracePeriodEndAt"),
580        customer_tier=customer_tier,
581        tier_warning=tier_warning,
582    )
583
584
585def register_organization_payment_config_tools(app: FastMCP) -> None:
586    """Register organization payment config tools with the FastMCP app."""
587    register_mcp_tools(app, mcp_module=__name__)