Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions backend/druks/mcp/oauth.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,8 @@ def _http() -> httpx.AsyncClient:


def get_connection(name: str, account_id: str) -> OauthConnection | None:
# One connection per (server, account) — MCP's policy over the shared table.
# One live connection per (server, account) — MCP's policy over the
# shared table. Revoked rows stay behind as history.
rows = OauthConnection.list_for_account(grant_provider(name), account_id)
return rows[0] if rows else None

Expand Down Expand Up @@ -276,10 +277,10 @@ async def evict_access_token(name: str, account_id: str) -> None:
await OauthClient(provider=grant_provider(name)).evict_access_token(connection.id)


async def disconnect(name: str, account_id: str) -> None:
async def disconnect(name: str, account_id: str, *, reason: str = "user") -> None:
connection = get_connection(name, account_id)
if connection:
await OauthClient(provider=grant_provider(name)).disconnect(connection)
await OauthClient(provider=grant_provider(name)).disconnect(connection, reason=reason)
registration = McpClientRegistration.get_for_account(name, account_id)
if registration:
registration.delete()
Expand Down
7 changes: 3 additions & 4 deletions backend/druks/mcp/routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -163,11 +163,10 @@ 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")
connections = oauth.list_connections(name)
# Revoke before the server row goes — the registration lookup needs it.
for connection in oauth.list_connections(name):
await oauth.disconnect(name, connection.account_id, reason="server_removed")
server.delete()
for connection in connections:
await oauth.evict_access_token(name, connection.account_id)
connection.delete()


@router.post("/{name}/connect", response_model=ConnectMcpServerResponse)
Expand Down
4 changes: 2 additions & 2 deletions backend/druks/services/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ async def get_access_token(self, scopes: tuple[str, ...] = (), cached: bool = Tr
)

async def disconnect(self) -> None:
await OauthClient(provider=self.service.name).disconnect(self.row)
await OauthClient(provider=self.service.name).disconnect(self.row, reason="user")


class ScopedService:
Expand Down Expand Up @@ -72,7 +72,7 @@ def list_for_account(self, account_id: str) -> list[Connection]:

def get(self, connection_id: str) -> Connection | None:
row = OauthConnection.get(connection_id)
if row and row.provider == self.service.name:
if row and row.provider == self.service.name and not row.revoked_at:
return Connection(self.service, row)


Expand Down
64 changes: 47 additions & 17 deletions backend/druks/services/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,11 @@ 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."""
engine rotates the refresh token on mint; nothing else writes here.

Revoking is a state, never a deletion: the consent happened, and the row
keeps its owner, identity, scopes, and dates forever. Only the refresh
token is cleared. ``reconnect`` returns a revoked row to life."""

__tablename__ = "oauth_connections"

Expand All @@ -67,6 +71,10 @@ class OauthConnection(Base, Uuid7Pk):
# consent; {} when the provider gave none.
identity: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict)
connected_at: Mapped[datetime] = mapped_column(default=Base.utc_now)
# NULL = live. Stamped with what revoked it: "user", "client_replaced",
# or "server_removed".
revoked_at: Mapped[datetime | None] = mapped_column(default=None)
revoked_reason: Mapped[str] = mapped_column(default="")

@classmethod
def get(cls, connection_id: str) -> "OauthConnection | None":
Expand Down Expand Up @@ -98,17 +106,28 @@ def list_for_account(cls, provider: str, account_id: str) -> "list[OauthConnecti
return list(
db_session().scalars(
select(cls)
.where(cls.provider == provider, cls.account_id == account_id)
.where(
cls.provider == provider,
cls.account_id == account_id,
cls.revoked_at.is_(None),
)
.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)))
def list_for_provider(
cls, provider: str, *, include_revoked: bool = False
) -> "list[OauthConnection]":
query = select(cls).where(cls.provider == provider)
if not include_revoked:
query = query.where(cls.revoked_at.is_(None))
return list(db_session().scalars(query))

