Agent key rotation

Added button to rotate api key. This is required for better UX, now user has to remake agent if key is lost.
A migration is needed to keep track of proper usage + to not leave any orphan entries.

tests/api/user/agents/test_routes_happy.py

@@ -1252,6 +1252,172 @@ class TestPinAgentMore:
         assert response.status_code == 500
 
 
+class TestRegenerateAgentKey:
+    def _seed_published_with_key(self, pg_conn, user, key):
+        from application.storage.db.repositories.agents import AgentsRepository
+
+        agent = _seed_agent(pg_conn, user=user, status="published")
+        AgentsRepository(pg_conn).update(str(agent["id"]), user, {"key": key})
+        return agent
+
+    def test_returns_401_unauthenticated(self, app):
+        from application.api.user.agents.routes import RegenerateAgentKey
+
+        with app.test_request_context(
+            "/api/regenerate_agent_key/abc", method="POST"
+        ):
+            from flask import request
+            request.decoded_token = None
+            response = RegenerateAgentKey().post("abc")
+        status = (
+            response[1] if isinstance(response, tuple) else response.status_code
+        )
+        assert status == 401
+
+    def test_returns_404_when_missing(self, app, pg_conn):
+        from application.api.user.agents.routes import RegenerateAgentKey
+
+        with _patch_db(pg_conn), app.test_request_context(
+            "/api/regenerate_agent_key/00000000-0000-0000-0000-000000000000",
+            method="POST",
+        ):
+            from flask import request
+            request.decoded_token = {"sub": "u"}
+            response = RegenerateAgentKey().post(
+                "00000000-0000-0000-0000-000000000000"
+            )
+        assert response.status_code == 404
+
+    def test_returns_404_for_non_owner(self, app, pg_conn):
+        """Owner-only: another user cannot rotate someone else's key."""
+        from application.api.user.agents.routes import RegenerateAgentKey
+
+        owner = "u-owner-key"
+        agent = self._seed_published_with_key(pg_conn, owner, "owner-old-key")
+
+        with _patch_db(pg_conn), app.test_request_context(
+            f"/api/regenerate_agent_key/{agent['id']}", method="POST"
+        ):
+            from flask import request
+            request.decoded_token = {"sub": "u-intruder"}
+            response = RegenerateAgentKey().post(str(agent["id"]))
+        assert response.status_code == 404
+
+    def test_returns_400_for_draft_without_key(self, app, pg_conn):
+        from application.api.user.agents.routes import RegenerateAgentKey
+
+        user = "u-draft-key"
+        agent = _seed_agent(pg_conn, user=user, status="draft")
+
+        with _patch_db(pg_conn), app.test_request_context(
+            f"/api/regenerate_agent_key/{agent['id']}", method="POST"
+        ):
+            from flask import request
+            request.decoded_token = {"sub": user}
+            response = RegenerateAgentKey().post(str(agent["id"]))
+        assert response.status_code == 400
+
+    def test_regenerates_key_and_invalidates_old(self, app, pg_conn):
+        from application.api.user.agents.routes import RegenerateAgentKey
+        from application.storage.db.repositories.agents import AgentsRepository
+
+        user = "u-regen"
+        old_key = "regen-old-key"
+        agent = self._seed_published_with_key(pg_conn, user, old_key)
+
+        with _patch_db(pg_conn), app.test_request_context(
+            f"/api/regenerate_agent_key/{agent['id']}", method="POST"
+        ):
+            from flask import request
+            request.decoded_token = {"sub": user}
+            response = RegenerateAgentKey().post(str(agent["id"]))
+
+        assert response.status_code == 200
+        new_key = response.json["key"]
+        assert new_key and new_key != old_key
+
+        repo = AgentsRepository(pg_conn)
+        # Persisted key is the new one; the old key no longer resolves.
+        assert repo.get(str(agent["id"]), user)["key"] == new_key
+        assert repo.find_by_key(old_key) is None
+        assert repo.find_by_key(new_key)["id"] == agent["id"]
+
+    def test_migrates_historical_logs_and_usage(self, app, pg_conn):
+        """stack_logs + token_usage rows follow the key so analytics and the
+        rate-limit window are not orphaned. The stack_logs row is stamped with
+        a *different* user_id (a caller, not the owner) to prove the migration
+        is not owner-scoped."""
+        from application.api.user.agents.routes import RegenerateAgentKey
+        from application.storage.db.repositories.stack_logs import (
+            StackLogsRepository,
+        )
+        from application.storage.db.repositories.token_usage import (
+            TokenUsageRepository,
+        )
+
+        user = "u-regen-hist"
+        old_key = "regen-hist-old-key"
+        agent = self._seed_published_with_key(pg_conn, user, old_key)
+
+        StackLogsRepository(pg_conn).insert(
+            activity_id="regen-act-1",
+            endpoint="webhook",
+            level="info",
+            user_id="external-caller",
+            api_key=old_key,
+        )
+        TokenUsageRepository(pg_conn).insert(
+            api_key=old_key, prompt_tokens=7, generated_tokens=3,
+        )
+
+        with _patch_db(pg_conn), app.test_request_context(
+            f"/api/regenerate_agent_key/{agent['id']}", method="POST"
+        ):
+            from flask import request
+            request.decoded_token = {"sub": user}
+            response = RegenerateAgentKey().post(str(agent["id"]))
+        assert response.status_code == 200
+        new_key = response.json["key"]
+
+        from sqlalchemy import text
+        stack_old = pg_conn.execute(
+            text("SELECT COUNT(*) FROM stack_logs WHERE api_key = :k"),
+            {"k": old_key},
+        ).scalar()
+        stack_new = pg_conn.execute(
+            text("SELECT COUNT(*) FROM stack_logs WHERE api_key = :k"),
+            {"k": new_key},
+        ).scalar()
+        usage_old = pg_conn.execute(
+            text("SELECT COUNT(*) FROM token_usage WHERE api_key = :k"),
+            {"k": old_key},
+        ).scalar()
+        usage_new = pg_conn.execute(
+            text("SELECT COUNT(*) FROM token_usage WHERE api_key = :k"),
+            {"k": new_key},
+        ).scalar()
+        assert stack_old == 0 and stack_new == 1
+        assert usage_old == 0 and usage_new == 1
+
+    def test_db_error_returns_500(self, app):
+        from application.api.user.agents.routes import RegenerateAgentKey
+
+        @contextmanager
+        def _broken():
+            raise RuntimeError("boom")
+            yield
+
+        with patch(
+            "application.api.user.agents.routes.db_session", _broken
+        ), app.test_request_context(
+            "/api/regenerate_agent_key/abc", method="POST"
+        ):
+            from flask import request
+            request.decoded_token = {"sub": "u"}
+            response = RegenerateAgentKey().post("abc")
+        assert response.status_code == 500
+
+
 class TestPinnedAgentsListing:
     def test_returns_pinned_after_pinning(self, app, pg_conn):
         from application.api.user.agents.routes import PinAgent, PinnedAgents

