The mint seam takes a grant, not two callables (#290)

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.

Author: czpython

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)

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

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:

backend/tests/test_oauth_client.py

@@ -48,35 +48,47 @@ 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
 
 
 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 ")