diff --git a/backend/druks/mcp/models.py b/backend/druks/mcp/models.py index 02f945fa..e6b5f46e 100644 --- a/backend/druks/mcp/models.py +++ b/backend/druks/mcp/models.py @@ -2,18 +2,22 @@ from datetime import datetime from typing import Any -from sqlalchemy import Boolean, ForeignKey, String, UniqueConstraint, select +from sqlalchemy import Boolean, ForeignKey, String, UniqueConstraint, select, update 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 +from druks.database import db_session, get_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, UnresolvedGrantAccountError +from druks.mcp.exceptions import ( + InvalidServerNameError, + MissingGrantError, + UnresolvedGrantAccountError, +) from druks.models import Base from druks.secrets.fields import EncryptedJsonField, EncryptedTextField, Secret @@ -258,6 +262,40 @@ def store( ).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 2b6a5468..23eceeff 100644 --- a/backend/druks/mcp/oauth.py +++ b/backend/druks/mcp/oauth.py @@ -4,7 +4,7 @@ from sqlalchemy import select, update from sqlalchemy.dialects.postgresql import insert as pg_insert -from druks.database import db_session, get_session +from druks.database import db_session from druks.mcp.constants import OAUTH_CALLBACK_PATH, OAUTH_PROVIDER from druks.mcp.enums import IdentityMode from druks.mcp.exceptions import GrantRefreshError, MissingGrantError, OauthConnectError @@ -244,41 +244,6 @@ async def mint_access_token(name: str, account_id: str) -> str: grant = McpOauthGrant.get_for_account(name, account_id) if not grant: raise MissingGrantError(name, account_id) - - def load_refresh_token() -> 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 == grant.id) - .execution_options(populate_existing=True) - ) - .one_or_none() - ) - if not fresh: - raise MissingGrantError(name, 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(refresh_token: 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 == grant.id) - .values(refresh_token=refresh_token) - ) - session.commit() - # Keep the enclosing transaction's copy true as well. - grant.refresh_token = refresh_token - db_session().flush() - client = OauthClient( provider=OAUTH_PROVIDER, token_endpoint=grant.token_endpoint, @@ -292,10 +257,6 @@ def save_refresh_token(refresh_token: str) -> None: http_factory=_http, ) try: - return await client.mint_access_token( - key=f"{name}:{account_id}", - load_refresh_token=load_refresh_token, - save_refresh_token=save_refresh_token, - ) + return await client.mint_access_token(key=f"{name}:{account_id}", grant=grant) except OauthRefreshError as error: raise GrantRefreshError(name, error.reason) from error diff --git a/backend/druks/services/oauth.py b/backend/druks/services/oauth.py index ea3b1fe0..dba3a525 100644 --- a/backend/druks/services/oauth.py +++ b/backend/druks/services/oauth.py @@ -61,8 +61,8 @@ class OauthClient: ``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, reached - through the two callables mint takes. A ``Service`` with declared OAuth + 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. @@ -213,13 +213,7 @@ async def complete_connect(self, *, state: str, code: str) -> tuple[dict, dict]: ) return tokens, pending - async def mint_access_token( - self, - *, - key: str, - load_refresh_token: Callable[[], str], - save_refresh_token: Callable[[str], None], - ) -> str: + 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 @@ -228,16 +222,19 @@ async def mint_access_token( backstop a live refresh cannot outlive). Losers poll for the winner's cache fill, for about one token-endpoint round trip, then fail loudly. - ``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 + ``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. - ``save_refresh_token(token)`` 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.""" + ``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.""" redis = get_client() token_key = f"{self.provider}:access_token:{key}" lock_key = f"{self.provider}:refresh_lock:{key}" @@ -255,7 +252,7 @@ async def mint_access_token( try: data = { "grant_type": "refresh_token", - "refresh_token": load_refresh_token(), + "refresh_token": grant.load_refresh_token(), **self.extra_token_params, } async with self._http() as http: @@ -286,7 +283,7 @@ async def mint_access_token( self.provider, "the token endpoint returned no access token" ) if tokens.get("refresh_token"): - save_refresh_token(tokens["refresh_token"]) + grant.save_refresh_token(tokens["refresh_token"]) try: ttl = int(tokens.get("expires_in", 3600)) - OAUTH_TOKEN_TTL_SKEW_SECONDS except (TypeError, ValueError) as error: diff --git a/backend/tests/test_oauth_client.py b/backend/tests/test_oauth_client.py index ef59eff1..09828a3f 100644 --- a/backend/tests/test_oauth_client.py +++ b/backend/tests/test_oauth_client.py @@ -48,18 +48,34 @@ def _client(token_endpoint: FakeTokenEndpoint, **overrides) -> OauthClient: return OauthClient(**kwargs) -def _fail_save(token: str) -> None: - pytest.fail("nothing rotated") +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) + + +class UntouchedGrant(GrantStub): + """A grant the mint under test must never read or rotate.""" + + 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") async def test_mint_serves_the_cache_without_a_refresh(token_endpoint): await get_client().set(_TOKEN_KEY, "at-cached") - token = await _client(token_endpoint).mint_access_token( - key="grant-1", - load_refresh_token=lambda: pytest.fail("cache hit reads no grant"), - save_refresh_token=_fail_save, - ) + token = await _client(token_endpoint).mint_access_token(key="grant-1", grant=UntouchedGrant()) assert token == "at-cached" assert not token_endpoint.requests @@ -67,16 +83,12 @@ 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} - saved: list[str] = [] + grant = GrantStub() - token = await _client(token_endpoint).mint_access_token( - key="grant-1", - load_refresh_token=lambda: "rt-old", - save_refresh_token=saved.append, - ) + token = await _client(token_endpoint).mint_access_token(key="grant-1", grant=grant) assert token == "at-2" - assert saved == ["rt-new"] + assert grant.saved == ["rt-new"] refresh = token_endpoint.requests[0] assert refresh["grant_type"] == "refresh_token" assert refresh["refresh_token"] == "rt-old" @@ -92,21 +104,30 @@ async def test_mint_refreshes_persists_rotation_and_fills_with_skewed_ttl(token_ async def test_mint_fills_the_cache_only_after_the_rotation_is_saved(token_endpoint): token_endpoint.response = {"access_token": "at-2", "refresh_token": "rt-new", "expires_in": 300} - def save_refresh_token(token: str) -> None: - raise RuntimeError("rotation write failed") + class UnsavableGrant(GrantStub): + def save_refresh_token(self, rotated: str) -> None: + raise RuntimeError("rotation write failed") with pytest.raises(RuntimeError, match="rotation write failed"): - await _client(token_endpoint).mint_access_token( - key="grant-1", - load_refresh_token=lambda: "rt-old", - save_refresh_token=save_refresh_token, - ) + await _client(token_endpoint).mint_access_token(key="grant-1", grant=UnsavableGrant()) 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) + + async def test_mint_losing_the_lock_polls_for_the_winners_token(token_endpoint): redis = get_client() await redis.set(_LOCK_KEY, "1") @@ -116,11 +137,7 @@ async def _winner_finishes(): await redis.delete(_LOCK_KEY) winner = asyncio.create_task(_winner_finishes()) - token = await _client(token_endpoint).mint_access_token( - key="grant-1", - load_refresh_token=lambda: "rt-old", - save_refresh_token=_fail_save, - ) + token = await _client(token_endpoint).mint_access_token(key="grant-1", grant=UntouchedGrant()) await winner assert token == "at-winner" @@ -132,21 +149,18 @@ async def test_mint_times_out_loudly_when_the_lock_never_frees(token_endpoint): with pytest.raises(OauthRefreshError, match="concurrent refresh"): await _client(token_endpoint, mint_wait_attempts=3).mint_access_token( - key="grant-1", - load_refresh_token=lambda: "rt-old", - save_refresh_token=_fail_save, + key="grant-1", grant=UntouchedGrant() ) async def test_mint_refresh_rejection_evicts_and_raises(token_endpoint): token_endpoint.status = 400 + grant = GrantStub() with pytest.raises(OauthRefreshError, match="HTTP 400"): - await _client(token_endpoint).mint_access_token( - key="grant-1", - load_refresh_token=lambda: "rt-old", - save_refresh_token=_fail_save, - ) + await _client(token_endpoint).mint_access_token(key="grant-1", grant=grant) + + assert not grant.saved redis = get_client() assert not await redis.get(_TOKEN_KEY) @@ -155,9 +169,7 @@ async def test_mint_refresh_rejection_evicts_and_raises(token_endpoint): async def test_mint_refresh_uses_basic_auth(token_endpoint): await _client(token_endpoint, basic_auth=True).mint_access_token( - key="grant-1", - load_refresh_token=lambda: "rt-old", - save_refresh_token=lambda token: None, + key="grant-1", grant=GrantStub() ) assert token_endpoint.authorizations[0].startswith("Basic ") diff --git a/docs/writing-an-extension.md b/docs/writing-an-extension.md index 1c34ac54..fb3ff5c9 100644 --- a/docs/writing-an-extension.md +++ b/docs/writing-an-extension.md @@ -792,37 +792,37 @@ whole grant. The engine raises `OauthRefreshError` when the grant does not refresh. Then ask the operator to connect again. ```python -token = await Acme.get_oauth_client().mint_access_token( - key=account_id, - load_refresh_token=load_refresh_token, - save_refresh_token=save_refresh_token, -) +grant = AcmeGrant.get_for_account(account_id) +token = await Acme.get_oauth_client().mint_access_token(key=account_id, grant=grant) ``` -The two callables connect the engine to your storage: +`grant` is your own grant row. It carries the two verbs that connect the +engine to your storage: ```python -def load_refresh_token() -> str: - # Runs under the refresh lock. Another process may have rotated and - # committed. Re-read the row; do not trust the identity map. - grant = db_session().scalars( - select(AcmeGrant) - .where(AcmeGrant.account_id == account_id) - .execution_options(populate_existing=True) - ).one() - return grant.secrets["refresh_token"] - +class AcmeGrant(Base): + ... -def save_refresh_token(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.account_id == account_id) + 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() - grant.secrets["refresh_token"] = rotated - session.commit() + 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