diff --git a/tests/test_agent_registry_store.py b/tests/test_agent_registry_store.py index d1c926869..df3756dc6 100644 --- a/tests/test_agent_registry_store.py +++ b/tests/test_agent_registry_store.py @@ -1354,3 +1354,119 @@ async def test_get_by_slug_rejects_glob_metacharacters(store): assert await store.get_by_slug("alpha?") is None assert await store.get_by_slug("") is None assert await store.get_by_slug("Alpha") is None + + +# --------------------------------------------------------------------------- +# Migration v5: token_min_iat — existing-DB pathway +# --------------------------------------------------------------------------- + + +class TestMigrationV5ExistingDB: + """Verify the guarded ALTER works on a DB created *before* token_min_iat existed.""" + + @pytest.mark.asyncio + async def test_existing_db_gains_column_and_tokens_still_work(self, tmp_path, signing_keypair): + """Build the pre-change agent_registry schema by hand, insert a row, + then init the store. The row must gain token_min_iat = 0 and old + tokens must still authenticate — a fresh-schema test would pass + vacuously because the column is already in SCHEMA. + """ + import sqlite3 + + db_path = tmp_path / "pre_v5.db" + priv, pub = signing_keypair + + # --- Build the pre-v5 schema by hand (no token_min_iat column) --- + conn = sqlite3.connect(str(db_path)) + conn.execute(""" + CREATE TABLE agent_registry ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + canonical_id TEXT NOT NULL UNIQUE, + display_name TEXT NOT NULL DEFAULT '', + framework TEXT NOT NULL DEFAULT '', + user_id TEXT NOT NULL DEFAULT '', + origin TEXT NOT NULL DEFAULT 'taos-deployed', + handle TEXT NOT NULL DEFAULT '', + role TEXT, + title TEXT, + reports_to TEXT, + capabilities TEXT NOT NULL DEFAULT '[]', + created_ts TEXT NOT NULL, + revoked_at TEXT, + status TEXT NOT NULL DEFAULT 'active' + ); + """) + # Insert a pre-existing row + conn.execute( + "INSERT INTO agent_registry " + "(canonical_id, display_name, framework, user_id, origin, " + "capabilities, created_ts, status) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + ( + "pre-v5-agent-20260101-000000", + "Pre V5 Agent", + "test", + "user-1", + "taos-deployed", + "[]", + "2026-01-01T00:00:00", + "active", + ), + ) + conn.commit() + conn.close() + + # --- Now init the store — this runs _migration_v5_add_token_min_iat --- + store = AgentRegistryStore(db_path) + await store.init() + + # The pre-existing row should have gained token_min_iat = 0 + row = await store.get("pre-v5-agent-20260101-000000") + assert row is not None + assert row["token_min_iat"] == 0 + assert row["display_name"] == "Pre V5 Agent" + + # Old tokens should still authenticate (token_min_iat=0 means no cutoff) + token = mint_registry_token( + "pre-v5-agent-20260101-000000", priv, + user_id="user-1", framework="test", + ) + payload = verify_registry_token(token, pub) + assert payload["sub"] == "pre-v5-agent-20260101-000000" + + await store.close() + + +# --------------------------------------------------------------------------- +# token_min_iat — per-identity token rotation +# --------------------------------------------------------------------------- + + +class TestBumpTokenMinIat: + @pytest.mark.asyncio + async def test_bump_sets_cutoff(self, store): + """bump_token_min_iat persists the cutoff and the new token_min_iat is + readable on the record.""" + row = await store.register(framework="test", display_name="test-agent") + cid = row["canonical_id"] + + ts = 1700000000 + updated = await store.bump_token_min_iat(cid, ts) + assert updated is not None + assert updated["token_min_iat"] == ts + + reread = await store.get(cid) + assert reread["token_min_iat"] == ts + + @pytest.mark.asyncio + async def test_default_zero_on_new_registration(self, store): + """Freshly registered agents have token_min_iat = 0 so all tokens + are valid (migration-safe default).""" + row = await store.register(framework="test", display_name="test-agent") + assert row["token_min_iat"] == 0 + + @pytest.mark.asyncio + async def test_bump_nonexistent_returns_none(self, store): + """bump_token_min_iat on an unknown canonical_id returns None.""" + result = await store.bump_token_min_iat("no-such-agent-20260101-000000", 1700000000) + assert result is None diff --git a/tests/test_token_rotation.py b/tests/test_token_rotation.py new file mode 100644 index 000000000..20c99c38b --- /dev/null +++ b/tests/test_token_rotation.py @@ -0,0 +1,295 @@ +"""Integration tests for per-identity token rotation (token_min_iat). + +Covers the full auth chain: store → auth path → route. +""" + +import time + +import pytest +import pytest_asyncio +from fastapi import HTTPException +from httpx import ASGITransport, AsyncClient + +from tinyagentos.agent_registry_store import mint_registry_token +from tinyagentos.agent_token_auth import check_agent_scope + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest_asyncio.fixture +async def agent_app(app, client): + """AsyncClient logged in as admin + agent_registry / agent_grants initialised. + + Reuses the shared ``client`` fixture (which initialises every store and + sets up the admin session) and only adds the agent_registry and + agent_grants stores on top. This prevents breakage when new stores get + added to conftest — there is no copy-paste to drift. + """ + for attr in ("agent_registry", "agent_grants"): + store = getattr(app.state, attr, None) + if store is not None and store._db is None: + await store.init() + yield client + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _FakeRequest: + """Minimal stand-in for a starlette Request used by check_agent_scope.""" + + def __init__(self, app, token: str | None = None): + self.app = app + self.headers = {} + if token is not None: + self.headers["Authorization"] = f"Bearer {token}" + + +async def _register_and_mint(app, *, user_id="u", owner_user_id=None, scopes=("a2a_receive",)): + """Register an active agent, add grants, and mint a signed JWT. + + If *owner_user_id* is given it is passed to ``register(user_id=...)`` + so the DB row's owner matches. Otherwise the DB row uses the default + (empty string, admin-only) while the JWT claim gets the separate + *user_id* value (legacy behaviour for auth-path tests). + + Returns (canonical_id, token). + """ + registry = app.state.agent_registry + grants = app.state.agent_grants + priv, _pub = app.state.agent_registry_keypair + rec = await registry.register( + framework="test", + display_name="TestAgent", + origin="external-selfjoin", + handle="@test", + user_id=owner_user_id if owner_user_id is not None else "", + ) + cid = rec["canonical_id"] + await registry.set_status(cid, "active") + for scope in scopes: + await grants.add_grant(cid, scope) + token = mint_registry_token(cid, priv, user_id=user_id, framework="test") + return cid, token + + +def _make_nonadmin_client( + app, auth, *, username: str, full_name: str, password: str +) -> tuple[AsyncClient, str]: + """Create an AsyncClient authenticated as a new non-admin user. + + Uses the invite flow (``add_user_invite`` + ``complete_invite``) because + ``setup_user`` only works for the first user. + Returns (AsyncClient, user_id). + """ + code = auth.add_user_invite(username, "admin") + auth.complete_invite(username, code, full_name, "", password) + record = auth.find_user(username) + uid = record["id"] + token = auth.create_session(user_id=uid) + transport = ASGITransport(app=app) + return AsyncClient( + transport=transport, + base_url="http://test", + cookies={"taos_session": token}, + ), uid + + +# --------------------------------------------------------------------------- +# Auth-path tests +# --------------------------------------------------------------------------- + + +class TestTokenMinIatAuth: + """Verify the token_min_iat check inside _verify_agent_scope.""" + + @pytest.mark.asyncio + async def test_old_token_rejected_after_bump(self, app): + """A token minted before the bump is rejected after bump_token_min_iat.""" + for attr in ("agent_registry", "agent_grants"): + store = getattr(app.state, attr, None) + if store is not None and store._db is None: + await store.init() + + registry = app.state.agent_registry + grants = app.state.agent_grants + priv, _pub = app.state.agent_registry_keypair + + rec = await registry.register( + framework="test", display_name="TestAgent", + origin="external-selfjoin", handle="@test", + ) + cid = rec["canonical_id"] + await registry.set_status(cid, "active") + await grants.add_grant(cid, "a2a_receive") + + # Mint old token + old_token = mint_registry_token(cid, priv, user_id="u", framework="test") + assert (await registry.get(cid)) is not None # token_min_iat is 0 + + # Bump the cutoff to a future timestamp so the old token's iat is + # strictly less (both happen in sub-second time in tests). + await registry.bump_token_min_iat(cid, int(time.time()) + 3600) + + # The old token should now be rejected + req = _FakeRequest(app, old_token) + with pytest.raises(HTTPException) as exc: + await check_agent_scope(req, "a2a_receive") + assert exc.value.status_code == 401 + assert exc.value.detail == "token superseded" + + @pytest.mark.asyncio + async def test_new_token_passes_after_bump(self, app): + """A token minted AFTER the bump passes the cutoff check.""" + for attr in ("agent_registry", "agent_grants"): + store = getattr(app.state, attr, None) + if store is not None and store._db is None: + await store.init() + + registry = app.state.agent_registry + grants = app.state.agent_grants + priv, _pub = app.state.agent_registry_keypair + + rec = await registry.register( + framework="test", display_name="TestAgent", + origin="external-selfjoin", handle="@test", + ) + cid = rec["canonical_id"] + await registry.set_status(cid, "active") + await grants.add_grant(cid, "a2a_receive") + + # Bump the cutoff + await registry.bump_token_min_iat(cid, int(time.time())) + + # Mint a new token after the bump + new_token = mint_registry_token(cid, priv, user_id="u", framework="test") + + req = _FakeRequest(app, new_token) + result = await check_agent_scope(req, "a2a_receive") + assert result == cid + + @pytest.mark.asyncio + async def test_default_zero_keeps_existing_tokens_valid(self, app): + """Default token_min_iat=0 means all tokens pass (no lockout on migration).""" + for attr in ("agent_registry", "agent_grants"): + store = getattr(app.state, attr, None) + if store is not None and store._db is None: + await store.init() + + registry = app.state.agent_registry + grants = app.state.agent_grants + priv, _pub = app.state.agent_registry_keypair + + rec = await registry.register( + framework="test", display_name="TestAgent", + origin="external-selfjoin", handle="@test", + ) + cid = rec["canonical_id"] + await registry.set_status(cid, "active") + await grants.add_grant(cid, "a2a_receive") + + # token_min_iat should be 0 by default + reread = await registry.get(cid) + assert reread["token_min_iat"] == 0 + + token = mint_registry_token(cid, priv, user_id="u", framework="test") + req = _FakeRequest(app, token) + result = await check_agent_scope(req, "a2a_receive") + assert result == cid + + +# --------------------------------------------------------------------------- +# Route-level tests +# --------------------------------------------------------------------------- + + +class TestRotateTokensRoute: + """Test POST /api/agents/registry/{id}/rotate-tokens.""" + + @pytest.mark.asyncio + async def test_admin_can_rotate(self, agent_app, app): + """An admin can bump token_min_iat on any identity.""" + cid, _token = await _register_and_mint(app, user_id="admin") + resp = await agent_app.post(f"/api/agents/registry/{cid}/rotate-tokens") + assert resp.status_code == 200 + body = resp.json() + assert body["token_min_iat"] > 0 + + @pytest.mark.asyncio + async def test_nonadmin_owner_can_rotate_own_identity(self, agent_app, app): + """A non-admin session owner can rotate their OWN identity (200).""" + # Create a non-admin user and register an agent they own. + client, uid = _make_nonadmin_client( + app, app.state.auth, username="owner1", full_name="Owner One", + password="password123", + ) + async with client: + cid, _token = await _register_and_mint( + app, user_id=uid, owner_user_id=uid, + ) + resp = await client.post( + f"/api/agents/registry/{cid}/rotate-tokens" + ) + assert resp.status_code == 200 + body = resp.json() + assert body["token_min_iat"] > 0 + + @pytest.mark.asyncio + async def test_nonadmin_cannot_rotate_others_agent(self, agent_app, app): + """A non-admin, non-owner rotating someone ELSE'S agent → 403.""" + # First user owns an agent. + owner_client, owner_uid = _make_nonadmin_client( + app, app.state.auth, username="owner2", full_name="Owner Two", + password="password123", + ) + # Second user tries to rotate the first user's agent. + intruder_client, intruder_uid = _make_nonadmin_client( + app, app.state.auth, username="intruder", full_name="Intruder", + password="password123", + ) + async with owner_client: + cid, _token = await _register_and_mint( + app, user_id=owner_uid, owner_user_id=owner_uid, + ) + async with intruder_client: + resp = await intruder_client.post( + f"/api/agents/registry/{cid}/rotate-tokens" + ) + assert resp.status_code == 403 + + @pytest.mark.asyncio + async def test_empty_userid_agent_is_admin_only(self, agent_app, app): + """An agent_registry row with user_id='' can ONLY be rotated by admin. + + Non-admin sessions get 403 because require_owner_or_admin compares + the session's user_id against an empty string (owner match fails) + and the session is not admin. + """ + # Register an agent with the default user_id="" (admin-only). + cid, _token = await _register_and_mint(app, user_id="admin") + r = await app.state.agent_registry.get(cid) + assert r["user_id"] == "" + + # A non-admin session trying to rotate it must get 403. + nonadmin_client, _uid = _make_nonadmin_client( + app, app.state.auth, username="randouser", full_name="Rando", + password="password123", + ) + async with nonadmin_client: + resp = await nonadmin_client.post( + f"/api/agents/registry/{cid}/rotate-tokens" + ) + assert resp.status_code == 403 + + @pytest.mark.asyncio + async def test_rotate_nonexistent_returns_404(self, agent_app): + """Rotating a nonexistent identity returns 404.""" + resp = await agent_app.post( + "/api/agents/registry/no-such-agent-20260101-000000/rotate-tokens" + ) + assert resp.status_code == 404 \ No newline at end of file diff --git a/tinyagentos/agent_registry_store.py b/tinyagentos/agent_registry_store.py index 4b5cbc84d..4957178e6 100644 --- a/tinyagentos/agent_registry_store.py +++ b/tinyagentos/agent_registry_store.py @@ -205,6 +205,26 @@ async def _migration_v4_dedupe_active_handles(conn) -> None: await conn.commit() +async def _migration_v5_add_token_min_iat(conn) -> None: + """Add token_min_iat column (idempotent) for per-identity token rotation. + + Default value 0 means every existing row's tokens remain valid so the + migration cannot lock the fleet out. Future bumps store a Unix timestamp; + any token whose ``iat`` claim is strictly less than this value is rejected + during auth, allowing per-identity credential rotation without revoking the + entire identity. + """ + existing_cols = { + row[1] + for row in await ( + await conn.execute("PRAGMA table_info(agent_registry)") + ).fetchall() + } + if "token_min_iat" not in existing_cols: + await conn.execute( + "ALTER TABLE agent_registry ADD COLUMN token_min_iat INTEGER NOT NULL DEFAULT 0" + ) + await conn.commit() # --------------------------------------------------------------------------- # Signing-key helpers (Ed25519, persisted to disk) # --------------------------------------------------------------------------- @@ -448,6 +468,7 @@ async def _post_init(self) -> None: # Dedupe BEFORE the index so a pre-invariant DB with duplicate active # handles cannot make the CREATE UNIQUE INDEX (hence boot) fail. await _migration_v4_dedupe_active_handles(self._db) + await _migration_v5_add_token_min_iat(self._db) # Created after the status migration so the partial index's WHERE clause # can reference the status column on the pre-status migration path. # Guard the index creation too: if some path we did not anticipate still @@ -850,6 +871,25 @@ async def revoke(self, canonical_id: str) -> Optional[dict]: await self._db.commit() return await self.get(canonical_id) + async def bump_token_min_iat(self, canonical_id: str, ts: int) -> Optional[dict]: + """Set *canonical_id*'s ``token_min_iat`` to *ts*, invalidating every + token minted before that Unix timestamp. + + The caller is responsible for authorisation (admin/session-owner checks). + Returns the updated record, or ``None`` if *canonical_id* does not exist. + """ + if self._db is None: + raise RuntimeError("AgentRegistryStore not initialised") + record = await self.get(canonical_id) + if record is None: + return None + await self._db.execute( + "UPDATE agent_registry SET token_min_iat = MAX(token_min_iat, ?) WHERE canonical_id = ?", + (ts, canonical_id), + ) + await self._db.commit() + return await self.get(canonical_id) + # ------------------------------------------------------------------ # Org model (#161): reporting lines, roles/titles, org tree # ------------------------------------------------------------------ diff --git a/tinyagentos/agent_token_auth.py b/tinyagentos/agent_token_auth.py index 54bfdcb6c..5db7b0748 100644 --- a/tinyagentos/agent_token_auth.py +++ b/tinyagentos/agent_token_auth.py @@ -17,8 +17,10 @@ - Valid signature but the agent is not active in the registry, the sub is unknown, or the required scope grant is missing/expired -> 403. -The registry JWT itself carries no exp claim; revocation is achieved by -suspending the agent (status != 'active') or by setting expires_at on the grant. +The registry JWT itself carries no exp claim; per-identity token rotation is +achieved by bumping ``token_min_iat`` on the registry record — any token whose +``iat`` is strictly less than the cutoff is rejected as superseded. Per-agent +revocation (suspending the agent or expiring the grant) remains available. """ from datetime import datetime, timezone @@ -110,6 +112,12 @@ async def _verify_agent_scope( if record is None or record.get("status") != "active": raise HTTPException(status_code=403, detail="agent is not active in the registry") + # Reject tokens issued before the identity's token_min_iat cutoff (rotation). + token_min_iat = record.get("token_min_iat") or 0 + token_iat = payload.get("iat") or 0 + if token_iat < token_min_iat: + raise HTTPException(status_code=401, detail="token superseded") + # Must hold an active grant for the required scope. grants_store = _get_grants_store(request) grants = await grants_store.list_grants(canonical_id) diff --git a/tinyagentos/routes/agent_registry.py b/tinyagentos/routes/agent_registry.py index a5b0d8c8f..a175c1a7e 100644 --- a/tinyagentos/routes/agent_registry.py +++ b/tinyagentos/routes/agent_registry.py @@ -15,6 +15,7 @@ POST /api/agents/registry/{id}/reject - lifecycle: pending → rejected (admin only) POST /api/agents/registry/{id}/suspend - lifecycle: active → suspended (admin only) POST /api/agents/registry/{id}/reactivate - lifecycle: suspended → active (admin only) +POST /api/agents/registry/{id}/rotate-tokens - bump token_min_iat, invalidate old tokens (owner/admin) Route ordering matters: /pubkey, /revoked, and /inactive are declared before /{canonical_id} so the literal strings are not captured as a path parameter. @@ -22,6 +23,7 @@ import asyncio import logging +import time from typing import Optional import aiosqlite @@ -218,23 +220,27 @@ async def _audit_governance( actor_user_id: str, before_status: str, after_status: str, + **extras, ) -> None: - """Write a governance audit event to the trace store (best-effort, non-fatal).""" + """Write a governance audit event to the trace store (best-effort, non-fatal). + + Extra keyword arguments are merged into the payload (e.g. before/after + token_min_iat for rotation events). + """ try: trace_registry = getattr(request.app.state, "trace_registry", None) if trace_registry is None: return ts = await trace_registry.get(_GOVERNANCE_SLUG) - await ts.record( - "governance", - payload={ - "action": action, - "canonical_id": canonical_id, - "actor_user_id": actor_user_id, - "before_status": before_status, - "after_status": after_status, - }, - ) + payload = { + "action": action, + "canonical_id": canonical_id, + "actor_user_id": actor_user_id, + "before_status": before_status, + "after_status": after_status, + **extras, + } + await ts.record("governance", payload=payload) except Exception: logger.exception("governance audit write failed (non-fatal)") @@ -755,6 +761,44 @@ async def reactivate_agent( return await _transition(request, canonical_id, "reactivate", "active", user) +@router.post("/api/agents/registry/{canonical_id}/rotate-tokens") +async def rotate_tokens( + request: Request, + canonical_id: str, + user: CurrentUser = Depends(current_user), +): + """Bump ``token_min_iat`` to the current Unix timestamp, invalidating every + token minted before now for this identity. + + Session owner or admin only. The rotation is a single-write DB bump (no + new token is minted — the caller re-mints after). Leaves a forensic + audit-log entry so every rotation is traceable to an actor. + """ + store = _get_store(request) + record = await store.get(canonical_id) + if record is None: + return JSONResponse({"error": "not found"}, status_code=404) + require_owner_or_admin(user, record["user_id"]) + + ts = int(time.time()) + before_iat = record.get("token_min_iat") or 0 + updated = await store.bump_token_min_iat(canonical_id, ts) + if updated is None: + return JSONResponse({"error": "not found"}, status_code=404) + + await _audit_governance( + request, + action="rotate-tokens", + canonical_id=canonical_id, + actor_user_id=user.user_id, + before_status=record.get("status") or "active", + after_status=updated.get("status") or "active", + before_token_min_iat=before_iat, + after_token_min_iat=ts, + ) + return updated + + # --------------------------------------------------------------------------- # Org model (#161): reporting lines, roles/titles, org tree # diff --git a/uv.lock b/uv.lock index e2e90aac8..c5f3ee940 100644 --- a/uv.lock +++ b/uv.lock @@ -3063,7 +3063,7 @@ wheels = [ [[package]] name = "tinyagentos" -version = "1.0.0b46" +version = "1.0.0b47" source = { editable = "." } dependencies = [ { name = "aiosqlite" },