airbyte_ops_mcp.mcp.connection_medic

  1# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
  2"""MCP tools for connection state and catalog: read-only inspection plus emergency (medic) writes.
  3
  4## MCP reference
  5
  6.. include:: ../../../docs/mcp-generated/connection_medic.md
  7    :start-line: 2
  8"""
  9
 10# NOTE: We intentionally do NOT use `from __future__ import annotations` here.
 11# FastMCP has issues resolving forward references when PEP 563 deferred annotations
 12# are used. See: https://github.com/jlowin/fastmcp/issues/905
 13# Python 3.12+ supports modern type hint syntax natively, so this is not needed.
 14
 15__all__: list[str] = []
 16
 17import json
 18import os
 19from typing import Annotated, Any
 20
 21from airbyte import constants
 22from airbyte.cloud.connections import CloudConnection
 23from airbyte.cloud.workspaces import CloudWorkspace
 24from airbyte.secrets.base import SecretString
 25from fastmcp import Context, FastMCP
 26from fastmcp_extensions import get_mcp_config, mcp_tool, register_mcp_tools
 27from pydantic import Field
 28
 29import airbyte_ops_mcp.cloud_admin.connection_state as cloud_connection_state
 30from airbyte_ops_mcp.cloud_admin.auth import CloudAuthError
 31from airbyte_ops_mcp.constants import (
 32    MEDIC_MODE_ENV_VAR,
 33    ServerConfigKey,
 34    WorkspaceAliasEnum,
 35)
 36
 37
 38def _is_medic_mode_enabled() -> bool:
 39    return os.getenv(MEDIC_MODE_ENV_VAR, "").lower() in ("1", "true")
 40
 41
 42def _resolve_cloud_auth(
 43    ctx: Context,
 44) -> tuple[str | None, str | None, str | None]:
 45    """Resolve authentication credentials for API calls.
 46
 47    Returns:
 48        Tuple of (bearer_token, client_id, client_secret).
 49    """
 50    bearer_token = get_mcp_config(ctx, ServerConfigKey.BEARER_TOKEN)
 51    if bearer_token:
 52        return bearer_token, None, None
 53
 54    try:
 55        client_id = get_mcp_config(ctx, ServerConfigKey.CLIENT_ID)
 56        client_secret = get_mcp_config(ctx, ServerConfigKey.CLIENT_SECRET)
 57        return None, client_id, client_secret
 58    except ValueError as e:
 59        raise CloudAuthError(
 60            f"Failed to resolve credentials. Ensure credentials are provided "
 61            f"via Authorization header (Bearer token), "
 62            f"HTTP headers (X-Airbyte-Cloud-Client-Id, X-Airbyte-Cloud-Client-Secret), "
 63            f"or environment variables. Error: {e}"
 64        ) from e
 65
 66
 67@mcp_tool(
 68    destructive=True,
 69    idempotent=False,
 70    open_world=True,
 71)
 72def update_connection_state(
 73    workspace_id: Annotated[
 74        str | WorkspaceAliasEnum,
 75        Field(
 76            description="The Airbyte Cloud workspace ID (UUID) or alias. "
 77            "Aliases: '@devin-ai-sandbox'.",
 78        ),
 79    ],
 80    connection_id: Annotated[
 81        str,
 82        Field(description="The connection ID (UUID) to update state for."),
 83    ],
 84    connection_state_json: Annotated[
 85        str,
 86        Field(
 87            description="The connection state as a JSON string. Must include: "
 88            "'stateType' (one of: 'global', 'stream', 'legacy'), "
 89            "and one of: 'state' (for legacy), 'streamState' (for stream), "
 90            "'globalState' (for global). "
 91            "Tip: Use get_connection_state first to see the current format."
 92        ),
 93    ],
 94    config_api_root: Annotated[
 95        str | None,
 96        Field(
 97            description="Optional API root URL override. "
 98            "Defaults to Airbyte Cloud. "
 99            "Use this to target local or self-hosted deployments.",
100            default=None,
101        ),
102    ] = None,
103    *,
104    ctx: Context,
105) -> list[dict[str, Any]]:
106    """Update the full state for an Airbyte connection.
107
108    WARNING: This is a destructive emergency operation. Use with caution.
109    Uses the safe variant that prevents updates while a sync is running (HTTP 423).
110    Get the current state first with get_connection_state, modify it, then pass
111    the full state object back here.
112    """
113    resolved_workspace_id = WorkspaceAliasEnum.resolve(workspace_id)
114    if resolved_workspace_id is None:
115        raise ValueError(
116            f"Unable to resolve workspace_id {workspace_id!r} to a concrete workspace ID. "
117            "Ensure you provided a valid UUID or supported alias."
118        )
119    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
120
121    conn = _get_cloud_connection(
122        workspace_id=resolved_workspace_id,
123        connection_id=connection_id,
124        api_root=config_api_root or constants.CLOUD_API_ROOT,
125        bearer_token=bearer_token,
126        client_id=client_id,
127        client_secret=client_secret,
128    )
129
130    try:
131        connection_state = json.loads(connection_state_json)
132    except json.JSONDecodeError as e:
133        raise ValueError(
134            "Invalid JSON for parameter 'connection_state_json'. "
135            "Expected a JSON object representing Airbyte connection state "
136            "(for example, an object with a 'stateType' field and related state data)."
137        ) from e
138    conn.import_raw_state(connection_state)
139    return conn.dump_raw_state()
140
141
142@mcp_tool(
143    destructive=True,
144    idempotent=False,
145    open_world=True,
146)
147def update_stream_state(
148    workspace_id: Annotated[
149        str | WorkspaceAliasEnum,
150        Field(
151            description="The Airbyte Cloud workspace ID (UUID) or alias. "
152            "Aliases: '@devin-ai-sandbox'.",
153        ),
154    ],
155    connection_id: Annotated[
156        str,
157        Field(description="The connection ID (UUID) to update state for."),
158    ],
159    stream_name: Annotated[
160        str,
161        Field(description="The name of the stream to update state for."),
162    ],
163    stream_state_json: Annotated[
164        str,
165        Field(
166            description="The state blob for this stream as a JSON string. "
167            "This is the inner state object (e.g., {'cursor': '2024-01-01'}), "
168            "not the full connection state. "
169            "Tip: Use get_connection_state with stream_name first to see the current format."
170        ),
171    ],
172    stream_namespace: Annotated[
173        str | None,
174        Field(
175            description="Optional stream namespace to identify the stream.",
176            default=None,
177        ),
178    ] = None,
179    config_api_root: Annotated[
180        str | None,
181        Field(
182            description="Optional API root URL override. "
183            "Defaults to Airbyte Cloud. "
184            "Use this to target local or self-hosted deployments.",
185            default=None,
186        ),
187    ] = None,
188    *,
189    ctx: Context,
190) -> list[dict[str, Any]]:
191    """Update the state for a single stream within a connection.
192
193    WARNING: This is a destructive emergency operation. Use with caution.
194    Fetches the current full state, replaces only the specified stream's state,
195    then sends the updated state back. If the stream doesn't exist, it is appended.
196    Uses the safe variant that prevents updates while a sync is running (HTTP 423).
197    """
198    resolved_workspace_id = WorkspaceAliasEnum.resolve(workspace_id)
199    if resolved_workspace_id is None:
200        raise ValueError(
201            f"Failed to resolve workspace_id '{workspace_id}' to a concrete workspace ID. "
202            "Use a valid workspace UUID or a supported alias."
203        )
204    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
205
206    conn = _get_cloud_connection(
207        workspace_id=resolved_workspace_id,
208        connection_id=connection_id,
209        api_root=config_api_root or constants.CLOUD_API_ROOT,
210        bearer_token=bearer_token,
211        client_id=client_id,
212        client_secret=client_secret,
213    )
214
215    try:
216        stream_state_dict = json.loads(stream_state_json)
217    except json.JSONDecodeError as exc:
218        raise ValueError(
219            "Invalid JSON provided for 'stream_state_json'. "
220            "Please supply a valid JSON object representing the stream state."
221        ) from exc
222    conn.set_stream_state(
223        stream_name=stream_name,
224        state_blob_dict=stream_state_dict,
225        stream_namespace=stream_namespace,
226    )
227    return conn.dump_raw_state()
228
229
230@mcp_tool(
231    destructive=True,
232    idempotent=False,
233    open_world=True,
234)
235def reset_stream_state(
236    workspace_id: Annotated[
237        str | WorkspaceAliasEnum,
238        Field(
239            description="The Airbyte Cloud workspace ID (UUID) or alias. "
240            "Aliases: '@devin-ai-sandbox'.",
241        ),
242    ],
243    connection_id: Annotated[
244        str,
245        Field(description="The connection ID (UUID) to update state for."),
246    ],
247    stream_name: Annotated[
248        str,
249        Field(description="The configured stream name whose state should be reset."),
250    ],
251    stream_namespace: Annotated[
252        str | None,
253        Field(
254            description="Optional stream namespace to identify the stream.",
255            default=None,
256        ),
257    ] = None,
258    config_api_root: Annotated[
259        str | None,
260        Field(
261            description="Optional API root URL override. "
262            "Defaults to Airbyte Cloud. "
263            "Use this to target local or self-hosted deployments.",
264            default=None,
265        ),
266    ] = None,
267    *,
268    ctx: Context,
269) -> cloud_connection_state.ResetStreamResult:
270    """Reset a configured stream's state so the next sync full-refreshes it.
271
272    WARNING: This is a destructive emergency operation. Use with caution.
273    Uses the safe variant that prevents updates while a sync is running (HTTP 423).
274    Returns `previous_state_backup` in raw Config API format so the state can be restored.
275    """
276    resolved_workspace_id = WorkspaceAliasEnum.resolve(workspace_id)
277    if resolved_workspace_id is None:
278        raise ValueError(
279            f"Failed to resolve workspace_id '{workspace_id}' to a concrete workspace ID. "
280            "Use a valid workspace UUID or a supported alias."
281        )
282    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
283
284    conn = _get_cloud_connection(
285        workspace_id=resolved_workspace_id,
286        connection_id=connection_id,
287        api_root=config_api_root or constants.CLOUD_API_ROOT,
288        bearer_token=bearer_token,
289        client_id=client_id,
290        client_secret=client_secret,
291    )
292    return cloud_connection_state.reset_stream_state(
293        conn,
294        stream_name=stream_name,
295        stream_namespace=stream_namespace,
296    )
297
298
299@mcp_tool(
300    destructive=True,
301    idempotent=False,
302    open_world=True,
303)
304def update_connection_catalog(
305    workspace_id: Annotated[
306        str | WorkspaceAliasEnum,
307        Field(
308            description="The Airbyte Cloud workspace ID (UUID) or alias. "
309            "Aliases: '@devin-ai-sandbox'.",
310        ),
311    ],
312    connection_id: Annotated[
313        str,
314        Field(description="The connection ID (UUID) to update catalog for."),
315    ],
316    configured_catalog_json: Annotated[
317        str,
318        Field(
319            description="The configured catalog as a JSON string. "
320            "This replaces the entire configured catalog for the connection. "
321            "Tip: Use get_connection_catalog first to see the current format."
322        ),
323    ],
324    config_api_root: Annotated[
325        str | None,
326        Field(
327            description="Optional API root URL override. "
328            "Defaults to Airbyte Cloud. "
329            "Use this to target local or self-hosted deployments.",
330            default=None,
331        ),
332    ] = None,
333    *,
334    ctx: Context,
335) -> dict[str, Any]:
336    """Replace the configured catalog for an Airbyte connection.
337
338    WARNING: This is a destructive emergency operation. Use with extreme caution.
339    This replaces the entire configured catalog, which controls which streams
340    are synced, their sync modes, primary keys, and cursor fields.
341    Get the current catalog first, modify it, then pass the full catalog back here.
342    """
343    resolved_workspace_id = WorkspaceAliasEnum.resolve(workspace_id)
344    if resolved_workspace_id is None:
345        raise ValueError(
346            f"Unable to resolve workspace_id {workspace_id!r} to a concrete workspace ID. "
347            "Ensure you provided a valid UUID or supported alias."
348        )
349    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
350
351    conn = _get_cloud_connection(
352        workspace_id=resolved_workspace_id,
353        connection_id=connection_id,
354        api_root=config_api_root or constants.CLOUD_API_ROOT,
355        bearer_token=bearer_token,
356        client_id=client_id,
357        client_secret=client_secret,
358    )
359
360    try:
361        configured_catalog = json.loads(configured_catalog_json)
362    except json.JSONDecodeError as e:
363        raise ValueError(
364            "Invalid JSON for parameter 'configured_catalog_json'. "
365            "Expected a JSON object representing an Airbyte configured catalog."
366        ) from e
367    conn.import_raw_catalog(configured_catalog)
368    result = conn.dump_raw_catalog()
369    if result is None:
370        raise RuntimeError(
371            "Failed to retrieve catalog after import. "
372            f"workspace_id={resolved_workspace_id!r}, connection_id={connection_id!r}"
373        )
374    return result
375
376
377def _get_cloud_connection(
378    workspace_id: str,
379    connection_id: str,
380    api_root: str,
381    bearer_token: str | None = None,
382    client_id: str | None = None,
383    client_secret: str | None = None,
384) -> CloudConnection:
385    """Create a CloudConnection from credentials."""
386    workspace = CloudWorkspace(
387        workspace_id=workspace_id,
388        api_root=api_root,
389        client_id=SecretString(client_id) if client_id else None,
390        client_secret=SecretString(client_secret) if client_secret else None,
391        bearer_token=SecretString(bearer_token) if bearer_token else None,
392    )
393    return workspace.get_connection(connection_id)
394
395
396@mcp_tool(
397    read_only=True,
398    idempotent=True,
399    open_world=True,
400)
401def get_connection_state(
402    workspace_id: Annotated[
403        str | WorkspaceAliasEnum,
404        Field(
405            description="The Airbyte Cloud workspace ID (UUID) or alias. "
406            "Aliases: '@devin-ai-sandbox'.",
407        ),
408    ],
409    connection_id: Annotated[
410        str,
411        Field(description="The connection ID (UUID) to fetch state for."),
412    ],
413    stream_name: Annotated[
414        str | None,
415        Field(
416            description="Optional stream name to filter state for a single stream. "
417            "When provided, only the matching stream's inner state blob is returned.",
418            default=None,
419        ),
420    ] = None,
421    stream_namespace: Annotated[
422        str | None,
423        Field(
424            description="Optional stream namespace to narrow the stream filter. "
425            "Only used when stream_name is also provided.",
426            default=None,
427        ),
428    ] = None,
429    config_api_root: Annotated[
430        str | None,
431        Field(
432            description="Optional API root URL override. "
433            "Defaults to Airbyte Cloud. "
434            "Use this to target local or self-hosted deployments.",
435            default=None,
436        ),
437    ] = None,
438    *,
439    ctx: Context,
440) -> dict[str, Any] | list[dict[str, Any]]:
441    """Get the current state for an Airbyte connection.
442
443    Returns the connection's sync state in Airbyte protocol format (snake_case).
444    The state can be one of: stream (per-stream), global, legacy, or not_set.
445
446    When `stream_name` is provided, returns a dict with `connection_id`,
447    `stream_name`, `stream_namespace`, and `stream_state` keys.
448    Otherwise returns the full state as a list of `AirbyteStateMessage` dicts.
449    """
450    resolved_workspace_id = WorkspaceAliasEnum.resolve(workspace_id)
451    if resolved_workspace_id is None:
452        raise ValueError(
453            f"Invalid workspace ID or alias: {workspace_id!r}. "
454            f"Supported aliases: {', '.join(sorted(alias.name.lower().replace('_', '-') for alias in WorkspaceAliasEnum))}"
455        )
456    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
457
458    conn = _get_cloud_connection(
459        workspace_id=resolved_workspace_id,
460        connection_id=connection_id,
461        api_root=config_api_root or constants.CLOUD_API_ROOT,
462        bearer_token=bearer_token,
463        client_id=client_id,
464        client_secret=client_secret,
465    )
466
467    if stream_name is not None:
468        stream_state = conn.get_stream_state(
469            stream_name=stream_name,
470            stream_namespace=stream_namespace,
471        )
472        return {
473            "connection_id": connection_id,
474            "stream_name": stream_name,
475            "stream_namespace": stream_namespace,
476            "stream_state": stream_state,
477        }
478
479    return conn.dump_raw_state()
480
481
482@mcp_tool(
483    read_only=True,
484    idempotent=True,
485    open_world=True,
486)
487def get_connection_catalog(
488    workspace_id: Annotated[
489        str | WorkspaceAliasEnum,
490        Field(
491            description="The Airbyte Cloud workspace ID (UUID) or alias. "
492            "Aliases: '@devin-ai-sandbox'.",
493        ),
494    ],
495    connection_id: Annotated[
496        str,
497        Field(description="The connection ID (UUID) to fetch catalog for."),
498    ],
499    config_api_root: Annotated[
500        str | None,
501        Field(
502            description="Optional API root URL override. "
503            "Defaults to Airbyte Cloud. "
504            "Use this to target local or self-hosted deployments.",
505            default=None,
506        ),
507    ] = None,
508    *,
509    ctx: Context,
510) -> dict[str, Any]:
511    """Get the configured catalog for an Airbyte connection.
512
513    Returns the connection's configured catalog, which defines which streams
514    are synced, their sync modes, primary keys, and cursor fields.
515    """
516    resolved_workspace_id = WorkspaceAliasEnum.resolve(workspace_id)
517    if resolved_workspace_id is None:
518        raise ValueError(
519            f"Invalid workspace ID or alias: {workspace_id!r}. "
520            f"Supported aliases: {', '.join(sorted(alias.name.lower().replace('_', '-') for alias in WorkspaceAliasEnum))}"
521        )
522    bearer_token, client_id, client_secret = _resolve_cloud_auth(ctx)
523
524    conn = _get_cloud_connection(
525        workspace_id=resolved_workspace_id,
526        connection_id=connection_id,
527        api_root=config_api_root or constants.CLOUD_API_ROOT,
528        bearer_token=bearer_token,
529        client_id=client_id,
530        client_secret=client_secret,
531    )
532
533    result = conn.dump_raw_catalog()
534    if result is None:
535        raise ValueError("No configured catalog found for this connection.")
536    return result
537
538
539# Destructive "break glass" tools that mutate production connection state.
540# These are gated behind AIRBYTE_OPS_MEDIC_MODE so they are only exposed when
541# an operator has explicitly opted into medic mode.
542_MEDIC_WRITE_TOOLS = (
543    "update_connection_state",
544    "update_stream_state",
545    "reset_stream_state",
546    "update_connection_catalog",
547)
548
549
550def register_connection_medic_tools(app: FastMCP) -> None:
551    """Register connection_medic tools with the FastMCP app.
552
553    The read-only inspection tools (`get_connection_state`,
554    `get_connection_catalog`) are always registered. The destructive
555    "break glass" write tools are only exposed when medic mode
556    (`AIRBYTE_OPS_MEDIC_MODE`) is enabled.
557    """
558    register_mcp_tools(app, mcp_module=__name__)
559    if not _is_medic_mode_enabled():
560        for tool_name in _MEDIC_WRITE_TOOLS:
561            app.local_provider.remove_tool(tool_name)