Bind a worker's name and labels into its token; list and revoke tokens by jti

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013GAeg4Jo6xjZzjwy9QDWMQ

backend/fluksio/alembic/versions/b7d2e9c4a1f3_worker_token.py

@@ -0,0 +1,37 @@
+"""worker_token
+
+Every worker credential a superuser mints, by the `jti` it carries, so one can
+be listed and revoked without rotating the key that signs them all.
+
+Revision ID: b7d2e9c4a1f3
+Revises: f8c3e5a91d27
+Create Date: 2026-10-06
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b7d2e9c4a1f3"
+down_revision = "f8c3e5a91d27"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table(
+        "worker_token",
+        sa.Column("jti", sqlmodel.sql.sqltypes.AutoString(length=32), nullable=False),
+        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(length=255), nullable=False),
+        sa.Column("labels", sa.JSON(), nullable=True),
+        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("revoked_at", sa.DateTime(timezone=True), nullable=True),
+        sa.PrimaryKeyConstraint("jti"),
+    )
+
+
+def downgrade():
+    op.drop_table("worker_token")

backend/fluksio/api/deps.py

@@ -20,7 +20,7 @@ from fluksio.flow.artifacts import is_reference
 from fluksio.flow.controller import FlowController
 from fluksio.flow.dashboards import DashboardDef, DashboardStore
 from fluksio.flow.workers import PythonWorkerPool
-from fluksio.models import OAuthClient, TokenPayload, User
+from fluksio.models import OAuthClient, TokenPayload, User, WorkerToken
 
 reusable_oauth2 = OAuth2PasswordBearer(
     tokenUrl=f"{settings.API_V1_STR}/login/access-token"
@@ -300,6 +300,25 @@ def oauth_client_lives(client_id: str) -> bool:
         return session.get(OAuthClient, key) is not None
 
 
+def worker_claims(token: str) -> dict[str, Any]:
+    """A worker credential's claims, if it is still good. Blocking.
+
+    Raises ``InvalidTokenError`` for one that does not verify, one minted
+    before tokens carried a `jti` and their labels, and one a superuser has
+    revoked. A token whose `jti` has no row is one the engine minted for
+    itself, a provisioned machine's, and is good until it expires.
+    """
+    claims = security.decode_worker_token(token)
+    jti = claims.get("jti")
+    if not jti:
+        raise InvalidTokenError("this worker token predates revocation; mint a new one")
+    with Session(engine) as session:
+        row = session.get(WorkerToken, str(jti))
+    if row is not None and row.revoked_at is not None:
+        raise InvalidTokenError("this worker token has been revoked")
+    return claims
+
+
 def _gate_agent(payload: dict[str, Any]) -> dict[str, Any]:
     """Refuse a token whose agent has been revoked."""
     if not oauth_client_lives(str(payload.get("client_id") or "")):

backend/fluksio/api/routes/artifacts.py

@@ -17,8 +17,7 @@ from pydantic import BaseModel
 from sqlmodel import Session
 from starlette.concurrency import run_in_threadpool
 
-from fluksio.api.deps import user_from_token
-from fluksio.core import security
+from fluksio.api.deps import user_from_token, worker_claims
 from fluksio.core.config import settings
 from fluksio.core.db import engine
 from fluksio.flow.artifacts import ArtifactStore
@@ -42,7 +41,7 @@ def artifact_caller(request: Request) -> str:
     if not token:
         raise HTTPException(status_code=401, detail="Not authenticated")
     try:
-        claims = security.decode_worker_token(token)
+        claims = worker_claims(token)
     except InvalidTokenError:
         pass
     else:

backend/fluksio/api/routes/workers.py

@@ -2,24 +2,39 @@
 
 A worker dials in rather than being dialled: the GPU box and the engine are
 usually on different networks, and only one of them can be reached. It presents
-a token minted here, says what it can do, and then answers calls on the socket
-it opened.
+a token minted here, which says who it is and which labels it carries, says
+what else it can do, and then answers calls on the socket it opened.
 """
 
 from __future__ import annotations
 
 import asyncio
 import logging
-from datetime import timedelta
+import uuid
+from datetime import UTC, datetime, timedelta
 from typing import Any
 
-from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket
+from fastapi import (
+    APIRouter,
+    Depends,
+    HTTPException,
+    Request,
+    WebSocket,
+    WebSocketDisconnect,
+)
 from fastapi.responses import PlainTextResponse
 from fluksio_worker import worker_main
 from jwt.exceptions import InvalidTokenError
 from pydantic import BaseModel, Field
