From 8d164a1fe103bdff2869f4ec9b7460b788d79c96 Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:41:06 +0200 Subject: [PATCH 1/2] registry: per-identity token rotation via token_min_iat (fixes #2205) Add token_min_iat INTEGER column (default 0) to agent_registry so tokens can be invalidated per-identity without revoking the whole identity. - migration _migration_v6_add_token_min_iat: idempotent ALTER TABLE with default 0 so the migration cannot lock the fleet out - AgentRegistryStore.bump_token_min_iat(canonical_id, ts): persist cutoff - Auth path (_verify_agent_scope): after the status check, reject tokens whose iat claim is < the record's token_min_iat with 401 'token superseded' - POST /api/agents/registry/{id}/rotate-tokens: session-owner/admin route that bumps the cutoff and writes a forensic audit-log entry - Tests: store-level (bump sets cutoff, default zero, nonexistent nil), auth-path (old token rejected, new token passes, default-zero safe), route (admin rotates, owner rotates, nonexistent 404) --- tests/test_agent_registry_store.py | 35 +++ tests/test_token_rotation.py | 411 +++++++++++++++++++++++++++ tinyagentos/agent_registry_store.py | 40 +++ tinyagentos/agent_token_auth.py | 12 +- tinyagentos/routes/agent_registry.py | 35 +++ 5 files changed, 531 insertions(+), 2 deletions(-) create mode 100644 tests/test_token_rotation.py diff --git a/tests/test_agent_registry_store.py b/tests/test_agent_registry_store.py index 69c695833..0df1fcaed 100644 --- a/tests/test_agent_registry_store.py +++ b/tests/test_agent_registry_store.py @@ -1183,3 +1183,38 @@ async def test_org_tree_dangling_reports_to_becomes_root(self, store): tree = await store.get_org_tree() ids = {node["canonical_id"] for node in tree} assert row["canonical_id"] in ids + + +# --------------------------------------------------------------------------- +# 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..8dd5ecc2a --- /dev/null +++ b/tests/test_token_rotation.py @@ -0,0 +1,411 @@ +"""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): + """AsyncClient logged in as admin + agent_registry / agent_grants initialised.""" + 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() + # Reuse the same init + admin-session logic as the main client fixture + store = app.state.metrics + if store._db is not None: + await store.close() + await store.init() + notif_store = app.state.notifications + if notif_store._db is not None: + await notif_store.close() + await notif_store.init() + await app.state.qmd_client.init() + secrets_store = app.state.secrets + if secrets_store._db is not None: + await secrets_store.close() + await secrets_store.init() + broker_store = app.state.broker_store + if broker_store._db is not None: + await broker_store.close() + await broker_store.init() + scheduler = app.state.scheduler + if scheduler._db is not None: + await scheduler.close() + await scheduler.init() + channel_store = app.state.channels + if channel_store._db is not None: + await channel_store.close() + await channel_store.init() + relationship_mgr = app.state.relationships + if relationship_mgr._db is not None: + await relationship_mgr.close() + await relationship_mgr.init() + conversion_mgr = app.state.conversion + if conversion_mgr._db is not None: + await conversion_mgr.close() + await conversion_mgr.init() + training_mgr = app.state.training + if training_mgr._db is not None: + await training_mgr.close() + await training_mgr.init() + agent_messages = app.state.agent_messages + if agent_messages._db is not None: + await agent_messages.close() + await agent_messages.init() + shared_folders = app.state.shared_folders + if shared_folders._db is not None: + await shared_folders.close() + await shared_folders.init() + streaming_sessions = app.state.streaming_sessions + if streaming_sessions._db is not None: + await streaming_sessions.close() + await streaming_sessions.init() + expert_agents = app.state.expert_agents + if expert_agents._db is not None: + await expert_agents.close() + await expert_agents.init() + chat_messages = app.state.chat_messages + if chat_messages._db is not None: + await chat_messages.close() + await chat_messages.init() + chat_channels = app.state.chat_channels + if chat_channels._db is not None: + await chat_channels.close() + await chat_channels.init() + project_store = app.state.project_store + if project_store._db is not None: + await project_store.close() + await project_store.init() + project_invites = app.state.project_invites + if project_invites._db is not None: + await project_invites.close() + await project_invites.init() + board_audit = app.state.board_audit + if board_audit._db is not None: + await board_audit.close() + await board_audit.init() + receipt_store = app.state.receipt_store + if receipt_store._db is not None: + await receipt_store.close() + await receipt_store.init() + project_task_store = app.state.project_task_store + if project_task_store._db is not None: + await project_task_store.close() + await project_task_store.init() + project_element_store = app.state.project_element_store + if project_element_store._db is not None: + await project_element_store.close() + await project_element_store.init() + routine_store = app.state.routine_store + if routine_store._db is not None: + await routine_store.close() + await routine_store.init() + decision_store = app.state.decision_store + if decision_store._db is not None: + await decision_store.close() + await decision_store.init() + execution_policies = app.state.execution_policies + if execution_policies._db is not None: + await execution_policies.close() + await execution_policies.init() + coding_session_store = app.state.coding_session_store + if coding_session_store._db is not None: + await coding_session_store.close() + await coding_session_store.init() + app.state.projects_root.mkdir(parents=True, exist_ok=True) + canvas_store = app.state.canvas_store + if canvas_store._db is not None: + await canvas_store.close() + await canvas_store.init() + themes = app.state.themes + if themes._db is not None: + await themes.close() + await themes.init() + office_docs = app.state.office_docs + if office_docs._db is not None: + await office_docs.close() + await office_docs.init() + web_sites = app.state.web_sites + if web_sites._db is not None: + await web_sites.close() + await web_sites.init() + song_store = app.state.song_store + if song_store._db is not None: + await song_store.close() + await song_store.init() + design_docs = app.state.design_docs + if design_docs._db is not None: + await design_docs.close() + await design_docs.init() + await app.state.app_grants.init() + await app.state.license_acceptances.init() + feedback_store = app.state.feedback_store + if feedback_store._db is not None: + await feedback_store.close() + await feedback_store.init() + client_log_store = app.state.client_log_store + if client_log_store._db is not None: + await client_log_store.close() + await client_log_store.init() + device_store = app.state.device_store + if device_store._db is not None: + await device_store.close() + await device_store.init() + council_roles = app.state.council_roles + if council_roles._db is not None: + await council_roles.close() + await council_roles.init() + council_members = app.state.council_members + if council_members._db is not None: + await council_members.close() + await council_members.init() + from tinyagentos.routes.desktop_browser.store import BrowserStore, BrowserCookieStore + _browser_store = BrowserStore(app.state.data_dir / "browser.sqlite3") + await _browser_store.init() + app.state.browser_store = _browser_store + _browser_cookie_store = BrowserCookieStore( + app.state.data_dir / "browser_cookies.sqlite3", key_hex="0" * 64 + ) + await _browser_cookie_store.init() + app.state.browser_cookie_store = _browser_cookie_store + from tinyagentos.routes.desktop_browser.copilot_ws import CopilotTicketStore, CopilotHub + app.state.copilot_ticket_store = CopilotTicketStore() + app.state.copilot_hub = CopilotHub() + app.state.auth.setup_user("admin", "Test Admin", "", "testpass") + _record = app.state.auth.find_user("admin") + _uid = _record["id"] if _record else "" + _token = app.state.auth.create_session(user_id=_uid, long_lived=True) + app.state._startup_complete = True + transport = ASGITransport(app=app) + async with AsyncClient( + transport=transport, + base_url="http://test", + cookies={"taos_session": _token}, + ) as c: + yield c + await canvas_store.close() + await project_task_store.close() + await routine_store.close() + await board_audit.close() + await project_store.close() + await project_invites.close() + await chat_channels.close() + await chat_messages.close() + await expert_agents.close() + await streaming_sessions.close() + await shared_folders.close() + await agent_messages.close() + await conversion_mgr.close() + await training_mgr.close() + await relationship_mgr.close() + await channel_store.close() + await scheduler.close() + await secrets_store.close() + await broker_store.close() + await notif_store.close() + await store.close() + await office_docs.close() + await web_sites.close() + await song_store.close() + await design_docs.close() + await coding_session_store.close() + await feedback_store.close() + await client_log_store.close() + await project_element_store.close() + await app.state.qmd_client.close() + await app.state.http_client.aclose() + await _browser_store.close() + await _browser_cookie_store.close() + await council_roles.close() + await council_members.close() + + +# --------------------------------------------------------------------------- +# 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", scopes=("a2a_receive",)): + """Register an active agent, add grants, and mint a signed JWT. + + 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", + ) + 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 + + +# --------------------------------------------------------------------------- +# 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 via the admin client.""" + + @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_owner_can_rotate(self, agent_app, app): + """A session owner (non-admin) can rotate their own 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_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 diff --git a/tinyagentos/agent_registry_store.py b/tinyagentos/agent_registry_store.py index 96f837337..879876009 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_v6_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) # --------------------------------------------------------------------------- @@ -417,6 +437,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_v6_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 @@ -745,6 +766,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 = ? 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 635a8db23..31755001e 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 d29d4e65a..00484ac1c 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 @@ -743,6 +745,39 @@ 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()) + updated = await store.bump_token_min_iat(canonical_id, ts) + + 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", + ) + return updated + + # --------------------------------------------------------------------------- # Org model (#161): reporting lines, roles/titles, org tree # From ee7a0a5a215e1e17b5b374e6515448e4edbba36f Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:55:18 +0200 Subject: [PATCH 2/2] fix(registry): address CodeRabbit findings on token rotation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - bump_token_min_iat: use MAX(token_min_iat, ?) so the cutoff is monotonically non-decreasing — a stale/lower timestamp cannot silently un-supersede tokens - rotate-tokens route: handle None from bump_token_min_iat (race condition) with 404 instead of AttributeError on updated.get() - audit entry: capture before_token_min_iat and after_token_min_iat in the governance audit payload so rotation values are traceable - tests: remove unused token bindings in route-level tests --- tests/test_token_rotation.py | 4 ++-- tinyagentos/agent_registry_store.py | 2 +- tinyagentos/routes/agent_registry.py | 31 ++++++++++++++++++---------- 3 files changed, 23 insertions(+), 14 deletions(-) diff --git a/tests/test_token_rotation.py b/tests/test_token_rotation.py index 8dd5ecc2a..c9c72cb98 100644 --- a/tests/test_token_rotation.py +++ b/tests/test_token_rotation.py @@ -387,7 +387,7 @@ class TestRotateTokensRoute: @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") + 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() @@ -396,7 +396,7 @@ async def test_admin_can_rotate(self, agent_app, app): @pytest.mark.asyncio async def test_owner_can_rotate(self, agent_app, app): """A session owner (non-admin) can rotate their own identity.""" - cid, token = await _register_and_mint(app, user_id="admin") + 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() diff --git a/tinyagentos/agent_registry_store.py b/tinyagentos/agent_registry_store.py index 879876009..1a47e573a 100644 --- a/tinyagentos/agent_registry_store.py +++ b/tinyagentos/agent_registry_store.py @@ -779,7 +779,7 @@ async def bump_token_min_iat(self, canonical_id: str, ts: int) -> Optional[dict] if record is None: return None await self._db.execute( - "UPDATE agent_registry SET token_min_iat = ? WHERE canonical_id = ?", + "UPDATE agent_registry SET token_min_iat = MAX(token_min_iat, ?) WHERE canonical_id = ?", (ts, canonical_id), ) await self._db.commit() diff --git a/tinyagentos/routes/agent_registry.py b/tinyagentos/routes/agent_registry.py index 00484ac1c..b958eb4c8 100644 --- a/tinyagentos/routes/agent_registry.py +++ b/tinyagentos/routes/agent_registry.py @@ -215,23 +215,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)") @@ -765,7 +769,10 @@ async def rotate_tokens( 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, @@ -774,6 +781,8 @@ async def rotate_tokens( 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