From 39b96333d7ca4b39030e22ffa0fa77d1ba31c3f0 Mon Sep 17 00:00:00 2001 From: Paulo Date: Wed, 19 Aug 2026 16:36:03 +0200 Subject: [PATCH] The mint seam takes a grant, not two callables mint_access_token(key=..., grant=...) now takes one object carrying the two grant verbs: load_refresh_token() and save_refresh_token(rotated). The engine's real dependency is a grant it can read fresh under the refresh lock and rotate durably; two loose callables atomized that noun and let a loader for one row pair silently with a saver for another. The lock and commit contracts now live on the grant parameter's documentation. McpOauthGrant carries the two verbs itself; the MCP mint wrapper passes the row. Redis keys, lock and cache semantics, and all observable MCP behavior are unchanged. The engine tests use small grant stubs, keep every behavioral assertion, and add one the seam now affords: an error raised by a grant's own verb surfaces unmasked. --- backend/druks/mcp/models.py | 44 +++++++++++++-- backend/druks/mcp/oauth.py | 43 +-------------- backend/druks/services/oauth.py | 35 ++++++------ backend/tests/test_oauth_client.py | 86 +++++++++++++++++------------- docs/writing-an-extension.md | 50 ++++++++--------- 5 files changed, 133 insertions(+), 125 deletions(-) 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