+from sqlmodel import col, select
+from starlette.concurrency import run_in_threadpool
 
-from fluksio.api.deps import get_current_active_superuser, get_current_user
+from fluksio.api.deps import (
+    SessionDep,
+    get_current_active_superuser,
+    get_current_user,
+    worker_claims,
+)
 from fluksio.core import security
 from fluksio.flow.remote import (
     PROTOCOL,
@@ -27,6 +42,7 @@ from fluksio.flow.remote import (
     RemoteWorkerHub,
     WorkerInventory,
 )
+from fluksio.models import WorkerToken
 
 logger = logging.getLogger(__name__)
 
@@ -35,6 +51,9 @@ router = APIRouter(prefix="/workers", tags=["workers"])
 #: Long, because a worker is a machine somebody set up once and left running.
 TOKEN_DAYS = 365
 
+#: The sockets attached with each token, by `jti`, so revoking one drops them.
+_attached: dict[str, set[WebSocket]] = {}
+
 
 class WorkerInfo(BaseModel):
     name: str
@@ -51,12 +70,16 @@ class WorkerInfo(BaseModel):
 
 
 class TokenRequest(BaseModel):
-    name: str
+    name: str = Field(min_length=1, max_length=255)
+    #: What a node's `device` may match on this worker; it cannot add to them.
+    labels: list[str] = Field(default_factory=list)
 
 
 class TokenIssued(BaseModel):
     name: str
     token: str
+    jti: str
+    labels: list[str] = Field(default_factory=list)
     expires_days: int = TOKEN_DAYS
 
 
@@ -145,14 +168,58 @@ def read_resources(request: Request) -> Any:
     response_model=TokenIssued,
     dependencies=[Depends(get_current_active_superuser)],
 )
-def issue_token(body: TokenRequest) -> Any:
+def issue_token(session: SessionDep, body: TokenRequest) -> Any:
     """Mint the credential a worker presents when it dials in.
 
-    Shown once. It is signed with the same keypair the agent tokens use, so
-    rotating that key revokes every worker along with them.
+    Shown once. The worker's name and labels are signed into it, so a worker is
+    what it was minted as whatever it says about itself, and only its `jti` is
+    kept here, to revoke it by.
     """
