kamidso-back/kamidso/native_connectors/executor.py

"""Call a predefined native connector function. Credentials stay inside this module."""

from __future__ import annotations

import json
import logging
import time
from typing import Any, Callable, Dict, Optional
from urllib.parse import urlparse

import httpx

from kamidso.native_connectors.errors import (
    AuthenticationFailed,
    CallTimeout,
    ConnectorNotFound,
    FunctionDisabled,
    FunctionNotFound,
    InvalidJSONResponse,
    NativeConnectorError,
    OperationNotAllowed,
    ProviderNotFound,
    RemoteCallFailed,
    ResponseTooLarge,
    ResultMaterializationError,
    TransportFailed,
)
from kamidso.native_connectors.models import ConnectorFunction, NativeProvider
from kamidso.native_connectors.redact import redact, redact_text, strip_sensitive_fields
from kamidso.native_connectors.registry import NativeConnectorRegistry, get_registry
from kamidso.native_connectors.result_store import (
    materialize_connector_result,
    render_success_observation,
)
from kamidso.native_connectors.validator import validate_body, validation_payload

logger = logging.getLogger(__name__)

_MAX_RESPONSE_BYTES = 5 * 1024 * 1024
_ATTEMPT_TIMEOUT = 30.0
_TOTAL_BUDGET = 60.0
_RETRY_STATUSES = {500, 502, 503}
RowLoader = Callable[[str], Optional[Dict[str, Any]]]


def _status_error_type(status_code: int) -> str:
    if status_code == 401:
        return "authentication_error"
    if status_code == 403:
        return "permission_error"
    if status_code == 429:
        return "rate_limit_error"
    if status_code == 400:
        return "remote_validation_error"
    if status_code >= 500:
        return "remote_server_error"
    return "http_error"


def suggest_functions(provider: NativeProvider, name: str) -> list:
    import difflib

    enabled = [item.name for item in provider.enabled_reads()]
    suggestions = list(provider.aliases.get(name) or [])
    if "lead" in name.lower():
        for candidate in ("search_persons", "search_deals"):
            if candidate in enabled and candidate not in suggestions:
                suggestions.append(candidate)
    close = difflib.get_close_matches(name, enabled, n=3, cutoff=0.5)
    for candidate in close:
        if candidate not in suggestions:
            suggestions.append(candidate)
    return [item for item in suggestions if item in enabled][:3]


def default_row_loader(connector_id: str) -> Optional[Dict[str, Any]]:
    from kamidso.engine.serve.connector.config import ServeConfig
    from kamidso.engine.serve.connector.models.models import (
        ConnectorInstanceDao,
        ConnectorInstanceEntity,
    )

    dao = ConnectorInstanceDao(ServeConfig())
    with dao.session() as session:
        entity = (
            session.query(ConnectorInstanceEntity)
            .filter(ConnectorInstanceEntity.connector_id == connector_id)
            .first()
        )
        if entity is None:
            return None
        config: Dict[str, Any] = {}
        raw_config = entity.config_json
        if isinstance(raw_config, str) and raw_config.strip():
            try:
                parsed = json.loads(raw_config)
            except json.JSONDecodeError:
                parsed = None
            if isinstance(parsed, dict):
                config = parsed
        return {
            "connector_id": entity.connector_id,
            "connector_type": entity.connector_type,
            "display_name": entity.display_name or "",
            "status": entity.status or "",
            "encrypted_credentials": entity.encrypted_credentials,
            "encryption_salt": entity.encryption_salt,
            "user_name": entity.user_name,
            "sys_code": entity.sys_code,
            "config": config,
        }


