From 1a7921a92143b60c3c0056d9fd4d0bce91bc82b5 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Tue, 29 Sep 2026 15:07:11 +0100 Subject: [PATCH] Keep non-object token_info values when encrypting connector sessions Migration 0040 dropped token_info values that were not objects before clearing the plaintext column. Objects stored as JSON strings are now decoded, and other values are kept in the envelope under legacy_token_info so downgrade restores them. --- docsgpt/alembic/versions/0040_connections.py | 50 ++++++++++++++++-- tests/storage/db/test_migration_0040.py | 54 ++++++++++++++++++++ 2 files changed, 100 insertions(+), 4 deletions(-) diff --git a/docsgpt/alembic/versions/0040_connections.py b/docsgpt/alembic/versions/0040_connections.py index beacde9d..33a18cc6 100644 --- a/docsgpt/alembic/versions/0040_connections.py +++ b/docsgpt/alembic/versions/0040_connections.py @@ -242,7 +242,8 @@ def _encrypt_session_secrets(bind) -> None: rows = bind.execute( text( """ - SELECT id, user_id, token_info, session_data FROM connector_sessions + SELECT id, user_id, token_info, token_info IS NOT NULL AS has_token_info, session_data + FROM connector_sessions WHERE encrypted_credentials IS NULL AND (token_info IS NOT NULL OR session_data ? 'tokens' OR session_data ? 'client_info') LIMIT :batch @@ -253,11 +254,16 @@ def _encrypt_session_secrets(bind) -> None: if not rows: return for row in rows: - token_info = row.token_info if isinstance(row.token_info, dict) else None + token_info = _token_info_object(row.token_info) session_data = dict(row.session_data or {}) secrets = {} - if token_info: + if token_info is not None: secrets["token_info"] = token_info + elif row.has_token_info: + # A value the app never read as a token (a bare string, a number, + # JSON ``null``): kept apart from ``token_info`` so downgrade can + # restore it, and the row is not selected again. + secrets["legacy_token_info"] = row.token_info for key in ("tokens", "client_info"): if key in session_data: secrets[key] = session_data.pop(key) @@ -288,6 +294,25 @@ def _encrypt_session_secrets(bind) -> None: ) +def _token_info_object(value): + """``token_info`` as an object, decoding one stored as a JSON string. + + Args: + value: The decoded ``token_info`` column value. + + Returns: + The object, or ``None`` when the value is not (and does not encode) one. + """ + import json + + if isinstance(value, str): + try: + value = json.loads(value) + except ValueError: + return None + return value if isinstance(value, dict) else None + + def _link_api_key_tools(bind) -> None: """5: one connection per distinct API credential, linked from its tools.""" from urllib.parse import urlparse @@ -459,13 +484,30 @@ def _decrypt_back(bind) -> None: "session_data = CAST(:session_data AS jsonb) WHERE id = :id" ), { - "token_info": json.dumps(secrets["token_info"]) if "token_info" in secrets else None, + "token_info": _restored_token_info(secrets), "session_data": json.dumps(session_data), "id": row.id, }, ) +def _restored_token_info(secrets: dict): + """The ``token_info`` JSON downgrade writes back, ``None`` for SQL NULL. + + Args: + secrets: The decrypted envelope. + + Returns: + The JSON text of ``token_info`` or of a preserved non-object value. + """ + import json + + for key in ("token_info", "legacy_token_info"): + if key in secrets: + return json.dumps(secrets[key]) + return None + + def _restore_account_index() -> None: # Several accounts per provider cannot survive the old unique index: keep # the most recently updated one. diff --git a/tests/storage/db/test_migration_0040.py b/tests/storage/db/test_migration_0040.py index a287f447..8cf924c3 100644 --- a/tests/storage/db/test_migration_0040.py +++ b/tests/storage/db/test_migration_0040.py @@ -298,6 +298,60 @@ class TestMigration0040Credentials: ).one()._mapping) assert read_secrets(row) == {"credentials": {"token": token}} + def test_non_object_token_info_is_kept_in_the_envelope(self, pg_engine): + from docsgpt.connectors.service import read_secrets + from docsgpt.security.encryption import decrypt_json + + values = { + # An object stored as a JSON string by an older writer. + "gina": json.dumps(json.dumps({"access_token": "dbl-at", "refresh_token": "dbl-rt", "scopes": "a b"})), + "hank": json.dumps("opaque-token"), + "ivan": "null", + } + + def seed(conn): + return { + user: conn.execute( + text( + "INSERT INTO connector_sessions (user_id, provider, session_token, status, token_info) " + "VALUES (:u, 'google_drive', :t, 'authorized', CAST(:ti AS jsonb)) RETURNING id" + ), + {"u": user, "t": f"tok-{user}", "ti": raw}, + ).scalar() + for user, raw in values.items() + } + + url, ids = self._upgrade_with(pg_engine, seed) + with pg_engine.connect() as conn: + rows = { + user: dict(conn.execute( + text("SELECT * FROM connector_sessions WHERE id = :i"), {"i": row_id} + ).one()._mapping) + for user, row_id in ids.items() + } + plaintext = conn.execute( + text("SELECT count(*) FROM connector_sessions WHERE id = ANY(:ids) AND token_info IS NOT NULL"), + {"ids": list(ids.values())}, + ).scalar() + assert plaintext == 0 + gina = rows["gina"] + assert read_secrets(gina)["token_info"]["refresh_token"] == "dbl-rt" + assert gina["has_refresh_token"] is True + assert gina["scopes"] == ["a", "b"] + assert decrypt_json(rows["hank"]["encrypted_credentials"], "hank")["legacy_token_info"] == "opaque-token" + assert "token_info" not in read_secrets(rows["hank"]) + + _run_alembic(url, "downgrade", _0037) + with pg_engine.connect() as conn: + restored = dict(conn.execute( + text("SELECT id, token_info::text FROM connector_sessions WHERE id = ANY(:ids)"), + {"ids": list(ids.values())}, + ).fetchall()) + assert json.loads(restored[ids["gina"]])["access_token"] == "dbl-at" + assert restored[ids["hank"]] == '"opaque-token"' + assert restored[ids["ivan"]] == "null" + _run_alembic(url, "upgrade", "head") + def test_oauth_mcp_tools_keep_member_credentials(self, pg_engine): _, ids = self._upgrade_with(pg_engine, _seed_secrets) with pg_engine.connect() as conn: