Connected {{ name }}. + You can close this tab and return to druks.
+ +{% endblock %} diff --git a/backend/druks/mcp/constants.py b/backend/druks/mcp/constants.py index 974ca98f..54ac9756 100644 --- a/backend/druks/mcp/constants.py +++ b/backend/druks/mcp/constants.py @@ -17,16 +17,6 @@ REGISTRY_SEARCH_CACHE_PREFIX = "mcp:registry:search:" REGISTRY_CACHE_TTL_SECONDS = 300 -# OAuth connect + mint plumbing rides the shared engine (druks.services' -# OauthClient) under this provider namespace. The namespace keys the engine's -# connect-state, token-cache, and refresh-lock Redis entries, so it is pinned: -# a rolling deploy's old and new processes must elect one refresher per grant. -# The prefixes spell out the derived keys; the token and lock keys append -# {name}:{account_id}. The callback path is public API surface — the -# authorization server redirects the operator's browser to -# {urls.endpoint}{OAUTH_CALLBACK_PATH} after consent. -OAUTH_PROVIDER = "mcp:oauth" +# Every dynamically-registered client pins this path as a redirect_uri, so +# renaming it orphans existing registrations. OAUTH_CALLBACK_PATH = "/api/mcp-servers/oauth/callback" -OAUTH_CONNECT_STATE_PREFIX = f"{OAUTH_PROVIDER}:connect:" -OAUTH_ACCESS_TOKEN_PREFIX = f"{OAUTH_PROVIDER}:access_token:" -OAUTH_REFRESH_LOCK_PREFIX = f"{OAUTH_PROVIDER}:refresh_lock:" diff --git a/backend/druks/mcp/helpers.py b/backend/druks/mcp/helpers.py index d2fb291d..2d0d637d 100644 --- a/backend/druks/mcp/helpers.py +++ b/backend/druks/mcp/helpers.py @@ -1,5 +1,26 @@ +from druks.accounts.constants import SYSTEM_ACCOUNT_ID from druks.mcp.constants import TOKEN_ENV_PREFIX, TOKEN_ENV_SUFFIX +from druks.mcp.enums import IdentityMode +from druks.mcp.exceptions import UnresolvedGrantAccountError def get_bearer_token_env_var(name: str) -> str: return f"{TOKEN_ENV_PREFIX}{name.upper()}{TOKEN_ENV_SUFFIX}" + + +def grant_provider(name: str) -> str: + # Namespaced so a server name can never collide with a Service name in + # the shared grant table and Redis keys. + return f"mcp:{name}" + + +def get_grant_account(identity_mode: str | None, run_account_id: str | None) -> str: + # Whose grant serves this caller: a shared server's grant lives under + # the system account whoever asks; a per-user server's under the asker. + if identity_mode == IdentityMode.PER_USER and run_account_id: + return run_account_id + if identity_mode == IdentityMode.PER_USER: + raise UnresolvedGrantAccountError(identity_mode, run_account_id) + if identity_mode == IdentityMode.SHARED: + return SYSTEM_ACCOUNT_ID + raise UnresolvedGrantAccountError(identity_mode, run_account_id) diff --git a/backend/druks/mcp/models.py b/backend/druks/mcp/models.py index e6b5f46e..e0d26423 100644 --- a/backend/druks/mcp/models.py +++ b/backend/druks/mcp/models.py @@ -2,24 +2,22 @@ from datetime import datetime from typing import Any -from sqlalchemy import Boolean, ForeignKey, String, UniqueConstraint, select, update +from sqlalchemy import Boolean, ForeignKey, String, UniqueConstraint, select from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.orm import Mapped, mapped_column from druks.accounts.constants import SYSTEM_ACCOUNT_ID from druks.core.models import Uuid7Pk -from druks.database import db_session, get_session +from druks.database import db_session from druks.extensions.registry import mcp_servers from druks.mcp.constants import NAME_PATTERN -from druks.mcp.enums import IdentityMode, TokenSource -from druks.mcp.exceptions import ( - InvalidServerNameError, - MissingGrantError, - UnresolvedGrantAccountError, -) +from druks.mcp.enums import TokenSource +from druks.mcp.exceptions import InvalidServerNameError +from druks.mcp.helpers import get_grant_account, grant_provider from druks.models import Base from druks.secrets.fields import EncryptedJsonField, EncryptedTextField, Secret +from druks.services.models import OauthConnection class McpServer(Base, Uuid7Pk): @@ -109,11 +107,11 @@ def get_resolved(cls, account_id: str | None) -> dict[str, dict]: elif source == TokenSource.OAUTH: server["has_token"] = False if server["identity_mode"]: - grant_account = McpOauthGrant.get_grant_account( - server["identity_mode"], account_id - ) + grant_account = get_grant_account(server["identity_mode"], account_id) server["has_token"] = bool( - McpOauthGrant.get_for_account(server["name"], grant_account) + OauthConnection.list_for_account( + grant_provider(server["name"]), grant_account + ) ) else: server["has_token"] = bool(server["token"]) @@ -172,130 +170,63 @@ def delete(self) -> None: session.flush() -class McpOauthGrant(Base, Uuid7Pk): - __tablename__ = "mcp_oauth_grants" - __table_args__ = (UniqueConstraint("server_name", "account_id"),) +class McpClientRegistration(Base, Uuid7Pk): + __tablename__ = "mcp_client_registrations" + __table_args__ = (UniqueConstraint("server_id", "account_id"),) - # One grant per (server, account): the durable outcome of an OAuth - # connect flow — exactly what mint needs to refresh an access token. - # Connect-time material (authorization endpoint, PKCE verifier, state) is - # transient and lives in Redis, never here. The refresh token never leaves - # the backend; the API exposes only that a grant exists. - server_name: Mapped[str] = mapped_column(String) + # One RFC 7591 registration per grant: druks registers a fresh client on + # every connect, so each account's grant refreshes as the client it + # consented through. The refresh token lives on the platform's OauthConnection. + server_id: Mapped[str] = mapped_column(ForeignKey("mcp_servers.id", ondelete="CASCADE")) account_id: Mapped[str] = mapped_column( ForeignKey("accounts.id", ondelete="RESTRICT"), default=SYSTEM_ACCOUNT_ID ) - # Ciphertext at rest; decrypted only into the refresh request body. - refresh_token = EncryptedTextField() token_endpoint: Mapped[str] = mapped_column(String) - # The MCP server url the grant is bound to (RFC 8707): an audience-binding - # authorization server rejects a refresh that doesn't carry the same - # ``resource`` the code exchange did. - resource: Mapped[str] = mapped_column(String) client_id: Mapped[str] = mapped_column(String) # "" for public clients (PKCE-only); some authorization servers issue one # even for token_endpoint_auth_method "none" and then expect it on refresh. client_secret = EncryptedTextField(default="") - # When the operator last completed consent. Stamped on every store — the - # row is upserted on re-connect, so row-creation time would lie. - connected_at: Mapped[datetime] = mapped_column(default=Base.utc_now) - - @staticmethod - def get_grant_account(identity_mode: str | None, run_account_id: str | None) -> str: - # Whose grant serves this caller: a shared server's grant lives under - # the system account whoever asks; a per-user server's under the asker. - if identity_mode == IdentityMode.PER_USER and run_account_id: - return run_account_id - if identity_mode == IdentityMode.PER_USER: - raise UnresolvedGrantAccountError(identity_mode, run_account_id) - if identity_mode == IdentityMode.SHARED: - return SYSTEM_ACCOUNT_ID - raise UnresolvedGrantAccountError(identity_mode, run_account_id) @classmethod - def get_for_account(cls, server_name: str, account_id: str) -> "McpOauthGrant | None": + def get_for_account(cls, server_name: str, account_id: str) -> "McpClientRegistration | None": return ( db_session() .execute( - select(cls).where(cls.server_name == server_name, cls.account_id == account_id) + select(cls) + .join(McpServer, McpServer.id == cls.server_id) + .where(McpServer.name == server_name, cls.account_id == account_id) ) .scalar_one_or_none() ) - @classmethod - def list_for_server(cls, server_name: str) -> list["McpOauthGrant"]: - return list(db_session().scalars(select(cls).where(cls.server_name == server_name))) - @classmethod def store( cls, *, - server_name: str, + server_id: str, account_id: str, - refresh_token: str, token_endpoint: str, - resource: str, client_id: str, client_secret: str = "", - ) -> "McpOauthGrant": + ) -> "McpClientRegistration": session = db_session() statement = pg_insert(cls).values( - server_name=server_name, + server_id=server_id, account_id=account_id, - refresh_token=refresh_token, token_endpoint=token_endpoint, - resource=resource, client_id=client_id, client_secret=client_secret, - connected_at=cls.utc_now(), ) statement = statement.on_conflict_do_update( - index_elements=["server_name", "account_id"], + index_elements=["server_id", "account_id"], set_={ - "refresh_token": statement.excluded.refresh_token, "token_endpoint": statement.excluded.token_endpoint, - "resource": statement.excluded.resource, "client_id": statement.excluded.client_id, "client_secret": statement.excluded.client_secret, - "connected_at": statement.excluded.connected_at, }, ).returning(cls) return session.scalars(statement, execution_options={"populate_existing": True}).one() - def load_refresh_token(self) -> str: - # Under the refresh lock: another process may have rotated and - # committed, and this transaction may already hold the row — - # populate_existing re-reads it past the identity map. - fresh = ( - db_session() - .scalars( - select(McpOauthGrant) - .where(McpOauthGrant.id == self.id) - .execution_options(populate_existing=True) - ) - .one_or_none() - ) - if not fresh: - raise MissingGrantError(self.server_name, self.account_id) - # The grant's secret halves are ciphertext at rest; the plaintext - # exists only in the refresh request body. - return fresh.refresh_token.decrypt() - - def save_refresh_token(self, rotated: str) -> None: - # The provider invalidated the old token the moment it rotated, so - # the write commits on its own session, never the enclosing - # transaction — a step that rolls back later must not brick the grant. - with get_session(db_session().get_bind()) as session: - session.execute( - update(McpOauthGrant) - .where(McpOauthGrant.id == self.id) - .values(refresh_token=rotated) - ) - session.commit() - # Keep the enclosing transaction's copy true as well. - self.refresh_token = rotated - db_session().flush() - def delete(self) -> None: session = db_session() session.delete(self) diff --git a/backend/druks/mcp/oauth.py b/backend/druks/mcp/oauth.py index 23eceeff..3bf92240 100644 --- a/backend/druks/mcp/oauth.py +++ b/backend/druks/mcp/oauth.py @@ -5,12 +5,15 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert from druks.database import db_session -from druks.mcp.constants import OAUTH_CALLBACK_PATH, OAUTH_PROVIDER +from druks.mcp.constants import OAUTH_CALLBACK_PATH from druks.mcp.enums import IdentityMode from druks.mcp.exceptions import GrantRefreshError, MissingGrantError, OauthConnectError -from druks.mcp.models import McpOauthGrant, McpServer +from druks.mcp.helpers import get_grant_account, grant_provider +from druks.mcp.models import McpClientRegistration, McpServer from druks.services import OauthClient, OauthExchangeError, OauthRefreshError from druks.services.constants import OAUTH_MINT_WAIT_ATTEMPTS, OAUTH_MINT_WAIT_INTERVAL_SECONDS +from druks.services.models import OauthConnection +from druks.services.oauth import complete_connect as complete_oauth_exchange def _http() -> httpx.AsyncClient: @@ -18,6 +21,16 @@ def _http() -> httpx.AsyncClient: return httpx.AsyncClient(timeout=30.0, follow_redirects=True) +def get_connection(name: str, account_id: str) -> OauthConnection | None: + # One connection per (server, account) — MCP's policy over the shared table. + rows = OauthConnection.list_for_account(grant_provider(name), account_id) + return rows[0] if rows else None + + +def list_connections(name: str) -> list[OauthConnection]: + return OauthConnection.list_for_provider(grant_provider(name)) + + def _origin(url: str) -> str: parts = urlparse(url) return f"{parts.scheme}://{parts.netloc}" @@ -166,7 +179,7 @@ async def begin_connect( raise OauthConnectError(name, "the authorization server does not support PKCE S256") registration = await _register_client(client, name, metadata, redirect_uri) return await OauthClient( - provider=OAUTH_PROVIDER, + provider=grant_provider(name), authorization_endpoint=metadata["authorization_endpoint"], token_endpoint=metadata["token_endpoint"], client_id=registration["client_id"], @@ -186,16 +199,11 @@ async def begin_connect( async def complete_connect(*, state: str, code: str) -> str: - """The callback half: the shared exchange (single-use state, code + - verifier), then the durable outcome — claim the server's identity mode and - store the grant. Returns the server name. The grant is the only outcome — - nothing is cached here, because it becomes real only when this request's - transaction commits, and a cache filled ahead of that would outlive its - failure. The first delivery mints from the committed grant.""" + """The callback half: the shared exchange, then the durable outcome — + claim the server's identity mode and store the registration and the + grant. Returns the server name.""" try: - tokens, pending = await OauthClient( - provider=OAUTH_PROVIDER, http_factory=_http - ).complete_connect(state=state, code=code) + tokens, pending = await complete_oauth_exchange(state=state, code=code) except OauthExchangeError as error: raise OauthConnectError(error.context.get("name", "unknown"), error.reason) from error name = pending["name"] @@ -217,46 +225,66 @@ async def complete_connect(*, state: str, code: str) -> str: .where(McpServer.name == name, McpServer.identity_mode.is_(None)) .values(identity_mode=pending["identity_mode"]) ) - effective_identity_mode = session.scalar( - select(McpServer.identity_mode).where(McpServer.name == name) - ) - account_id = McpOauthGrant.get_grant_account(effective_identity_mode, pending["account_id"]) - McpOauthGrant.store( - server_name=name, + server = session.scalars(select(McpServer).where(McpServer.name == name)).one() + account_id = get_grant_account(server.identity_mode, pending["account_id"]) + McpClientRegistration.store( + server_id=server.id, account_id=account_id, - refresh_token=tokens["refresh_token"], token_endpoint=pending["token_endpoint"], - resource=pending["server_url"], client_id=pending["client_id"], client_secret=pending["client_secret"], ) + connection = get_connection(name, account_id) + if connection: + connection.reconnect(refresh_token=tokens["refresh_token"], scopes=[]) + # A reconsent's stale cached token must not serve until its TTL runs out. + await evict_access_token(name, account_id) + else: + OauthConnection.create( + provider=grant_provider(name), + account_id=account_id, + refresh_token=tokens["refresh_token"], + scopes=[], + ) return name async def evict_access_token(name: str, account_id: str) -> None: - await OauthClient(provider=OAUTH_PROVIDER).evict_access_token(f"{name}:{account_id}") + connection = get_connection(name, account_id) + if connection: + await OauthClient(provider=grant_provider(name)).evict_access_token(connection.id) + + +async def disconnect(name: str, account_id: str) -> None: + connection = get_connection(name, account_id) + if connection: + await OauthClient(provider=grant_provider(name)).disconnect(connection) + registration = McpClientRegistration.get_for_account(name, account_id) + if registration: + registration.delete() async def mint_access_token(name: str, account_id: str) -> str: """The delivery-side token for a connected server, minted by the shared engine from this server's grant — delivery never ships a server the agent can't authenticate to.""" - grant = McpOauthGrant.get_for_account(name, account_id) - if not grant: + connection = get_connection(name, account_id) + registration = McpClientRegistration.get_for_account(name, account_id) + server = McpServer.get_for_name(name) + if not connection or not registration or not server: raise MissingGrantError(name, account_id) client = OauthClient( - provider=OAUTH_PROVIDER, - token_endpoint=grant.token_endpoint, - client_id=grant.client_id, - client_secret=grant.client_secret.decrypt(), + provider=grant_provider(name), + token_endpoint=registration.token_endpoint, + client_id=registration.client_id, + client_secret=registration.client_secret.decrypt(), # RFC 8707: an audience-binding server expects the refresh to carry # the same resource the code exchange was bound to. - extra_token_params={"resource": grant.resource}, + extra_token_params={"resource": server.url}, mint_wait_interval_seconds=OAUTH_MINT_WAIT_INTERVAL_SECONDS, mint_wait_attempts=OAUTH_MINT_WAIT_ATTEMPTS, - http_factory=_http, ) try: - return await client.mint_access_token(key=f"{name}:{account_id}", grant=grant) + return await client.mint_access_token(connection=connection) except OauthRefreshError as error: raise GrantRefreshError(name, error.reason) from error diff --git a/backend/druks/mcp/routes.py b/backend/druks/mcp/routes.py index 74260441..111905f1 100644 --- a/backend/druks/mcp/routes.py +++ b/backend/druks/mcp/routes.py @@ -14,7 +14,8 @@ OauthConnectError, RegistryUnavailableError, ) -from druks.mcp.models import McpOauthGrant, McpServer +from druks.mcp.helpers import get_grant_account +from druks.mcp.models import McpServer from druks.mcp.schemas import ( ConnectMcpServerResponse, CreateMcpServerRequest, @@ -162,11 +163,11 @@ async def remove_mcp_server(name: str) -> None: server = McpServer.get_for_name(name) if not server: raise HTTPException(status_code=404, detail=f"MCP server {name!r} not found") - grants = McpOauthGrant.list_for_server(name) + connections = oauth.list_connections(name) server.delete() - for grant in grants: - grant.delete() - await oauth.evict_access_token(name, grant.account_id) + for connection in connections: + await oauth.evict_access_token(name, connection.account_id) + connection.delete() @router.post("/{name}/connect", response_model=ConnectMcpServerResponse) @@ -178,7 +179,7 @@ async def connect_mcp_server( server = McpServer.get_resolved(current_account_id.get()).get(name) if not server or server["token_source"] != TokenSource.OAUTH: raise HTTPException(status_code=404, detail=f"MCP server {name!r} is not an OAuth server.") - if McpOauthGrant.list_for_server(name) and server["identity_mode"] != identity_mode: + if oauth.list_connections(name) and server["identity_mode"] != identity_mode: raise HTTPException( status_code=409, detail=f"MCP server {name!r} already uses {server['identity_mode']!r} identity.", @@ -236,16 +237,15 @@ async def disconnect_mcp_server(name: str) -> None: raise HTTPException(status_code=404, detail=f"MCP server {name!r} is not an OAuth server.") if not server["identity_mode"]: raise HTTPException(status_code=404, detail=f"MCP server {name!r} has no grant.") - account_id = McpOauthGrant.get_grant_account(server["identity_mode"], current_account_id.get()) - grant = McpOauthGrant.get_for_account(name, account_id) - if not grant: + account_id = get_grant_account(server["identity_mode"], current_account_id.get()) + connection = oauth.get_connection(name, account_id) + if not connection: raise HTTPException( status_code=404, detail=f"MCP server {name!r} has no grant for account {account_id!r}.", ) - await oauth.evict_access_token(name, grant.account_id) - grant.delete() - if not McpOauthGrant.list_for_server(name): + await oauth.disconnect(name, account_id) + if not oauth.list_connections(name): # The last grant leaving reopens the mode choice: the next connect is # a first connect again. server_row = McpServer.get_for_name(name) diff --git a/backend/druks/sandbox/datastructures.py b/backend/druks/sandbox/datastructures.py index 84c9017b..9427d7a9 100644 --- a/backend/druks/sandbox/datastructures.py +++ b/backend/druks/sandbox/datastructures.py @@ -12,7 +12,7 @@ from druks.mcp.constants import TOKEN_ENV_PREFIX from druks.mcp.enums import TokenSource from druks.mcp.exceptions import MissingTokenError, SourceEnvVarUnsetError -from druks.mcp.helpers import get_bearer_token_env_var +from druks.mcp.helpers import get_bearer_token_env_var, get_grant_account from druks.user_settings.models import UserSettings if TYPE_CHECKING: @@ -205,9 +205,7 @@ async def with_mcp_servers(self, account_id: str | None, **kwargs: Any) -> dict[ if not token: raise SourceEnvVarUnsetError(server["name"], server["source_env_var"]) else: # oauth - grant_account = mcp_models.McpOauthGrant.get_grant_account( - server["identity_mode"], run_account - ) + grant_account = get_grant_account(server["identity_mode"], run_account) token = await oauth.mint_access_token(server["name"], grant_account) bearer_token_env_var = "" if token: diff --git a/backend/druks/services/base.py b/backend/druks/services/base.py index 6b0b2377..f59443c8 100644 --- a/backend/druks/services/base.py +++ b/backend/druks/services/base.py @@ -3,14 +3,72 @@ from pydantic import BaseModel, ValidationError from druks.extensions.base import NAME_RE +from druks.extensions.loader import iter_extensions from druks.extensions.registry import services from druks.extensions.settings import field_kind, field_multiline from .exceptions import ServiceConnectError, ServiceNotConnectedError -from .models import ServiceIdentity +from .models import OauthConnection, ServiceIdentity from .oauth import OauthClient +class Connection: + """One signed-in provider account, reached through the extension's + declared handle. Mint and disconnect act on this sign-in only.""" + + def __init__(self, service: "type[Service]", row: OauthConnection) -> None: + self.service = service + self.row = row + + @property + def id(self) -> str: + return self.row.id + + @property + def scopes(self) -> list[str]: + return self.row.scopes + + @property + def connected_at(self): + return self.row.connected_at + + async def mint_access_token(self) -> str: + return await self.service.get_oauth_client().mint_access_token(connection=self.row) + + async def disconnect(self) -> None: + await OauthClient(provider=self.service.name).disconnect(self.row) + + +class ScopedService: + """A service seen through one extension's declared scopes + (``gmail = Gmail.with_scopes("gmail.readonly")``). The declaration + feeds the consent union; the handle reads the connections that grant + it.""" + + def __init__(self, service: "type[Service]", scopes: tuple[str, ...]) -> None: + self.service = service + self.scopes = scopes + + def __set_name__(self, owner: type, name: str) -> None: + self.owner = owner + self.name = name + + @property + def label(self) -> str: + return f"{self.owner.name}.{self.name}" + + def list_for_account(self, account_id: str) -> list[Connection]: + return [ + Connection(self.service, row) + for row in OauthConnection.list_for_account(self.service.name, account_id) + ] + + def get(self, connection_id: str) -> Connection | None: + row = OauthConnection.get(connection_id) + if row and row.provider == self.service.name: + return Connection(self.service, row) + + class Service: """The appliance's own identity at an external provider — one per service, declared by the code that consumes it. Subclass in a ``services`` module, @@ -31,8 +89,9 @@ class Service: settings_model: ClassVar[type[BaseModel]] # Set both endpoints when the registered app is an OAuth client; # ``get_oauth_client()`` then hands back the connected identity as a - # configured ``OauthClient``. Scopes are not declared here — each - # ``begin_connect`` asks for its own. + # configured ``OauthClient``. Scopes are not declared here — the + # extensions that use the service declare them (``connection``), and + # the connect door asks for their union. authorization_endpoint: ClassVar[str] = "" token_endpoint: ClassVar[str] = "" # HTTP Basic on the token endpoint; False sends the secret in the body. @@ -90,6 +149,29 @@ def connect_fields(cls) -> list[dict[str, Any]]: def get(cls) -> ServiceIdentity: return ServiceIdentity.get(cls.name) + @classmethod + def with_scopes(cls, *scopes: str) -> ScopedService: + """Declare this extension's use of the service and the scopes its + calls need.""" + if not cls.token_endpoint: + raise TypeError(f"{cls.__name__} declares no OAuth endpoints") + return ScopedService(cls, scopes) + + @classmethod + def declarations(cls) -> "list[ScopedService]": + return [ + value + for extension in iter_extensions() + for value in vars(extension).values() + if isinstance(value, ScopedService) and value.service is cls + ] + + @classmethod + def required_scopes(cls) -> tuple[str, ...]: + """The union of every installed declaration's scopes — the consent ask.""" + scopes = {scope for declaration in cls.declarations() for scope in declaration.scopes} + return tuple(sorted(scopes)) + @classmethod def get_oauth_client(cls) -> OauthClient: """The connected identity as a configured ``OauthClient``, keyed by diff --git a/backend/druks/services/models.py b/backend/druks/services/models.py index 43873b8c..f9abc056 100644 --- a/backend/druks/services/models.py +++ b/backend/druks/services/models.py @@ -1,13 +1,15 @@ from datetime import datetime from typing import Any +from sqlalchemy import ForeignKey, String, select, update from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.orm import Mapped, mapped_column -from druks.database import db_session +from druks.core.models import Uuid7Pk +from druks.database import db_session, get_session from druks.models import Base -from druks.secrets.fields import EncryptedJsonField -from druks.services.exceptions import ServiceNotConnectedError +from druks.secrets.fields import EncryptedJsonField, EncryptedTextField +from druks.services.exceptions import OauthRefreshError, ServiceNotConnectedError class ServiceIdentity(Base): @@ -24,10 +26,9 @@ class ServiceIdentity(Base): @classmethod def get(cls, service: str) -> "ServiceIdentity": - identity = db_session().get(cls, service) - if identity is None: - raise ServiceNotConnectedError(service) - return identity + if identity := db_session().get(cls, service): + return identity + raise ServiceNotConnectedError(service) @classmethod def connect( @@ -36,7 +37,7 @@ def connect( # The caller verifies the credentials against the service first; this # trusts what it is given and overwrites whatever was connected. row = db_session().get(cls, service) - if row is None: + if not row: row = cls(service=service) db_session().add(row) row.identity = identity @@ -44,3 +45,103 @@ def connect( row.connected_at = Base.utc_now() db_session().flush() return row + + +class OauthConnection(Base, Uuid7Pk): + """One signed-in provider account: the durable outcome of an OAuth + consent, owned by the druks account that completed it. An account can + hold many per provider — one per mailbox, handle, or workspace. The + engine rotates the refresh token on mint; nothing else writes here.""" + + __tablename__ = "oauth_connections" + + provider: Mapped[str] = mapped_column(String) + # Who in druks connected it — every read scopes through the owner. + account_id: Mapped[str] = mapped_column(ForeignKey("accounts.id", ondelete="RESTRICT")) + # Ciphertext at rest; decrypted only into the refresh request body. + refresh_token = EncryptedTextField() + # The token response's ``scope`` when the provider echoes one, else the + # scopes the consent asked for. + scopes: Mapped[list[str]] = mapped_column(JSONB, default=list) + connected_at: Mapped[datetime] = mapped_column(default=Base.utc_now) + + @classmethod + def get(cls, connection_id: str) -> "OauthConnection | None": + return db_session().get(cls, connection_id) + + @classmethod + def create( + cls, *, provider: str, account_id: str, refresh_token: str, scopes: list[str] + ) -> "OauthConnection": + connection = cls( + provider=provider, account_id=account_id, refresh_token=refresh_token, scopes=scopes + ) + db_session().add(connection) + db_session().flush() + return connection + + @classmethod + def list_for_account(cls, provider: str, account_id: str) -> "list[OauthConnection]": + return list( + db_session().scalars( + select(cls) + .where(cls.provider == provider, cls.account_id == account_id) + .order_by(cls.connected_at) + ) + ) + + @classmethod + def list_for_provider(cls, provider: str) -> "list[OauthConnection]": + return list(db_session().scalars(select(cls).where(cls.provider == provider))) + + @classmethod + def list_owned_by(cls, account_id: str | None) -> "list[OauthConnection]": + return list( + db_session().scalars( + select(cls).where(cls.account_id == account_id).order_by(cls.connected_at) + ) + ) + + def reconnect(self, *, refresh_token: str, scopes: list[str]) -> None: + self.refresh_token = refresh_token + self.scopes = scopes + self.connected_at = Base.utc_now() + db_session().flush() + + def delete(self) -> None: + session = db_session() + session.delete(self) + session.flush() + + def _load_refresh_token(self) -> str: + # Under the refresh lock: another process may have rotated and + # committed, and this transaction may already hold the row — + # populate_existing re-reads it past the identity map. + fresh = ( + db_session() + .scalars( + select(OauthConnection) + .where(OauthConnection.id == self.id) + .execution_options(populate_existing=True) + ) + .one_or_none() + ) + if fresh: + return fresh.refresh_token.decrypt() + raise OauthRefreshError(self.provider, "the connection was removed mid-refresh") + + def _save_refresh_token(self, rotated: str) -> None: + # The provider invalidated the old token the moment it rotated, so + # the write commits on its own session, never the enclosing + # transaction — a step that rolls back later must not brick the + # connection. + with get_session(db_session().get_bind()) as session: + session.execute( + update(OauthConnection) + .where(OauthConnection.id == self.id) + .values(refresh_token=rotated) + ) + session.commit() + # Keep the enclosing transaction's copy true as well. + self.refresh_token = rotated + db_session().flush() diff --git a/backend/druks/services/oauth.py b/backend/druks/services/oauth.py index dba3a525..034f4f4b 100644 --- a/backend/druks/services/oauth.py +++ b/backend/druks/services/oauth.py @@ -3,7 +3,6 @@ import hashlib import json import secrets -from collections.abc import Callable from typing import Any, cast from urllib.parse import urlencode @@ -19,6 +18,7 @@ OAUTH_TOKEN_TTL_SKEW_SECONDS, ) from .exceptions import OauthExchangeError, OauthRefreshError +from .models import OauthConnection def _http() -> httpx.AsyncClient: @@ -58,19 +58,20 @@ class OauthClient: basic_auth=True, ) - ``begin_connect`` returns the consent URL to open; ``complete_connect`` - consumes the callback's single-use state and exchanges the code; - ``mint_access_token`` serves delivery from the Redis token cache, electing - one refresher per grant. Grants live on the caller's own rows — mint takes - the grant object that owns them. A ``Service`` with declared OAuth - endpoints hands back a configured client via ``get_oauth_client()`` — - construct directly only when no service holds the client credentials. + ``begin_connect`` returns the consent URL to open; the module-level + ``complete_connect`` consumes the callback's single-use state and + exchanges the code; ``mint_access_token`` serves delivery from the Redis + token cache, electing one refresher per connection. The caller stores an + ``OauthConnection`` from the completed exchange and hands it back to + mint. A ``Service`` with declared OAuth endpoints hands back a + configured client via ``get_oauth_client()`` — construct directly only + when no service holds the client credentials. - ``provider`` keys every Redis entry — connect state, token cache, refresh - lock — so all clients constructed with one provider name share them, and - across a rolling deploy old and new processes elect the same single - refresher. Completion needs only ``provider``: the begun flow's endpoints - and client identity ride the stashed state, pinned at begin time so a + The Redis token cache and refresh lock key on the connection id, so all + clients constructed for one provider share them, and across a rolling + deploy old and new processes elect the same single refresher. Connect + state is keyed by the state value alone: the begun flow's provider, + endpoints, and client identity ride the stash, pinned at begin time so a configuration change mid-consent cannot mismatch the PKCE verifier. ``basic_auth`` picks HTTP Basic on the token endpoint, for both the code @@ -92,7 +93,6 @@ def __init__( extra_token_params: dict[str, str] | None = None, mint_wait_interval_seconds: float = OAUTH_MINT_WAIT_INTERVAL_SECONDS, mint_wait_attempts: int = OAUTH_MINT_WAIT_ATTEMPTS, - http_factory: Callable[[], httpx.AsyncClient] | None = None, ) -> None: self.provider = provider self.authorization_endpoint = authorization_endpoint @@ -103,7 +103,6 @@ def __init__( self.extra_token_params = dict(extra_token_params or {}) self.mint_wait_interval_seconds = mint_wait_interval_seconds self.mint_wait_attempts = mint_wait_attempts - self._http = http_factory or _http async def begin_connect( self, @@ -129,6 +128,8 @@ async def begin_connect( ) pending = { **(context or {}), + "provider": self.provider, + "scopes": list(scopes), "code_verifier": code_verifier, "redirect_uri": redirect_uri, "token_endpoint": self.token_endpoint, @@ -138,7 +139,7 @@ async def begin_connect( "extra_token_params": self.extra_token_params, } await get_client().set( - f"{self.provider}:connect:{state}", + f"oauth:connect:{state}", json.dumps(pending), ex=OAUTH_CONNECT_STATE_TTL_SECONDS, ) @@ -155,89 +156,19 @@ async def begin_connect( query.update(extra_authorize_params or {}) return f"{self.authorization_endpoint}?{urlencode(query)}" - async def complete_connect(self, *, state: str, code: str) -> tuple[dict, dict]: - """The callback half: consume the pending state (single-use, GETDEL) - and exchange the code + verifier for tokens. Returns ``(tokens, - context)`` — the token response and the begun flow's stash, the - caller's begin-time context with the flow's ``token_endpoint``, - ``client_id``, and ``client_secret`` merged in. A response without a - ``refresh_token`` is rejected: a grant must survive offline. Nothing - is cached here — the caller's grant becomes real only when its own - write commits, and the first mint refreshes from it.""" - raw = await get_client().getdel(f"{self.provider}:connect:{state}") - if not raw: - raise OauthExchangeError( - self.provider, - "unknown or expired state; start the connect flow again", - context={}, - ) - pending = json.loads(raw) - data = { - "grant_type": "authorization_code", - "code": code, - "redirect_uri": pending["redirect_uri"], - "code_verifier": pending["code_verifier"], - **pending["extra_token_params"], - } - async with self._http() as http: - try: - response = await _post_token( - http, - pending["token_endpoint"], - data, - client_id=pending["client_id"], - client_secret=pending["client_secret"], - basic_auth=pending["basic_auth"], - ) - except httpx.HTTPError as error: - raise OauthExchangeError( - self.provider, f"code exchange failed: {error}", context=pending - ) from error - if response.status_code != 200: - raise OauthExchangeError( - self.provider, - f"code exchange failed: HTTP {response.status_code}", - context=pending, - ) - try: - tokens = response.json() - except ValueError as error: - raise OauthExchangeError( - self.provider, "the token endpoint returned malformed JSON", context=pending - ) from error - if not isinstance(tokens, dict) or not tokens.get("refresh_token"): - raise OauthExchangeError( - self.provider, - "the authorization server granted no refresh token; druks needs offline access", - context=pending, - ) - return tokens, pending - - async def mint_access_token(self, *, key: str, grant) -> str: - """The delivery-side token for one grant: the cached access token - while it lives, else one refreshed through the grant's refresh token. - The provider may rotate the refresh token on use — two concurrent - refreshes trip its reuse detection and can revoke the whole grant — - so Redis elects one refresher per ``key`` (SET NX; the TTL is a crash - backstop a live refresh cannot outlive). Losers poll for the winner's - cache fill, for about one token-endpoint round trip, then fail loudly. - - ``grant`` is the caller's own object — typically the row the grant - lives on — carrying two verbs: - - ``grant.load_refresh_token()`` runs under the refresh lock and must - observe rotations other processes committed — a naive re-select can - return a row this transaction already identity-mapped, so read with - ``populate_existing`` or on a fresh session. - - ``grant.save_refresh_token(rotated)`` receives a rotated refresh - token and must have committed it before returning: the provider has - already invalidated the old token, so the write cannot ride an - enclosing transaction that may later roll back. The cache fills only - after it returns.""" + async def mint_access_token(self, *, connection: OauthConnection) -> str: + """The delivery-side token for one connection: the cached access + token while it lives, else one refreshed through the stored refresh + token. The provider may rotate the refresh token on use — two + concurrent refreshes trip its reuse detection and can revoke the + whole connection — so Redis elects one refresher per connection (SET + NX; the TTL is a crash backstop a live refresh cannot outlive). + Losers poll for the winner's cache fill, for about one token-endpoint + round trip, then fail loudly. The engine reads the connection fresh + under the lock and commits a rotated token before the cache fills.""" redis = get_client() - token_key = f"{self.provider}:access_token:{key}" - lock_key = f"{self.provider}:refresh_lock:{key}" + token_key = f"{self.provider}:access_token:{connection.id}" + lock_key = f"{self.provider}:refresh_lock:{connection.id}" for _ in range(self.mint_wait_attempts): cached = await redis.get(token_key) if cached: @@ -252,10 +183,10 @@ async def mint_access_token(self, *, key: str, grant) -> str: try: data = { "grant_type": "refresh_token", - "refresh_token": grant.load_refresh_token(), + "refresh_token": connection._load_refresh_token(), **self.extra_token_params, } - async with self._http() as http: + async with _http() as http: try: response = await _post_token( http, @@ -283,7 +214,7 @@ async def mint_access_token(self, *, key: str, grant) -> str: self.provider, "the token endpoint returned no access token" ) if tokens.get("refresh_token"): - grant.save_refresh_token(tokens["refresh_token"]) + connection._save_refresh_token(tokens["refresh_token"]) try: ttl = int(tokens.get("expires_in", 3600)) - OAUTH_TOKEN_TTL_SKEW_SECONDS except (TypeError, ValueError) as error: @@ -296,5 +227,67 @@ async def mint_access_token(self, *, key: str, grant) -> str: finally: await redis.delete(lock_key) - async def evict_access_token(self, key: str) -> None: - await get_client().delete(f"{self.provider}:access_token:{key}") + async def evict_access_token(self, connection_id: str) -> None: + await get_client().delete(f"{self.provider}:access_token:{connection_id}") + + async def disconnect(self, connection: OauthConnection) -> None: + """Delete the connection and evict its cached access token.""" + connection.delete() + await self.evict_access_token(connection.id) + + +async def complete_connect(*, state: str, code: str) -> tuple[dict, dict]: + """Consume the pending state (single-use, GETDEL) and exchange the code + for tokens; a callback route knows only ``state`` and ``code``, so the + flow's provider and client identity ride the stash. Returns ``(tokens, + pending)`` — the caller stores the grant, because only it knows the + grant's account.""" + raw = await get_client().getdel(f"oauth:connect:{state}") + if not raw: + raise OauthExchangeError( + "oauth", + "unknown or expired state; start the connect flow again", + context={}, + ) + pending = json.loads(raw) + provider = pending["provider"] + data = { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": pending["redirect_uri"], + "code_verifier": pending["code_verifier"], + **pending["extra_token_params"], + } + async with _http() as http: + try: + response = await _post_token( + http, + pending["token_endpoint"], + data, + client_id=pending["client_id"], + client_secret=pending["client_secret"], + basic_auth=pending["basic_auth"], + ) + except httpx.HTTPError as error: + raise OauthExchangeError( + provider, f"code exchange failed: {error}", context=pending + ) from error + if response.status_code != 200: + raise OauthExchangeError( + provider, + f"code exchange failed: HTTP {response.status_code}", + context=pending, + ) + try: + tokens = response.json() + except ValueError as error: + raise OauthExchangeError( + provider, "the token endpoint returned malformed JSON", context=pending + ) from error + if not isinstance(tokens, dict) or not tokens.get("refresh_token"): + raise OauthExchangeError( + provider, + "the authorization server granted no refresh token; druks needs offline access", + context=pending, + ) + return tokens, pending diff --git a/backend/druks/services/routes.py b/backend/druks/services/routes.py index 1dc4ff2f..b052c827 100644 --- a/backend/druks/services/routes.py +++ b/backend/druks/services/routes.py @@ -1,12 +1,21 @@ -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi.responses import HTMLResponse, RedirectResponse +from druks.accounts.context import current_account_id from druks.accounts.dependencies import current_session_account +from druks.core.templates import render_page from druks.extensions.registry import services -from druks.services.exceptions import ServiceConnectError, ServiceNotConnectedError -from druks.services.models import ServiceIdentity -from druks.services.schemas import ServiceResponse +from druks.services.exceptions import ( + OauthExchangeError, + ServiceConnectError, + ServiceNotConnectedError, +) +from druks.services.models import OauthConnection, ServiceIdentity +from druks.services.oauth import OauthClient, complete_connect +from druks.services.schemas import ConnectionResponse, ServiceResponse router = APIRouter(prefix="/api/services", tags=["services"]) +oauth_router = APIRouter(prefix="/api/oauth", tags=["oauth"]) @router.get("", response_model=list[ServiceResponse], response_model_by_alias=True) @@ -17,7 +26,10 @@ async def list_services() -> list[ServiceResponse]: row = ServiceIdentity.get(service.name) except ServiceNotConnectedError: row = None - entries.append(ServiceResponse.from_row(service, row)) + connections = [] + if service.token_endpoint: + connections = OauthConnection.list_for_provider(service.name) + entries.append(ServiceResponse.from_row(service, row, connections)) return entries @@ -37,4 +49,101 @@ async def connect_service(name: str, payload: dict[str, str]) -> ServiceResponse row = await service.connect(payload) except ServiceConnectError as error: raise HTTPException(status_code=422, detail=str(error)) from error + if service.token_endpoint: + # A replaced client can never refresh the old client's connections. + client = OauthClient(provider=name) + for connection in OauthConnection.list_for_provider(name): + await client.disconnect(connection) return ServiceResponse.from_row(service, row) + + +def _get_oauth_service(name: str): + service = services.get(name) + if not service or not service.token_endpoint: + raise HTTPException(status_code=404, detail=f"No OAuth service {name!r}.") + return service + + +@oauth_router.get("/{name}/connect", dependencies=[Depends(current_session_account)]) +async def connect_oauth_service( + name: str, request: Request, connection: str = "" +) -> RedirectResponse: + service = _get_oauth_service(name) + account_id = current_account_id.get() + if connection: + row = OauthConnection.get(connection) + if not row or row.provider != name: + raise HTTPException( + status_code=404, detail=f"No connection {connection!r} on {name!r}." + ) + endpoint = request.app.state.settings.urls.endpoint + if not endpoint: + raise HTTPException( + status_code=409, + detail="The provider redirects the operator's browser back to druks. " + "Set urls.endpoint to the address druks has in that browser.", + ) + try: + client = service.get_oauth_client() + except ServiceNotConnectedError as error: + raise HTTPException(status_code=409, detail=str(error)) from error + url = await client.begin_connect( + redirect_uri=f"{endpoint.rstrip('/')}/api/oauth/callback", + scopes=service.required_scopes(), + context={"account_id": account_id, "connection_id": connection}, + ) + return RedirectResponse(url) + + +@oauth_router.get("/callback", response_class=HTMLResponse) +async def oauth_callback(state: str = "", code: str = "", error: str = "") -> HTMLResponse: + if error: + raise HTTPException( + status_code=400, detail=f"The authorization server denied the request: {error}" + ) + if not state or not code: + raise HTTPException(status_code=400, detail="Missing state or code in the callback.") + try: + tokens, pending = await complete_connect(state=state, code=code) + except OauthExchangeError as exchange_error: + raise HTTPException(status_code=400, detail=str(exchange_error)) from exchange_error + provider = pending["provider"] + if not services.get(provider): + # A state begun by another door (an MCP connect) finishes at its own callback. + raise HTTPException(status_code=400, detail=f"No OAuth service {provider!r}.") + granted = tokens.get("scope", "").split() or pending["scopes"] + if pending["connection_id"]: + row = OauthConnection.get(pending["connection_id"]) + if not row: + raise HTTPException( + status_code=400, detail="The connection was removed while consent was open." + ) + row.reconnect(refresh_token=tokens["refresh_token"], scopes=granted) + # A reconsent's narrower cached token must not serve until its TTL runs out. + await OauthClient(provider=provider).evict_access_token(row.id) + else: + OauthConnection.create( + provider=provider, + account_id=pending["account_id"], + refresh_token=tokens["refresh_token"], + scopes=granted, + ) + return render_page("service_oauth_callback.html", name=provider) + + +@oauth_router.get("/connections", dependencies=[Depends(current_session_account)]) +async def list_connections() -> list[ConnectionResponse]: + rows = OauthConnection.list_owned_by(current_account_id.get()) + return [ConnectionResponse.model_validate(row) for row in rows] + + +@oauth_router.delete( + "/connections/{connection_id}", + status_code=204, + dependencies=[Depends(current_session_account)], +) +async def disconnect_connection(connection_id: str) -> None: + row = OauthConnection.get(connection_id) + if not row: + raise HTTPException(status_code=404, detail=f"No connection {connection_id!r}.") + await OauthClient(provider=row.provider).disconnect(row) diff --git a/backend/druks/services/schemas.py b/backend/druks/services/schemas.py index 847a6b81..1ad45a18 100644 --- a/backend/druks/services/schemas.py +++ b/backend/druks/services/schemas.py @@ -1,11 +1,13 @@ from datetime import datetime from typing import TYPE_CHECKING, Any +from pydantic import ConfigDict + from druks.schemas import BaseResponse if TYPE_CHECKING: from druks.services.base import Service - from druks.services.models import ServiceIdentity + from druks.services.models import OauthConnection, ServiceIdentity # The settings-form field vocabulary (label/help/type), minus everything a @@ -18,6 +20,15 @@ class ServiceFieldSpec(BaseResponse): multiline: bool +class ConnectionResponse(BaseResponse): + model_config = ConfigDict(from_attributes=True) + + id: str + provider: str + scopes: list[str] + connected_at: datetime + + class ServiceResponse(BaseResponse): # Connection state and identity facts only — never a stored secret. name: str @@ -28,9 +39,18 @@ class ServiceResponse(BaseResponse): facts: dict[str, Any] connected_at: datetime | None fields: list[ServiceFieldSpec] + is_oauth: bool + required_scopes: list[str] + used_by: list[str] + connections: list[ConnectionResponse] @classmethod - def from_row(cls, service: "type[Service]", row: "ServiceIdentity | None") -> "ServiceResponse": + def from_row( + cls, + service: "type[Service]", + row: "ServiceIdentity | None", + connections: "list[OauthConnection] | None" = None, + ) -> "ServiceResponse": return cls( name=service.name, title=service.title, @@ -40,4 +60,8 @@ def from_row(cls, service: "type[Service]", row: "ServiceIdentity | None") -> "S facts=row.identity if row else {}, connected_at=row.connected_at if row else None, fields=[ServiceFieldSpec(**spec) for spec in service.connect_fields()], + is_oauth=bool(service.token_endpoint), + required_scopes=list(service.required_scopes()), + used_by=[declaration.label for declaration in service.declarations()], + connections=[ConnectionResponse.model_validate(c) for c in connections or []], ) diff --git a/backend/migrations/versions/a7c2e9f14b38_oauth_connections_are_platform_rows.py b/backend/migrations/versions/a7c2e9f14b38_oauth_connections_are_platform_rows.py new file mode 100644 index 00000000..9549eec1 --- /dev/null +++ b/backend/migrations/versions/a7c2e9f14b38_oauth_connections_are_platform_rows.py @@ -0,0 +1,112 @@ +"""oauth connections are platform rows + +Revision ID: a7c2e9f14b38 +Revises: e3a1c8f92d74 +Create Date: 2026-08-20 00:00:00.000000 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +from druks.core.models import uuid7_str +from druks.secrets import utils + +# revision identifiers, used by Alembic. +revision: str = "a7c2e9f14b38" +down_revision: str | Sequence[str] | None = "e3a1c8f92d74" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _reencrypt(envelope: bytes, old_aad: str, new_aad: str) -> bytes: + # Envelopes bind to their table.column as AAD, so a moved secret must be + # decrypted under the old identity and sealed under the new one. + if not envelope: + return b"" + return utils.encrypt(utils.decrypt(envelope, old_aad), new_aad) + + +def upgrade() -> None: + op.create_table( + "mcp_client_registrations", + sa.Column("id", sa.String(), nullable=False), + sa.Column("server_id", sa.String(), nullable=False), + sa.Column("account_id", sa.String(), server_default="system", nullable=False), + sa.Column("token_endpoint", sa.String(), nullable=False), + sa.Column("client_id", sa.String(), nullable=False), + sa.Column("client_secret", sa.LargeBinary(), nullable=False), + sa.ForeignKeyConstraint(["server_id"], ["mcp_servers.id"], ondelete="CASCADE"), + sa.ForeignKeyConstraint(["account_id"], ["accounts.id"], ondelete="RESTRICT"), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("server_id", "account_id"), + ) + op.create_table( + "oauth_connections", + sa.Column("id", sa.String(), nullable=False), + sa.Column("provider", sa.String(), nullable=False), + sa.Column("account_id", sa.String(), nullable=False), + sa.Column("refresh_token", sa.LargeBinary(), nullable=False), + sa.Column("scopes", postgresql.JSONB(), nullable=False), + sa.Column("connected_at", sa.DateTime(timezone=True), nullable=False), + sa.ForeignKeyConstraint(["account_id"], ["accounts.id"], ondelete="RESTRICT"), + sa.PrimaryKeyConstraint("id"), + ) + bind = op.get_bind() + grants = bind.execute( + sa.text( + "SELECT grants.account_id, grants.refresh_token, grants.token_endpoint," + " grants.client_id, grants.client_secret, grants.connected_at," + " grants.server_name, servers.id AS server_id" + " FROM mcp_oauth_grants grants" + " JOIN mcp_servers servers ON servers.name = grants.server_name" + ) + ).mappings() + for grant in grants: + bind.execute( + sa.text( + "INSERT INTO mcp_client_registrations" + " (id, server_id, account_id, token_endpoint, client_id, client_secret)" + " VALUES (:id, :server_id, :account_id, :token_endpoint, :client_id," + " :client_secret)" + ), + { + "id": uuid7_str(), + "server_id": grant["server_id"], + "account_id": grant["account_id"], + "token_endpoint": grant["token_endpoint"], + "client_id": grant["client_id"], + "client_secret": _reencrypt( + bytes(grant["client_secret"]), + "mcp_oauth_grants.client_secret", + "mcp_client_registrations.client_secret", + ), + }, + ) + bind.execute( + sa.text( + "INSERT INTO oauth_connections" + " (id, provider, account_id, refresh_token, scopes, connected_at)" + " VALUES (:id, :provider, :account_id, :refresh_token, '[]'::jsonb," + " :connected_at)" + ), + { + "id": uuid7_str(), + "provider": f"mcp:{grant['server_name']}", + "account_id": grant["account_id"], + "refresh_token": _reencrypt( + bytes(grant["refresh_token"]), + "mcp_oauth_grants.refresh_token", + "oauth_connections.refresh_token", + ), + "connected_at": grant["connected_at"], + }, + ) + op.drop_table("mcp_oauth_grants") + + +def downgrade() -> None: + raise NotImplementedError("oauth rows moved; restore from backup instead") diff --git a/backend/tests/test_auth_boundary.py b/backend/tests/test_auth_boundary.py index 455dd7b4..34aec68c 100644 --- a/backend/tests/test_auth_boundary.py +++ b/backend/tests/test_auth_boundary.py @@ -30,6 +30,9 @@ ("DELETE", "/api/harnesses/{name}/connection"), ("PATCH", "/api/settings/extensions"), ("POST", "/api/services/{name}"), + ("GET", "/api/oauth/{name}/connect"), + ("GET", "/api/oauth/connections"), + ("DELETE", "/api/oauth/connections/{connection_id}"), } # Session-gated routes that also sit behind their router's identity gate — @@ -37,6 +40,9 @@ DUAL_GATED_API_PATHS = { "/api/settings/extensions", "/api/services/{name}", + "/api/oauth/{name}/connect", + "/api/oauth/connections", + "/api/oauth/connections/{connection_id}", } # The connection flow must answer during none/zero setup, before any account diff --git a/backend/tests/test_mcp_oauth.py b/backend/tests/test_mcp_oauth.py index 96e923b0..e483762f 100644 --- a/backend/tests/test_mcp_oauth.py +++ b/backend/tests/test_mcp_oauth.py @@ -11,11 +11,6 @@ from druks.accounts.models import Account from druks.extensions.registry import mcp_servers from druks.mcp import oauth -from druks.mcp.constants import ( - OAUTH_ACCESS_TOKEN_PREFIX, - OAUTH_CONNECT_STATE_PREFIX, - OAUTH_REFRESH_LOCK_PREFIX, -) from druks.mcp.enums import IdentityMode, TokenSource from druks.mcp.exceptions import ( GrantRefreshError, @@ -23,10 +18,11 @@ OauthConnectError, UnresolvedGrantAccountError, ) -from druks.mcp.helpers import get_bearer_token_env_var -from druks.mcp.models import McpOauthGrant, McpServer +from druks.mcp.helpers import get_bearer_token_env_var, get_grant_account +from druks.mcp.models import McpClientRegistration, McpServer from druks.redis import close_client, get_client from druks.sandbox.datastructures import Workspace +from druks.services.models import OauthConnection from druks.testing import configure_app_for_test, make_settings from druks.user_settings.models import UserSettings from fastapi.testclient import TestClient @@ -97,9 +93,12 @@ def handler(self, request: httpx.Request) -> httpx.Response: @pytest.fixture def auth_server(monkeypatch): fake = FakeAuthServer() - monkeypatch.setattr( - oauth, "_http", lambda: httpx.AsyncClient(transport=httpx.MockTransport(fake.handler)) - ) + + def http(): + return httpx.AsyncClient(transport=httpx.MockTransport(fake.handler)) + + monkeypatch.setattr(oauth, "_http", http) + monkeypatch.setattr("druks.services.oauth._http", http) return fake @@ -120,7 +119,7 @@ def _store_grant( *, account_id: str = SYSTEM_ACCOUNT_ID, identity_mode: IdentityMode = IdentityMode.SHARED, -) -> McpOauthGrant: +) -> OauthConnection: server = McpServer.get_for_name(_NAME) if not server: server = McpServer.create( @@ -129,14 +128,30 @@ def _store_grant( token_source=TokenSource.OAUTH, ) server.identity_mode = identity_mode - return McpOauthGrant.store( - server_name=_NAME, + McpClientRegistration.store( + server_id=server.id, account_id=account_id, - refresh_token=refresh_token, token_endpoint=f"{_AUTH_BASE}/token", - resource=_SERVER_URL, client_id="client-123", ) + return OauthConnection.create( + provider=f"mcp:{_NAME}", + account_id=account_id, + refresh_token=refresh_token, + scopes=[], + ) + + +def _state_key(state: str) -> str: + return f"oauth:connect:{state}" + + +def _token_key(account_id: str) -> str: + return f"mcp:{_NAME}:access_token:{oauth.get_connection(_NAME, account_id).id}" + + +def _lock_key(account_id: str) -> str: + return f"mcp:{_NAME}:refresh_lock:{oauth.get_connection(_NAME, account_id).id}" @pytest.mark.parametrize( @@ -149,7 +164,7 @@ def _store_grant( ) def test_get_grant_account_rejects_unresolved_modes(identity_mode, account_id): with pytest.raises(UnresolvedGrantAccountError): - McpOauthGrant.get_grant_account(identity_mode, account_id) + get_grant_account(identity_mode, account_id) # --- connect: discovery + DCR + PKCE --------------------------------------- @@ -172,7 +187,7 @@ async def test_begin_connect_builds_consent_url_and_stashes_pkce_state(auth_serv assert params["code_challenge_method"] == "S256" assert params["resource"] == _SERVER_URL - raw = await get_client().get(f"{OAUTH_CONNECT_STATE_PREFIX}{params['state']}") + raw = await get_client().get(_state_key(params["state"])) pending = json.loads(raw) # The challenge in the consent URL is the S256 hash of the stashed verifier. expected = ( @@ -291,10 +306,11 @@ async def test_complete_connect_exchanges_code_and_stores_the_grant(auth_server, name = await oauth.complete_connect(state=state, code="code-1") assert name == _NAME - grant = McpOauthGrant.get_for_account(_NAME, SYSTEM_ACCOUNT_ID) + grant = oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) assert grant.refresh_token.decrypt() == "rt-1" - assert grant.resource == _SERVER_URL - assert grant.client_id == "client-123" + registration = McpClientRegistration.get_for_account(_NAME, SYSTEM_ACCOUNT_ID) + assert registration.client_id == "client-123" + assert registration.token_endpoint == f"{_AUTH_BASE}/token" exchange = auth_server.token_requests[0] assert exchange["grant_type"] == "authorization_code" assert exchange["code"] == "code-1" @@ -303,7 +319,7 @@ async def test_complete_connect_exchanges_code_and_stores_the_grant(auth_server, # Nothing is cached at connect (the grant is real only once this commits); # the first delivery mints from it, carrying the grant's resource binding. - assert not await get_client().get(f"{OAUTH_ACCESS_TOKEN_PREFIX}{_NAME}:{SYSTEM_ACCOUNT_ID}") + assert not await get_client().get(_token_key(SYSTEM_ACCOUNT_ID)) assert await oauth.mint_access_token(_NAME, SYSTEM_ACCOUNT_ID) == "at-1" refresh = auth_server.token_requests[1] assert refresh["grant_type"] == "refresh_token" @@ -327,7 +343,27 @@ async def test_complete_connect_without_refresh_token_stores_nothing(auth_server with pytest.raises(OauthConnectError, match="no refresh token"): await oauth.complete_connect(state=state, code="code-1") - assert not McpOauthGrant.list_for_server(_NAME) + assert not oauth.list_connections(_NAME) + + +async def test_reconsent_replaces_the_grant_and_evicts_the_stale_token(auth_server, druks_db): + _store_grant(refresh_token="rt-stale") + await get_client().set(_token_key(SYSTEM_ACCOUNT_ID), "at-stale") + + url = await oauth.begin_connect( + _NAME, + _SERVER_URL, + _ENDPOINT, + account_id=SYSTEM_ACCOUNT_ID, + identity_mode=IdentityMode.SHARED, + ) + state = dict(parse_qsl(urlparse(url).query))["state"] + await oauth.complete_connect(state=state, code="code-1") + + grant = oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) + assert grant.refresh_token.decrypt() == "rt-1" + # The stale narrow token must not keep serving until its TTL runs out. + assert not await get_client().get(_token_key(SYSTEM_ACCOUNT_ID)) async def test_two_shared_connects_converge_on_one_grant(auth_server, druks_db): @@ -345,7 +381,7 @@ async def test_two_shared_connects_converge_on_one_grant(auth_server, druks_db): state = dict(parse_qsl(urlparse(url).query))["state"] await oauth.complete_connect(state=state, code=account.id) - grants = McpOauthGrant.list_for_server(_NAME) + grants = oauth.list_connections(_NAME) assert McpServer.get_for_name(_NAME).identity_mode == IdentityMode.SHARED assert [grant.account_id for grant in grants] == [SYSTEM_ACCOUNT_ID] @@ -365,7 +401,7 @@ async def test_two_per_user_connects_store_two_grants(auth_server, druks_db): state = dict(parse_qsl(urlparse(url).query))["state"] await oauth.complete_connect(state=state, code=account.id) - grants = McpOauthGrant.list_for_server(_NAME) + grants = oauth.list_connections(_NAME) assert McpServer.get_for_name(_NAME).identity_mode == IdentityMode.PER_USER assert {grant.account_id for grant in grants} == {first.id, second.id} @@ -394,7 +430,7 @@ async def test_a_later_connect_stores_under_the_claimed_mode(auth_server, druks_ await oauth.complete_connect(state=second_state, code="second") assert McpServer.get_for_name(_NAME).identity_mode == IdentityMode.SHARED - grant_accounts = {grant.account_id for grant in McpOauthGrant.list_for_server(_NAME)} + grant_accounts = {grant.account_id for grant in oauth.list_connections(_NAME)} assert grant_accounts == {SYSTEM_ACCOUNT_ID} @@ -418,7 +454,7 @@ async def test_mint_refreshes_on_cache_miss_and_persists_rotation(auth_server, d assert refresh["refresh_token"] == "rt-old" assert refresh["resource"] == _SERVER_URL # Rotation: the provider's new refresh token replaced the stored one. - stored = McpOauthGrant.get_for_account(_NAME, SYSTEM_ACCOUNT_ID) + stored = oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) assert stored.refresh_token.decrypt() == "rt-new" # A second mint within the TTL reuses the cache — no second refresh. @@ -437,7 +473,7 @@ async def test_mint_refresh_rejection_fails_loudly_and_evicts_the_cache(auth_ser with pytest.raises(GrantRefreshError, match=_NAME): await oauth.mint_access_token(_NAME, SYSTEM_ACCOUNT_ID) - assert not await get_client().get(f"{OAUTH_ACCESS_TOKEN_PREFIX}{_NAME}:{SYSTEM_ACCOUNT_ID}") + assert not await get_client().get(_token_key(SYSTEM_ACCOUNT_ID)) async def test_mint_rejects_a_malformed_token_response(auth_server, druks_db): @@ -465,11 +501,11 @@ async def test_mint_losing_the_refresh_lock_polls_for_the_winners_token( _store_grant() redis = get_client() monkeypatch.setattr(oauth, "OAUTH_MINT_WAIT_INTERVAL_SECONDS", 0) - await redis.set(f"{OAUTH_REFRESH_LOCK_PREFIX}{_NAME}:{SYSTEM_ACCOUNT_ID}", "1") + await redis.set(_lock_key(SYSTEM_ACCOUNT_ID), "1") async def _winner_finishes(): - await redis.set(f"{OAUTH_ACCESS_TOKEN_PREFIX}{_NAME}:{SYSTEM_ACCOUNT_ID}", "at-winner") - await redis.delete(f"{OAUTH_REFRESH_LOCK_PREFIX}{_NAME}:{SYSTEM_ACCOUNT_ID}") + await redis.set(_token_key(SYSTEM_ACCOUNT_ID), "at-winner") + await redis.delete(_lock_key(SYSTEM_ACCOUNT_ID)) winner = asyncio.create_task(_winner_finishes()) assert await oauth.mint_access_token(_NAME, SYSTEM_ACCOUNT_ID) == "at-winner" @@ -481,7 +517,7 @@ async def test_mint_times_out_loudly_when_the_refresh_lock_never_frees(druks_db, _store_grant() monkeypatch.setattr(oauth, "OAUTH_MINT_WAIT_INTERVAL_SECONDS", 0) monkeypatch.setattr(oauth, "OAUTH_MINT_WAIT_ATTEMPTS", 3) - await get_client().set(f"{OAUTH_REFRESH_LOCK_PREFIX}{_NAME}:{SYSTEM_ACCOUNT_ID}", "1") + await get_client().set(_lock_key(SYSTEM_ACCOUNT_ID), "1") with pytest.raises(GrantRefreshError, match="concurrent refresh"): await oauth.mint_access_token(_NAME, SYSTEM_ACCOUNT_ID) @@ -493,11 +529,11 @@ async def test_mint_cache_and_refresh_lock_are_per_account(auth_server, druks_db _store_grant(account_id=first.id, identity_mode=IdentityMode.PER_USER) _store_grant(account_id=second.id, identity_mode=IdentityMode.PER_USER) redis = get_client() - await redis.set(f"{OAUTH_REFRESH_LOCK_PREFIX}{_NAME}:{first.id}", "1") + await redis.set(_lock_key(first.id), "1") assert await oauth.mint_access_token(_NAME, second.id) == "at-1" - assert not await redis.get(f"{OAUTH_ACCESS_TOKEN_PREFIX}{_NAME}:{first.id}") - assert await redis.get(f"{OAUTH_ACCESS_TOKEN_PREFIX}{_NAME}:{second.id}") == b"at-1" + assert not await redis.get(_token_key(first.id)) + assert await redis.get(_token_key(second.id)) == b"at-1" # --- delivery: the oauth branch of the fold --------------------------------- @@ -627,7 +663,7 @@ def test_callback_route_completes_the_connect(tmp_path, registry_state, auth_ser # The page notifies the opener tab, then closes itself. assert "BroadcastChannel('druks-mcp-connect')" in page.text assert "window.close()" in page.text - assert McpOauthGrant.get_for_account(_NAME, SYSTEM_ACCOUNT_ID) + assert oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) # Connecting is the explicit "use this server" — it enables too. assert McpServer.get_for_name(_NAME).is_enabled is True @@ -653,13 +689,15 @@ async def test_disconnect_route_drops_grant_and_cache( _register_oauth_server() _store_grant() await oauth.mint_access_token(_NAME, SYSTEM_ACCOUNT_ID) + token_key = _token_key(SYSTEM_ACCOUNT_ID) # The mint's Redis client is bound to this test's loop; close it so the # route dials its own — the cached token lives in Redis either way. await close_client() with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: assert client.delete(f"/api/mcp-servers/{_NAME}/grant").status_code == 204 - assert not McpOauthGrant.get_for_account(_NAME, SYSTEM_ACCOUNT_ID) + assert not oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID) + assert not McpClientRegistration.get_for_account(_NAME, SYSTEM_ACCOUNT_ID) # The mirror of connect-enables: no grant, no calls, so no dead entry # riding into VMs. assert McpServer.get_for_name(_NAME).is_enabled is False @@ -670,7 +708,7 @@ async def test_disconnect_route_drops_grant_and_cache( # shutdown to have nulled it — a portal-bound client awaited from here # parks forever, it doesn't raise. druks.redis._client = None - assert not await get_client().get(f"{OAUTH_ACCESS_TOKEN_PREFIX}{_NAME}:{SYSTEM_ACCOUNT_ID}") + assert not await get_client().get(token_key) def test_shared_disconnect_allows_per_user_reconnect( @@ -709,7 +747,7 @@ def test_shared_disconnect_allows_per_user_reconnect( ) assert McpServer.get_for_name(_NAME).identity_mode == IdentityMode.PER_USER - assert {grant.account_id for grant in McpOauthGrant.list_for_server(_NAME)} == {operator.id} + assert {grant.account_id for grant in oauth.list_connections(_NAME)} == {operator.id} def test_api_has_token_reflects_the_grant_and_leaks_no_secret(tmp_path, registry_state, druks_db): @@ -774,8 +812,8 @@ async def test_per_user_disconnect_preserves_other_accounts_grant_and_cache(tmp_ _store_grant(account_id=disconnected.id, identity_mode=IdentityMode.PER_USER) _store_grant(account_id=connected.id, identity_mode=IdentityMode.PER_USER) redis = get_client() - disconnected_key = f"{OAUTH_ACCESS_TOKEN_PREFIX}{_NAME}:{disconnected.id}" - connected_key = f"{OAUTH_ACCESS_TOKEN_PREFIX}{_NAME}:{connected.id}" + disconnected_key = _token_key(disconnected.id) + connected_key = _token_key(connected.id) await redis.set(disconnected_key, "disconnect-token") await redis.set(connected_key, "connected-token") await close_client() @@ -791,8 +829,8 @@ async def test_per_user_disconnect_preserves_other_accounts_grant_and_cache(tmp_ ) assert response.status_code == 204 - assert not McpOauthGrant.get_for_account(_NAME, disconnected.id) - assert McpOauthGrant.get_for_account(_NAME, connected.id) + assert not oauth.get_connection(_NAME, disconnected.id) + assert oauth.get_connection(_NAME, connected.id) assert McpServer.get_for_name(_NAME).is_enabled is True druks.redis._client = None assert not await get_client().get(disconnected_key) @@ -805,8 +843,8 @@ async def test_removal_drops_every_grant_and_cached_token(tmp_path, druks_db): _store_grant(account_id=first.id, identity_mode=IdentityMode.PER_USER) _store_grant(account_id=second.id, identity_mode=IdentityMode.PER_USER) redis = get_client() - first_key = f"{OAUTH_ACCESS_TOKEN_PREFIX}{_NAME}:{first.id}" - second_key = f"{OAUTH_ACCESS_TOKEN_PREFIX}{_NAME}:{second.id}" + first_key = _token_key(first.id) + second_key = _token_key(second.id) await redis.set(first_key, "first-token") await redis.set(second_key, "second-token") await close_client() @@ -815,7 +853,7 @@ async def test_removal_drops_every_grant_and_cached_token(tmp_path, druks_db): assert client.delete(f"/api/mcp-servers/{_NAME}").status_code == 204 assert not McpServer.get_for_name(_NAME) - assert not McpOauthGrant.list_for_server(_NAME) + assert not oauth.list_connections(_NAME) druks.redis._client = None assert not await get_client().get(first_key) assert not await get_client().get(second_key) diff --git a/backend/tests/test_mcp_registry.py b/backend/tests/test_mcp_registry.py index 4c9a0e43..12c3e34c 100644 --- a/backend/tests/test_mcp_registry.py +++ b/backend/tests/test_mcp_registry.py @@ -6,8 +6,9 @@ from druks.mcp import registry from druks.mcp.enums import IdentityMode from druks.mcp.exceptions import RegistryUnavailableError -from druks.mcp.models import McpOauthGrant, McpServer +from druks.mcp.models import McpServer from druks.mcp.registry import derive_server_name, resolve_candidates, search_registry +from druks.services.models import OauthConnection from druks.settings import PACKAGED_MCP_TRUSTED from druks.testing import configure_app_for_test, make_settings from fastapi.testclient import TestClient @@ -442,17 +443,12 @@ def test_removing_a_connected_row_drops_its_grant(tmp_path, monkeypatch, druks_d "/api/mcp-servers/registry", json={"name": "grafana", "registry": "io.github.grafana/mcp-grafana", "headers": {}}, ) - McpOauthGrant.store( - server_name="grafana", - account_id=SYSTEM_ACCOUNT_ID, - refresh_token="rt", - token_endpoint="https://as.example/token", - resource="https://mcp.grafana.com/mcp", - client_id="cid", + OauthConnection.create( + provider="mcp:grafana", account_id=SYSTEM_ACCOUNT_ID, refresh_token="rt", scopes=[] ) assert client.delete("/api/mcp-servers/grafana").status_code == 204 # An orphan grant would revive as this name's credential on re-add. assert not McpServer.get_for_name("grafana") - assert not McpOauthGrant.list_for_server("grafana") + assert not OauthConnection.list_for_provider("mcp:grafana") diff --git a/backend/tests/test_oauth_client.py b/backend/tests/test_oauth_client.py index 09828a3f..51127fb3 100644 --- a/backend/tests/test_oauth_client.py +++ b/backend/tests/test_oauth_client.py @@ -3,15 +3,17 @@ import httpx import pytest +from druks.accounts.constants import SYSTEM_ACCOUNT_ID +from druks.database import db_session from druks.redis import get_client from druks.services import OauthClient, OauthExchangeError, OauthRefreshError +from druks.services.models import OauthConnection +from druks.services.oauth import complete_connect _PROVIDER = "acme" _AUTHORIZATION_ENDPOINT = "https://auth.acme.test/authorize" _TOKEN_ENDPOINT = "https://auth.acme.test/token" -_REDIRECT_URI = "https://druks.example/api/acme/oauth/callback" -_TOKEN_KEY = f"{_PROVIDER}:access_token:grant-1" -_LOCK_KEY = f"{_PROVIDER}:refresh_lock:grant-1" +_REDIRECT_URI = "https://druks.example/api/oauth/callback" class FakeTokenEndpoint: @@ -28,11 +30,16 @@ def handler(self, request: httpx.Request) -> httpx.Response: @pytest.fixture -def token_endpoint(): - return FakeTokenEndpoint() +def token_endpoint(monkeypatch): + fake = FakeTokenEndpoint() + monkeypatch.setattr( + "druks.services.oauth._http", + lambda: httpx.AsyncClient(transport=httpx.MockTransport(fake.handler)), + ) + return fake -def _client(token_endpoint: FakeTokenEndpoint, **overrides) -> OauthClient: +def _client(**overrides) -> OauthClient: kwargs = dict( provider=_PROVIDER, authorization_endpoint=_AUTHORIZATION_ENDPOINT, @@ -40,42 +47,33 @@ def _client(token_endpoint: FakeTokenEndpoint, **overrides) -> OauthClient: client_id="client-123", client_secret="secret-123", mint_wait_interval_seconds=0, - http_factory=lambda: httpx.AsyncClient( - transport=httpx.MockTransport(token_endpoint.handler) - ), ) kwargs.update(overrides) return OauthClient(**kwargs) -class GrantStub: - """A grant with the two mint verbs over an in-memory refresh token.""" - - def __init__(self) -> None: - self.refresh_token = "rt-old" - self.saved: list[str] = [] - - def load_refresh_token(self) -> str: - return self.refresh_token - - def save_refresh_token(self, rotated: str) -> None: - self.saved.append(rotated) +def _connection(refresh_token: str = "rt-old") -> OauthConnection: + return OauthConnection.create( + provider=_PROVIDER, + account_id=SYSTEM_ACCOUNT_ID, + refresh_token=refresh_token, + scopes=[], + ) -class UntouchedGrant(GrantStub): - """A grant the mint under test must never read or rotate.""" +def _token_key(connection: OauthConnection) -> str: + return f"{_PROVIDER}:access_token:{connection.id}" - def load_refresh_token(self) -> str: - pytest.fail("this mint reads no grant") - def save_refresh_token(self, rotated: str) -> None: - pytest.fail("nothing rotated") +def _lock_key(connection: OauthConnection) -> str: + return f"{_PROVIDER}:refresh_lock:{connection.id}" async def test_mint_serves_the_cache_without_a_refresh(token_endpoint): - await get_client().set(_TOKEN_KEY, "at-cached") + connection = _connection() + await get_client().set(_token_key(connection), "at-cached") - token = await _client(token_endpoint).mint_access_token(key="grant-1", grant=UntouchedGrant()) + token = await _client().mint_access_token(connection=connection) assert token == "at-cached" assert not token_endpoint.requests @@ -83,12 +81,13 @@ async def test_mint_serves_the_cache_without_a_refresh(token_endpoint): async def test_mint_refreshes_persists_rotation_and_fills_with_skewed_ttl(token_endpoint): token_endpoint.response = {"access_token": "at-2", "refresh_token": "rt-new", "expires_in": 300} - grant = GrantStub() + connection = _connection() - token = await _client(token_endpoint).mint_access_token(key="grant-1", grant=grant) + token = await _client().mint_access_token(connection=connection) assert token == "at-2" - assert grant.saved == ["rt-new"] + db_session().expire_all() + assert OauthConnection.get(connection.id).refresh_token.decrypt() == "rt-new" refresh = token_endpoint.requests[0] assert refresh["grant_type"] == "refresh_token" assert refresh["refresh_token"] == "rt-old" @@ -96,48 +95,38 @@ async def test_mint_refreshes_persists_rotation_and_fills_with_skewed_ttl(token_ assert refresh["client_id"] == "client-123" assert refresh["client_secret"] == "secret-123" redis = get_client() - assert await redis.get(_TOKEN_KEY) == b"at-2" - assert 0 < await redis.ttl(_TOKEN_KEY) <= 240 - assert not await redis.get(_LOCK_KEY) + assert await redis.get(_token_key(connection)) == b"at-2" + assert 0 < await redis.ttl(_token_key(connection)) <= 240 + assert not await redis.get(_lock_key(connection)) -async def test_mint_fills_the_cache_only_after_the_rotation_is_saved(token_endpoint): +async def test_mint_fills_the_cache_only_after_the_rotation_is_saved(token_endpoint, monkeypatch): token_endpoint.response = {"access_token": "at-2", "refresh_token": "rt-new", "expires_in": 300} + connection = _connection() - class UnsavableGrant(GrantStub): - def save_refresh_token(self, rotated: str) -> None: - raise RuntimeError("rotation write failed") + def _unsavable(self, rotated: str) -> None: + raise RuntimeError("rotation write failed") + monkeypatch.setattr(OauthConnection, "_save_refresh_token", _unsavable) with pytest.raises(RuntimeError, match="rotation write failed"): - await _client(token_endpoint).mint_access_token(key="grant-1", grant=UnsavableGrant()) + await _client().mint_access_token(connection=connection) redis = get_client() - assert not await redis.get(_TOKEN_KEY) - assert not await redis.get(_LOCK_KEY) - - -async def test_mint_surfaces_the_grants_own_load_error(token_endpoint): - class GoneGrant(GrantStub): - def load_refresh_token(self) -> str: - raise LookupError("the grant row is gone") - - with pytest.raises(LookupError, match="the grant row is gone"): - await _client(token_endpoint).mint_access_token(key="grant-1", grant=GoneGrant()) - - assert not token_endpoint.requests - assert not await get_client().get(_LOCK_KEY) + assert not await redis.get(_token_key(connection)) + assert not await redis.get(_lock_key(connection)) async def test_mint_losing_the_lock_polls_for_the_winners_token(token_endpoint): + connection = _connection() redis = get_client() - await redis.set(_LOCK_KEY, "1") + await redis.set(_lock_key(connection), "1") async def _winner_finishes(): - await redis.set(_TOKEN_KEY, "at-winner") - await redis.delete(_LOCK_KEY) + await redis.set(_token_key(connection), "at-winner") + await redis.delete(_lock_key(connection)) winner = asyncio.create_task(_winner_finishes()) - token = await _client(token_endpoint).mint_access_token(key="grant-1", grant=UntouchedGrant()) + token = await _client().mint_access_token(connection=connection) await winner assert token == "at-winner" @@ -145,32 +134,30 @@ async def _winner_finishes(): async def test_mint_times_out_loudly_when_the_lock_never_frees(token_endpoint): - await get_client().set(_LOCK_KEY, "1") + connection = _connection() + await get_client().set(_lock_key(connection), "1") with pytest.raises(OauthRefreshError, match="concurrent refresh"): - await _client(token_endpoint, mint_wait_attempts=3).mint_access_token( - key="grant-1", grant=UntouchedGrant() - ) + await _client(mint_wait_attempts=3).mint_access_token(connection=connection) async def test_mint_refresh_rejection_evicts_and_raises(token_endpoint): token_endpoint.status = 400 + connection = _connection() - grant = GrantStub() with pytest.raises(OauthRefreshError, match="HTTP 400"): - await _client(token_endpoint).mint_access_token(key="grant-1", grant=grant) - - assert not grant.saved + await _client().mint_access_token(connection=connection) + assert OauthConnection.get(connection.id).refresh_token.decrypt() == "rt-old" redis = get_client() - assert not await redis.get(_TOKEN_KEY) - assert not await redis.get(_LOCK_KEY) + assert not await redis.get(_token_key(connection)) + assert not await redis.get(_lock_key(connection)) async def test_mint_refresh_uses_basic_auth(token_endpoint): - await _client(token_endpoint, basic_auth=True).mint_access_token( - key="grant-1", grant=GrantStub() - ) + connection = _connection() + + await _client(basic_auth=True).mint_access_token(connection=connection) assert token_endpoint.authorizations[0].startswith("Basic ") # Basic auth keeps the client credentials out of the form body. @@ -178,8 +165,18 @@ async def test_mint_refresh_uses_basic_auth(token_endpoint): assert "client_secret" not in token_endpoint.requests[0] +async def test_disconnect_drops_the_connection_and_the_cached_token(token_endpoint): + connection = _connection() + await get_client().set(_token_key(connection), "at-cached") + + await _client().disconnect(connection) + + assert not OauthConnection.get(connection.id) + assert not await get_client().get(_token_key(connection)) + + async def test_connect_roundtrip_exchanges_with_basic_auth(token_endpoint): - url = await _client(token_endpoint, basic_auth=True).begin_connect( + url = await _client(basic_auth=True).begin_connect( redirect_uri=_REDIRECT_URI, scopes=("profile.read", "posts.write"), context={"account": "a-1"}, @@ -192,18 +189,15 @@ async def test_connect_roundtrip_exchanges_with_basic_auth(token_endpoint): assert params["audience"] == "api" assert params["code_challenge_method"] == "S256" - # Completion needs only the provider: the begun flow's client identity - # rides the stashed state. - tokens, context = await OauthClient( - provider=_PROVIDER, - http_factory=lambda: httpx.AsyncClient( - transport=httpx.MockTransport(token_endpoint.handler) - ), - ).complete_connect(state=params["state"], code="code-1") + # Completion needs only the state: the begun flow's provider and client + # identity ride the stash. + tokens, pending = await complete_connect(state=params["state"], code="code-1") assert tokens["refresh_token"] == "rt-1" - assert context["account"] == "a-1" - assert context["client_id"] == "client-123" + assert pending["account"] == "a-1" + assert pending["provider"] == _PROVIDER + assert pending["scopes"] == ["profile.read", "posts.write"] + assert pending["client_id"] == "client-123" exchange = token_endpoint.requests[0] assert exchange["grant_type"] == "authorization_code" assert exchange["code"] == "code-1" @@ -215,8 +209,8 @@ async def test_connect_roundtrip_exchanges_with_basic_auth(token_endpoint): async def test_complete_connect_requires_a_refresh_token(token_endpoint): token_endpoint.response = {"access_token": "at-1", "expires_in": 3600} - url = await _client(token_endpoint).begin_connect(redirect_uri=_REDIRECT_URI) + url = await _client().begin_connect(redirect_uri=_REDIRECT_URI) state = dict(parse_qsl(urlparse(url).query))["state"] with pytest.raises(OauthExchangeError, match="no refresh token"): - await _client(token_endpoint).complete_connect(state=state, code="code-1") + await complete_connect(state=state, code="code-1") diff --git a/backend/tests/test_secrets.py b/backend/tests/test_secrets.py index 8b31bcf3..bdfe3650 100644 --- a/backend/tests/test_secrets.py +++ b/backend/tests/test_secrets.py @@ -4,10 +4,11 @@ import pytest from druks.accounts.constants import SYSTEM_ACCOUNT_ID from druks.core.models import Uuid7Pk -from druks.mcp.models import McpOauthGrant, McpServer +from druks.mcp.models import McpClientRegistration, McpServer from druks.models import Base from druks.secrets.exceptions import SecretDecryptError from druks.secrets.fields import EncryptedJsonField +from druks.services.models import OauthConnection from druks.settings import load_settings from pydantic import ValidationError from sqlalchemy import text @@ -35,16 +36,23 @@ def _set_key(monkeypatch, tmp_path, value: str) -> None: monkeypatch.setenv("DRUKS_CONFIG", str(config_path)) -def _store_grant(refresh_token: str = "rt-secret", client_secret: str = "") -> McpOauthGrant: - return McpOauthGrant.store( - server_name="notion", +def _store_grant(refresh_token: str = "rt-secret", client_secret: str = "") -> OauthConnection: + server = McpServer.get_for_name("notion") or McpServer.create( + name="notion", url="https://mcp.notion.test/sse" + ) + McpClientRegistration.store( + server_id=server.id, account_id=SYSTEM_ACCOUNT_ID, - refresh_token=refresh_token, token_endpoint="https://auth.test/token", - resource="https://mcp.notion.test/sse", client_id="client-123", client_secret=client_secret, ) + return OauthConnection.create( + provider="mcp:notion", + account_id=SYSTEM_ACCOUNT_ID, + refresh_token=refresh_token, + scopes=[], + ) def test_stored_secrets_are_ciphertext_and_reads_restore_them(druks_db): @@ -65,9 +73,10 @@ def test_grant_secret_halves_round_trip(druks_db): _store_grant(refresh_token="rt-secret", client_secret="cs-secret") druks_db.expire_all() - grant = McpOauthGrant.get_for_account("notion", SYSTEM_ACCOUNT_ID) + grant = OauthConnection.list_for_account("mcp:notion", SYSTEM_ACCOUNT_ID)[0] + registration = McpClientRegistration.get_for_account("notion", SYSTEM_ACCOUNT_ID) assert grant.refresh_token.decrypt() == "rt-secret" - assert grant.client_secret.decrypt() == "cs-secret" + assert registration.client_secret.decrypt() == "cs-secret" def test_loaded_secrets_are_lazy_and_redacted(monkeypatch, tmp_path, druks_db): @@ -161,16 +170,25 @@ def test_ciphertext_is_bound_to_its_column(druks_db): McpServer.create(name="linear", url="https://mcp.linear.app/sse", token=_TOKEN) _store_grant(refresh_token="rt-secret", client_secret="cs-secret") druks_db.execute( - text("UPDATE mcp_oauth_grants SET refresh_token = (SELECT token FROM mcp_servers)") + text( + "UPDATE oauth_connections SET refresh_token =" + " (SELECT token FROM mcp_servers WHERE name = 'linear')" + ) + ) + druks_db.execute( + text( + "UPDATE mcp_client_registrations SET client_secret =" + " (SELECT refresh_token FROM oauth_connections WHERE provider = 'mcp:notion')" + ) ) - druks_db.execute(text("UPDATE mcp_oauth_grants SET client_secret = refresh_token")) druks_db.expire_all() - grant = McpOauthGrant.get_for_account("notion", SYSTEM_ACCOUNT_ID) + grant = OauthConnection.list_for_account("mcp:notion", SYSTEM_ACCOUNT_ID)[0] + registration = McpClientRegistration.get_for_account("notion", SYSTEM_ACCOUNT_ID) with pytest.raises(SecretDecryptError): grant.refresh_token.decrypt() with pytest.raises(SecretDecryptError): - grant.client_secret.decrypt() + registration.client_secret.decrypt() def test_prepended_key_still_decrypts(monkeypatch, tmp_path, druks_db): @@ -185,7 +203,7 @@ def test_prepended_key_still_decrypts(monkeypatch, tmp_path, druks_db): druks_db.expire_all() assert McpServer.get_for_name("linear").token.decrypt() == _TOKEN assert ( - McpOauthGrant.get_for_account("notion", SYSTEM_ACCOUNT_ID).refresh_token.decrypt() + OauthConnection.list_for_account("mcp:notion", SYSTEM_ACCOUNT_ID)[0].refresh_token.decrypt() == "rt-secret" ) diff --git a/backend/tests/test_services.py b/backend/tests/test_services.py index 1aa0f2fb..50d3dd9e 100644 --- a/backend/tests/test_services.py +++ b/backend/tests/test_services.py @@ -448,3 +448,269 @@ class Settings(BaseModel): with pytest.raises(TypeError, match="no OAuth endpoints"): Plain.get_oauth_client() + + +def test_with_scopes_declares_the_union_and_reads_connections(declared_services, monkeypatch): + from druks.services import Service + from pydantic import BaseModel, SecretStr + + class Acme(Service): + name = "acme" + title = "Acme OAuth app" + authorization_endpoint = "https://acme.test/authorize" + token_endpoint = "https://acme.test/token" + + class Settings(BaseModel): + client_id: str + client_secret: SecretStr + + class NightWatch: + name = "night_watch" + acme = Acme.with_scopes("profile.read", "posts.write") + + class Digest: + name = "digest" + acme = Acme.with_scopes("profile.read") + + monkeypatch.setattr("druks.services.base.iter_extensions", lambda: [NightWatch, Digest]) + + assert NightWatch.acme.scopes == ("profile.read", "posts.write") + assert Acme.required_scopes() == ("posts.write", "profile.read") + assert [declaration.label for declaration in Acme.declarations()] == [ + "night_watch.acme", + "digest.acme", + ] + + from druks.accounts.constants import SYSTEM_ACCOUNT_ID + from druks.services.models import OauthConnection + + row = OauthConnection.create( + provider="acme", account_id=SYSTEM_ACCOUNT_ID, refresh_token="rt-1", scopes=["profile.read"] + ) + connections = NightWatch.acme.list_for_account(SYSTEM_ACCOUNT_ID) + assert [connection.id for connection in connections] == [row.id] + assert connections[0].scopes == ["profile.read"] + assert NightWatch.acme.get(row.id).id == row.id + assert not NightWatch.acme.get("missing") + + +def test_with_scopes_requires_oauth_endpoints(declared_services): + from druks.services import Service + from pydantic import BaseModel, SecretStr + + class Plain(Service): + name = "plain_no_oauth" + title = "Plain" + + class Settings(BaseModel): + api_key: SecretStr + + with pytest.raises(TypeError, match="no OAuth endpoints"): + Plain.with_scopes("profile.read") + + +# --- The connect door: /api/oauth ------------------------------------------ + + +@pytest.fixture +def acme(declared_services, monkeypatch): + from druks.services import Service + from pydantic import BaseModel, SecretStr + + class Acme(Service): + name = "acme" + title = "Acme OAuth app" + authorization_endpoint = "https://acme.test/authorize" + token_endpoint = "https://acme.test/token" + + class Settings(BaseModel): + client_id: str + client_secret: SecretStr + + class NightWatch: + name = "night_watch" + acme = Acme.with_scopes("profile.read", "posts.write") + + monkeypatch.setattr("druks.services.base.iter_extensions", lambda: [NightWatch]) + tokens = { + "access_token": "at-1", + "refresh_token": "rt-1", + "expires_in": 3600, + "scope": "profile.read posts.write", + } + monkeypatch.setattr( + "druks.services.oauth._http", + lambda: httpx.AsyncClient( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json=tokens)) + ), + ) + return Acme + + +def test_oauth_connect_redirects_to_consent_with_the_scope_union( + tmp_path, acme, druks_db, monkeypatch +): + from urllib.parse import parse_qsl, urlparse + + from druks.testing import configure_app_for_test + + ServiceIdentity.connect( + "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} + ) + settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) + with TestClient(configure_app_for_test(settings=settings)) as client: + response = client.get("/api/oauth/acme/connect", follow_redirects=False) + + assert response.status_code == 307 + consent = urlparse(response.headers["location"]) + params = dict(parse_qsl(consent.query)) + assert response.headers["location"].startswith("https://acme.test/authorize?") + assert params["scope"] == "posts.write profile.read" + assert params["redirect_uri"] == "https://druks.example/api/oauth/callback" + + +def test_oauth_connect_guards(tmp_path, acme, druks_db): + from druks.testing import configure_app_for_test + + with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: + assert client.get("/api/oauth/github/connect").status_code == 404 + # No urls.endpoint configured. + assert client.get("/api/oauth/acme/connect").status_code == 409 + + settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) + with TestClient(configure_app_for_test(settings=settings)) as client: + # The client credentials are not connected yet. + response = client.get("/api/oauth/acme/connect", follow_redirects=False) + assert response.status_code == 409 + assert "not connected" in response.json()["detail"] + + +async def test_oauth_callback_creates_and_reconnects_a_connection( + tmp_path, acme, druks_db, monkeypatch +): + from urllib.parse import parse_qsl, urlparse + + import druks.redis + from druks.redis import close_client, get_client + from druks.services.models import OauthConnection + from druks.testing import configure_app_for_test + + ServiceIdentity.connect( + "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} + ) + await close_client() + settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) + with TestClient(configure_app_for_test(settings=settings)) as client: + consent = client.get("/api/oauth/acme/connect", follow_redirects=False).headers["location"] + state = dict(parse_qsl(urlparse(consent).query))["state"] + page = client.get("/api/oauth/callback", params={"state": state, "code": "c-1"}) + assert page.status_code == 200 + assert "BroadcastChannel('druks-service-connect')" in page.text + # The state is single-use; a denied consent lands loudly. + assert ( + client.get("/api/oauth/callback", params={"state": state, "code": "c-1"}).status_code + == 400 + ) + assert ( + client.get( + "/api/oauth/callback", params={"state": "s", "code": "c", "error": "denied"} + ).status_code + == 400 + ) + + [connection] = OauthConnection.list_for_provider("acme") + assert connection.refresh_token.decrypt() == "rt-1" + assert connection.scopes == ["profile.read", "posts.write"] + + # Reconsent through the same connection replaces its tokens and + # evicts the stale cached access token. + stale_key = f"acme:access_token:{connection.id}" + client.get("/api/oauth/acme/connect?connection=zzz", follow_redirects=False) + reconnect = client.get( + f"/api/oauth/acme/connect?connection={connection.id}", follow_redirects=False + ) + state = dict(parse_qsl(urlparse(reconnect.headers["location"]).query))["state"] + finish = client.get("/api/oauth/callback", params={"state": state, "code": "c-2"}) + assert finish.status_code == 200 + assert len(OauthConnection.list_for_provider("acme")) == 1 + + druks.redis._client = None + assert not await get_client().get(stale_key) + + +def test_oauth_connect_rejects_an_unknown_reconnect_target(tmp_path, acme, druks_db): + from druks.testing import configure_app_for_test + + ServiceIdentity.connect( + "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} + ) + settings = make_settings(tmp_path, urls={"endpoint": "https://druks.example"}) + with TestClient(configure_app_for_test(settings=settings)) as client: + assert client.get("/api/oauth/acme/connect?connection=zzz").status_code == 404 + + +def test_connections_list_and_revoke(tmp_path, acme, druks_db): + from druks.accounts.models import Account + from druks.services.models import OauthConnection + from druks.testing import configure_app_for_test + + me = Account.get_or_create("op@example.com") + with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: + row = OauthConnection.create( + provider="acme", account_id=me.id, refresh_token="rt-1", scopes=["profile.read"] + ) + [listed] = client.get("/api/oauth/connections").json() + assert listed["id"] == row.id + assert listed["provider"] == "acme" + assert listed["scopes"] == ["profile.read"] + + assert client.delete(f"/api/oauth/connections/{row.id}").status_code == 204 + assert not OauthConnection.list_for_provider("acme") + assert client.delete(f"/api/oauth/connections/{row.id}").status_code == 404 + + +def test_replacing_the_client_credentials_deletes_its_connections(tmp_path, acme, druks_db): + from druks.accounts.constants import SYSTEM_ACCOUNT_ID + from druks.services.models import OauthConnection + from druks.testing import configure_app_for_test + + OauthConnection.create( + provider="acme", account_id=SYSTEM_ACCOUNT_ID, refresh_token="rt-old", scopes=[] + ) + + with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: + response = client.post( + "/api/services/acme", json={"client_id": "id-2", "client_secret": "sec-2"} + ) + assert response.status_code == 200 + + # The new client can never refresh the old client's connections. + assert not OauthConnection.list_for_provider("acme") + + +def test_list_serves_the_connections_beside_the_declared_union(tmp_path, acme, druks_db): + from druks.accounts.constants import SYSTEM_ACCOUNT_ID + from druks.services.models import OauthConnection + from druks.testing import configure_app_for_test + + def entry(client, name="acme"): + return next(e for e in client.get("/api/services").json() if e["name"] == name) + + with TestClient(configure_app_for_test(settings=make_settings(tmp_path))) as client: + assert entry(client, "github")["isOauth"] is False + before = entry(client) + assert before["isOauth"] is True + assert before["connections"] == [] + assert before["requiredScopes"] == ["posts.write", "profile.read"] + assert before["usedBy"] == ["night_watch.acme"] + + row = OauthConnection.create( + provider="acme", + account_id=SYSTEM_ACCOUNT_ID, + refresh_token="rt-1", + scopes=["profile.read"], + ) + [connection] = entry(client)["connections"] + assert connection["id"] == row.id + assert connection["scopes"] == ["profile.read"] + assert connection["connectedAt"] diff --git a/docs/writing-an-extension.md b/docs/writing-an-extension.md index fb3ff5c9..ef3daa36 100644 --- a/docs/writing-an-extension.md +++ b/docs/writing-an-extension.md @@ -318,9 +318,9 @@ from druks.browser import BrowserSession from druks.extensions import Extension -class XMe(Extension): - name = "x_me" - x = BrowserSession(site="x.com", persist=True) +class NightWatch(Extension): + name = "night_watch" + acme = BrowserSession(site="acme.example", persist=True) ``` A workflow borrows the logged-in browser as a playwright handle — the @@ -330,13 +330,13 @@ dies with the block; a ``persist`` session is exported and stored back first): ```python -async with XMe.x.playwright() as browser: +async with NightWatch.acme.playwright() as browser: page = await browser.new_page() # opened on the logged-in context - await page.goto("https://x.com/home") + await page.goto("https://acme.example/home") ``` ``playwright()`` yields the logged-in browser context; pages you open on -it carry the session. ``XMe.x.cdp()`` is the same borrow yielding the raw CDP +it carry the session. ``NightWatch.acme.cdp()`` is the same borrow yielding the raw CDP url, for any other client — an existing test suite, raw CDP, your own wrapper. @@ -681,9 +681,10 @@ router for its own resource and the question never comes up. A service identity is the appliance's own registered app at an external provider — one per deployment, keyed by a service string; the platform's -GitHub App is the first one. Per-user OAuth grants are not service identities -(keep those on your own rows), and a credential only your extension posts with -belongs in your extension settings instead. +GitHub App is the first one. OAuth grants are not service identities — the +platform stores those when the operator connects (see "Connect provider +accounts") — and a credential only your extension posts with belongs in your +extension settings instead. Declare one class in `services.py` and the platform does the rest: it renders the connect card in Settings, verifies and stores the paste (`SecretStr` @@ -737,13 +738,9 @@ blast-radius control. ## Connect provider accounts (OAuth) -`OauthClient` runs the OAuth 2.0 authorization-code + PKCE flow. Use it for a -provider with fixed endpoints and a registered client. It mints access tokens -and keeps refresh-token rotation safe. Your extension stores each grant on its -own rows. The platform stores no grants. - -Declare the endpoints on the service that holds the client credentials. The -`Settings` model must have `client_id` and `client_secret` fields: +Declare the OAuth endpoints on the service that holds the client +credentials. The `Settings` model must have `client_id` and `client_secret` +fields: ```python class Acme(Service): @@ -759,76 +756,45 @@ class Acme(Service): client_secret: SecretStr = Field(title="Client secret") ``` -`Acme.get_oauth_client()` returns a configured client for the connected -identity. Call `begin_connect` from your connect route. Call `complete_connect` -from your callback route: +Declare your extension's use of the service, with the scopes your calls +need: ```python -url = await Acme.get_oauth_client().begin_connect( - redirect_uri="https://druks.example/api/acme/oauth/callback", - scopes=("profile.read", "posts.write"), - context={"account_id": account_id}, -) -# ... the operator consents; the provider redirects back with state + code ... -tokens, context = await Acme.get_oauth_client().complete_connect(state=state, code=code) -AcmeGrant.store(account_id=context["account_id"], refresh_token=tokens["refresh_token"]) -``` - -Scopes belong to one authorization, not to the service. Each `begin_connect` -call asks for its own scopes. The provider registration sets the ceiling. The -grant keeps the scopes the user approved. - -`begin_connect` stores the pending exchange in Redis and returns the consent -URL. The state is single-use and expires after a short time. `complete_connect` -consumes the state and exchanges the code for tokens. It rejects a token -response without a `refresh_token`, because a grant must work offline. It -raises `OauthExchangeError` when the flow is denied or expired. It returns your -`context` from begin time, with the flow's client identity merged in. - -When a run needs the provider, call `mint_access_token`. The engine serves -tokens from a Redis cache. It lets only one refresher run for each `key`. This -is necessary: two refreshes at the same time can make the provider revoke the -whole grant. The engine raises `OauthRefreshError` when the grant does not -refresh. Then ask the operator to connect again. - -```python -grant = AcmeGrant.get_for_account(account_id) -token = await Acme.get_oauth_client().mint_access_token(key=account_id, grant=grant) -``` - -`grant` is your own grant row. It carries the two verbs that connect the -engine to your storage: - -```python -class AcmeGrant(Base): - ... - - def load_refresh_token(self) -> str: - # Runs under the refresh lock. Another process may have rotated and - # committed. Re-read the row; do not trust the identity map. - fresh = db_session().scalars( - select(AcmeGrant) - .where(AcmeGrant.id == self.id) - .execution_options(populate_existing=True) - ).one() - return fresh.secrets["refresh_token"] - - def save_refresh_token(self, rotated: str) -> None: - # The provider has already invalidated the old token. Commit on an own - # session, never on the enclosing step transaction. A later rollback - # must not lose the new token. - with Session(db_session().get_bind()) as session: - grant = session.scalars( - select(AcmeGrant).where(AcmeGrant.id == self.id) - ).one() - grant.secrets["refresh_token"] = rotated - session.commit() -``` - -`secrets` is an `EncryptedJsonField` column on your grant row (see -[models](#models-and-migrations)). The refresh token is ciphertext at rest. -When the operator disconnects, delete your row and evict the cached access -token: `await Acme.get_oauth_client().evict_access_token(account_id)`. +class NightWatch(Extension): + name = "night_watch" + acme = Acme.with_scopes("profile.read", "posts.write") +``` + +A *connection* is one signed-in provider account. A user can hold many per +provider — one per mailbox, handle, or workspace — and the platform stores +each one: the refresh token, the granted scopes, the owner. Your workflow +code reads them through the declaration and mints per connection: + +```python +for connection in NightWatch.acme.list_for_account(account_id): + token = await connection.mint_access_token() +``` + +`NightWatch.acme.get(connection_id)` returns one connection when your own +row stored its id. + +Your UI starts a sign-in by opening `/api/oauth/acme/connect` — the +platform runs the consent with the union of every installed extension's +declared scopes and stores the connection for the signed-in user. To widen +an existing connection's scopes, open +`/api/oauth/acme/connect?connection=Connected {new Date(service.connectedAt).toLocaleString()}
)} + {service.isOauth &&The accounts you have signed in to. Revoke one here.
+No connections yet.
} + {connections.length > 0 && ( +