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.
This commit is contained in:
Pavel committed 2026-07-25 18:15:56 +02:00
1 parent 2b1cd2d404
commit 1394a7a0b4
24 files changed
+727 -13

No files matched your search

@@ -0,0 +1,52 @@
"""0026 stack_logs.agent_id — stable per-agent attribution for activity logs.
``stack_logs`` (webhook + system-error activity logs) previously carried only
the agent's ``api_key`` string, so it could be joined to an agent only through
that key. Now that agents can rotate their key, a key-only join would orphan
history on rotation. This adds a nullable ``agent_id`` UUID (no FK, mirroring
``token_usage.agent_id`` so a stray/legacy id can never block a log write),
indexes it, and backfills existing rows by matching the stored ``api_key`` to
``agents.key``. The backfill is exhaustive today because no key has been
rotated yet. Idempotent both ways.
Revision ID: 0026_stack_logs_agent_id
Revises: 0025_artifacts
"""
from typing import Sequence, Union
from alembic import op
revision: str = "0026_stack_logs_agent_id"
down_revision: Union[str, None] = "0025_artifacts"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
op.execute("ALTER TABLE stack_logs ADD COLUMN IF NOT EXISTS agent_id uuid;")
op.execute(
"""
CREATE INDEX IF NOT EXISTS ix_stack_logs_agent_id
ON stack_logs (agent_id)
WHERE agent_id IS NOT NULL;
"""
)
# Backfill from the still-current key mapping. ``agents.key`` is CITEXT and
# unique, so each api_key resolves to at most one agent.
op.execute(
"""
UPDATE stack_logs s
SET agent_id = a.id
FROM agents a
WHERE s.agent_id IS NULL
AND s.api_key IS NOT NULL
AND a.key = s.api_key;
"""
)
def downgrade() -> None:
op.execute("DROP INDEX IF EXISTS ix_stack_logs_agent_id;")
op.execute("ALTER TABLE stack_logs DROP COLUMN IF EXISTS agent_id;")
+91
View File
@@ -28,6 +28,8 @@ from application.api.user.team_sharing import (
)
from application.storage.db.repositories.agent_folders import AgentFoldersRepository
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.stack_logs import StackLogsRepository
from application.storage.db.repositories.token_usage import TokenUsageRepository
from application.storage.db.repositories.users import UsersRepository
from application.storage.db.repositories.workflow_edges import WorkflowEdgesRepository
from application.storage.db.repositories.workflow_nodes import WorkflowNodesRepository
@@ -1293,6 +1295,95 @@ class UpdateAgent(Resource):
return make_response(jsonify(response_data), 200)
@agents_ns.route("/regenerate_agent_key/<string:agent_id>")
class RegenerateAgentKey(Resource):
@api.doc(
params={"agent_id": "ID of the agent"},
description=(
"Rotate an agent's API key. The previous key is invalidated "
"immediately and a fresh key is returned once. Historical logging "
"and usage records are re-pointed to the new key so analytics stay "
"intact. Owner only."
),
)
def post(self, agent_id):
if not (decoded_token := request.decoded_token):
return make_response(
jsonify({"success": False, "message": "Unauthorized"}), 401
)
user = decoded_token.get("sub")
try:
with db_session() as conn:
agents_repo = AgentsRepository(conn)
# Owner-only: rotating a credential is destructive to live
# integrations, so this is intentionally stricter than
# update_agent (which also allows team editors).
existing_agent = agents_repo.get_any(agent_id, user)
if not existing_agent:
return make_response(
jsonify(
{
"success": False,
"message": "Agent not found or not authorized",
}
),
404,
)
pg_agent_id = str(existing_agent["id"])
old_key = existing_agent.get("key")
if not old_key:
# Draft agents have no key; the first key is minted on
# publish. Nothing to rotate.
return make_response(
jsonify(
{
"success": False,
"message": "Publish the agent before resetting its key",
}
),
400,
)
new_key = str(uuid.uuid4())
updated = agents_repo.update(pg_agent_id, user, {"key": new_key})
if not updated:
return make_response(
jsonify(
{
"success": False,
"message": "Failed to regenerate key",
}
),
500,
)
# Re-point key-string-only history to the new key so it is not
# orphaned. conversations/token analytics also match on
# agent_id, but stack_logs (webhook + system-error logs) has no
# agent_id column, and the 24h rate-limit window is keyed by the
# api-key string only. Runs in the same transaction as the swap.
StackLogsRepository(conn).reassign_api_key(
old_key=old_key, new_key=new_key
)
TokenUsageRepository(conn).reassign_api_key(
old_key=old_key, new_key=new_key
)
except Exception as err:
current_app.logger.error(
f"Error regenerating agent key: {err}", exc_info=True
)
return make_response(
jsonify(
{"success": False, "message": "Error regenerating agent key"}
),
500,
)
return make_response(jsonify({"success": True, "key": new_key}), 200)
@agents_ns.route("/delete_agent")
class DeleteAgent(Resource):
@api.doc(params={"id": "ID of the agent"}, description="Delete an agent by ID")
+10 -2
View File
@@ -801,15 +801,23 @@ class GetUserLogs(Resource):
"(l.data->>'api_key' = :api_key"
" OR l.data->>'agent_id' = :agent_pg_id)"
]
# ``stack_logs.agent_id`` (0026) is the stable join key;
# ``api_key`` is mutable (agents rotate keys), so match the
# id first and keep the key arm for rows predating the
# backfill / any null agent_id.
stack_agent_match = (
"(s.agent_id = CAST(:agent_pg_id AS uuid)"
" OR s.api_key = :api_key)"
)
webhook_where = [
"COALESCE(s.endpoint, '') = 'webhook'",
"s.api_key = :api_key",
stack_agent_match,
webhook_dedupe,
]
system_where = [
"s.level = 'error'",
"COALESCE(s.endpoint, '') NOT IN ('webhook', 'schedule')",
"s.api_key = :api_key",
stack_agent_match,
]
# Owner-gated agent match: drop the user clause so a
# shared agent's runs (stamped with the caller's
+9 -2
View File
@@ -17,11 +17,12 @@ logging.basicConfig(
class LogContext:
def __init__(self, endpoint, activity_id, user, api_key, query):
def __init__(self, endpoint, activity_id, user, api_key, query, agent_id=None):
self.endpoint = endpoint
self.activity_id = activity_id
self.user = user
self.api_key = api_key
self.agent_id = agent_id
self.query = query
self.stacks = []
# Per-activity response aggregates populated by ``_consume_and_log``
@@ -100,7 +101,9 @@ def log_activity() -> Callable:
# so nested activities record the parent → child link.
parent_activity_id = log_context.snapshot().get("activity_id")
context = LogContext(endpoint, activity_id, user, api_key, query)
context = LogContext(
endpoint, activity_id, user, api_key, query, agent_id=agent_id
)
kwargs["log_context"] = context
ctx_token = log_context.bind(
@@ -217,6 +220,7 @@ def _consume_and_log(generator: Generator, context: "LogContext"):
activity_id=context.activity_id,
user=context.user,
api_key=context.api_key,
agent_id=context.agent_id,
query=context.query,
stacks=context.stacks,
level="error",
@@ -228,6 +232,7 @@ def _consume_and_log(generator: Generator, context: "LogContext"):
activity_id=context.activity_id,
user=context.user,
api_key=context.api_key,
agent_id=context.agent_id,
query=context.query,
stacks=context.stacks,
level="info",
@@ -242,6 +247,7 @@ def _log_activity_to_db(
query: str,
stacks: List[Dict],
level: str,
agent_id: str = None,
) -> None:
"""Append a per-request activity log row to Postgres (``stack_logs``)."""
try:
@@ -259,6 +265,7 @@ def _log_activity_to_db(
level=_truncate(level),
user_id=_truncate(user),
api_key=_truncate(api_key),
agent_id=agent_id,
query=_truncate(query),
stacks=stacks,
timestamp=datetime.datetime.now(datetime.timezone.utc),
+5
View File
@@ -255,6 +255,11 @@ stack_logs_table = Table(
Column("level", Text),
Column("user_id", Text),
Column("api_key", Text),
# Stable per-agent attribution. ``api_key`` is mutable (agents can rotate
# their key), so analytics join on this UUID first and fall back to the
# key string. No FK (mirrors ``token_usage.agent_id``) so a stray/legacy
# id can never block an activity-log write. Added in ``0026``.
Column("agent_id", UUID(as_uuid=True)),
Column("query", Text),
Column("stacks", JSONB, nullable=False, server_default="[]"),
Column("timestamp", DateTime(timezone=True), nullable=False, server_default=func.now()),
@@ -52,16 +52,26 @@ class StackLogsRepository:
level: Optional[str] = None,
user_id: Optional[str] = None,
api_key: Optional[str] = None,
agent_id: Optional[str] = None,
query: Optional[str] = None,
stacks: Optional[list] = None,
timestamp: Optional[datetime] = None,
) -> None:
# ``agent_id`` is a UUID column. Coerce anything not shaped like a UUID
# (36 chars with hyphens) to NULL so a stray/legacy id never breaks the
# activity-log write. Mirrors TokenUsageRepository.insert.
agent_id_uuid: Optional[str] = None
if agent_id:
s = str(agent_id)
if len(s) == 36 and "-" in s:
agent_id_uuid = s
self._conn.execute(
text(
"""
INSERT INTO stack_logs (activity_id, endpoint, level, user_id, api_key, query, stacks, timestamp)
INSERT INTO stack_logs (activity_id, endpoint, level, user_id, api_key, agent_id, query, stacks, timestamp)
VALUES (
:activity_id, :endpoint, :level, :user_id, :api_key, :query,
:activity_id, :endpoint, :level, :user_id, :api_key,
CAST(:agent_id AS uuid), :query,
CAST(:stacks AS jsonb),
COALESCE(:timestamp, now())
)
@@ -73,6 +83,7 @@ class StackLogsRepository:
"level": strip_null_bytes(level),
"user_id": strip_null_bytes(user_id),
"api_key": strip_null_bytes(api_key),
"agent_id": agent_id_uuid,
"query": strip_null_bytes(query),
"stacks": json.dumps(
redact_secrets(
@@ -83,3 +94,27 @@ class StackLogsRepository:
"timestamp": timestamp,
},
)
def reassign_api_key(self, *, old_key: str, new_key: str) -> int:
"""Re-point historical rows from ``old_key`` to ``new_key``.
``stack_logs`` has no ``agent_id`` column, so webhook / system-error
logs are matched by the api-key string alone (see the analytics
``webhook_where`` / ``system_where`` clauses). When an agent's key is
rotated this rewrite keeps that history attached to the agent instead
of orphaning it.
Matched by ``api_key`` only — deliberately NOT scoped by ``user_id``:
``agents.key`` is globally unique, so a key maps to exactly one agent,
and rows are stamped with the *caller's* user_id (which is not the
owner for webhook / external-api-key traffic). Scoping by owner would
skip precisely the rows this migration exists to preserve.
Returns the number of rows updated.
"""
if not old_key or not new_key:
return 0
result = self._conn.execute(
text("UPDATE stack_logs SET api_key = :new_key WHERE api_key = :old_key"),
{"old_key": old_key, "new_key": new_key},
)
return result.rowcount
@@ -83,6 +83,25 @@ class TokenUsageRepository:
},
)
def reassign_api_key(self, *, old_key: str, new_key: str) -> int:
"""Re-point historical rows from ``old_key`` to ``new_key``.
Called when an agent's key is rotated. Although per-agent analytics
also match on ``agent_id``, the 24h rate-limit window
(``sum_tokens_in_range`` / ``count_in_range``) and any legacy rows with
a NULL ``agent_id`` are keyed by the api-key string only. Rewriting the
key here preserves the running rate-limit window across a rotation (so
rotating cannot be used to reset it) and keeps legacy rows attributed.
Returns the number of rows updated.
"""
if not old_key or not new_key:
return 0
result = self._conn.execute(
text("UPDATE token_usage SET api_key = :new_key WHERE api_key = :old_key"),
{"old_key": old_key, "new_key": new_key},
)
return result.rowcount
def sum_tokens_in_range(
self,
*,
+1
View File
@@ -1455,6 +1455,7 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) {
mode={effectiveMode}
modalState={agentDetails}
setModalState={setAgentDetails}
onKeyRegenerated={(key) => setAgent((prev) => ({ ...prev, key }))}
/>
{shareModalOpen && agent.id && (
<ShareToTeamModal
@@ -3077,6 +3077,9 @@ function WorkflowBuilderInner() {
mode="edit"
modalState={agentDetails}
setModalState={setAgentDetails}
onKeyRegenerated={(key) =>
setCurrentAgent((prev) => ({ ...prev, key }))
}
/>
)}
</div>
+1
View File
@@ -27,6 +27,7 @@ const endpoints = {
TEMPLATE_AGENTS: '/api/template_agents',
ADOPT_AGENT: (id: string) => `/api/adopt_agent?id=${id}`,
AGENT_WEBHOOK: (id: string) => `/api/agent_webhook?id=${id}`,
REGENERATE_AGENT_KEY: (id: string) => `/api/regenerate_agent_key/${id}`,
EXPORT_AGENT: (id: string) => `/api/export_agent?id=${id}`,
IMPORT_AGENT_PLAN: '/api/import_agent/plan',
IMPORT_AGENT: '/api/import_agent',
+2
View File
@@ -66,6 +66,8 @@ const userService = {
apiClient.post(endpoints.USER.ADOPT_AGENT(id), {}, token),
getAgentWebhook: (id: string, token: string | null): Promise<any> =>
apiClient.get(endpoints.USER.AGENT_WEBHOOK(id), token),
regenerateAgentKey: (id: string, token: string | null): Promise<any> =>
apiClient.post(endpoints.USER.REGENERATE_AGENT_KEY(id), {}, token),
getPrompts: (token: string | null): Promise<any> =>
apiClient.get(endpoints.USER.PROMPTS, token),
createPrompt: (data: any, token: string | null): Promise<any> =>
+3 -1
View File
@@ -1001,7 +1001,9 @@
"webhookUrl": "Webhook-URL",
"generate": "Generieren",
"test": "Testen",
"learnMore": "Mehr erfahren"
"learnMore": "Mehr erfahren",
"resetKey": "Schlüssel zurücksetzen",
"resetKeyConfirm": "Möchten Sie den API-Schlüssel wirklich zurücksetzen? Der aktuelle Schlüssel funktioniert sofort nicht mehr und diese Aktion kann nicht rückgängig gemacht werden."
},
"importSpec": {
"title": "API-Spezifikation importieren",
+3 -1
View File
@@ -1002,7 +1002,9 @@
"webhookUrl": "Webhook URL",
"generate": "Generate",
"test": "Test",
"learnMore": "Learn more"
"learnMore": "Learn more",
"resetKey": "Reset key",
"resetKeyConfirm": "Are you sure you want to reset the API key? The current key will stop working immediately and this action cannot be undone."
},
"importSpec": {
"title": "Import API Specification",
+3 -1
View File
@@ -1001,7 +1001,9 @@
"webhookUrl": "Webhook URL",
"generate": "Generate",
"test": "Test",
"learnMore": "Learn more"
"learnMore": "Learn more",
"resetKey": "Restablecer clave",
"resetKeyConfirm": "¿Seguro que quieres restablecer la clave de API? La clave actual dejará de funcionar de inmediato y esta acción no se puede deshacer."
},
"importSpec": {
"title": "Importar especificación de API",
+3 -1
View File
@@ -1001,7 +1001,9 @@
"webhookUrl": "Webhook URL",
"generate": "Generate",
"test": "Test",
"learnMore": "Learn more"
"learnMore": "Learn more",
"resetKey": "キーをリセット",
"resetKeyConfirm": "APIキーをリセットしてもよろしいですか?現在のキーは直ちに無効になり、この操作は元に戻せません。"
},
"importSpec": {
"title": "API仕様のインポート",
+3 -1
View File
@@ -1001,7 +1001,9 @@
"webhookUrl": "Webhook URL",
"generate": "Generate",
"test": "Test",
"learnMore": "Learn more"
"learnMore": "Learn more",
"resetKey": "Сбросить ключ",
"resetKeyConfirm": "Вы уверены, что хотите сбросить API-ключ? Текущий ключ немедленно перестанет работать, и это действие нельзя отменить."
},
"importSpec": {
"title": "Импорт спецификации API",
+3 -1
View File
@@ -1001,7 +1001,9 @@
"webhookUrl": "Webhook URL",
"generate": "Generate",
"test": "Test",
"learnMore": "Learn more"
"learnMore": "Learn more",
"resetKey": "重設金鑰",
"resetKeyConfirm": "確定要重設 API 金鑰嗎?目前的金鑰將立即停止運作,此操作無法復原。"
},
"importSpec": {
"title": "匯入 API 規格",
+3 -1
View File
@@ -1001,7 +1001,9 @@
"webhookUrl": "Webhook URL",
"generate": "Generate",
"test": "Test",
"learnMore": "Learn more"
"learnMore": "Learn more",
"resetKey": "重置密钥",
"resetKeyConfirm": "确定要重置 API 密钥吗?当前密钥将立即停止工作,此操作无法撤销。"
},
"importSpec": {
"title": "导入 API 规范",
+42
View File
@@ -11,6 +11,7 @@ import { Button } from '../components/ui/button';
import { Modal } from '../components/ui/modal';
import { ActiveState } from '../models/misc';
import { selectToken } from '../preferences/preferenceSlice';
import ConfirmationModal from './ConfirmationModal';
const baseURL = import.meta.env.VITE_BASE_URL;
@@ -19,6 +20,7 @@ type AgentDetailsModalProps = {
mode: 'new' | 'edit' | 'draft';
modalState: ActiveState;
setModalState: (state: ActiveState) => void;
onKeyRegenerated?: (key: string) => void;
};
export default function AgentDetailsModal({
@@ -26,6 +28,7 @@ export default function AgentDetailsModal({
mode,
modalState,
setModalState,
onKeyRegenerated,
}: AgentDetailsModalProps) {
const { t } = useTranslation();
const token = useSelector(selectToken);
@@ -35,6 +38,8 @@ export default function AgentDetailsModal({
);
const [apiKey, setApiKey] = useState<string | null>(null);
const [webhookUrl, setWebhookUrl] = useState<string | null>(null);
const [resetKeyConfirmState, setResetKeyConfirmState] =
useState<ActiveState>('INACTIVE');
const [loadingStates, setLoadingStates] = useState({
publicLink: false,
apiKey: false,
@@ -75,12 +80,26 @@ export default function AgentDetailsModal({
setLoading('webhook', false);
};
const handleRegenerateKey = async () => {
setLoading('apiKey', true);
const response = await userService.regenerateAgentKey(agent.id ?? '', token);
if (!response.ok) {
setLoading('apiKey', false);
return;
}
const data = await response.json();
setApiKey(data.key);
onKeyRegenerated?.(data.key);
setLoading('apiKey', false);
};
useEffect(() => {
setSharedToken(agent.shared_token ?? null);
setApiKey(agent.key ?? null);
}, [agent]);
return (
<>
<Modal
open={modalState === 'ACTIVE'}
onOpenChange={(o) => !o && setModalState('INACTIVE')}
@@ -174,6 +193,20 @@ export default function AgentDetailsModal({
/>
</a>
)}
{apiKey.includes('...') && (
<button
type="button"
onClick={() => setResetKeyConfirmState('ACTIVE')}
disabled={loadingStates.apiKey}
className="border-primary text-primary hover:bg-primary/90 ml-8 flex items-center justify-center gap-1 rounded-full border px-5 py-1.5 text-sm font-medium transition-colors hover:text-white disabled:opacity-60"
>
{loadingStates.apiKey ? (
<Spinner size="small" />
) : (
t('modals.agentDetails.resetKey')
)}
</button>
)}
</div>
</div>
) : (
@@ -238,5 +271,14 @@ export default function AgentDetailsModal({
</div>
</div>
</Modal>
<ConfirmationModal
message={t('modals.agentDetails.resetKeyConfirm')}
modalState={resetKeyConfirmState}
setModalState={setResetKeyConfirmState}
submitLabel={t('modals.agentDetails.resetKey')}
handleSubmit={handleRegenerateKey}
variant="danger"
/>
</>
);
}
+166
View File
@@ -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
+49
View File
@@ -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
@@ -33,6 +33,32 @@ class TestInsert:
assert mapping["user_id"] == "u1"
assert mapping["stacks"] == [{"component": "retriever", "data": {"docs": 3}}]
def test_inserts_with_agent_id(self, pg_conn):
repo = _repo(pg_conn)
agent_uuid = "11111111-1111-1111-1111-111111111111"
repo.insert(
activity_id="act-agent",
api_key="k-agent",
agent_id=agent_uuid,
)
row = pg_conn.execute(
text("SELECT agent_id FROM stack_logs WHERE activity_id = 'act-agent'")
).fetchone()
assert str(dict(row._mapping)["agent_id"]) == agent_uuid
def test_non_uuid_agent_id_coerced_to_null(self, pg_conn):
# A stray/legacy (non-UUID) id must not break the activity-log write.
repo = _repo(pg_conn)
repo.insert(
activity_id="act-badid",
api_key="k-badid",
agent_id="507f1f77bcf86cd799439011", # 24-hex Mongo ObjectId
)
row = pg_conn.execute(
text("SELECT agent_id FROM stack_logs WHERE activity_id = 'act-badid'")
).fetchone()
assert dict(row._mapping)["agent_id"] is None
def test_inserts_with_empty_stacks(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(activity_id="act-2", level="error")
@@ -171,3 +197,33 @@ class TestInsert:
text("SELECT stacks FROM stack_logs WHERE activity_id = 'act-small'")
).fetchone()
assert dict(row._mapping)["stacks"][0]["data"]["result"] == "ok"
class TestReassignApiKey:
def test_rewrites_matching_rows_regardless_of_user(self, pg_conn):
# stack_logs has no agent_id; on key rotation, rows must follow the
# key. Rows are stamped with the caller's user_id (not the owner),
# so the migration must NOT be scoped by user_id.
repo = _repo(pg_conn)
repo.insert(activity_id="ra-1", api_key="old-k", user_id="owner")
repo.insert(activity_id="ra-2", api_key="old-k", user_id="a-caller")
repo.insert(activity_id="ra-3", api_key="other-k", user_id="owner")
moved = repo.reassign_api_key(old_key="old-k", new_key="new-k")
assert moved == 2
assert pg_conn.execute(
text("SELECT COUNT(*) FROM stack_logs WHERE api_key = 'old-k'")
).scalar() == 0
assert pg_conn.execute(
text("SELECT COUNT(*) FROM stack_logs WHERE api_key = 'new-k'")
).scalar() == 2
# Unrelated key untouched.
assert pg_conn.execute(
text("SELECT COUNT(*) FROM stack_logs WHERE api_key = 'other-k'")
).scalar() == 1
def test_noop_on_blank_keys(self, pg_conn):
repo = _repo(pg_conn)
assert repo.reassign_api_key(old_key="", new_key="x") == 0
assert repo.reassign_api_key(old_key="x", new_key="") == 0
@@ -36,6 +36,31 @@ class TestInsert:
assert total == 30
class TestReassignApiKey:
def test_rewrites_and_preserves_rate_limit_window(self, pg_conn):
# Rotating an agent key must carry the running 24h usage window over
# to the new key so rotation cannot be used to reset the limit.
repo = _repo(pg_conn)
repo.insert(api_key="tu-old", prompt_tokens=20, generated_tokens=10)
repo.insert(api_key="tu-other", prompt_tokens=1, generated_tokens=1)
moved = repo.reassign_api_key(old_key="tu-old", new_key="tu-new")
assert moved == 1
window = dict(
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1)
)
assert repo.sum_tokens_in_range(api_key="tu-old", **window) == 0
assert repo.sum_tokens_in_range(api_key="tu-new", **window) == 30
# Unrelated key untouched.
assert repo.sum_tokens_in_range(api_key="tu-other", **window) == 2
def test_noop_on_blank_keys(self, pg_conn):
repo = _repo(pg_conn)
assert repo.reassign_api_key(old_key="", new_key="x") == 0
assert repo.reassign_api_key(old_key="x", new_key="") == 0
class TestModelId:
def _latest_model_id(self, conn, user_id):
return conn.execute(
+138
View File
@@ -0,0 +1,138 @@
"""Migration round-trip + backfill test for 0026_stack_logs_agent_id."""
from __future__ import annotations
import os
import subprocess
import sys
from pathlib import Path
import pytest
from sqlalchemy import text
pytestmark = pytest.mark.integration
def _alembic_ini() -> Path:
return Path(__file__).resolve().parents[3] / "application" / "alembic.ini"
def _run_alembic(url: str, *args: str) -> None:
subprocess.check_call(
[sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args],
timeout=60,
env={**os.environ, "POSTGRES_URI": url},
)
def _alembic_heads(url: str) -> list[str]:
out = subprocess.check_output(
[sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "heads"],
timeout=60,
env={**os.environ, "POSTGRES_URI": url},
text=True,
)
return [line for line in out.splitlines() if line.strip()]
def _alembic_version(conn) -> str:
return conn.execute(text("SELECT version_num FROM alembic_version")).scalar()
def _column_exists(conn, table: str, column: str) -> bool:
row = conn.execute(
text(
"SELECT 1 FROM information_schema.columns "
"WHERE table_name = :t AND column_name = :c AND table_schema = 'public'"
),
{"t": table, "c": column},
).fetchone()
return row is not None
def _index_exists(conn, name: str) -> bool:
row = conn.execute(
text("SELECT 1 FROM pg_indexes WHERE indexname = :n"),
{"n": name},
).fetchone()
return row is not None
_0026 = "0026_stack_logs_agent_id"
_0025 = "0025_artifacts"
class TestMigration0026RoundTrip:
def test_single_head(self, pg_engine):
url = pg_engine.url.render_as_string(hide_password=False)
assert len(_alembic_heads(url)) == 1
def test_head_has_agent_id_column_and_index(self, pg_engine):
with pg_engine.connect() as conn:
assert _alembic_version(conn) >= _0026
assert _column_exists(conn, "stack_logs", "agent_id")
assert _index_exists(conn, "ix_stack_logs_agent_id")
def test_downgrade_drops_then_upgrade_restores(self, pg_engine):
url = pg_engine.url.render_as_string(hide_password=False)
_run_alembic(url, "downgrade", _0025)
with pg_engine.connect() as conn:
assert _alembic_version(conn) == _0025
assert not _column_exists(conn, "stack_logs", "agent_id")
assert not _index_exists(conn, "ix_stack_logs_agent_id")
_run_alembic(url, "upgrade", "head")
with pg_engine.connect() as conn:
assert _alembic_version(conn) >= _0026
assert _column_exists(conn, "stack_logs", "agent_id")
def test_upgrade_backfills_agent_id_from_api_key(self, pg_engine):
"""A legacy stack_logs row (written before the column existed) is
attributed to its agent by matching the stored api_key to
``agents.key`` during the upgrade."""
from application.storage.db.repositories.agents import AgentsRepository
url = pg_engine.url.render_as_string(hide_password=False)
_run_alembic(url, "downgrade", _0025)
with pg_engine.begin() as conn:
agent = AgentsRepository(conn).create(
"u-mig26", "a", "published", key="mig26-key",
)
agent_id = str(agent["id"])
# Pre-0026 shape: no agent_id column.
conn.execute(
text(
"INSERT INTO stack_logs "
"(activity_id, endpoint, level, api_key, stacks) "
"VALUES ('mig26-act', 'webhook', 'info', 'mig26-key', "
"'[]'::jsonb)"
)
)
# A row whose key matches no agent must stay NULL.
conn.execute(
text(
"INSERT INTO stack_logs "
"(activity_id, endpoint, level, api_key, stacks) "
"VALUES ('mig26-orphan', 'webhook', 'info', 'no-such-key', "
"'[]'::jsonb)"
)
)
_run_alembic(url, "upgrade", "head")
with pg_engine.connect() as conn:
matched = conn.execute(
text(
"SELECT agent_id FROM stack_logs "
"WHERE activity_id = 'mig26-act'"
)
).fetchone()
orphan = conn.execute(
text(
"SELECT agent_id FROM stack_logs "
"WHERE activity_id = 'mig26-orphan'"
)
).fetchone()
assert str(matched._mapping["agent_id"]) == agent_id
assert orphan._mapping["agent_id"] is None