1# Copyright (c) 2026 Airbyte, Inc., all rights reserved.
  2"""Airbyte Agents MCP operations.
  3
  4.. include:: ../../docs/mcp-generated/agents.md
  5"""
  6
  7# No public Python API — MCP primitives are registered via decorators and
  8# documented via the generated Markdown include above. Setting `__all__` to an
  9# empty list tells pdoc (and other doc tools) not to surface the individual
 10# tool / helper definitions as a redundant "API Documentation" list.
 11__all__: list[str] = []
 12
 13import json
 14from typing import Annotated, Any, Literal, get_args
 15
 16from fastmcp import Context, FastMCP
 17from fastmcp_extensions import get_mcp_config, mcp_tool, register_mcp_tools
 18from pydantic import BaseModel, Field
 19
 20from airbyte.agents.connectors import AgentConnector
 21from airbyte.agents.organizations import AgentOrganization
 22from airbyte.agents.workspaces import AgentWorkspace
 23from airbyte.constants import (
 24    CLOUD_BEARER_TOKEN_ENV_VAR,
 25    CLOUD_CLIENT_ID_ENV_VAR,
 26    CLOUD_CLIENT_SECRET_ENV_VAR,
 27    CLOUD_WORKSPACE_ID_ENV_VAR,
 28    MCP_BEARER_TOKEN_HEADER,
 29    MCP_CONFIG_BEARER_TOKEN,
 30    MCP_CONFIG_CLIENT_ID,
 31    MCP_CONFIG_CLIENT_SECRET,
 32    MCP_CONFIG_WORKSPACE_ID,
 33    MCP_WORKSPACE_ID_HEADER,
 34)
 35from airbyte.exceptions import PyAirbyteInputError
 36from airbyte.mcp._arg_resolvers import resolve_list_of_strings
 37from airbyte.mcp._tool_utils import AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET
 38from airbyte.mcp.cloud import _add_defaults_for_exclude_args
 39
 40
 41AgentReadAction = Literal["list", "get", "search", "api_search"]
 42"""The connector actions that only read data.
 43
 44The `download` action is deliberately absent even though it reads: it returns a binary
 45stream rather than JSON, which PyAirbyte does not yet support.
 46"""
 47
 48AgentAction = Literal["list", "get", "search", "api_search", "create", "update", "delete"]
 49"""Every connector action callable through the MCP layer, including writes."""
 50
 51AGENTS_AUTH_TIP_TEXT = (
 52    f"The Airbyte Agents API authenticates with Airbyte Cloud credentials. When connecting "
 53    f"to a hosted MCP server, provide a bearer token via the `{MCP_BEARER_TOKEN_HEADER}` "
 54    f"header, or client credentials via the transport `Client-Id` and `Client-Secret` "
 55    f"headers. For local or stdio connections, set the `{CLOUD_BEARER_TOKEN_ENV_VAR}` "
 56    f"environment variable, or both `{CLOUD_CLIENT_ID_ENV_VAR}` and "
 57    f"`{CLOUD_CLIENT_SECRET_ENV_VAR}`. Call `list_agent_connectors` to discover connector "
 58    f"IDs, then `describe_agent_connector` to learn which entities a connector supports, "
 59    f"before calling `execute_agent_connector`."
 60)
 61WORKSPACE_ID_TIP_TEXT = (
 62    f"Workspace ID. Hosted MCP connections pass it via the `{MCP_WORKSPACE_ID_HEADER}` "
 63    f"header; local or stdio connections use the `{CLOUD_WORKSPACE_ID_ENV_VAR}` "
 64    f"environment variable."
 65)
 66
 67
 68class AgentWorkspaceResult(BaseModel):
 69    """Information about a workspace on the Airbyte Agents platform."""
 70
 71    workspace_id: str
 72    """The workspace ID."""
 73
 74    workspace_name: str | None = None
 75    """Display name of the workspace."""
 76
 77    organization_id: str | None = None
 78    """The organization that owns the workspace, when reported."""
 79
 80
 81class AgentWorkspaceListResult(BaseModel):
 82    """Result of listing workspaces on the Airbyte Agents platform."""
 83
 84    workspaces: list[AgentWorkspaceResult]
 85    """Workspaces reachable through the Agents API with these credentials."""
 86
 87
 88class AgentConnectorResult(BaseModel):
 89    """Information about a connector configured on the Airbyte Agents platform."""
 90
 91    connector_id: str
 92    """The connector ID, used as `connector_id` in the other Agents tools."""
 93
 94    connector_name: str | None = None
 95    """Display name of the connector."""
 96
 97
 98class AgentConnectorListResult(BaseModel):
 99    """Result of listing connectors in an Airbyte Agents workspace."""
