DocsGPT/docsgpt/connectors/resolve.py

"""Which connection a shared tool or source uses at runtime.

A resource that points at a connection runs either with its owner's account
(``owner`` mode, the default for new shares) or with the invoking member's
own account for the same service (``member`` mode). Resolution never returns
credentials; callers read them from the resolved row through
``docsgpt.connectors.service``.
"""

from __future__ import annotations

from datetime import datetime, timezone
import logging
from dataclasses import dataclass
from typing import Optional

from docsgpt.connectors import catalog, service
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
from docsgpt.storage.db.session import db_readonly

logger = logging.getLogger(__name__)

MODE_OWNER = "owner"
MODE_MEMBER = "member"

# Why a connection-backed tool can't run (``connection_stop_reason``).
CONNECTION_NEEDS_RECONNECT = "connection_needs_reconnect"
CONNECTION_REMOVED = "connection_removed"
CONNECTOR_DISABLED = "connector_disabled"


@dataclass(frozen=True)
class ResolvedConnection:
    """The connection a call runs with, or why there is none.

    Attributes:
        row: The ``connector_sessions`` row, None when missing.
        available: Whether it can be used right now.
        connector_key: Catalog key of the service.
        connector_name: Name shown to the user ("Notion", or a custom label).
        delegated: The row belongs to someone other than the invoker.
        writes_allowed: Whether an admin lets agents make changes through a
            connector that offers them as an opt-in (GitHub); True elsewhere.
        enabled: Whether the connector is switched on (an admin can turn it off).
    """

    row: Optional[dict]
    available: bool
    connector_key: Optional[str]
    connector_name: Optional[str]
    delegated: bool = False
    writes_allowed: bool = True
    enabled: bool = True

    @property
    def connection_id(self) -> Optional[str]:
        return str(self.row["id"]) if self.row else None


def _name_for(row: Optional[dict], fallback_key: Optional[str]) -> Optional[str]:
    if row is not None:
        return service.serialize_connection(row)["name"]
    definition = catalog.get_definition(fallback_key)
    return definition.name if definition else None


def resolve_connection(
    resource: dict,
    invoker_user_id: Optional[str],
    *,
    conn=None,
    policies: Optional[dict] = None,
) -> Optional[ResolvedConnection]:
    """Pick the connection a tool or source uses for ``invoker_user_id``.

    ``owner`` mode uses ``resource.connection_id``. ``member`` mode uses the
    invoker's own connection to the same service (and, for MCP, the same
    server), falling back to the owner's when the invoker is the owner.

    Args:
        resource: A ``user_tools`` or ``sources`` row.
        invoker_user_id: Who is running it.
        conn: An open connection to reuse; a read-only one is opened when None.
        policies: Connector policies already loaded with ``service.load_policies``,
            so a caller resolving many resources loads them once.

    Returns:
        None when the resource has no connection at all; otherwise the
        resolution, possibly with ``available=False``.
    """
    if not resource.get("connection_id"):
        return None
    if conn is None:
        with db_readonly() as own_conn:
            return _resolve(own_conn, resource, invoker_user_id, policies)
    return _resolve(conn, resource, invoker_user_id, policies)


def _resolve(conn, resource: dict, invoker_user_id: Optional[str], policies: Optional[dict]) -> ResolvedConnection:
    connection_id = resource.get("connection_id")
    mode = resource.get("credential_mode") or MODE_OWNER
    owner = resource.get("user_id")
    repo = ConnectorSessionsRepository(conn)
    owned = repo.get(str(connection_id))
    owned_key = catalog.connector_key_for_row(owned) if owned else None
    if policies is None:
        policies = service.load_policies(conn)
    policy = (policies.get(owned_key) or {}) if owned_key else {}
    if policy.get("credential_mode") in (MODE_OWNER, MODE_MEMBER):
        # An admin forces whose account every share of this connector uses.
        mode = policy["credential_mode"]
    if owned is not None and owner and owned.get("user_id") != owner:
        # A resource may only point at its own owner's connection.
        logger.warning(
            "resource %s points at a connection it does not own", resource.get("id"),
        )
        owned = None
    row = owned
    if mode == MODE_MEMBER and invoker_user_id and invoker_user_id != owner:
        row = _member_connection(repo, owned, invoker_user_id)
    key = catalog.connector_key_for_row(row or owned or {})
    enabled = service.connector_is_enabled(policies, key)
    available = (
        row is not None
        and service.normalize_status(row) == service.STATUS_CONNECTED
        and enabled
    )
    return ResolvedConnection(
        row=row,
        available=available,
        connector_key=key,
        connector_name=_name_for(row or owned, key),
        delegated=bool(row and invoker_user_id and row.get("user_id") != invoker_user_id),
        writes_allowed=_writes_allowed(policies, key),
        enabled=enabled,
    )