class NativeConnectorExecutor:
    def __init__(
        self,
        registry: Optional[NativeConnectorRegistry] = None,
        *,
        client: Optional[httpx.AsyncClient] = None,
        credential_store: Any = None,
        row_loader: Optional[RowLoader] = None,
    ) -> None:
        self._registry = registry
        self._client = client
        self._credential_store = credential_store
        self._row_loader = row_loader or default_row_loader

    def load_row(self, connector_id: str) -> Optional[Dict[str, Any]]:
        return self._row_loader(connector_id)

    @property
    def registry(self) -> NativeConnectorRegistry:
        if self._registry is None:
            self._registry = get_registry()
        return self._registry

    async def probe(self, provider_id: str, credentials: Dict[str, str]) -> Dict[str, Any]:
        try:
            provider = self.registry.get(provider_id)
        except ProviderNotFound as exc:
            return {"success": False, "message": exc.message}
        try:
            result = await self._call(
                provider=provider,
                function_name=provider.health_function,
                body=dict(provider.health_body),
                credentials=credentials,
                react_state=None,
                materialize=False,
                display_name=provider.display_name,
            )
        except NativeConnectorError as exc:
            payload = exc.to_payload()
            status_code = payload.get("status_code")
            return {
                "success": False,
                "message": redact_text(exc.message),
                "status_code": status_code if isinstance(status_code, int) else None,
                "error_type": exc.error_type,
            }
        if result.get("ok"):
            return {"success": True, "message": "Connection successful"}
        return {
            "success": False,
            "message": redact_text(str(result.get("message") or "Connection test failed")),
        }

    async def probe_stored(self, connector_id: str) -> Dict[str, Any]:
        row = self._row_loader(connector_id)
        if row is None:
            return {"success": False, "message": f"Connector '{connector_id}' not found"}
        try:
            credentials = self._decrypt(row)
        except NativeConnectorError as exc:
            return {"success": False, "message": exc.message}
        return await self.probe(str(row.get("connector_type") or ""), credentials)

    async def execute(
        self,
        *,
        connector_id: str,
        function_name: str,
        body: Any,
        react_state: Optional[Dict[str, Any]],
        display_name: Optional[str] = None,
    ) -> Dict[str, Any]:
        row = self._row_loader(connector_id)
        if row is None or str(row.get("status") or "") != "active":
            raise ConnectorNotFound(f"Connector '{connector_id}' is not available.")
        provider = self.registry.get(str(row.get("connector_type") or ""))
        label = display_name or str(row.get("display_name") or provider.display_name)
        return await self._call(
            provider=provider,
            function_name=function_name,
            body=body,
            credentials=None,
            react_state=react_state,
            materialize=True,
            display_name=label,
            audit_user=str((react_state or {}).get("user_name") or row.get("user_name") or ""),
            connector_id=connector_id,
            encrypted_row=row,
        )

    def _decrypt(self, row: Dict[str, Any]) -> Dict[str, str]:
        if self._credential_store is None:
            raise AuthenticationFailed("Credential store is not configured.")
        try:
            decrypted = self._credential_store.decrypt(
                row.get("encrypted_credentials"),
                row.get("encryption_salt"),
            )
        except Exception as exc:
            raise AuthenticationFailed("Stored credentials could not be decrypted.") from exc
        if not isinstance(decrypted, dict):
            raise AuthenticationFailed("Stored credentials could not be decrypted.")
        return {str(key): str(value) for key, value in decrypted.items()}

    async def _call(
        self,
        *,
        provider: NativeProvider,
        function_name: str,
        body: Any,
        credentials: Optional[Dict[str, str]],
        react_state: Optional[Dict[str, Any]],
        materialize: bool,
        display_name: str,
        audit_user: str = "",
        connector_id: str = "",
        encrypted_row: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        function = provider.function(function_name)
        if function is None:
            raise FunctionNotFound(
                function_name,
                provider.display_name,
                suggest_functions(provider, function_name),
            )
        if function.effect != "read":
            raise OperationNotAllowed(function.name, function.effect)
        if not function.enabled:
            raise FunctionDisabled(
                f"Function '{function.name}' is disabled.",
                function=function.name,
            )
        normalized, issues = validate_body(function, body)
        if issues:
            return validation_payload(
                connector_name=display_name,
                function=function,
                issues=issues,
            )
        assert normalized is not None
        if credentials is None:
            if encrypted_row is None:
                raise AuthenticationFailed(
                    "The connector credential is missing.",
                    function=function.name,
                )
            credentials = self._decrypt(encrypted_row)
        secret = str(credentials.get(provider.credential_field) or "").strip()
        if not secret:
            raise AuthenticationFailed(
                "The connector credential is missing.",
                function=function.name,
            )
        started = time.perf_counter()
        response, request_size = await self._send(
            provider, function, strip_sensitive_fields(normalized), secret
        )
        duration_ms = int((time.perf_counter() - started) * 1000)
        response_size = len(response.content or b"")
        self._audit(
            react_state=react_state,
            user_name=audit_user,
            connector_id=connector_id,
            connector_type=provider.id,
            function=function,
            status_code=response.status_code,
            duration_ms=duration_ms,
            request_size=request_size,
            response_size=response_size,
        )
        if response_size > _MAX_RESPONSE_BYTES:
            raise ResponseTooLarge(
                "Use narrower filters or pagination.",
                function=function.name,
                details={"status_code": response.status_code},
            )
        if response.status_code in {301, 302, 303, 307, 308}:
            raise RemoteCallFailed(
                "The connector refused a redirect.",
                error_type="http_error",
                function=function.name,
                details={
                    "provider": provider.id,
                    "status_code": response.status_code,
                    "body": None,
                },
            )
        payload = self._decode_json(response, function.name)
        if response.status_code >= 400:
            details: Dict[str, Any] = {
                "provider": provider.id,
                "status_code": response.status_code,
                "body": redact(payload),
            }
            retry_after = response.headers.get("Retry-After")
            if response.status_code == 429 and retry_after:
                details["retry_after"] = retry_after
            raise RemoteCallFailed(
                f"The connector returned HTTP {response.status_code}.",
                error_type=_status_error_type(response.status_code),
                function=function.name,
                details=details,
            )
        if not materialize or react_state is None:
            return {
                "ok": True,
                "http_sent": True,
                "status_code": response.status_code,
                "body": payload,
            }
        try:
            meta = materialize_connector_result(react_state, payload)
        except OSError as exc:
            raise ResultMaterializationError(
                "The connector response could not be saved.",
                function=function.name,
            ) from exc
        observation = render_success_observation(
            provider=provider,
            function=function,
            body=normalized,
            status_code=response.status_code,
            payload=payload,
            relative_path=meta["relative_path"],
        )
        return {
            "ok": True,
            "http_sent": True,
            "status_code": response.status_code,
            "relative_path": meta["relative_path"],
            "observation": observation,
        }

    def _decode_json(self, response: httpx.Response, function_name: str) -> Any:
        if not response.content:
            return {}
        try:
            return response.json()
        except (json.JSONDecodeError, ValueError) as exc:
            raise InvalidJSONResponse(
                "The connector returned a non-JSON response.",
                function=function_name,
                details={"status_code": response.status_code},
            ) from exc

    async def _send(
        self,
        provider: NativeProvider,
        function: ConnectorFunction,
        body: Dict[str, Any],
        secret: str,
    ):
        url = _safe_url(provider, function.path)
        params = {provider.auth_parameter: secret}
        headers = dict(provider.default_headers)
        request_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
        deadline = time.monotonic() + _TOTAL_BUDGET
        attempt = 0
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise CallTimeout(
                    "The connector request exceeded the time budget.",
                    function=function.name,
                )
            timeout = httpx.Timeout(min(_ATTEMPT_TIMEOUT, remaining))
            try:
                response = await self._request(
                    function.method,
                    url,
                    params=params,
                    headers=headers,
                    body=body,
                    timeout=timeout,
                )
            except httpx.TimeoutException as exc:
                raise CallTimeout(
                    "The connector request timed out.",
                    function=function.name,
                ) from exc
            except httpx.TransportError as exc:
                if attempt == 0 and _retryable_transport(exc) and deadline - time.monotonic() > 0:
                    attempt += 1
                    continue
                raise TransportFailed(
                    redact_text(str(exc)) or "The connector request failed.",
                    function=function.name,
                ) from exc
            if response.status_code in _RETRY_STATUSES and attempt == 0 and deadline - time.monotonic() > 0:
                attempt += 1
                continue
            return response, request_size

    async def _request(
        self,
        method: str,
        url: str,
        *,
        params: Dict[str, str],
        headers: Dict[str, str],
        body: Dict[str, Any],
        timeout: httpx.Timeout,
    ) -> httpx.Response:
        kwargs: Dict[str, Any] = {
            "params": params,
            "headers": headers,
            "timeout": timeout,
            "follow_redirects": False,
        }
        if method.upper() != "GET":
            kwargs["json"] = body

        async def _send(client: httpx.AsyncClient) -> httpx.Response:
            return await client.request(method, url, **kwargs)

        if self._client is not None:
            return await _send(self._client)
        async with httpx.AsyncClient(follow_redirects=False) as client:
            return await _send(client)

    def _audit(
        self,
        *,
        react_state: Optional[Dict[str, Any]],
        user_name: str,
        connector_id: str,
        connector_type: str,
        function: ConnectorFunction,
        status_code: int,
        duration_ms: int,
        request_size: int,
        response_size: int,
    ) -> None:
        logger.info(
            "native_connector conversation_id=%s user_name=%s connector_id=%s "
            "connector_type=%s function=%s effect=%s http_status=%s duration_ms=%s "
            "request_size=%s response_size=%s cache_hit=%s",
            (react_state or {}).get("conv_id") or "",
            user_name,
            connector_id,
            connector_type,
            function.name,
            function.effect,
            status_code,
            duration_ms,
            request_size,
            response_size,
            False,
        )


def _retryable_transport(exc: httpx.TransportError) -> bool:
    if isinstance(exc, httpx.TimeoutException):
        return False
    return isinstance(exc, (httpx.NetworkError, httpx.RemoteProtocolError))


def _safe_url(provider: NativeProvider, path: str) -> str:
    base = urlparse(provider.base_url)
    if base.scheme != "https" or not base.hostname:
        raise TransportFailed("The connector host is not allowed.")
    if not path.startswith("/") or "://" in path or ".." in path.split("/"):
        raise TransportFailed("The connector path is not allowed.")
    url = provider.base_url.rstrip("/") + path
    parsed = urlparse(url)
    if parsed.scheme != "https" or parsed.hostname != base.hostname:
        raise TransportFailed("The connector host is not allowed.")
    return url