From 6c63c91643f9eabc9368a7c841fbaf8b7acee420 Mon Sep 17 00:00:00 2001 From: Pavel Date: Mon, 27 Jul 2026 22:10:04 +0200 Subject: [PATCH] The conversations fix --- application/api/user/agents/routes.py | 18 ++++--- .../storage/db/repositories/conversations.py | 21 ++++++++ tests/api/user/agents/test_routes_happy.py | 48 ++++++++++--------- .../db/repositories/test_conversations.py | 29 +++++++++++ 4 files changed, 87 insertions(+), 29 deletions(-) diff --git a/application/api/user/agents/routes.py b/application/api/user/agents/routes.py index b8c6a6f0..9792febf 100644 --- a/application/api/user/agents/routes.py +++ b/application/api/user/agents/routes.py @@ -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 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/tests/api/user/agents/test_routes_happy.py b/tests/api/user/agents/test_routes_happy.py index 28e4177d..6c71c381 100644 --- a/tests/api/user/agents/test_routes_happy.py +++ b/tests/api/user/agents/test_routes_happy.py @@ -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 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