def connection_stop_reason(tool: dict, resolved: Optional[ResolvedConnection]) -> Optional[str]:
    """Why a tool's connection keeps it from running, or None when it can run.

    Args:
        tool: The ``user_tools`` row.
        resolved: What :func:`resolve_connection` returned for it.

    Returns:
        :data:`CONNECTION_REMOVED` when the connection is gone (its row was
        deleted, or a built-in service's tool lost its connection and has no
        credentials of its own), :data:`CONNECTOR_DISABLED` when an admin
        turned the service off, :data:`CONNECTION_NEEDS_RECONNECT` when the
        account must sign in again; else None.
    """
    if resolved is None:
        if tool.get("connection_id") or not catalog.definition_for_tool(tool.get("name") or ""):
            return None
        # Removing a connection but keeping its tools nulls their link; a
        # built-in service's tool has no secrets of its own to fall back to.
        config = tool.get("config") or {}
        return None if config.get("encrypted_credentials") else CONNECTION_REMOVED
    if resolved.row is None:
        return CONNECTION_REMOVED
    if not resolved.enabled:
        return CONNECTOR_DISABLED
    if not resolved.available:
        return CONNECTION_NEEDS_RECONNECT
    return None


def _writes_allowed(policies: dict, key: Optional[str]) -> bool:
    definition = catalog.get_definition(key) if key else None
    if definition is None or not definition.mcp_write_url:
        return True
    return service.writes_allowed(policies, key)


def _member_connection(repo: ConnectorSessionsRepository, owned: Optional[dict], invoker: str) -> Optional[dict]:
    """The invoker's own connection to the service the owner's connection is for.

    A member with several connected accounts of that service gets the one
    most recently used or connected, whichever is later: the account they
    are working in, or the one a "Connect to continue" just added (never
    used yet). Ties go to the last used, then the last updated.
    """
    if owned is None:
        return None
    candidates = [
        row for row in repo.list_for_user(invoker)
        if row.get("provider") == owned.get("provider")
        and (row.get("server_url") or "") == (owned.get("server_url") or "")
        and service.normalize_status(row) == service.STATUS_CONNECTED
    ]
    if not candidates:
        return None

    def when(value) -> datetime:
        if isinstance(value, str):
            value = datetime.fromisoformat(value)
        if not isinstance(value, datetime):
            return datetime.min.replace(tzinfo=timezone.utc)
        return value if value.tzinfo else value.replace(tzinfo=timezone.utc)

    def recency(row: dict) -> tuple:
        used, updated, created = (when(row.get(field)) for field in ("last_used_at", "updated_at", "created_at"))
        return max(used, created), used, updated, created

    return max(candidates, key=recency)


def audit_delegation(resolved: ResolvedConnection, *, invoker: Optional[str], resource_type: str,
                     resource_id: Optional[str], agent_id: Optional[str] = None) -> None:
    """Log a call that runs with someone else's account (``owner`` mode)."""
    if not resolved.delegated or resolved.row is None:
        return
    logger.info(
        "tool_credential_delegation",
        extra={
            "invoker": invoker,
            "tool_owner": resolved.row.get("user_id"),
            "connection_id": resolved.connection_id,
            "resource_type": resource_type,
            "resource_id": resource_id,
            "agent_id": agent_id,
        },
    )