Add OAuth client list and revoke endpoints
Superuser-only management for agents that registered themselves: list them with whether anyone approved them, and withdraw one without rotating the signing key and cutting off every other agent. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H7LwYgJfpkbLCTeiAf8U4A
This commit is contained in:
@@ -31,14 +31,21 @@ from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode, urlparse
|
||||
|
||||
from fastapi import APIRouter, Depends, Form, Request
|
||||
from fastapi import APIRouter, Depends, Form, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from sqlmodel import select
|
||||
from pydantic import BaseModel
|
||||
from sqlmodel import col, select
|
||||
|
||||
from app.api.deps import CurrentUser, SessionDep, get_current_user
|
||||
from app.api.deps import (
|
||||
CurrentUser,
|
||||
SessionDep,
|
||||
get_current_active_superuser,
|
||||
get_current_user,
|
||||
)
|
||||
from app.core import security
|
||||
from app.core.config import settings
|
||||
from app.models import (
|
||||
Message,
|
||||
OAuthAuthorizationCode,
|
||||
OAuthAuthorizeInfo,
|
||||
OAuthAuthorizeRequest,
|
||||
@@ -444,3 +451,86 @@ def _revoke_family(
|
||||
).all():
|
||||
token.revoked = True
|
||||
session.add(token)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Management
|
||||
#
|
||||
# Registration is open, so without these the only way to withdraw one agent's
|
||||
# access was rotating the signing key — which cuts off every agent at once.
|
||||
# They answer in FastAPI's error shape rather than OAuth's, because the
|
||||
# dashboard reads them and no OAuth client ever does, and they keep working
|
||||
# with MCP switched off: that is exactly when leftover clients want clearing.
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
class RegisteredClient(BaseModel):
|
||||
"""A registered agent, and whether anyone actually let it in."""
|
||||
|
||||
id: uuid.UUID
|
||||
client_name: str
|
||||
redirect_uris: list[str]
|
||||
created_at: datetime
|
||||
#: Live refresh tokens. Zero means it registered and was never approved.
|
||||
active_tokens: int
|
||||
last_authorized_at: datetime | None
|
||||
|
||||
|
||||
class RegisteredClients(BaseModel):
|
||||
data: list[RegisteredClient]
|
||||
count: int
|
||||
|
||||
|
||||
@router.get(
|
||||
"/clients",
|
||||
response_model=RegisteredClients,
|
||||
dependencies=[Depends(get_current_active_superuser)],
|
||||
)
|
||||
def read_clients(session: SessionDep) -> Any:
|
||||
"""Every agent that registered itself, newest first."""
|
||||
# ponytail: counts live tokens in Python; a GROUP BY if this ever grows.
|
||||
# It also keeps the expiry comparison off a naive Postgres column.
|
||||
now = _now()
|
||||
live: dict[uuid.UUID, list[datetime]] = defaultdict(list)
|
||||
for token in session.exec(
|
||||
select(OAuthRefreshToken).where(col(OAuthRefreshToken.revoked).is_(False))
|
||||
).all():
|
||||
if _aware(token.expires_at) > now:
|
||||
live[token.client_id].append(_aware(token.created_at))
|
||||
|
||||
data = [
|
||||
RegisteredClient(
|
||||
id=client.id,
|
||||
client_name=client.client_name,
|
||||
redirect_uris=client.redirect_uris,
|
||||
created_at=_aware(client.created_at),
|
||||
active_tokens=len(live[client.id]),
|
||||
last_authorized_at=max(live[client.id], default=None),
|
||||
)
|
||||
for client in session.exec(
|
||||
select(OAuthClient).order_by(col(OAuthClient.created_at).desc())
|
||||
).all()
|
||||
]
|
||||
return RegisteredClients(data=data, count=len(data))
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/clients/{client_id}",
|
||||
response_model=Message,
|
||||
dependencies=[Depends(get_current_active_superuser)],
|
||||
)
|
||||
def revoke_client(client_id: uuid.UUID, session: SessionDep) -> Any:
|
||||
"""Withdraw one agent's access, leaving every other agent alone.
|
||||
|
||||
Deleting the client cascades to its codes and refresh tokens, so it can
|
||||
get nothing new and cannot come back without registering again. An access
|
||||
token already in its hands keeps working until it expires
|
||||
(``MCP_TOKEN_EXPIRE_MINUTES``) — those are stateless by design.
|
||||
"""
|
||||
client = session.get(OAuthClient, client_id)
|
||||
if client is None:
|
||||
raise HTTPException(status_code=404, detail="No such client")
|
||||
name = client.client_name
|
||||
session.delete(client)
|
||||
session.commit()
|
||||
return Message(message=f"Revoked '{name}'")
|
||||
|
||||
@@ -224,3 +224,44 @@ def test_everything_is_refused_while_mcp_is_off(
|
||||
client.post(f"{PREFIX}/token", data={"grant_type": "authorization_code"}).status_code
|
||||
== 403
|
||||
)
|
||||
|
||||
|
||||
def test_one_agent_can_be_revoked_without_touching_the_others(
|
||||
client: TestClient,
|
||||
superuser_token_headers: dict[str, str],
|
||||
normal_user_token_headers: dict[str, str],
|
||||
) -> None:
|
||||
registered = register(client, client_name="Doomed agent")
|
||||
client_id = registered["client_id"]
|
||||
verifier, challenge = pkce()
|
||||
code = approve(client, superuser_token_headers, client_id, challenge)
|
||||
tokens = exchange(client, client_id, code, verifier).json()
|
||||
|
||||
def listed() -> list[dict]:
|
||||
response = client.get(f"{PREFIX}/clients", headers=superuser_token_headers)
|
||||
assert response.status_code == 200
|
||||
return response.json()["data"]
|
||||
|
||||
entry = next(c for c in listed() if c["id"] == client_id)
|
||||
assert entry["client_name"] == "Doomed agent"
|
||||
# The column that tells an approved agent from one that only registered.
|
||||
assert entry["active_tokens"] == 1
|
||||
|
||||
# Withdrawing access is a superuser's job.
|
||||
assert (
|
||||
client.get(f"{PREFIX}/clients", headers=normal_user_token_headers).status_code
|
||||
== 403
|
||||
)
|
||||
|
||||
revoked = client.delete(
|
||||
f"{PREFIX}/clients/{client_id}", headers=superuser_token_headers
|
||||
)
|
||||
assert revoked.status_code == 200, revoked.text
|
||||
assert all(c["id"] != client_id for c in listed())
|
||||
|
||||
# Its refresh token went with it, so it cannot mint itself a new one.
|
||||
refreshed = client.post(
|
||||
f"{PREFIX}/token",
|
||||
data={"grant_type": "refresh_token", "refresh_token": tokens["refresh_token"]},
|
||||
)
|
||||
assert refreshed.status_code == 400
|
||||
|
||||
Reference in New Issue
Block a user