The mint seam takes a grant, not two callables

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 ")