airbyte_ops_mcp.mcp.organization_agentic_flag

  1# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
  2"""MCP tools for organization `is_agentic` flag management.
  3
  4Provides tools to read and update the Airbyte Cloud organization `is_agentic`
  5flag.
  6
  7## MCP reference
  8
  9.. include:: ../../../docs/mcp-generated/organization_agentic_flag.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
 19from typing import Annotated
 20
 21from airbyte import constants
 22from fastmcp import Context, FastMCP
 23from fastmcp_extensions import get_mcp_config, mcp_tool, register_mcp_tools
 24from pydantic import Field
 25
 26from airbyte_ops_mcp.approval_resolution import (
 27    ApprovalResolutionError,
 28    resolve_admin_email_from_approval,
 29)
 30from airbyte_ops_mcp.cloud_admin.auth import (
 31    CloudAuthError,
 32    require_internal_admin_flag_only,
 33)
 34from airbyte_ops_mcp.cloud_admin.models import (
 35    OrganizationAgenticFlagBatchInfo,
 36    OrganizationAgenticFlagBatchUpdateResult,
 37    OrganizationAgenticFlagInfo,
 38    OrganizationAgenticFlagUpdateResult,
 39    OrganizationInfo,
 40)
 41from airbyte_ops_mcp.cloud_admin.organization_agentic_flag import (
 42    OrganizationAgenticFlagAPIError,
 43)
 44from airbyte_ops_mcp.cloud_admin.organization_agentic_flag import (
 45    get_organization_info as _get_organization_info,
 46)
 47from airbyte_ops_mcp.cloud_admin.organization_agentic_flag import (
 48    set_organization_agentic_status as _set_organization_agentic_status,
 49)
 50from airbyte_ops_mcp.cloud_admin.version_overrides import validate_tier_filter
 51from airbyte_ops_mcp.constants import ServerConfigKey
 52from airbyte_ops_mcp.mcp.organization_payment_config import (
 53    _build_tier_warning,
 54    _validate_organization_name,
 55)
 56from airbyte_ops_mcp.prod_db_access.queries import (
 57    query_organization_agentic_flags,
 58)
 59from airbyte_ops_mcp.tier_cache import TierFilter, get_org_tier
 60
 61
 62def _normalize_organization_ids(organization_ids: str | list[str]) -> list[str]:
 63    """Normalize one organization ID or a list of organization IDs."""
 64    ids = [organization_ids] if isinstance(organization_ids, str) else organization_ids
 65
 66    normalized = [org_id.strip() for org_id in ids if org_id and org_id.strip()]
 67    return list(dict.fromkeys(normalized))
 68
 69
 70def _build_update_failure_results(
 71    organization_ids: list[str],
 72    message: str,
 73) -> list[OrganizationAgenticFlagUpdateResult]:
 74    """Build per-org failure results for request-level validation failures."""
 75    return [
 76        OrganizationAgenticFlagUpdateResult(
 77            success=False,
 78            message=message,
 79            organization_id=organization_id,
 80        )
 81        for organization_id in organization_ids
 82    ]
 83
 84
 85def _resolve_cloud_auth(ctx: Context) -> tuple[str | None, str | None, str | None]:
 86    """Resolve auth credentials, returning `(bearer_token, client_id, client_secret)`."""
 87    bearer_token = get_mcp_config(ctx, ServerConfigKey.BEARER_TOKEN)
 88    if bearer_token:
 89        return bearer_token, None, None
 90
 91    client_id = get_mcp_config(ctx, ServerConfigKey.CLIENT_ID)
 92    client_secret = get_mcp_config(ctx, ServerConfigKey.CLIENT_SECRET)
 93    return None, client_id, client_secret
 94
 95
 96def _org_info_to_flag_info(
 97    org_info: OrganizationInfo,
 98    customer_tier: str | None = None,
 99    tier_warning: str | None = None,
100) -> OrganizationAgenticFlagInfo:
101    """Convert Config API organization info into agentic flag info."""
102    return OrganizationAgenticFlagInfo(
103        organization_id=org_info.organization_id,
104        organization_name=org_info.organization_name,
105        email=org_info.email,
106        tombstone=False,
107        is_agentic=org_info.is_agentic or False,
108        customer_tier=customer_tier,
109        tier_warning=tier_warning,
110    )
111
112
113def _row_to_flag_info(
114    row: dict[str, object],
115    customer_tier: str | None = None,
116    tier_warning: str | None = None,
117) -> OrganizationAgenticFlagInfo:
118    """Convert a DB row into agentic flag info."""
119    return OrganizationAgenticFlagInfo(
120        organization_id=str(row["organization_id"]),
121        organization_name=str(row["organization_name"])
122        if row["organization_name"] is not None
123        else None,
124        email=str(row["email"]) if row["email"] is not None else None,
125        tombstone=bool(row["tombstone"]),
126        is_agentic=bool(row["is_agentic"]),
127        customer_tier=customer_tier,
128        tier_warning=tier_warning,
129    )
130
131
132@mcp_tool(
133    destructive=False,
134    idempotent=True,
135    open_world=True,
136)
137def get_organization_agentic_flag(
138    organization_ids: Annotated[
139        str | list[str],
140        Field(
141            description="One organization UUID, or a list of organization UUIDs.",
142        ),
143    ],
144    config_api_root: Annotated[
145        str | None,
146        Field(
147            description="Optional Config API root URL override. When omitted, "
148            "the tool reads from the Prod DB replica.",
149            default=None,
150        ),
151    ] = None,
152    *,
153    ctx: Context,
154) -> OrganizationAgenticFlagBatchInfo:
155    """Get the current `is_agentic` flag for one or more organizations."""
156    normalized_ids = _normalize_organization_ids(organization_ids)
157    if not normalized_ids:
158        return OrganizationAgenticFlagBatchInfo(
159            organizations=[],
160            missing_organization_ids=[],
161        )
162
163    if config_api_root is None:
164        rows = query_organization_agentic_flags(normalized_ids)
165        rows_by_org_id = {str(row["organization_id"]): row for row in rows}
166        organizations: list[OrganizationAgenticFlagInfo] = []
167        for organization_id in normalized_ids:
168            row = rows_by_org_id.get(organization_id)
169            if row is None:
170                continue
171            tier_result = get_org_tier(organization_id)
172            organizations.append(
173                _row_to_flag_info(
174                    row,
175                    customer_tier=tier_result.customer_tier,
176                    tier_warning=_build_tier_warning(tier_result.customer_tier),
177                )
178            )
179        return OrganizationAgenticFlagBatchInfo(
180            organizations=organizations,
181            missing_organization_ids=[
182                organization_id
183                for organization_id in normalized_ids
184                if organization_id not in rows_by_org_id
185            ],
186        )
187
188    resolved_api_root = config_api_root or constants.CLOUD_CONFIG_API_ROOT
189    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
190
191    organizations: list[OrganizationAgenticFlagInfo] = []
192    missing_organization_ids: list[str] = []
193    for organization_id in normalized_ids:
194        try:
195            org_info = _get_organization_info(
196                organization_id=organization_id,
197                config_api_root=resolved_api_root,
198                client_id=client_id,
199                client_secret=client_secret,
200                bearer_token=bearer_token,
201            )
202        except OrganizationAgenticFlagAPIError:
203            missing_organization_ids.append(organization_id)
204            continue
205        if org_info is None:
206            missing_organization_ids.append(organization_id)
207            continue
208
209        tier_result = get_org_tier(organization_id)
210        organizations.append(
211            _org_info_to_flag_info(
212                org_info,
213                customer_tier=tier_result.customer_tier,
214                tier_warning=_build_tier_warning(tier_result.customer_tier),
215            )
216        )
217
218    return OrganizationAgenticFlagBatchInfo(
219        organizations=organizations,
220        missing_organization_ids=missing_organization_ids,
221    )
222
223
224@mcp_tool(
225    destructive=True,
226    idempotent=True,
227    open_world=True,
228)
229def update_organization_agentic_flag(
230    organization_ids: Annotated[
231        str | list[str],
232        Field(
233            description="One organization UUID, or a list of organization UUIDs.",
234        ),
235    ],
236    is_agentic: Annotated[
237        bool,
238        Field(description="The desired `is_agentic` value."),
239    ],
240    approval_comment_url: Annotated[
241        str,
242        Field(
243            description="URL to the Slack approval record. Obtain this by calling "
244            "`escalate_to_human` with `approval_requested=True`; the backend "
245            "delivers the approval record URL when a human clicks Approve.",
246        ),
247    ],
248    organization_name: Annotated[
249        str | None,
250        Field(
251            description="Confirmation of the target organization. Accepts the "
252            "organization name, email address, or email domain. Required when "
253            "updating a single organization. Ignored for multi-org updates; use "
254            "`organization_names` instead.",
255            default=None,
256        ),
257    ] = None,
258    organization_names: Annotated[
259        dict[str, str] | None,
260        Field(
261            description="Per-organization confirmation for multi-org updates. "
262            "Keys are organization UUIDs; values may be organization names, "
263            "email addresses, or email domains.",
264            default=None,
265        ),
266    ] = None,
267    customer_tier_filter: Annotated[
268        TierFilter,
269        Field(
270            description=(
271                "Required tier filter: 'TIER_0', 'TIER_1', 'TIER_2', or 'ALL'. "
272                "The operation is rejected if the actual customer tier does not match. "
273                "Use 'ALL' to proceed regardless of tier after human approval."
274            ),
275        ),
276    ] = "TIER_2",
277    *,
278    ctx: Context,
279) -> OrganizationAgenticFlagBatchUpdateResult:
280    """Update the `is_agentic` flag for one or more organizations.
281
282    All updates require human-in-the-loop approval via `escalate_to_human`.
283    The tool validates each organization identity before writing the flag.
284    """
285    normalized_ids = _normalize_organization_ids(organization_ids)
286    if not normalized_ids:
287        return OrganizationAgenticFlagBatchUpdateResult(
288            success=False,
289            message="No organization IDs provided.",
290            results=[],
291        )
292
293    try:
294        require_internal_admin_flag_only()
295    except CloudAuthError as e:
296        return OrganizationAgenticFlagBatchUpdateResult(
297            success=False,
298            message=f"Admin authentication failed: {e}",
299            results=_build_update_failure_results(
300                normalized_ids, f"Admin authentication failed: {e}"
301            ),
302        )
303
304    try:
305        resolve_admin_email_from_approval(
306            approval_comment_url=approval_comment_url,
307        )
308    except ApprovalResolutionError as e:
309        return OrganizationAgenticFlagBatchUpdateResult(
310            success=False,
311            message=str(e),
312            results=_build_update_failure_results(normalized_ids, str(e)),
313        )
314
315    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
316
317    results: list[OrganizationAgenticFlagUpdateResult] = []
318    for organization_id in normalized_ids:
319        rows = query_organization_agentic_flags([organization_id])
320        if not rows:
321            results.append(
322                OrganizationAgenticFlagUpdateResult(
323                    success=False,
324                    message=f"Organization {organization_id} not found.",
325                    organization_id=organization_id,
326                )
327            )
328            continue
329
330        current_info = _row_to_flag_info(rows[0])
331        if current_info.tombstone:
332            results.append(
333                OrganizationAgenticFlagUpdateResult(
334                    success=False,
335                    message=f"Organization {organization_id} is tombstoned.",
336                    organization_id=organization_id,
337                    organization_name=current_info.organization_name,
338                    email=current_info.email,
339                    previous_is_agentic=current_info.is_agentic,
340                )
341            )
342            continue
343
344        org_info = OrganizationInfo.model_validate(
345            {
346                "organizationId": current_info.organization_id,
347                "organizationName": current_info.organization_name or "",
348                "email": current_info.email,
349                "isAgentic": current_info.is_agentic,
350            }
351        )
352        expected_name = (
353            organization_names.get(organization_id)
354            if organization_names is not None
355            else organization_name
356        )
357        name_ok, name_error = _validate_organization_name(
358            organization_id, expected_name, org_info
359        )
360        if not name_ok:
361            results.append(
362                OrganizationAgenticFlagUpdateResult(
363                    success=False,
364                    message=name_error or "Organization name validation failed.",
365                    organization_id=organization_id,
366                    organization_name=current_info.organization_name,
367                    email=current_info.email,
368                    previous_is_agentic=current_info.is_agentic,
369                )
370            )
371            continue
372
373        tier_result = get_org_tier(organization_id)
374        customer_tier = tier_result.customer_tier
375        tier_warning = _build_tier_warning(customer_tier)
376        tier_ok, tier_error = validate_tier_filter(customer_tier, customer_tier_filter)
377        if not tier_ok:
378            results.append(
379                OrganizationAgenticFlagUpdateResult(
380                    success=False,
381                    message=tier_error or "Customer tier validation failed.",
382                    organization_id=organization_id,
383                    organization_name=current_info.organization_name,
384                    email=current_info.email,
385                    previous_is_agentic=current_info.is_agentic,
386                    customer_tier=customer_tier,
387                    tier_warning=tier_warning,
388                )
389            )
390            continue
391        previous_is_agentic = current_info.is_agentic
392
393        if previous_is_agentic == is_agentic:
394            results.append(
395                OrganizationAgenticFlagUpdateResult(
396                    success=True,
397                    message=(
398                        f"Organization {organization_id} already has managed agentic "
399                        f"org status {is_agentic}."
400                    ),
401                    organization_id=organization_id,
402                    organization_name=current_info.organization_name,
403                    email=current_info.email,
404                    previous_is_agentic=previous_is_agentic,
405                    new_is_agentic=is_agentic,
406                    customer_tier=customer_tier,
407                    tier_warning=tier_warning,
408                )
409            )
410            continue
411
412        try:
413            updated = _set_organization_agentic_status(
414                organization_id=organization_id,
415                is_agentic=is_agentic,
416                config_api_root=constants.CLOUD_CONFIG_API_ROOT,
417                client_id=client_id,
418                client_secret=client_secret,
419                bearer_token=bearer_token,
420            )
421        except Exception as e:
422            results.append(
423                OrganizationAgenticFlagUpdateResult(
424                    success=False,
425                    message=str(e),
426                    organization_id=organization_id,
427                    organization_name=current_info.organization_name,
428                    email=current_info.email,
429                    previous_is_agentic=previous_is_agentic,
430                    customer_tier=customer_tier,
431                    tier_warning=tier_warning,
432                )
433            )
434            continue
435
436        updated_info = _org_info_to_flag_info(
437            updated,
438            customer_tier=customer_tier,
439            tier_warning=tier_warning,
440        )
441        results.append(
442            OrganizationAgenticFlagUpdateResult(
443                success=True,
444                message=(
445                    f"Organization {organization_id} updated managed agentic org "
446                    f"status to {updated_info.is_agentic}."
447                ),
448                organization_id=organization_id,
449                organization_name=updated_info.organization_name,
450                email=updated_info.email,
451                previous_is_agentic=previous_is_agentic,
452                new_is_agentic=updated_info.is_agentic,
453                customer_tier=customer_tier,
454                tier_warning=tier_warning,
455            )
456        )
457
458    success = all(result.success for result in results)
459    return OrganizationAgenticFlagBatchUpdateResult(
460        success=success,
461        message=(
462            f"Updated {sum(result.success for result in results)} of "
463            f"{len(results)} organization agentic flag request(s)."
464        ),
465        results=results,
466    )
467
468
469def register_organization_agentic_flag_tools(app: FastMCP) -> None:
470    """Register organization agentic flag MCP tools."""
471    register_mcp_tools(app, mcp_module=__name__)