airbyte_ops_mcp.mcp.connection_resources

  1"""MCP tools for connection-level worker resource requirements."""
  2
  3from __future__ import annotations
  4
  5from typing import Annotated
  6
  7from fastmcp import Context, FastMCP
  8from fastmcp_extensions import mcp_tool, register_mcp_tools
  9from pydantic import Field
 10
 11from airbyte_ops_mcp.cloud_admin.connection_resources import (
 12    CpuRung,
 13    DiskRung,
 14    MemoryRung,
 15)
 16from airbyte_ops_mcp.cloud_admin.connection_resources import (
 17    get_connection_resource_requirements as get_connection_resource_requirements_core,
 18)
 19from airbyte_ops_mcp.cloud_admin.connection_resources import (
 20    set_connection_resource_requirements as set_connection_resource_requirements_core,
 21)
 22from airbyte_ops_mcp.cloud_admin.models import (
 23    ConnectionResourceRequirementsInfo,
 24    ConnectionResourceRequirementsOperationResult,
 25)
 26from airbyte_ops_mcp.mcp.cloud_auth import resolve_cloud_auth
 27from airbyte_ops_mcp.tier_cache import TierFilter
 28
 29
 30@mcp_tool(read_only=True, idempotent=True, open_world=True)
 31def get_connection_resource_requirements(
 32    connection_id: Annotated[str, Field(description="The Airbyte Cloud connection ID")],
 33    config_api_root: Annotated[
 34        str | None,
 35        Field(description="Optional Config API root override."),
 36    ] = None,
 37    *,
 38    ctx: Context,
 39) -> ConnectionResourceRequirementsInfo:
 40    """Get connection-level resource requirements.
 41
 42    Returns the explicit CPU and memory values, plus whether the connection
 43    inherits platform defaults.
 44    """
 45    return get_connection_resource_requirements_core(
 46        auth=resolve_cloud_auth(ctx),
 47        connection_id=connection_id,
 48        config_api_root=config_api_root,
 49    )
 50
 51
 52@mcp_tool(destructive=True, idempotent=False, open_world=True)
 53def set_connection_resource_requirements(
 54    connection_id: Annotated[str, Field(description="The Airbyte Cloud connection ID")],
 55    workspace_id: Annotated[
 56        str,
 57        Field(description="The workspace that must own the connection"),
 58    ],
 59    approval_comment_url: Annotated[
 60        str,
 61        Field(description="Slack or GitHub approval record URL"),
 62    ],
 63    issue_url: Annotated[
 64        str,
 65        Field(description="GitHub issue URL providing context for this operation"),
 66    ],
 67    override_reason: Annotated[
 68        str,
 69        Field(description="Reason for the change; at least 10 characters"),
 70    ],
 71    cpu_rung: Annotated[
 72        CpuRung | None,
 73        Field(
 74            description=(
 75                "Absolute CPU limit rung. Omit to preserve the current dimension; "
 76                "`DEFAULT` clears it."
 77            )
 78        ),
 79    ] = None,
 80    memory_rung: Annotated[
 81        MemoryRung | None,
 82        Field(
 83            description=(
 84                "Absolute memory limit rung. Omit to preserve the current dimension; "
 85                "`DEFAULT` clears it."
 86            )
 87        ),
 88    ] = None,
 89    disk_rung: Annotated[
 90        DiskRung | None,
 91        Field(
 92            description=(
 93                "Absolute ephemeral-storage limit rung. Omit to preserve the current "
 94                "dimension; `DEFAULT` clears it."
 95            )
 96        ),
 97    ] = None,
 98    unset: Annotated[
 99        bool,
100        Field(
101            description="Clear the connection override and inherit platform defaults."
102        ),
103    ] = False,
104    customer_tier_filter: Annotated[
105        TierFilter,
106        Field(description="Required customer tier filter; defaults to `TIER_2`."),
107    ] = "TIER_2",
108    cpu_impact_acknowledged: Annotated[
109        bool,
110        Field(
111            description=(
112                "Acknowledge shared worker capacity impact when increasing CPU."
113            )
114        ),
115    ] = False,
116    ai_agent_session_url: Annotated[
117        str | None,
118        Field(
119            description="Optional URL for the AI agent session performing the change."
120        ),
121    ] = None,
122    config_api_root: Annotated[
123        str | None,
124        Field(description="Optional Config API root override."),
125    ] = None,
126    *,
127    ctx: Context,
128) -> ConnectionResourceRequirementsOperationResult:
129    """Set or clear connection-level worker resource requirements.
130
131    CPU, memory, and ephemeral-storage use bounded absolute rungs. Omit a dimension
132    to preserve its current values, or select `DEFAULT` to inherit platform
133    defaults for that dimension. Setting CPU above the current value requires
134    `cpu_impact_acknowledged=True` because each data worker consumes shared
135    capacity. The production profile called `Boosted` corresponds to selecting
136    4 CPU and 4Gi memory. Changes apply to the next sync attempt; an in-flight
137    attempt keeps its current pod sizing. Clearing sends an empty
138    resource-requirements object.
139    """
140    return set_connection_resource_requirements_core(
141        auth=resolve_cloud_auth(ctx),
142        connection_id=connection_id,
143        workspace_id=workspace_id,
144        cpu_rung=cpu_rung,
145        memory_rung=memory_rung,
146        disk_rung=disk_rung,
147        unset=unset,
148        override_reason=override_reason,
149        issue_url=issue_url,
150        approval_comment_url=approval_comment_url,
151        customer_tier_filter=customer_tier_filter,
152        cpu_impact_acknowledged=cpu_impact_acknowledged,
153        ai_agent_session_url=ai_agent_session_url,
154        config_api_root=config_api_root,
155    )
156
157
158def register_connection_resource_tools(app: FastMCP) -> None:
159    """Register connection resource requirement tools with the FastMCP app."""
160    register_mcp_tools(app, mcp_module=__name__)