@classmethod
def list_owned_by(cls, account_id: str | None) -> "list[OauthConnection]":
# The audit read: everything this account ever authorized, revoked
# rows included.
return list(
db_session().scalars(
select(cls).where(cls.account_id == account_id).order_by(cls.connected_at)
Expand All @@ -123,12 +142,17 @@ def reconnect(
if identity:
self.identity = identity
self.connected_at = Base.utc_now()
self.revoked_at = None
self.revoked_reason = ""
db_session().flush()

def delete(self) -> None:
session = db_session()
session.delete(self)
session.flush()
def revoke(self, reason: str) -> None:
# The consent's facts survive; only the secret is cleared. A second
# revoke keeps the first stamp.
self.revoked_at = self.revoked_at or Base.utc_now()
self.revoked_reason = self.revoked_reason or reason
self.refresh_token = ""
db_session().flush()

def _load_refresh_token(self) -> str:
# Under the refresh lock: another process may have rotated and
Expand All @@ -141,24 +165,30 @@ def _load_refresh_token(self) -> str:
.where(OauthConnection.id == self.id)
.execution_options(populate_existing=True)
)
.one_or_none()
.one()
)
if fresh:
return fresh.refresh_token.decrypt()
raise OauthRefreshError(self.provider, "the connection was removed mid-refresh")
if fresh.revoked_at:
raise OauthRefreshError(self.provider, "the connection was revoked mid-refresh")
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
# connection.
with get_session(db_session().get_bind()) as session:
session.execute(
stored = session.execute(
update(OauthConnection)
.where(OauthConnection.id == self.id)
.where(OauthConnection.id == self.id, OauthConnection.revoked_at.is_(None))
.values(refresh_token=rotated)
)
session.commit()
# Keep the enclosing transaction's copy true as well.
self.refresh_token = rotated
db_session().flush()
if not stored.rowcount:
# A revoke landed mid-refresh. Nothing secret outlives the
# consent at rest, so the rotated token is not stored.
raise OauthRefreshError(self.provider, "the connection was revoked mid-refresh")
# Expire the stale copy instead of assigning it. The next read loads
# the rotated value, and the enclosing transaction never writes this
# column at commit — a revoke that lands between the two commits
# must keep its cleared token.
db_session().expire(self, ["refresh_token"])
11 changes: 8 additions & 3 deletions backend/druks/services/oauth.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,10 @@ async def get_access_token(
scopes. ``cached=False`` skips the cache read for a full-lifetime
token, still electing one refresher and filling the cache for later
callers."""
if connection.revoked_at:
raise OauthRefreshError(
self.provider, "the connection is revoked; sign in again to restore it"
)
requested = tuple(sorted(scopes))
if requested and not set(requested) <= set(connection.scopes):
missing = ", ".join(sorted(set(requested) - set(connection.scopes)))
Expand Down Expand Up @@ -273,9 +277,10 @@ async def evict_access_token(self, connection_id: str) -> None:
async for key in redis.scan_iter(match=f"{self.provider}:access_token:{connection_id}*"):
await redis.delete(key)

async def disconnect(self, connection: OauthConnection) -> None:
"""Delete the connection and evict its cached access token."""
connection.delete()
async def disconnect(self, connection: OauthConnection, *, reason: str) -> None:
"""Revoke the connection and evict its cached access token. The row
and its facts survive; the refresh token dies with the consent."""
connection.revoke(reason)
await self.evict_access_token(connection.id)


Expand Down
29 changes: 17 additions & 12 deletions backend/druks/services/routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,8 @@ async def list_services() -> list[ServiceResponse]:
row = None
connections = []
if service.token_endpoint:
connections = OauthConnection.list_for_provider(service.name)
# The detail shows revoked connections as history beside the live.
connections = OauthConnection.list_for_provider(service.name, include_revoked=True)
entries.append(ServiceResponse.from_row(service, row, connections))
return entries

Expand All @@ -51,16 +52,16 @@ async def connect_service(name: str, payload: dict[str, str]) -> ServiceResponse
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.
# A replaced client can never refresh the old client's connections —
# revoke every live one; the consents stay on record.
client = OauthClient(provider=name)
for connection in OauthConnection.list_for_provider(name):
connection_id, account_id = connection.id, connection.account_id
await client.disconnect(connection)
await client.disconnect(connection, reason="client_replaced")
await publish(
"oauth.disconnected",
provider=name,
connection_id=connection_id,
account_id=account_id,
connection_id=connection.id,
account_id=connection.account_id,
)
return ServiceResponse.from_row(service, row)

Expand Down Expand Up @@ -132,8 +133,10 @@ async def oauth_callback(state: str = "", code: str = "", error: str = "") -> Re
raise HTTPException(
status_code=400, detail="The connection was removed while consent was open."
)
# Reconsent names the row, so it also returns a revoked one to life.
# A fresh sign-in never does — it always creates a new connection.
row.reconnect(refresh_token=tokens["refresh_token"], scopes=granted, identity=identity)
# A reconsent's narrower cached token must not serve until its TTL runs out.
# A token cached before this consent must not serve the new one.
await OauthClient(provider=provider).evict_access_token(row.id)
else:
row = OauthConnection.create(
Expand Down Expand Up @@ -170,11 +173,13 @@ 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}.")
provider, account_id = row.provider, row.account_id
await OauthClient(provider=provider).disconnect(row)
if row.revoked_at:
# Revoking is idempotent — the second delete finds the state true.
return
await OauthClient(provider=row.provider).disconnect(row, reason="user")
await publish(
"oauth.disconnected",
provider=provider,
connection_id=connection_id,
account_id=account_id,
provider=row.provider,
connection_id=row.id,
account_id=row.account_id,
)
2 changes: 2 additions & 0 deletions backend/druks/services/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@ class ConnectionResponse(BaseResponse):
scopes: list[str]
identity: dict[str, Any]
connected_at: datetime
revoked_at: datetime | None
revoked_reason: str


class ServiceResponse(BaseResponse):
Expand Down
34 changes: 34 additions & 0 deletions backend/migrations/versions/b6d4f2a81c93_revoking_is_a_state.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
"""revoking is a state

Revision ID: b6d4f2a81c93
Revises: c4f8a2d97e15
Create Date: 2026-08-22 00:00:00.000000

"""

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision: str = "b6d4f2a81c93"
down_revision: str | Sequence[str] | None = "c4f8a2d97e15"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
op.add_column(
"oauth_connections",
sa.Column("revoked_at", sa.DateTime(timezone=True), nullable=True),
)
op.add_column(
"oauth_connections",
sa.Column("revoked_reason", sa.String(), server_default="", nullable=False),
)


def downgrade() -> None:
op.drop_column("oauth_connections", "revoked_reason")
op.drop_column("oauth_connections", "revoked_at")
27 changes: 27 additions & 0 deletions backend/tests/test_mcp_oauth.py
Original file line number Diff line number Diff line change
Expand Up @@ -389,6 +389,33 @@ async def test_reconsent_replaces_the_grant_and_evicts_the_stale_token(auth_serv
assert not await get_client().get(_token_key(SYSTEM_ACCOUNT_ID))


async def test_reconnect_after_disconnect_creates_a_new_grant(auth_server, druks_db):
_store_grant(refresh_token="rt-stale")
revoked = oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID)
await oauth.disconnect(_NAME, SYSTEM_ACCOUNT_ID)
assert not oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID)

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

from druks.database import db_session

db_session().expire_all()
# A re-connect creates a new grant. The revoked row stays as history,
# so at most one live connection holds the (server, account) slot.
grant = oauth.get_connection(_NAME, SYSTEM_ACCOUNT_ID)
assert grant.id != revoked.id
assert grant.refresh_token.decrypt() == "rt-1"
assert revoked.revoked_at


async def test_two_shared_connects_converge_on_one_grant(auth_server, druks_db):
first = Account.get_or_create("first@example.com")
second = Account.get_or_create("second@example.com")
Expand Down
Loading