diff --git a/docsgpt/alembic/versions/0038_connections.py b/docsgpt/alembic/versions/0038_connections.py index 962e5c4a..052a40db 100644 --- a/docsgpt/alembic/versions/0038_connections.py +++ b/docsgpt/alembic/versions/0038_connections.py @@ -31,7 +31,9 @@ Backfill (idempotent, only fills NULLs or unconverted rows): 5. API-key tools (Brave, Telegram, ntfy, PostgreSQL, custom MCP with a key, bearer token or basic auth) get one connection per distinct credential, re-encrypted into v2. The tool keeps its v1 copy for one release so a - rollback still works; the executor prefers the connection. + rollback still works; the executor prefers the connection. Downgrade + writes a v1 copy back to every tool on an API-key connection, including + tools added after the upgrade. 6. ``token_info`` and the secret parts of ``session_data`` (``tokens``, ``client_info``) are encrypted into ``encrypted_credentials`` and removed from the plaintext columns. This needs ``ENCRYPTION_SECRET_KEY`` set to @@ -386,9 +388,32 @@ def _decrypt_back(bind) -> None: from sqlalchemy import text - from docsgpt.security.encryption import CredentialDecryptionError, decrypt_json + from docsgpt.security.encryption import CredentialDecryptionError, decrypt_json, encrypt_credentials - # Tools keep their v1 secrets; connections made only for them go away. + # API-key connections go away, so their tools get a v1 copy back: tools + # added after the upgrade never had one, and a reconnect may have changed + # the key since the backfill. + linked = bind.execute( + text( + "SELECT t.id, t.user_id, c.encrypted_credentials FROM user_tools t " + "JOIN connector_sessions c ON c.id = t.connection_id " + "WHERE c.auth_kind = 'api_key' AND c.encrypted_credentials IS NOT NULL" + ) + ).fetchall() + for row in linked: + try: + credentials = decrypt_json(row.encrypted_credentials, row.user_id).get("credentials") + except CredentialDecryptionError: + continue + if not credentials: + continue + bind.execute( + text( + "UPDATE user_tools SET config = COALESCE(config, '{}'::jsonb) " + "|| jsonb_build_object('encrypted_credentials', CAST(:blob AS text)) WHERE id = :id" + ), + {"blob": encrypt_credentials(credentials, row.user_id), "id": row.id}, + ) bind.execute(text("UPDATE user_tools SET connection_id = NULL WHERE connection_id IN " "(SELECT id FROM connector_sessions WHERE auth_kind = 'api_key')")) bind.execute(text("DELETE FROM connector_sessions WHERE auth_kind = 'api_key'")) diff --git a/tests/storage/db/test_migration_0038.py b/tests/storage/db/test_migration_0038.py index 5b84179b..c99884c3 100644 --- a/tests/storage/db/test_migration_0038.py +++ b/tests/storage/db/test_migration_0038.py @@ -312,6 +312,37 @@ class TestMigration0038Credentials: assert api_rows == 0 _run_alembic(url, "upgrade", "head") + def test_downgrade_keeps_secrets_of_tools_made_after_upgrade(self, pg_engine): + from docsgpt.security.encryption import decrypt_credentials, encrypt_json + + url, _ = self._upgrade_with(pg_engine, _seed_secrets) + # A tool added through the wizard keeps its key only on the connection. + with pg_engine.begin() as conn: + connection_id = conn.execute( + text( + "INSERT INTO connector_sessions (user_id, provider, connector_key, auth_kind, status, " + "account_label, encrypted_credentials, session_data) VALUES ('erin', 'ntfy', 'ntfy', " + "'api_key', 'connected', '…9999', :blob, '{}'::jsonb) RETURNING id" + ), + {"blob": encrypt_json({"credentials": {"token": "ntfy-secret-9999"}}, "erin")}, + ).scalar() + tool_id = conn.execute( + text("INSERT INTO user_tools (user_id, name, config, connection_id) " + "VALUES ('erin', 'ntfy', CAST(:c AS jsonb), :cid) RETURNING id"), + {"c": json.dumps({"server_url": "https://ntfy.sh"}), "cid": connection_id}, + ).scalar() + _run_alembic(url, "downgrade", _0037) + with pg_engine.connect() as conn: + config = conn.execute(text("SELECT config FROM user_tools WHERE id = :i"), {"i": tool_id}).scalar() + assert config["server_url"] == "https://ntfy.sh" + assert decrypt_credentials(config["encrypted_credentials"], "erin") == {"token": "ntfy-secret-9999"} + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + relinked = conn.execute( + text("SELECT connection_id FROM user_tools WHERE id = :i"), {"i": tool_id} + ).scalar() + assert relinked is not None + def test_second_upgrade_is_a_no_op(self, pg_engine): from docsgpt.connectors.service import read_secrets