airbyte_ops_mcp.mcp.feature_flags

  1# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
  2"""MCP tools for LaunchDarkly flag organization targeting.
  3
  4This module exposes tools for reading LaunchDarkly feature flags and for
  5adding or removing individual organization targets on allowlisted flags
  6(those tagged `ops-mcp`), gated by human-in-the-loop approval.
  7
  8## MCP reference
  9
 10.. include:: ../../../docs/mcp-generated/feature_flags.md
 11    :start-line: 2
 12"""
 13
 14# NOTE: We intentionally do NOT use `from __future__ import annotations` here.
 15# FastMCP has issues resolving forward references when PEP 563 deferred
 16# annotations are used. See: https://github.com/jlowin/fastmcp/issues/905
 17
 18__all__: list[str] = []
 19
 20from typing import Annotated, Any, Literal
 21
 22from fastmcp import Context, FastMCP
 23from fastmcp_extensions import 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 import launchdarkly_flags
 31from airbyte_ops_mcp.cloud_admin.auth import (
 32    CloudAuthError,
 33    require_internal_admin_flag_only,
 34)
 35from airbyte_ops_mcp.cloud_admin.launchdarkly_flags import (
 36    LD_OPS_MCP_TAG,
 37    LaunchDarklyAPIError,
 38    LaunchDarklyVerificationError,
 39    apply_org_target_change,
 40    build_change_comment,
 41    flag_is_ops_targetable,
 42    list_flags,
 43    select_variation,
 44    variation_id_of,
 45)
 46from airbyte_ops_mcp.cloud_admin.models import (
 47    LDFeatureFlagInfo,
 48    LDFeatureFlagList,
 49    LDFlagOrgTargetUpdateResult,
 50    LDFlagVariationInfo,
 51    OrganizationAgenticFlagInfo,
 52    OrganizationInfo,
 53)
 54from airbyte_ops_mcp.mcp.organization_admin import (
 55    _row_to_flag_info,
 56    _validate_organization_name,
 57)
 58from airbyte_ops_mcp.prod_db_access.queries import query_organization_agentic_flags
 59from airbyte_ops_mcp.tier_cache import (
 60    get_org_tier,
 61    tier_source_warnings,
 62)
 63
 64
 65def _flag_to_targets_info(flag: dict[str, Any]) -> LDFeatureFlagInfo:
 66    """Convert a raw LaunchDarkly flag dict into org-targeting info."""
 67    env = (flag.get("environments") or {}).get("production") or {}
 68    variations: list[LDFlagVariationInfo] = []
 69    for index, variation in enumerate(flag.get("variations") or []):
 70        variations.append(
 71            LDFlagVariationInfo(
 72                index=index,
 73                variation_id=variation_id_of(variation),
 74                name=variation.get("name"),
 75                value=variation.get("value"),
 76                organization_targets=launchdarkly_flags.get_org_targets(flag, index),
 77            )
 78        )
 79    return LDFeatureFlagInfo(
 80        flag_key=str(flag.get("key")),
 81        flag_name=flag.get("name"),
 82        kind=flag.get("kind"),
 83        tags=list(flag.get("tags") or []),
 84        ops_mcp_addressable=flag_is_ops_targetable(flag),
 85        on=bool(env.get("on")),
 86        variations=variations,
 87    )
 88
 89
 90def _failure_result(
 91    message: str,
 92    flag_key: str,
 93    organization_id: str,
 94    action: Literal["add", "remove"],
 95    **kwargs: Any,
 96) -> LDFlagOrgTargetUpdateResult:
 97    """Build a `success=False` update result for a rejected request."""
 98    return LDFlagOrgTargetUpdateResult(
 99        success=False,
100        message=message,
101        flag_key=flag_key,
102        organization_id=organization_id,
103        action=action,
104        **kwargs,
105    )
106
107
108def _success_result(
109    message: str,
110    flag_key: str,
111    organization_id: str,
112    action: Literal["add", "remove"],
113    **kwargs: Any,
114) -> LDFlagOrgTargetUpdateResult:
115    """Build a `success=True` update result."""
116    return LDFlagOrgTargetUpdateResult(
117        success=True,
118        message=message,
119        flag_key=flag_key,
120        organization_id=organization_id,
121        action=action,
122        **kwargs,
123    )
124
125
126def _load_target_organization(
127    organization_id: str,
128    organization_name: str | None,
129) -> tuple[OrganizationAgenticFlagInfo | None, str | None]:
130    """Look up the target organization and validate `organization_name`.
131
132    Returns `(info, error_message)`; `error_message` is `None` on success.
133    `info` is returned alongside the error when the organization exists but
134    failed the tombstone or name check, so callers can report identity.
135    """
136    rows = query_organization_agentic_flags([organization_id])
137    if not rows:
138        return None, f"Organization {organization_id} not found."
139    current_info = _row_to_flag_info(rows[0])
140    if current_info.tombstone:
141        return current_info, f"Organization {organization_id} is tombstoned."
142    org_info = OrganizationInfo.model_validate(
143        {
144            "organizationId": current_info.organization_id,
145            "organizationName": current_info.organization_name or "",
146            "email": current_info.email,
147            "isAgentic": current_info.is_agentic,
148        }
149    )
150    name_ok, name_error = _validate_organization_name(
151        organization_id, organization_name, org_info
152    )
153    if not name_ok:
154        return current_info, name_error or "Organization name validation failed."
155    return current_info, None
156
157
158@mcp_tool(
159    destructive=False,
160    idempotent=True,
161    open_world=True,
162)
163def list_launchdarkly_feature_flags(
164    tag: Annotated[
165        str | None,
166        Field(
167            description="Optional LaunchDarkly tag filter, passed through to "
168            "the API. Use 'ops-mcp' to list only flags addressable by "
169            "`update_launchdarkly_feature_flag_targeting`.",
170            default=None,
171        ),
172    ] = None,
173    flag_contains: Annotated[
174        str | None,
175        Field(
176            description="Optional case-insensitive substring matched against "
177            "the flag key or display name.",
178            default=None,
179        ),
180    ] = None,
181    limit: Annotated[
182        int,
183        Field(
184            description="Maximum number of flags to return.",
185            ge=1,
186            le=500,
187        ),
188    ] = 50,
189    *,
190    ctx: Context,
191) -> LDFeatureFlagList:
192    """List LaunchDarkly feature flags with their organization targets.
193
194    Returns flags in the `default` project with each variation's individual
195    `organization` targets in the `production` environment. Results can be
196    narrowed by the LaunchDarkly `tag` filter and by a case-insensitive
197    `flag_contains` substring on the flag key or display name, and are
198    truncated to `limit`.
199
200    The `ops_mcp_addressable` field marks flags tagged `ops-mcp` — only
201    those may be modified with `update_launchdarkly_feature_flag_targeting`.
202    """
203    flags = list_flags(tag=tag)
204    if flag_contains is not None:
205        needle = flag_contains.lower()
206        flags = [
207            flag
208            for flag in flags
209            if needle in str(flag.get("key") or "").lower()
210            or needle in str(flag.get("name") or "").lower()
211        ]
212    truncated = len(flags) > limit
213    return LDFeatureFlagList(
214        flags=[_flag_to_targets_info(flag) for flag in flags[:limit]],
215        total_returned=min(len(flags), limit),
216        truncated=truncated,
217    )
218
219
220@mcp_tool(
221    destructive=False,
222    idempotent=True,
223    open_world=True,
224)
225def describe_launchdarkly_feature_flag(
226    flag_key: Annotated[
227        str,
228        Field(description="The LaunchDarkly flag key."),
229    ],
230    *,
231    ctx: Context,
232) -> LDFeatureFlagInfo:
233    """Get a LaunchDarkly flag's production organization targets.
234
235    Works for any flag in the `default` project, including untagged flags;
236    the result's `ops_mcp_addressable` field reports whether the flag carries
237    the `ops-mcp` tag required by `update_launchdarkly_feature_flag_targeting`.
238    """
239    flag = launchdarkly_flags.get_flag(flag_key)
240    return _flag_to_targets_info(flag)
241
242
243@mcp_tool(
244    destructive=True,
245    idempotent=True,
246    open_world=True,
247)
248def update_launchdarkly_feature_flag_targeting(
249    flag_key: Annotated[
250        str,
251        Field(description="The LaunchDarkly flag key."),
252    ],
253    organization_id: Annotated[
254        str,
255        Field(description="The Airbyte Cloud organization UUID to target."),
256    ],
257    action: Annotated[
258        Literal["add", "remove"],
259        Field(
260            description="Whether to add or remove the organization as an "
261            "individual target."
262        ),
263    ],
264    approval_comment_url: Annotated[
265        str,
266        Field(
267            description="URL to the Slack approval record. Obtain this by calling "
268            "`escalate_to_human` with `approval_requested=True`; the backend "
269            "delivers the approval record URL when a human clicks Approve.",
270        ),
271    ],
272    organization_name: Annotated[
273        str | None,
274        Field(
275            description="Confirmation of the target organization. Accepts the "
276            "organization name, email address, or email domain. Required.",
277            default=None,
278        ),
279    ] = None,
280    variation: Annotated[
281        str | None,
282        Field(
283            description="Variation to target, matched case-insensitively on the "
284            "variation name or on `str(value)` (so 'true'/'false' work). "
285            "For `add`: optional for boolean flags (defaults to the `True` "
286            "variation), required for multivariate flags. For `remove`: "
287            "omitting it removes the organization from whichever variation "
288            "currently targets it.",
289            default=None,
290        ),
291    ] = None,
292    agent_session_url: Annotated[
293        str | None,
294        Field(
295            description="Optional agent session URL recorded in the "
296            "LaunchDarkly change comment for auditability.",
297            default=None,
298        ),
299    ] = None,
300    *,
301    ctx: Context,
302) -> LDFlagOrgTargetUpdateResult:
303    """Add or remove an organization as an individual target on a flag.
304
305    Guardrails:
306
307    - Only emits `addTargets`/`removeTargets` semantic-patch instructions on
308      `contextKind=organization` individual targets.
309    - Only writes to the `production` environment of the `default` project.
310    - Only flags tagged `ops-mcp` in LaunchDarkly are eligible; all other
311      flags are rejected before any write.
312    - Requires human-in-the-loop approval: call `escalate_to_human` with
313      `approval_requested=True` and pass the resulting approval record URL
314      as `approval_comment_url`.
315    """
316    try:
317        require_internal_admin_flag_only()
318    except CloudAuthError as e:
319        return _failure_result(
320            f"Admin authentication failed: {e}",
321            flag_key,
322            organization_id,
323            action,
324        )
325
326    organization_id = organization_id.strip().lower()
327
328    try:
329        flag = launchdarkly_flags.get_flag(flag_key)
330    except LaunchDarklyAPIError as e:
331        return _failure_result(str(e), flag_key, organization_id, action)
332    if not flag_is_ops_targetable(flag):
333        return _failure_result(
334            f"Flag {flag_key!r} is not addressable by Ops MCP. Add the "
335            f"{LD_OPS_MCP_TAG!r} tag to the flag in LaunchDarkly to opt it in.",
336            flag_key,
337            organization_id,
338            action,
339        )
340
341    try:
342        selection = select_variation(
343            flag,
344            action=action,
345            organization_id=organization_id,
346            variation=variation,
347        )
348    except LaunchDarklyAPIError as e:
349        return _failure_result(str(e), flag_key, organization_id, action)
350
351    try:
352        approver_email = resolve_admin_email_from_approval(
353            approval_comment_url=approval_comment_url,
354        )
355    except ApprovalResolutionError as e:
356        return _failure_result(str(e), flag_key, organization_id, action)
357
358    current_info, org_error = _load_target_organization(
359        organization_id, organization_name
360    )
361    if org_error is not None:
362        return _failure_result(
363            org_error,
364            flag_key,
365            organization_id,
366            action,
367            organization_name=(
368                current_info.organization_name if current_info else None
369            ),
370            email=current_info.email if current_info else None,
371        )
372    assert current_info is not None
373
374    tier_result = get_org_tier(
375        organization_id=organization_id,
376        allow_degraded=True,
377    )
378    customer_tier = tier_result.customer_tier
379    warnings = tier_source_warnings(tier_result.source_health)
380
381    previous_targets = selection.previous_targets
382    already_targeted = organization_id in previous_targets
383    result_kwargs: dict[str, Any] = {
384        "organization_name": current_info.organization_name,
385        "email": current_info.email,
386        "variation_name": selection.name,
387        "previous_targets": previous_targets,
388        "customer_tier": customer_tier,
389        "warnings": warnings,
390    }
391    if action == "add" and already_targeted:
392        return _success_result(
393            f"Organization {organization_id} is already individually "
394            f"targeted on flag {flag_key!r}.",
395            flag_key,
396            organization_id,
397            action,
398            new_targets=previous_targets,
399            **result_kwargs,
400        )
401    if action == "remove" and (selection.index is None or not already_targeted):
402        return _success_result(
403            f"Organization {organization_id} is not individually "
404            f"targeted on flag {flag_key!r}; nothing to remove.",
405            flag_key,
406            organization_id,
407            action,
408            new_targets=previous_targets,
409            **result_kwargs,
410        )
411
412    comment = build_change_comment(
413        action,
414        organization_id,
415        approver_email,
416        agent_session_url,
417    )
418    try:
419        new_targets, warnings = apply_org_target_change(
420            flag_key,
421            action=action,
422            selection=selection,
423            organization_id=organization_id,
424            comment=comment,
425            verify_all_variations=(action == "remove" and variation is None),
426        )
427    except LaunchDarklyVerificationError as e:
428        return _failure_result(
429            str(e),
430            flag_key,
431            organization_id,
432            action,
433            new_targets=e.new_targets,
434            **result_kwargs,
435        )
436    except LaunchDarklyAPIError as e:
437        return _failure_result(
438            str(e), flag_key, organization_id, action, **result_kwargs
439        )
440    result_kwargs["warnings"] = warnings
441
442    verb = "added to" if action == "add" else "removed from"
443    return _success_result(
444        f"Organization {organization_id} {verb} the individual targets "
445        f"of flag {flag_key!r}.",
446        flag_key,
447        organization_id,
448        action,
449        new_targets=new_targets,
450        **result_kwargs,
451    )
452
453
454def register_launchdarkly_ops_tools(app: FastMCP) -> None:
455    """Register launchdarkly_ops tools with the FastMCP app."""
456    register_mcp_tools(app, mcp_module=__name__)