Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion airbyte/_direct_connectors/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -376,6 +376,9 @@ class ExternalSearchResult(BaseModel):
"""Server warnings, plus per-index search errors collected from `metadata` so partial
failures stand out, and a notice when no index was searched at all."""

errors: list[str] = Field(default_factory=list)
"""Errors that stopped the search, for example search not being enabled."""

@field_validator("warnings", mode="before")
@classmethod
def _none_to_empty(cls, value: object) -> object:
Expand All @@ -389,7 +392,7 @@ def _collect_index_errors(self) -> ExternalSearchResult:
for entry in self.metadata
if entry.error
]
if not self.metadata:
if not self.metadata and not self.errors:
collected.append(
"No index matched this search type and stream selection, so empty results "
"do not mean there are no matches. Call `get_cloud_search_status` to see "
Expand Down Expand Up @@ -507,6 +510,9 @@ class ExternalSearchStatusResult(BaseModel):
warnings: list[str] = Field(default_factory=list)
"""Non-fatal issues, for example a stream filter that matched no indexed stream."""

errors: list[str] = Field(default_factory=list)
"""Errors that stopped the status check, for example search not being enabled."""

@field_validator("warnings", mode="before")
@classmethod
def _none_to_empty(cls, value: object) -> object:
Expand Down
46 changes: 37 additions & 9 deletions airbyte/cloud/connectors.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,15 @@
)


_GENERIC_CONFLICT_MESSAGES = frozenset(
{
"The request conflicts with the current resource state.",
"Airbyte Cloud returned HTTP 409.",
}
)
"""409 messages that carry no Cloud detail: the platform default and PyAirbyte's fallback."""


class _ConnectorDefinitionLike(Protocol):
"""The connector definition fields PyAirbyte reads from the Cloud public API."""

Expand Down Expand Up @@ -597,7 +606,10 @@ def execute_search_query( # noqa: PLR0913 # Explicit args are the point of thi
)
except exc.AirbyteCloudError as error:
self._raise_if_feature_not_enabled(
error, ConnectorFeature.SEARCH_INDEXING, guidance=_SEARCH_NOT_ENABLED_GUIDANCE
error,
ConnectorFeature.SEARCH_INDEXING,
guidance=_SEARCH_NOT_ENABLED_GUIDANCE,
enablement_url=self.connector_url,
)
raise

Expand All @@ -616,7 +628,10 @@ def get_search_status(self) -> ExternalSearchStatusResult:
)
except exc.AirbyteCloudError as error:
self._raise_if_feature_not_enabled(
error, ConnectorFeature.SEARCH_INDEXING, guidance=_SEARCH_NOT_ENABLED_GUIDANCE
error,
ConnectorFeature.SEARCH_INDEXING,
guidance=_SEARCH_NOT_ENABLED_GUIDANCE,
enablement_url=self.connector_url,
)
raise

