mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 16:14:10 +00:00
The conversations fix
This commit is contained in:
1 parent
10e1cbb359
commit
6c63c91643
4 files changed
+87
-29
No files matched your search
@@ -28,6 +28,7 @@ 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.stack_logs import StackLogsRepository
|
||||
from application.storage.db.repositories.token_usage import TokenUsageRepository
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
@@ -1359,19 +1360,22 @@ class RegenerateAgentKey(Resource):
|
||||
500,
|
||||
)
|
||||
|
||||
# Re-point key-string history to the new key. conversations,
|
||||
# token_usage and stack_logs (since 0026) all carry agent_id, so
|
||||
# per-agent analytics survive rotation via agent_id regardless.
|
||||
# These rewrites still run to (a) preserve the 24h rate-limit
|
||||
# window, which is keyed by api_key only (token_usage), and (b)
|
||||
# keep the api_key column consistent / re-attach any legacy rows
|
||||
# whose agent_id is NULL. Runs in the same transaction as the swap.
|
||||
# Re-point key-string history to the new key so per-agent
|
||||
# analytics survive rotation. 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).
|
||||
# All three rewrites run 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
|
||||
)
|
||||
ConversationsRepository(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
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -1343,11 +1343,16 @@ class TestRegenerateAgentKey:
|
||||
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."""
|
||||
"""stack_logs + token_usage + conversations 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, 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.stack_logs import (
|
||||
StackLogsRepository,
|
||||
)
|
||||
@@ -1369,6 +1374,10 @@ class TestRegenerateAgentKey:
|
||||
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.
|
||||
ConversationsRepository(pg_conn).create(
|
||||
user, "conv", api_key=old_key,
|
||||
)
|
||||
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/regenerate_agent_key/{agent['id']}", method="POST"
|
||||
@@ -1380,24 +1389,19 @@ class TestRegenerateAgentKey:
|
||||
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 _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
|
||||
|
||||
def test_db_error_returns_500(self, app):
|
||||
from application.api.user.agents.routes import RegenerateAgentKey
|
||||
|
||||
@@ -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
|
||||
Reference in new issue
Block a user