tests/api/user/test_analytics.py

@@ -413,6 +413,55 @@ class TestGetUserLogs:
         assert response.status_code == 200
         assert len(response.json["logs"]) == 1
 
+    def test_webhook_logs_match_by_agent_id_after_key_rotation(
+        self, app, pg_conn,
+    ):
+        """A webhook activity-log row stamped with the agent's id but an
+        *old* api_key (the state after a key rotation) must still surface on
+        the agent's timeline via the stable agent_id join."""
+        from application.api.user.analytics.routes import GetUserLogs
+        from application.storage.db.repositories.agents import AgentsRepository
+        from application.storage.db.repositories.stack_logs import (
+            StackLogsRepository,
+        )
+
+        user = "u-logs-rotate"
+        agent = AgentsRepository(pg_conn).create(
+            user, "a", "published", key="current-key",
+        )
+        stack_repo = StackLogsRepository(pg_conn)
+        # Written under the previous key, but carries the stable agent_id.
+        stack_repo.insert(
+            activity_id="wh-rotated",
+            endpoint="webhook",
+            level="info",
+            user_id="external-caller",
+            api_key="rotated-away-key",
+            agent_id=str(agent["id"]),
+        )
+        # A different agent's webhook log must NOT leak in.
+        stack_repo.insert(
+            activity_id="wh-other",
+            endpoint="webhook",
+            level="info",
+            user_id="external-caller",
+            api_key="unrelated-key",
+            agent_id="99999999-9999-9999-9999-999999999999",
+        )
+
+        with _patch_analytics_db(pg_conn), app.test_request_context(
+            "/api/get_user_logs",
+            method="POST",
+            json={"api_key_id": str(agent["id"]), "event_type": "webhook"},
+        ):
+            from flask import request
+            request.decoded_token = {"sub": user}
+            response = GetUserLogs().post()
+        assert response.status_code == 200
+        logs = response.json["logs"]
+        assert len(logs) == 1
+        assert logs[0]["id"]  # the rotated-away row surfaced via agent_id
+
     def test_db_error_returns_400(self, app):
         from application.api.user.analytics.routes import GetUserLogs