100
101    connectors: list[AgentConnectorResult]
102    """Connectors configured in the workspace."""
103
104
105class AgentConnectorDetailsResult(BaseModel):
106    """Details about a single Airbyte Agents connector."""
107
108    connector_id: str
109    """The connector ID."""
110
111    connector_name: str | None = None
112    """Display name of the connector."""
113
114    workspace_id: str | None = None
115    """The workspace that owns the connector."""
116
117    source_definition_name: str | None = None
118    """The name of the underlying source definition, for example `GitHub`."""
119
120    context_store_entities: list[str]
121    """Entities this connector can cache in the Context Store.
122
123    This is not an exhaustive list of executable entities: an entity may be executable via
124    `execute_agent_connector` without appearing here.
125    """
126
127    warnings: list[str]
128    """Warnings the Agents API reported about this connector."""
129
130
131class AgentExecuteToolResult(BaseModel):
132    """Result of executing a single action against an Airbyte Agents connector."""
133
134    status: str
135    """The execution status reported by the Agents API, for example `success`."""
136
137    result: Any = None
138    """The action's payload. Entity-returning actions put a list of entities here."""
139
140    has_next_page: bool = False
141    """Whether the connector reported more entities after this page."""
142
143    end_cursor: str | None = None
144    """The cursor to pass as `cursor` to fetch the next page, when one is available."""
145
146    execution_time_ms: int | None = None
147    """How long the connector took to execute the action, when reported."""
148
149    warning: dict[str, Any] | None = None
150    """A warning reported alongside an otherwise successful result."""
151
152
153def _resolve_api_args(api_args: dict[str, Any] | str | None) -> dict[str, Any] | None:
154    """Resolve `api_args` from a dictionary or a JSON object string."""
155    if api_args is None or isinstance(api_args, dict):
156        return api_args
157
158    try:
159        parsed: Any = json.loads(api_args)
160    except json.JSONDecodeError as ex:
161        raise PyAirbyteInputError(
162            message="The `api_args` string is not valid JSON.",
163            guidance="Pass `api_args` as an object, or as a JSON object string.",
164        ) from ex
165
166    if not isinstance(parsed, dict):
167        raise PyAirbyteInputError(
168            message="The `api_args` string is not a JSON object.",
169            guidance="Pass `api_args` as an object, or as a JSON object string.",
170            context={"parsed_type": type(parsed).__name__},
171        )
172    return parsed
173
174
175def _get_agent_organization(ctx: Context, organization_id: str | None) -> AgentOrganization:
176    """Build an `AgentOrganization` from MCP config."""
177    return AgentOrganization(
178        organization_id=organization_id,
179        client_id=get_mcp_config(ctx, MCP_CONFIG_CLIENT_ID),
180        client_secret=get_mcp_config(ctx, MCP_CONFIG_CLIENT_SECRET),
181        bearer_token=get_mcp_config(ctx, MCP_CONFIG_BEARER_TOKEN),
182    )
183
184
185def _get_agent_workspace(ctx: Context, workspace_id: str | None) -> AgentWorkspace:
186    """Build an `AgentWorkspace` from MCP config."""
187    return AgentWorkspace(
188        workspace_id=workspace_id or get_mcp_config(ctx, MCP_CONFIG_WORKSPACE_ID),
189        client_id=get_mcp_config(ctx, MCP_CONFIG_CLIENT_ID),
190        client_secret=get_mcp_config(ctx, MCP_CONFIG_CLIENT_SECRET),
191        bearer_token=get_mcp_config(ctx, MCP_CONFIG_BEARER_TOKEN),
192    )
193
194
195def _get_agent_connector(
196    ctx: Context,
197    connector_id: str,
198    workspace_id: str | None = None,
199) -> AgentConnector:
200    """Get an `AgentConnector` from its workspace, using MCP config.
201
202    The Agents API addresses a connector by ID alone, but the connector is fetched through
203    its workspace anyway, so a connector ID belonging to another workspace raises before
204    any action runs.
205    """
206    return _get_agent_workspace(ctx, workspace_id).get_connector(connector_id)
207
208
209def _execute(  # noqa: PLR0913  # Mirrors the tool signatures it serves.
210    ctx: Context,
211    *,
212    connector_id: str,
213    workspace_id: str | None,
214    entity_type: str,
215    action: str,
216    api_args: dict[str, Any] | str | None,
217    select_fields: list[str] | str | None,
218    exclude_fields: list[str] | str | None,
219    page_size: int | None,
220    cursor: str | None,
221    intent: str | None,
222    read_only: bool | None = None,
223) -> AgentExecuteToolResult:
224    """Execute one connector action and shape it into an `AgentExecuteToolResult`.
225
226    When `read_only` is `True`, write actions are rejected before any request is sent.
227    """
228    if read_only and action not in get_args(AgentReadAction):
229        raise PyAirbyteInputError(
230            message="This action writes data and cannot run in read-only mode.",
231            guidance=f"Read-only actions are: {', '.join(get_args(AgentReadAction))}.",
232            context={"action": action},
233        )
234
235    result = _get_agent_connector(ctx, connector_id, workspace_id).execute(
236        entity_type,
237        action,
238        _resolve_api_args(api_args),
239        select_fields=resolve_list_of_strings(select_fields),
240        exclude_fields=resolve_list_of_strings(exclude_fields),
241        page_size=page_size,
242        cursor=cursor,
243        intent=intent,
244    )
245    return AgentExecuteToolResult(
246        status=result.status,
247        result=result.result,
248        has_next_page=result.has_next_page,
249        end_cursor=result.end_cursor,
250        execution_time_ms=result.execution_metadata.execution_time_ms,
251        warning=result.warning,
252    )
253
254
255@mcp_tool(
256    read_only=True,
257    idempotent=True,
258    open_world=True,
259    extra_help_text=AGENTS_AUTH_TIP_TEXT,
260)
261def list_agent_workspaces(
262    ctx: Context,
263    *,
264    organization_id: Annotated[
265        str | None,
266        Field(
267            description=(
268                "Organization ID to scope the listing to. Omit it when the credentials "
269                "belong to exactly one organization."
270            ),
271            default=None,
272        ),
273    ],
274) -> AgentWorkspaceListResult:
275    """List the workspaces reachable through the Airbyte Agents API."""
276    organization = _get_agent_organization(ctx, organization_id)
277    return AgentWorkspaceListResult(
278        workspaces=[
279            AgentWorkspaceResult(
280                workspace_id=workspace.workspace_id,
281                workspace_name=workspace.name,
282                organization_id=workspace.organization_id,
283            )
284            for workspace in organization.list_workspaces()
285        ]
286    )
287
288
289@mcp_tool(
290    read_only=True,
291    idempotent=True,
292    open_world=True,
293    extra_help_text=AGENTS_AUTH_TIP_TEXT,
294)
295def list_agent_connectors(
296    ctx: Context,
297    *,
298    workspace_id: Annotated[
299        str | None,
300        Field(
301            description=WORKSPACE_ID_TIP_TEXT,
302            default=None,
303        ),
304    ],
305) -> AgentConnectorListResult:
306    """List the connectors configured in an Airbyte Agents workspace."""
307    workspace = _get_agent_workspace(ctx, workspace_id)
308    return AgentConnectorListResult(
309        connectors=[
310            AgentConnectorResult(
311                connector_id=connector.connector_id,
312                connector_name=connector.name,
313            )
314            for connector in workspace.list_connectors()
315        ]
316    )
317
318
319@mcp_tool(
320    read_only=True,
321    idempotent=True,
322    open_world=True,
323    extra_help_text=AGENTS_AUTH_TIP_TEXT,
324)
325def describe_agent_connector(
326    ctx: Context,
327    connector_id: Annotated[
328        str,
329        Field(description="The ID of the Airbyte Agents connector."),
330    ],
331    *,
332    workspace_id: Annotated[
333        str | None,
334        Field(
335            description=WORKSPACE_ID_TIP_TEXT,
336            default=None,
337        ),
338    ],
339) -> AgentConnectorDetailsResult:
340    """Describe an Airbyte Agents connector, including its Context Store entities.
341
342    Call this before `execute_agent_connector` to learn what the connector exposes. The
343    connector must belong to the given workspace.
344    """
345    details = _get_agent_connector(ctx, connector_id, workspace_id).describe()
346    return AgentConnectorDetailsResult(
347        connector_id=details.connector_id,
348        connector_name=details.name,
349        workspace_id=details.workspace_id,
350        source_definition_name=details.source_definition_name,
351        context_store_entities=details.context_store_entities,
352        warnings=[str(warning) for warning in details.warnings],
353    )
354
355
356@mcp_tool(
357    read_only=True,
358    idempotent=True,
359    open_world=True,
360    extra_help_text=AGENTS_AUTH_TIP_TEXT,
361)
362def execute_agent_connector_ro(  # noqa: PLR0913  # Explicit args are the point of this tool.
363    ctx: Context,
364    connector_id: Annotated[
365        str,
366        Field(description="The ID of the Airbyte Agents connector."),
367    ],
368    entity_type: Annotated[
369        str,
370        Field(
371            description=(
372                "The type of entity to act on, for example 'issues'. Call "
373                "`describe_agent_connector` to see the entity types a connector supports."
374            ),
375        ),
376    ],
377    action: Annotated[
378        AgentReadAction,
379        Field(description="The read action to run against the entity type."),
380    ],
381    api_args: Annotated[
382        dict[str, Any] | str | None,
383        Field(
384            description=(
385                "Connector-specific arguments for the action, as an object or a JSON "
386                "object string. For example {'repository': 'airbytehq/PyAirbyte'}."
387            ),
388            default=None,
389        ),
390    ],
391    *,
392    select_fields: Annotated[
393        list[str] | str | None,
394        Field(
395            description="Fields to keep in the response, as a list or a CSV string.",
396            default=None,
397        ),
398    ],
399    exclude_fields: Annotated[
400        list[str] | str | None,
401        Field(
402            description="Fields to drop from the response, as a list or a CSV string.",
403            default=None,
404        ),
405    ],
406    page_size: Annotated[
407        int | None,
408        Field(description="Maximum number of entities to return in this page.", default=None),
409    ],
410    cursor: Annotated[
411        str | None,
412        Field(
413            description="Pagination cursor, taken from `end_cursor` of a previous result.",
414            default=None,
415        ),
416    ],
417    intent: Annotated[
418        str | None,
419        Field(
420            description="A short description of why the action is being run.",
421            default=None,
422        ),
423    ],
424    workspace_id: Annotated[
425        str | None,
426        Field(
427            description=WORKSPACE_ID_TIP_TEXT,
428            default=None,
429        ),
430    ],
431) -> AgentExecuteToolResult:
432    """Read data from an Airbyte Agents connector, without modifying anything.
433
434    This tool only accepts read actions, so it stays available in read-only mode. Use
435    `execute_agent_connector` for actions that create, update, or delete data. Entity types
436    are connector-specific, so call `describe_agent_connector` first. The connector must
437    belong to the given workspace.
438    """
439    return _execute(
440        ctx,
441        connector_id=connector_id,
442        workspace_id=workspace_id,
443        entity_type=entity_type,
444        action=action,
445        api_args=api_args,
446        select_fields=select_fields,
447        exclude_fields=exclude_fields,
448        page_size=page_size,
449        cursor=cursor,
450        intent=intent,
451        read_only=True,
452    )
453
454
455@mcp_tool(
456    open_world=True,
457    extra_help_text=AGENTS_AUTH_TIP_TEXT,
458)
459def execute_agent_connector(  # noqa: PLR0913  # Explicit args are the point of this tool.
460    ctx: Context,
461    connector_id: Annotated[
462        str,
463        Field(description="The ID of the Airbyte Agents connector."),
464    ],
465    entity_type: Annotated[
466        str,
467        Field(
468            description=(
469                "The type of entity to act on, for example 'issues'. Call "
470                "`describe_agent_connector` to see the entity types a connector supports."
471            ),
472        ),
473    ],
474    action: Annotated[
475        AgentAction,
476        Field(description="The action to run against the entity type."),
477    ],
478    api_args: Annotated[
479        dict[str, Any] | str | None,
480        Field(
481            description=(
482                "Connector-specific arguments for the action, as an object or a JSON "
483                "object string. For example {'repository': 'airbytehq/PyAirbyte'}."
484            ),
485            default=None,
486        ),
487    ],
488    *,
489    select_fields: Annotated[
490        list[str] | str | None,
491        Field(
492            description="Fields to keep in the response, as a list or a CSV string.",
493            default=None,
494        ),
495    ],
496    exclude_fields: Annotated[
497        list[str] | str | None,
498        Field(
499            description="Fields to drop from the response, as a list or a CSV string.",
500            default=None,
501        ),
502    ],
503    page_size: Annotated[
504        int | None,
505        Field(description="Maximum number of entities to return in this page.", default=None),
506    ],
507    cursor: Annotated[
508        str | None,
509        Field(
510            description="Pagination cursor, taken from `end_cursor` of a previous result.",
511            default=None,
512        ),
513    ],
514    intent: Annotated[
515        str | None,
516        Field(
517            description="A short description of why the action is being run.",
518            default=None,
519        ),
520    ],
521    read_only: Annotated[
522        bool | None,
523        Field(
524            description=(
525                "Set to `true` to reject write actions before any request is sent, when the "
526                "caller wants a read guarantee from this tool."
527            ),
528            default=None,
529        ),
530    ],
531    workspace_id: Annotated[
532        str | None,
533        Field(
534            description=WORKSPACE_ID_TIP_TEXT,
535            default=None,
536        ),
537    ],
538) -> AgentExecuteToolResult:
539    """Execute a single action against an Airbyte Agents connector, including writes.
540
541    Prefer `execute_agent_connector_ro` when only reading, since it is available in
542    read-only mode. Entity types and actions are connector-specific, so call
543    `describe_agent_connector` first. The connector must belong to the given workspace.
544    """
545    return _execute(
546        ctx,
547        connector_id=connector_id,
548        workspace_id=workspace_id,
549        entity_type=entity_type,
550        action=action,
551        api_args=api_args,
552        select_fields=select_fields,
553        exclude_fields=exclude_fields,
554        page_size=page_size,
555        cursor=cursor,
556        intent=intent,
557        read_only=read_only,
558    )
559
560
561def register_agents_tools(app: FastMCP) -> None:
562    """Register the Airbyte Agents tools with the FastMCP app."""
563    exclude_args = ["workspace_id"] if AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET else None
564    if exclude_args:
565        _add_defaults_for_exclude_args(exclude_args)
566    register_mcp_tools(
567        app,
568        mcp_module=__name__,
569        exclude_args=exclude_args,
570    )