mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 18:15:02 +00:00
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:
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;")
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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,
|
||||
*,
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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> =>
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -1001,7 +1001,9 @@
|
||||
"webhookUrl": "Webhook URL",
|
||||
"generate": "Generate",
|
||||
"test": "Test",
|
||||
"learnMore": "Learn more"
|
||||
"learnMore": "Learn more",
|
||||
"resetKey": "キーをリセット",
|
||||
"resetKeyConfirm": "APIキーをリセットしてもよろしいですか?現在のキーは直ちに無効になり、この操作は元に戻せません。"
|
||||
},
|
||||
"importSpec": {
|
||||
"title": "API仕様のインポート",
|
||||
|
||||
@@ -1001,7 +1001,9 @@
|
||||
"webhookUrl": "Webhook URL",
|
||||
"generate": "Generate",
|
||||
"test": "Test",
|
||||
"learnMore": "Learn more"
|
||||
"learnMore": "Learn more",
|
||||
"resetKey": "Сбросить ключ",
|
||||
"resetKeyConfirm": "Вы уверены, что хотите сбросить API-ключ? Текущий ключ немедленно перестанет работать, и это действие нельзя отменить."
|
||||
},
|
||||
"importSpec": {
|
||||
"title": "Импорт спецификации API",
|
||||
|
||||
@@ -1001,7 +1001,9 @@
|
||||
"webhookUrl": "Webhook URL",
|
||||
"generate": "Generate",
|
||||
"test": "Test",
|
||||
"learnMore": "Learn more"
|
||||
"learnMore": "Learn more",
|
||||
"resetKey": "重設金鑰",
|
||||
"resetKeyConfirm": "確定要重設 API 金鑰嗎?目前的金鑰將立即停止運作,此操作無法復原。"
|
||||
},
|
||||
"importSpec": {
|
||||
"title": "匯入 API 規格",
|
||||
|
||||
@@ -1001,7 +1001,9 @@
|
||||
"webhookUrl": "Webhook URL",
|
||||
"generate": "Generate",
|
||||
"test": "Test",
|
||||
"learnMore": "Learn more"
|
||||
"learnMore": "Learn more",
|
||||
"resetKey": "重置密钥",
|
||||
"resetKeyConfirm": "确定要重置 API 密钥吗?当前密钥将立即停止工作,此操作无法撤销。"
|
||||
},
|
||||
"importSpec": {
|
||||
"title": "导入 API 规范",
|
||||
|
||||
@@ -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"
|
||||
/>
|
||||
</>
|
||||
);
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
Reference in new issue
Block a user