-    token = security.create_worker_token(body.name, timedelta(days=TOKEN_DAYS))
-    return TokenIssued(name=body.name, token=token)
+    labels = sorted({label.strip() for label in body.labels if label.strip()})
+    row = WorkerToken(
+        jti=uuid.uuid4().hex,
+        name=body.name,
+        labels=labels,
+        expires_at=datetime.now(UTC) + timedelta(days=TOKEN_DAYS),
+    )
+    token = security.create_worker_token(
+        body.name, timedelta(days=TOKEN_DAYS), labels=labels, jti=row.jti
+    )
+    session.add(row)
+    session.commit()
+    return TokenIssued(name=body.name, token=token, jti=row.jti, labels=labels)
+
+
+@router.get(
+    "/tokens",
+    response_model=list[WorkerToken],
+    dependencies=[Depends(get_current_active_superuser)],
+)
+def read_tokens(session: SessionDep) -> Any:
+    """Every worker token minted here, newest first, without the tokens."""
+    return session.exec(
+        select(WorkerToken).order_by(col(WorkerToken.created_at).desc())
+    ).all()
+
+
+@router.post(
+    "/tokens/{jti}/revoke",
+    response_model=WorkerToken,
+    dependencies=[Depends(get_current_active_superuser)],
+)
+async def revoke_token(session: SessionDep, jti: str) -> Any:
+    """Refuse this token from now on, and drop a worker attached with it."""
+    row = session.get(WorkerToken, jti)
+    if row is None:
+        raise HTTPException(status_code=404, detail="No such worker token")
+    if row.revoked_at is None:
+        row.revoked_at = datetime.now(UTC)
+        session.add(row)
+        session.commit()
+        session.refresh(row)
+    for socket in list(_attached.get(jti, ())):
+        await socket.close(code=1008)
+    return row
 
 
 @router.get(
@@ -177,14 +244,29 @@ async def attach(websocket: WebSocket, token: str = "") -> None:
     does: a websocket handshake carries no headers of its own.
     """
     try:
-        claims = security.decode_worker_token(token)
+        claims = await run_in_threadpool(worker_claims, token)
     except InvalidTokenError:
         await websocket.close(code=1008)
         return
 
+    jti = str(claims["jti"])
     await websocket.accept()
+    _attached.setdefault(jti, set()).add(websocket)
+    try:
+        await _attached_session(websocket, claims)
+    finally:
+        sockets = _attached.get(jti, set())
+        sockets.discard(websocket)
+        if not sockets:
+            _attached.pop(jti, None)
+
+
+async def _attached_session(websocket: WebSocket, claims: dict[str, Any]) -> None:
+    """The hello, then calls, until the socket goes or the token is revoked."""
     try:
         hello = await asyncio.wait_for(websocket.receive_json(), timeout=30)
+    except WebSocketDisconnect:
+        return
     except (TimeoutError, ValueError):
         await websocket.close(code=1002)
         return
@@ -196,13 +278,27 @@ async def attach(websocket: WebSocket, token: str = "") -> None:
         await websocket.close(code=1002)
         return
 
-    # The token names the worker; what it calls itself is a suggestion, so two
-    # hosts cannot fight over one identity by claiming the same name.
+    # The token names the worker and its labels. What it calls itself is a
+    # suggestion, so two hosts cannot fight over one identity by claiming the
+    # same name; labels it may repeat but not change, since they decide which
+    # node code and secrets come its way.
+    labels = sorted({str(label) for label in claims.get("labels") or []})
+    claimed = sorted({str(label) for label in hello.get("labels") or []})
+    if claimed and claimed != labels:
+        await websocket.send_json(
+            {
+                "op": "refused",
+                "reason": f"this token carries the labels {labels}, not {claimed}",
+            }
+        )
+        await websocket.close(code=1008)
+        return
+
     name = str(claims.get("sub") or hello.get("name") or "worker")
     hub = _hub(websocket.app)
     worker = RemoteWorker(
         name=name,
-        labels=[str(label) for label in (hello.get("labels") or [])],
+        labels=labels,
         send=websocket.send_json,
         loop=asyncio.get_running_loop(),
         max_parallel=max(1, int(hello.get("max_parallel") or 1)),

backend/fluksio/core/security.py

@@ -132,13 +132,19 @@ def create_oauth_access_token(
     )
 
 
-def create_worker_token(name: str, expires_delta: timedelta) -> str:
+def create_worker_token(
+    name: str,
+    expires_delta: timedelta,
+    labels: list[str] | None = None,
+    jti: str | None = None,
+) -> str:
     """A credential a remote worker presents when it dials in.
 
-    Signed with the same keypair the agent tokens use, so the whole set can be
-    revoked by rotating one key, and told apart from them by its audience: a
-    worker's token grants no API access, and an agent's opens no worker
-    connection.
+    Signed with the same keypair the agent tokens use, and told apart from them
+    by its audience: a worker's token grants no API access, and an agent's
+    opens no worker connection. It names the worker and the labels it may
+    carry, so a worker cannot claim work meant for another machine, and its
+    `jti` is what revoking one token rather than the whole key goes by.
     """
     now = datetime.now(UTC)
     payload = {
@@ -148,6 +154,8 @@ def create_worker_token(name: str, expires_delta: timedelta) -> str:
         "iat": now,
         "exp": now + expires_delta,
         "scope": WORKER_SCOPE,
+        "labels": sorted(set(labels or [])),
+        "jti": jti or uuid.uuid4().hex,
     }
     return jwt.encode(
         payload,

backend/fluksio/flow/provision.py

@@ -199,7 +199,9 @@ class SlurmProvisioner:
 
     def _submit(self, profile: SlurmProfile) -> None:
         worker = f"{self.name}-{profile.name}-{uuid.uuid4().hex[:8]}"
-        token = security.create_worker_token(worker, timedelta(hours=TOKEN_HOURS))
+        token = security.create_worker_token(
+            worker, timedelta(hours=TOKEN_HOURS), labels=profile.labels
+        )
         try:
             result = self._ssh(["sbatch", "--parsable"], self._payload(profile, token))
         except (OSError, subprocess.SubprocessError) as exc:

backend/fluksio/models.py

@@ -566,3 +566,25 @@ class ConnectorState(SQLModel, table=True):
     key: str = Field(primary_key=True, max_length=200)
     value: str
     updated_at: datetime = Field(default_factory=get_datetime_utc, sa_type=UTCDateTime)
+
+
+# -----------------------------------------------------------------------------
+# Worker tokens
+#
+# What a superuser minted for a remote worker, so it can be listed and revoked.
+# The token itself is not stored: it carries its own name and labels, and this
+# row is how one is refused before it expires.
+# -----------------------------------------------------------------------------
+
+
+class WorkerToken(SQLModel, table=True):
+    """One worker credential, by the id it carries as `jti`."""
+
+    __tablename__ = "worker_token"
+
+    jti: str = Field(primary_key=True, max_length=32)
+    name: str = Field(max_length=255)
+    labels: list[str] = Field(sa_column=Column(JSON), default_factory=list)
+    created_at: datetime = Field(default_factory=get_datetime_utc, sa_type=UTCDateTime)
+    expires_at: datetime = Field(sa_type=UTCDateTime)
+    revoked_at: datetime | None = Field(default=None, sa_type=UTCDateTime)

backend/tests/api/routes/test_workers.py

@@ -0,0 +1,99 @@
+"""A worker is what its token says, and a superuser can take the token back."""
+
+from datetime import timedelta
+
+import jwt
+import pytest
+from fastapi.testclient import TestClient
+from starlette.websockets import WebSocketDisconnect
+
+from fluksio.core import security
+from fluksio.core.config import settings
+from fluksio.flow.remote import PROTOCOL
+
+PREFIX = f"{settings.API_V1_STR}/workers"
+MISSING = "sha256:" + "0" * 64
+
+
+def mint(client: TestClient, headers: dict[str, str]) -> dict[str, str]:
+    answer = client.post(
+        f"{PREFIX}/tokens", headers=headers, json={"name": "gpu-box", "labels": ["gpu"]}
+    )
+    assert answer.status_code == 200
+    return answer.json()
+
+
+def hello(*labels: str) -> dict[str, object]:
+    return {"op": "hello", "protocol": PROTOCOL, "name": "elsewhere", "labels": labels}
+
+
+def test_the_token_names_the_worker_and_its_labels(
+    client: TestClient, superuser_token_headers: dict[str, str]
+) -> None:
+    issued = mint(client, superuser_token_headers)
+    with client.websocket_connect(f"{PREFIX}/attach?token={issued['token']}") as ws:
+        ws.send_json(hello())
+        assert ws.receive_json()["name"] == "gpu-box"
+        attached = client.get(PREFIX, headers=superuser_token_headers).json()
+    assert [(w["name"], w["labels"]) for w in attached] == [("gpu-box", ["gpu"])]
+
+
+def test_a_worker_cannot_claim_a_label_its_token_lacks(
+    client: TestClient, superuser_token_headers: dict[str, str]
+) -> None:
+    issued = mint(client, superuser_token_headers)
+    with client.websocket_connect(f"{PREFIX}/attach?token={issued['token']}") as ws:
+        ws.send_json(hello("gpu", "vault"))
+        assert ws.receive_json()["op"] == "refused"
+
+
+def test_a_revoked_token_is_dropped_and_refused(
+    client: TestClient,
+    superuser_token_headers: dict[str, str],
+    normal_user_token_headers: dict[str, str],
+) -> None:
+    issued = mint(client, superuser_token_headers)
+    url = f"{PREFIX}/attach?token={issued['token']}"
+    bearer = {"Authorization": f"Bearer {issued['token']}"}
+    assert (
+        client.get(
+            f"{settings.API_V1_STR}/artifacts/{MISSING}", headers=bearer
+        ).status_code
+        == 404
+    )
+    revoke = f"{PREFIX}/tokens/{issued['jti']}/revoke"
+    assert client.post(revoke, headers=normal_user_token_headers).status_code == 403
+
+    with client.websocket_connect(url) as ws:
+        ws.send_json(hello("gpu"))
+        ws.receive_json()
+        assert client.post(revoke, headers=superuser_token_headers).json()["revoked_at"]
+        with pytest.raises(WebSocketDisconnect):
+            ws.receive_json()
+
+    with pytest.raises(WebSocketDisconnect), client.websocket_connect(url):
+        pass
+    assert (
+        client.get(
+            f"{settings.API_V1_STR}/artifacts/{MISSING}", headers=bearer
+        ).status_code
+        == 401
+    )
+    listed = client.get(f"{PREFIX}/tokens", headers=superuser_token_headers).json()
+    assert [row for row in listed if row["jti"] == issued["jti"]][0]["revoked_at"]
+    assert all("token" not in row for row in listed)
+
+
+def test_a_token_minted_before_revocation_is_refused(client: TestClient) -> None:
+    claims = security.decode_worker_token(
+        security.create_worker_token("old", timedelta(days=1))
+    )
+    del claims["jti"], claims["labels"]
+    legacy = jwt.encode(
+        claims, security.oauth_key(), algorithm=security.OAUTH_ALGORITHM
+    )
+    with (
+        pytest.raises(WebSocketDisconnect),
+        client.websocket_connect(f"{PREFIX}/attach?token={legacy}"),
+    ):
+        pass