diff --git a/application/alembic/versions/0026_stack_logs_agent_id.py b/application/alembic/versions/0026_stack_logs_agent_id.py new file mode 100644 index 00000000..8e365d9a --- /dev/null +++ b/application/alembic/versions/0026_stack_logs_agent_id.py @@ -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;") diff --git a/application/api/user/agents/routes.py b/application/api/user/agents/routes.py index a58f526c..2104788b 100644 --- a/application/api/user/agents/routes.py +++ b/application/api/user/agents/routes.py @@ -28,6 +28,12 @@ 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.conversations import ConversationsRepository +from application.storage.db.repositories.shared_conversations import ( + SharedConversationsRepository, +) +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 +1299,106 @@ class UpdateAgent(Resource): return make_response(jsonify(response_data), 200) +@agents_ns.route("/regenerate_agent_key/") +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 every api_key reference to the new key so rotation + # doesn't break history or live links. token_usage/stack_logs + # (0026) carry agent_id, but rows can still be api_key-only: + # conversations set api_key without an agent_id unless the caller + # passes one, stack_logs rows may have a NULL agent_id, and the + # 24h rate-limit window is keyed by api_key only (token_usage). + # shared_conversations is functional, not just analytics: the + # public share endpoint hands its stored api_key to the widget, + # so a stale key would break promptable shared links. + # All rewrites run in the same transaction as the key 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 + ) + ConversationsRepository(conn).reassign_api_key( + old_key=old_key, new_key=new_key + ) + SharedConversationsRepository(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") diff --git a/application/api/user/analytics/routes.py b/application/api/user/analytics/routes.py index 9bf3ba1a..9a193068 100644 --- a/application/api/user/analytics/routes.py +++ b/application/api/user/analytics/routes.py @@ -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 diff --git a/application/logging.py b/application/logging.py index d9f5b768..5c03f2a2 100644 --- a/application/logging.py +++ b/application/logging.py @@ -5,7 +5,7 @@ import time import logging import uuid -from typing import Any, Callable, Dict, Generator, List +from typing import Any, Callable, Dict, Generator, List, Optional from application.core import log_context from application.storage.db.repositories.stack_logs import StackLogsRepository @@ -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: Optional[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), diff --git a/application/storage/db/models.py b/application/storage/db/models.py index 261e5c1b..bdde2263 100644 --- a/application/storage/db/models.py +++ b/application/storage/db/models.py @@ -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()), diff --git a/application/storage/db/repositories/conversations.py b/application/storage/db/repositories/conversations.py index dc652c99..97f8c56f 100644 --- a/application/storage/db/repositories/conversations.py +++ b/application/storage/db/repositories/conversations.py @@ -574,6 +574,27 @@ class ConversationsRepository: self._reap_artifact_storage(paths) return result.rowcount + def reassign_api_key(self, *, old_key: str, new_key: str) -> int: + """Re-point conversation rows from ``old_key`` to ``new_key``. + + Called when an agent's key is rotated. Per-agent analytics match + ``c.api_key = :api_key OR c.agent_id = :agent_pg_id``, but a + conversation created from api-key traffic may have ``api_key`` set and + ``agent_id`` NULL (``conversation_service`` only sets ``agent_id`` when + the caller passes one). Rewriting ``api_key`` keeps those rows attached + to the agent after rotation instead of orphaning them. + + Matched by ``api_key`` only — ``agents.key`` is globally unique so a key + maps to exactly one agent. Returns the number of rows updated. + """ + if not old_key or not new_key: + return 0 + result = self._conn.execute( + text("UPDATE conversations SET api_key = :new_key WHERE api_key = :old_key"), + {"old_key": old_key, "new_key": new_key}, + ) + return result.rowcount + @staticmethod def _reap_artifact_storage(paths: list) -> None: """Best-effort delete artifact bytes whose rows were cascade-removed; never raises.""" diff --git a/application/storage/db/repositories/shared_conversations.py b/application/storage/db/repositories/shared_conversations.py index 4d3bf40f..5681cbef 100644 --- a/application/storage/db/repositories/shared_conversations.py +++ b/application/storage/db/repositories/shared_conversations.py @@ -211,3 +211,31 @@ class SharedConversationsRepository: {"conv_id": conversation_id}, ) return [row_to_dict(r) for r in result.fetchall()] + + def reassign_api_key(self, *, old_key: str, new_key: str) -> int: + """Re-point promptable share rows from ``old_key`` to ``new_key``. + + A *promptable* shared-conversation link stores the backing agent's key + in ``api_key`` and the public share endpoint returns it so the shared + widget can keep prompting. If that agent's key is rotated without this + rewrite, the endpoint would hand out an invalidated key and break every + existing shared link. Rewriting keeps those links working with the new + key. + + Matched by ``api_key`` only — ``agents.key`` is globally unique, so a + key maps to exactly one agent. The uniform rewrite cannot violate the + ``(conversation_id, user_id, is_promptable, first_n_queries, + COALESCE(api_key, ''))`` partial unique index: rows already differ on + the non-key columns, and ``new_key`` is freshly generated so no row + already holds it. Returns the number of rows updated. + """ + if not old_key or not new_key: + return 0 + result = self._conn.execute( + text( + "UPDATE shared_conversations SET api_key = :new_key " + "WHERE api_key = :old_key" + ), + {"old_key": old_key, "new_key": new_key}, + ) + return result.rowcount diff --git a/application/storage/db/repositories/stack_logs.py b/application/storage/db/repositories/stack_logs.py index a5ac17fc..750376d6 100644 --- a/application/storage/db/repositories/stack_logs.py +++ b/application/storage/db/repositories/stack_logs.py @@ -10,6 +10,7 @@ Covers the single operation the legacy Mongo code performs: from __future__ import annotations import json +import uuid from datetime import datetime from typing import Optional @@ -52,16 +53,28 @@ 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. Parse it with ``uuid.UUID`` and coerce + # anything that isn't a valid UUID (e.g. a 24-hex legacy Mongo ObjectId) + # to NULL so a stray/legacy id never reaches ``CAST(... AS uuid)`` and + # breaks the activity-log write. + agent_id_uuid: Optional[str] = None + if agent_id: + try: + agent_id_uuid = str(uuid.UUID(str(agent_id))) + except (ValueError, AttributeError, TypeError): + agent_id_uuid = None 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 +86,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 +97,28 @@ 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``. + + Since migration 0026 ``stack_logs`` has an ``agent_id`` column and the + analytics ``webhook_where`` / ``system_where`` clauses match + ``agent_id`` first with ``api_key`` as a fallback, so most rows already + survive a key rotation via ``agent_id``. This rewrite still runs to keep + the ``api_key`` column consistent and to re-attach any rows whose + ``agent_id`` is NULL (e.g. a log written without agent context). + + 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 rewrite 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 diff --git a/application/storage/db/repositories/token_usage.py b/application/storage/db/repositories/token_usage.py index 127d213e..c23ec90b 100644 --- a/application/storage/db/repositories/token_usage.py +++ b/application/storage/db/repositories/token_usage.py @@ -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, *, diff --git a/frontend/src/agents/NewAgent.tsx b/frontend/src/agents/NewAgent.tsx index c12ba288..1f93efe2 100644 --- a/frontend/src/agents/NewAgent.tsx +++ b/frontend/src/agents/NewAgent.tsx @@ -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 && ( + setCurrentAgent((prev) => ({ ...prev, key })) + } /> )} diff --git a/frontend/src/api/endpoints.ts b/frontend/src/api/endpoints.ts index edc89b97..c32d62e9 100644 --- a/frontend/src/api/endpoints.ts +++ b/frontend/src/api/endpoints.ts @@ -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', diff --git a/frontend/src/api/services/userService.ts b/frontend/src/api/services/userService.ts index 563743b5..c4aef35b 100644 --- a/frontend/src/api/services/userService.ts +++ b/frontend/src/api/services/userService.ts @@ -66,6 +66,8 @@ const userService = { apiClient.post(endpoints.USER.ADOPT_AGENT(id), {}, token), getAgentWebhook: (id: string, token: string | null): Promise => apiClient.get(endpoints.USER.AGENT_WEBHOOK(id), token), + regenerateAgentKey: (id: string, token: string | null): Promise => + apiClient.post(endpoints.USER.REGENERATE_AGENT_KEY(id), {}, token), getPrompts: (token: string | null): Promise => apiClient.get(endpoints.USER.PROMPTS, token), createPrompt: (data: any, token: string | null): Promise => diff --git a/frontend/src/locale/de.json b/frontend/src/locale/de.json index 60e40eee..ebfdc15e 100644 --- a/frontend/src/locale/de.json +++ b/frontend/src/locale/de.json @@ -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", diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index 472961b0..b3bf62dc 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -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", diff --git a/frontend/src/locale/es.json b/frontend/src/locale/es.json index 700d8732..f3f577c9 100644 --- a/frontend/src/locale/es.json +++ b/frontend/src/locale/es.json @@ -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", diff --git a/frontend/src/locale/jp.json b/frontend/src/locale/jp.json index 02157c05..5fbae862 100644 --- a/frontend/src/locale/jp.json +++ b/frontend/src/locale/jp.json @@ -1001,7 +1001,9 @@ "webhookUrl": "Webhook URL", "generate": "Generate", "test": "Test", - "learnMore": "Learn more" + "learnMore": "Learn more", + "resetKey": "キーをリセット", + "resetKeyConfirm": "APIキーをリセットしてもよろしいですか?現在のキーは直ちに無効になり、この操作は元に戻せません。" }, "importSpec": { "title": "API仕様のインポート", diff --git a/frontend/src/locale/ru.json b/frontend/src/locale/ru.json index df5d395f..da2c418e 100644 --- a/frontend/src/locale/ru.json +++ b/frontend/src/locale/ru.json @@ -1001,7 +1001,9 @@ "webhookUrl": "Webhook URL", "generate": "Generate", "test": "Test", - "learnMore": "Learn more" + "learnMore": "Learn more", + "resetKey": "Сбросить ключ", + "resetKeyConfirm": "Вы уверены, что хотите сбросить API-ключ? Текущий ключ немедленно перестанет работать, и это действие нельзя отменить." }, "importSpec": { "title": "Импорт спецификации API", diff --git a/frontend/src/locale/zh-TW.json b/frontend/src/locale/zh-TW.json index 03087484..1a56a7b5 100644 --- a/frontend/src/locale/zh-TW.json +++ b/frontend/src/locale/zh-TW.json @@ -1001,7 +1001,9 @@ "webhookUrl": "Webhook URL", "generate": "Generate", "test": "Test", - "learnMore": "Learn more" + "learnMore": "Learn more", + "resetKey": "重設金鑰", + "resetKeyConfirm": "確定要重設 API 金鑰嗎?目前的金鑰將立即停止運作,此操作無法復原。" }, "importSpec": { "title": "匯入 API 規格", diff --git a/frontend/src/locale/zh.json b/frontend/src/locale/zh.json index 202ff839..c1a4db88 100644 --- a/frontend/src/locale/zh.json +++ b/frontend/src/locale/zh.json @@ -1001,7 +1001,9 @@ "webhookUrl": "Webhook URL", "generate": "Generate", "test": "Test", - "learnMore": "Learn more" + "learnMore": "Learn more", + "resetKey": "重置密钥", + "resetKeyConfirm": "确定要重置 API 密钥吗?当前密钥将立即停止工作,此操作无法撤销。" }, "importSpec": { "title": "导入 API 规范", diff --git a/frontend/src/modals/AgentDetailsModal.tsx b/frontend/src/modals/AgentDetailsModal.tsx index 5b95c950..b2bb8ee3 100644 --- a/frontend/src/modals/AgentDetailsModal.tsx +++ b/frontend/src/modals/AgentDetailsModal.tsx @@ -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(null); const [webhookUrl, setWebhookUrl] = useState(null); + const [resetKeyConfirmState, setResetKeyConfirmState] = + useState('INACTIVE'); const [loadingStates, setLoadingStates] = useState({ publicLink: false, apiKey: false, @@ -75,12 +80,29 @@ export default function AgentDetailsModal({ setLoading('webhook', false); }; + const handleRegenerateKey = async () => { + setLoading('apiKey', true); + try { + const response = await userService.regenerateAgentKey( + agent.id ?? '', + token, + ); + if (!response.ok) return; + const data = await response.json(); + setApiKey(data.key); + onKeyRegenerated?.(data.key); + } finally { + setLoading('apiKey', false); + } + }; + useEffect(() => { setSharedToken(agent.shared_token ?? null); setApiKey(agent.key ?? null); }, [agent]); return ( + <> !o && setModalState('INACTIVE')} @@ -174,6 +196,20 @@ export default function AgentDetailsModal({ /> )} + {apiKey.includes('...') && ( + + )} ) : ( @@ -238,5 +274,14 @@ export default function AgentDetailsModal({ + + ); } diff --git a/tests/api/user/agents/test_routes_happy.py b/tests/api/user/agents/test_routes_happy.py index c483da35..079fec1c 100644 --- a/tests/api/user/agents/test_routes_happy.py +++ b/tests/api/user/agents/test_routes_happy.py @@ -1252,6 +1252,187 @@ 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 + conversations + shared_conversations rows + follow the key so analytics, the rate-limit window, and promptable shared + links are not orphaned/broken. The stack_logs row is stamped with a + *different* user_id (a caller, not the owner) to prove the migration is + not owner-scoped, and the conversation is created with api_key set but + agent_id NULL to prove the api_key-only path is covered. + """ + from application.api.user.agents.routes import RegenerateAgentKey + from application.storage.db.repositories.conversations import ( + ConversationsRepository, + ) + from application.storage.db.repositories.shared_conversations import ( + SharedConversationsRepository, + ) + 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, + ) + # api_key set, agent_id intentionally omitted (NULL) — the orphan case. + conv = ConversationsRepository(pg_conn).create( + user, "conv", api_key=old_key, + ) + # A promptable shared link stores the agent key; rotation must not + # leave it pointing at an invalidated key. + SharedConversationsRepository(pg_conn).create( + conv["id"], user, is_promptable=True, api_key=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"] + + from sqlalchemy import text + + def _count(table, key): + return pg_conn.execute( + text(f"SELECT COUNT(*) FROM {table} WHERE api_key = :k"), + {"k": key}, + ).scalar() + + assert _count("stack_logs", old_key) == 0 + assert _count("stack_logs", new_key) == 1 + assert _count("token_usage", old_key) == 0 + assert _count("token_usage", new_key) == 1 + assert _count("conversations", old_key) == 0 + assert _count("conversations", new_key) == 1 + assert _count("shared_conversations", old_key) == 0 + assert _count("shared_conversations", new_key) == 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 diff --git a/tests/api/user/test_analytics.py b/tests/api/user/test_analytics.py index 4afd3aec..4c47490c 100644 --- a/tests/api/user/test_analytics.py +++ b/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 diff --git a/tests/storage/db/repositories/test_conversations.py b/tests/storage/db/repositories/test_conversations.py index 5ca3475d..21b1c544 100644 --- a/tests/storage/db/repositories/test_conversations.py +++ b/tests/storage/db/repositories/test_conversations.py @@ -885,3 +885,32 @@ class TestUuidShapeGate: msg = repo.get_message_at(conv["id"], 0) assert msg is not None assert msg["prompt"] == "q" + + +class TestReassignApiKey: + def test_rewrites_matching_rows_including_null_agent(self, pg_conn): + # A conversation created from api-key traffic can have api_key set and + # agent_id NULL; on key rotation it must follow the key so per-agent + # analytics (api_key OR agent_id) don't orphan it. + repo = _repo(pg_conn) + repo.create("owner", "c1", api_key="conv-old") + repo.create("caller", "c2", api_key="conv-old") + repo.create("owner", "c3", api_key="conv-other") + + moved = repo.reassign_api_key(old_key="conv-old", new_key="conv-new") + assert moved == 2 + + assert pg_conn.execute( + text("SELECT COUNT(*) FROM conversations WHERE api_key = 'conv-old'") + ).scalar() == 0 + assert pg_conn.execute( + text("SELECT COUNT(*) FROM conversations WHERE api_key = 'conv-new'") + ).scalar() == 2 + assert pg_conn.execute( + text("SELECT COUNT(*) FROM conversations WHERE api_key = 'conv-other'") + ).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 diff --git a/tests/storage/db/repositories/test_shared_conversations.py b/tests/storage/db/repositories/test_shared_conversations.py index 3bb12e96..b7fd939b 100644 --- a/tests/storage/db/repositories/test_shared_conversations.py +++ b/tests/storage/db/repositories/test_shared_conversations.py @@ -115,3 +115,34 @@ class TestFindByUuidShapeGate: share = repo.create(conv["id"], "user-1", first_n_queries=1) found = repo.find_by_uuid(str(share["uuid"])) assert found is not None + + +class TestReassignApiKey: + def test_rewrites_promptable_share_keys(self, pg_conn): + # A promptable share stores the backing agent's key; on rotation it must + # follow the key so the public share endpoint keeps returning a valid one. + from sqlalchemy import text + + repo = _repo(pg_conn) + c1, c2, c3 = _conv(pg_conn), _conv(pg_conn), _conv(pg_conn) + repo.create(c1["id"], "u", is_promptable=True, api_key="sc-old") + repo.create(c2["id"], "u", is_promptable=True, api_key="sc-old") + repo.create(c3["id"], "u", is_promptable=True, api_key="sc-other") + + moved = repo.reassign_api_key(old_key="sc-old", new_key="sc-new") + assert moved == 2 + + def _count(key): + return pg_conn.execute( + text("SELECT COUNT(*) FROM shared_conversations WHERE api_key = :k"), + {"k": key}, + ).scalar() + + assert _count("sc-old") == 0 + assert _count("sc-new") == 2 + assert _count("sc-other") == 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 diff --git a/tests/storage/db/repositories/test_stack_logs.py b/tests/storage/db/repositories/test_stack_logs.py index 300ad5e1..e77286dc 100644 --- a/tests/storage/db/repositories/test_stack_logs.py +++ b/tests/storage/db/repositories/test_stack_logs.py @@ -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,35 @@ 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): + # On key rotation the rewrite is intentionally NOT scoped by user_id: + # agents.key is unique so a key maps to one agent, and rows are stamped + # with the caller's user_id (not the owner), so scoping by owner would + # skip caller-logged rows. (Rows also carry agent_id since 0026, but + # this rewrite still covers api_key consistency and NULL-agent_id rows.) + 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 diff --git a/tests/storage/db/repositories/test_token_usage.py b/tests/storage/db/repositories/test_token_usage.py index 39975095..c4781b4c 100644 --- a/tests/storage/db/repositories/test_token_usage.py +++ b/tests/storage/db/repositories/test_token_usage.py @@ -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( diff --git a/tests/storage/db/test_migration_0026.py b/tests/storage/db/test_migration_0026.py new file mode 100644 index 00000000..4a3d68e2 --- /dev/null +++ b/tests/storage/db/test_migration_0026.py @@ -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