Expand All @@ -630,14 +645,22 @@ def _raise_if_feature_not_enabled(
) -> None:
"""Map a forbidden or not-found `error` to a clear error, when one applies.

Only a 403 or 404 is re-checked against the connector's enablement, and only an
enablement lookup reporting `feature` as disabled raises
`AirbyteExternalAccessNotEnabledError`. When the enablement lookup itself fails,
`error` is re-raised instead: with wrong-ID guidance on a 404 (no connector of
this kind), with access guidance on a 403, and as-is otherwise. When `feature` is
enabled, this returns.
Only a 403 or 404 (or a 409 without Cloud detail for `SEARCH_INDEXING`) is re-checked
against the connector's enablement, and only an enablement lookup reporting `feature`
as disabled raises `AirbyteExternalAccessNotEnabledError`. When the enablement lookup
itself fails, `error` is re-raised instead: with wrong-ID guidance on a 404 (no
connector of this kind), with access guidance on a 403, and as-is otherwise. When
`feature` is enabled, this returns.
"""
if not agents_api_util.is_not_enabled_error(error):
# Search also reports a connector with search turned off as a 409 conflict. A 409
# with its own Cloud detail is relayed as-is instead.
search_disabled = (
feature == ConnectorFeature.SEARCH_INDEXING
and isinstance(error, exc.AirbyteCloudApiError)
and error.status_code == HTTPStatus.CONFLICT
Comment thread
devin-ai-integration[bot] marked this conversation as resolved.
and (error.message is None or error.message in _GENERIC_CONFLICT_MESSAGES)
)
if not (agents_api_util.is_not_enabled_error(error) or search_disabled):
return
try:
enabled = self.is_feature_enabled(feature)
Expand Down Expand Up @@ -671,6 +694,11 @@ def _raise_if_feature_not_enabled(
not_enabled.guidance = guidance
if enablement_url is not None:
not_enabled.enablement_url = enablement_url
if self._enablement is not None and self._enablement.organization_id:
not_enabled.organization_enablement_url = (
f"{get_web_url_root(self.workspace.api_root)}/organization/"
f"{self._enablement.organization_id}/context-layer"
)
raise not_enabled from error

@property
Expand Down
3 changes: 3 additions & 0 deletions airbyte/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -626,6 +626,9 @@ class AirbyteExternalAccessNotEnabledError(AirbyteCloudError):
enablement_url: str | None = None
"""The web URL where the user can enable direct access."""

organization_enablement_url: str | None = None
"""The web URL where an organization admin can enable the Context layer, when known."""


@dataclass
class AirbyteCloudConnectionError(AirbyteCloudError):
Expand Down
69 changes: 55 additions & 14 deletions airbyte/mcp/cloud.py
Original file line number Diff line number Diff line change
Expand Up @@ -1892,13 +1892,41 @@ def execute_external_sql_query(
return _direct_access_disabled_result(error, connector)


def _search_disabled_errors(
error: AirbyteExternalAccessNotEnabledError,
connector: CloudConnector,
) -> list[str]:
"""Return user-facing search enablement guidance as MCP result errors.

TODO: Move this translation to the platform, which drops Sonar's specific search error
details today: https://linear.app/airbyteio/issue/PLAT-1557
"""
label = error.connector_name or f"{connector.connector_type} {connector.connector_id}"
if error.organization_enablement_url:
steps = (
f"If the Context layer isn't enabled for your organization yet, an organization "
Comment thread
girarda marked this conversation as resolved.
f"admin can [click here]({error.organization_enablement_url}) to enable it first. "
f"Then [click here]({error.enablement_url}) to turn on agent access and search "
f"for {label} specifically."
)
else:
steps = f"[Click here]({error.enablement_url}) to enable it."
return [
redact_agent_text(
f"Show the user this message as-is: Search is currently disabled for {label}. {steps}"
),
"Once search is enabled, the connector's synced data must be indexed before "
"searches return results. Call `get_cloud_search_status` to check which streams "
"are indexed and whether their backfill is complete.",
]


@mcp_tool(
read_only=True,
idempotent=True,
open_world=True,
external_access=True,
tracing=lambda args: agent_action_attributes("execute_external_search_query", args),
insiders_only=True,
)
def execute_external_search_query( # noqa: PLR0913 # Explicit args mirror the connector API.
ctx: Context,
Expand Down Expand Up @@ -2010,24 +2038,30 @@ def execute_external_search_query( # noqa: PLR0913 # Explicit args mirror the
connector = _get_cloud_connector(
_get_cloud_workspace(ctx, workspace_id), connector_id, connector_type
)
return connector.execute_search_query(
prompt,
search_type=search_type,
limit=limit,
streams=resolve_list_of_dicts(streams, arg_name="streams"),
lookback_seconds=lookback_seconds,
max_context_chars=max_context_chars,
min_similarity=min_similarity,
max_similarity_diff=max_similarity_diff,
destination_id=destination_id,
)
try:
return connector.execute_search_query(
prompt,
search_type=search_type,
limit=limit,
streams=resolve_list_of_dicts(streams, arg_name="streams"),
lookback_seconds=lookback_seconds,
max_context_chars=max_context_chars,
min_similarity=min_similarity,
max_similarity_diff=max_similarity_diff,
destination_id=destination_id,
)
except AirbyteExternalAccessNotEnabledError as error:
if error.enablement_url is None:
raise
return ExternalSearchResult(
hits=[], metadata=[], errors=_search_disabled_errors(error, connector)
)


@mcp_tool(
read_only=True,
idempotent=True,
external_access=True,
insiders_only=True,
)
def get_cloud_search_status(
ctx: Context,
Expand Down Expand Up @@ -2080,7 +2114,14 @@ def get_cloud_search_status(
connector = _get_cloud_connector(
_get_cloud_workspace(ctx, workspace_id), connector_id, connector_type
)
status = connector.get_search_status()
try:
status = connector.get_search_status()
except AirbyteExternalAccessNotEnabledError as error:
if error.enablement_url is None:
raise
return ExternalSearchStatusResult(
sources=[], errors=_search_disabled_errors(error, connector)
)
if stream_name is None and namespace is None:
return status

Expand Down
97 changes: 96 additions & 1 deletion tests/unit_tests/test_cloud_direct_execute.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
from airbyte._direct_connectors import api_util as agents_api_util
from airbyte._util import api_util
from airbyte._direct_connectors.models import (
ConnectorEnablement,
ExternalApiExecuteResult,
ExternalSearchHit,
ExternalSearchResult,
Expand Down Expand Up @@ -1177,6 +1178,43 @@ def test_search_input_validation(
assert calls == []


@pytest.mark.parametrize("method_name", ["execute_search_query", "get_search_status"])
@pytest.mark.parametrize(
("organization_id", "expected_url"),
[
pytest.param(
"org-1",
"https://cloud.airbyte.com/organization/org-1/context-layer",
id="organization_known",
),
pytest.param(None, None, id="organization_unknown"),
],
)
def test_search_not_enabled_sets_organization_enablement_url(
monkeypatch: pytest.MonkeyPatch,
method_name: str,
organization_id: str | None,
expected_url: str | None,
) -> None:
"""A search not-enabled error links to the org Context layer page when the org is known."""
workspace = _make_workspace(monkeypatch)
_patch_context_layer(monkeypatch)
_patch_search(monkeypatch, {}, error=AirbyteCloudApiError(status_code=403))
source = _seed_source(workspace, "source-1", "GitHub Issues")
source._enablement = ConnectorEnablement( # noqa: SLF001
enable_agent_access=False,
enable_indexing=False,
organizationId=organization_id,
)
args = ("refunds",) if method_name == "execute_search_query" else ()

with pytest.raises(AirbyteExternalAccessNotEnabledError) as exc_info:
getattr(source, method_name)(*args)

assert exc_info.value.enablement_url == source.connector_url
assert exc_info.value.organization_enablement_url == expected_url


@pytest.mark.parametrize("method_name", ["execute_search_query", "get_search_status"])
@pytest.mark.parametrize(
("status_code", "search_flag", "expected_exc", "expected_status"),
Expand All @@ -1188,6 +1226,14 @@ def test_search_input_validation(
None,
id="forbidden_disabled_raises_not_enabled",
),
pytest.param(
409,
False,
AirbyteExternalAccessNotEnabledError,
None,
id="conflict_disabled_raises_not_enabled",
),
pytest.param(409, True, AirbyteCloudError, 409, id="conflict_enabled_reraises"),
pytest.param(
404, True, AirbyteCloudError, 404, id="not_found_enabled_reraises"
),
Expand Down Expand Up @@ -1231,7 +1277,7 @@ def test_search_error_handling(
with flag_patch as flag_mock, pytest.raises(expected_exc) as exc_info:
getattr(source, method_name)(*args)

if status_code in {403, 404}:
if status_code in {403, 404, 409}:
flag_mock.assert_called_once_with(ConnectorFeature.SEARCH_INDEXING)
else:
flag_mock.assert_not_called()
Expand All @@ -1240,12 +1286,61 @@ def test_search_error_handling(
assert exc_info.value.connector_id == "source-1"
assert exc_info.value.connector_name == "GitHub Issues"
assert "search_indexing" in (exc_info.value.guidance or "")
assert exc_info.value.enablement_url == source.connector_url
else:
assert not isinstance(exc_info.value, AirbyteExternalAccessNotEnabledError)
assert isinstance(exc_info.value, AirbyteCloudApiError)
assert exc_info.value.status_code == expected_status


@pytest.mark.parametrize("method_name", ["execute_search_query", "get_search_status"])
@pytest.mark.parametrize(
("message", "expected_exc"),
[
pytest.param(
"The request conflicts with the current resource state.",
AirbyteExternalAccessNotEnabledError,
id="platform_default_detail_raises_not_enabled",
),
pytest.param(
"Airbyte Cloud returned HTTP 409.",
AirbyteExternalAccessNotEnabledError,
id="status_fallback_raises_not_enabled",
),
pytest.param(
"Search is turned off for GitHub Issues.",
AirbyteCloudApiError,
id="cloud_detail_relayed",
),
],
)
def test_search_conflict_relays_cloud_detail(
monkeypatch: pytest.MonkeyPatch,
method_name: str,
message: str,
expected_exc: type[Exception],
) -> None:
"""A search 409 with its own Cloud detail keeps that detail instead of local guidance."""
workspace = _make_workspace(monkeypatch)
_patch_context_layer(monkeypatch)
cloud_error = AirbyteCloudApiError(status_code=409, message=message)
_patch_search(monkeypatch, {}, error=cloud_error)
source = _seed_source(workspace, "source-1", "GitHub Issues")
args = ("refunds",) if method_name == "execute_search_query" else ()

with (
patch.object(CloudConnector, "is_feature_enabled", return_value=False),
pytest.raises(expected_exc) as exc_info,
):
getattr(source, method_name)(*args)

if expected_exc is AirbyteCloudApiError:
assert exc_info.value is cloud_error
assert exc_info.value.message == message
else:
assert exc_info.value.__cause__ is cloud_error


def test_search_methods_raise_without_context_layer(
monkeypatch: pytest.MonkeyPatch,
) -> None:
Expand Down
7 changes: 0 additions & 7 deletions tests/unit_tests/test_mcp_arg_trace_parity.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
from opentelemetry.instrumentation.requests import RequestsInstrumentor
from opentelemetry.sdk.trace import ReadableSpan

from airbyte.constants import MCP_INSIDERS_ENV_VAR
from airbyte.mcp import _otel as observability
from airbyte.mcp import cloud
from airbyte.mcp.server import app as server_app
Expand Down Expand Up @@ -60,12 +59,6 @@
}


@pytest.fixture(autouse=True)
def _insiders_mode(monkeypatch: pytest.MonkeyPatch) -> None:
"""Advertise the insiders-only search tools so their args are covered."""
monkeypatch.setenv(MCP_INSIDERS_ENV_VAR, "1")


def _plan() -> dict[str, dict[str, str]]:
plan = asyncio.run(trace_plan(server_app))
return {name: entry["args"] for name, entry in plan.items()}
Expand Down
Loading
Loading