From 44562aa2819a78a1fa4f203d2cf193cd4119ead7 Mon Sep 17 00:00:00 2001 From: xr843 Date: Wed, 6 May 2026 17:42:03 +0800 Subject: [PATCH 1/2] fix(commands): purge tenant tool credentials on reset-encrypt-key-pair (#35396) The 'reset-encrypt-key-pair' CLI command rotates every tenant's asymmetric key and purges Provider / ProviderModel rows that hold LLM credentials encrypted under the old key. However it does not touch the tool provider tables that encrypt credentials under the same key: - tool_builtin_providers (BuiltinToolProvider.encrypted_credentials) - tool_api_providers (ApiToolProvider.credentials_str) - tool_mcp_providers (MCPToolProvider.encrypted_credentials/headers/server_url) Stale ciphertext in these tables surfaces as a 500 from /console/api/workspaces/current/tool-providers because the new key cannot decrypt it (#35396). This change extends the existing per-tenant purge to also delete rows from the three tool provider tables. The help text and confirmation prompt are updated to reflect the broader scope. The existing UX guarantee ('all LLM credentials will become invalid, requiring re-entry') already implies tool credentials, so this is a behavior fix rather than a contract change. Tests added under tests/unit_tests/commands/ following the existing test_upgrade_db.py pattern (mocked sessionmaker, no DB required): - aborts when EDITION is not SELF_HOSTED - purges every encrypted-credential table for a single tenant - iterates correctly across multiple tenants --- api/commands/system.py | 17 ++- .../commands/test_reset_encrypt_key_pair.py | 108 ++++++++++++++++++ 2 files changed, 122 insertions(+), 3 deletions(-) create mode 100644 api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py diff --git a/api/commands/system.py b/api/commands/system.py index 39b2e991ed..7755d3b5bc 100644 --- a/api/commands/system.py +++ b/api/commands/system.py @@ -14,6 +14,7 @@ from libs.rsa import generate_key_pair from models import Tenant from models.model import App, AppMode, Conversation from models.provider import Provider, ProviderModel +from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider logger = logging.getLogger(__name__) @@ -23,13 +24,16 @@ DB_UPGRADE_LOCK_TTL_SECONDS = 60 @click.command( "reset-encrypt-key-pair", help="Reset the asymmetric key pair of workspace for encrypt LLM credentials. " - "After the reset, all LLM credentials will become invalid, " - "requiring re-entry." + "After the reset, all LLM credentials and tool provider credentials " + "(builtin / API / MCP) will be purged, requiring re-entry. " "Only support SELF_HOSTED mode.", ) @click.confirmation_option( prompt=click.style( - "Are you sure you want to reset encrypt key pair? This operation cannot be rolled back!", fg="red" + "Are you sure you want to reset encrypt key pair? " + "This will also purge builtin / API / MCP tool provider records for every tenant. " + "This operation cannot be rolled back!", + fg="red", ) ) def reset_encrypt_key_pair(): @@ -53,6 +57,13 @@ def reset_encrypt_key_pair(): session.execute(delete(Provider).where(Provider.provider_type == "custom", Provider.tenant_id == tenant.id)) session.execute(delete(ProviderModel).where(ProviderModel.tenant_id == tenant.id)) + # Purge tool provider records that hold credentials encrypted under the + # tenant key. Leaving them in place causes /console/api/workspaces/current/ + # tool-providers to 500 because decryption fails on stale ciphertext (#35396). + session.execute(delete(BuiltinToolProvider).where(BuiltinToolProvider.tenant_id == tenant.id)) + session.execute(delete(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant.id)) + session.execute(delete(MCPToolProvider).where(MCPToolProvider.tenant_id == tenant.id)) + click.echo( click.style( f"Congratulations! The asymmetric key pair of workspace {tenant.id} has been reset.", diff --git a/api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py b/api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py new file mode 100644 index 0000000000..31b4d71d0f --- /dev/null +++ b/api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py @@ -0,0 +1,108 @@ +"""Unit tests for the reset-encrypt-key-pair CLI command (#35396). + +The command must purge every table that stores ciphertext encrypted with the +tenant's asymmetric key, otherwise stale rows cause downstream API failures +such as `/console/api/workspaces/current/tool-providers` returning 500. +""" + +from unittest.mock import MagicMock, patch + +import commands +from commands import system as system_commands +from models.provider import Provider, ProviderModel +from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider + + +def _invoke_reset() -> int: + try: + commands.reset_encrypt_key_pair.callback() + except SystemExit as e: + return int(e.code or 0) + return 0 + + +def _delete_targets(session_mock: MagicMock) -> list: + """Extract the model class targeted by each `delete(...)` call on the session.""" + targets = [] + for call in session_mock.execute.call_args_list: + stmt = call.args[0] + # `delete(Foo)` constructs a `Delete` statement whose entity is `Foo`. + try: + targets.append(stmt.table.name) + except AttributeError: + targets.append(repr(stmt)) + return targets + + +def test_reset_aborts_when_not_self_hosted(monkeypatch, capsys): + monkeypatch.setattr(system_commands.dify_config, "EDITION", "CLOUD") + + exit_code = _invoke_reset() + captured = capsys.readouterr() + + assert exit_code == 0 + assert "only for SELF_HOSTED" in captured.out + + +def test_reset_purges_provider_and_tool_tables_for_each_tenant(monkeypatch, capsys): + """The command must purge LLM provider rows AND every tool provider table + that stores ciphertext encrypted under the tenant key (#35396).""" + monkeypatch.setattr(system_commands.dify_config, "EDITION", "SELF_HOSTED") + monkeypatch.setattr(system_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}") + + fake_tenant = MagicMock(id="tenant-abc", encrypt_public_key="old-key") + session = MagicMock() + session.scalars.return_value.all.return_value = [fake_tenant] + + fake_sessionmaker = MagicMock() + fake_sessionmaker.begin.return_value.__enter__.return_value = session + fake_sessionmaker.begin.return_value.__exit__.return_value = False + + with ( + patch.object(system_commands, "db", MagicMock()), + patch.object(system_commands, "sessionmaker", return_value=fake_sessionmaker), + ): + exit_code = _invoke_reset() + + captured = capsys.readouterr() + assert exit_code == 0 + assert "tenant-abc" in captured.out + + # New key pair generated and assigned. + assert fake_tenant.encrypt_public_key == "new-key-tenant-abc" + + # Every encrypted-credential table should have been purged for this tenant. + table_names = _delete_targets(session) + expected = { + Provider.__tablename__, + ProviderModel.__tablename__, + BuiltinToolProvider.__tablename__, + ApiToolProvider.__tablename__, + MCPToolProvider.__tablename__, + } + assert expected.issubset(set(table_names)), f"missing purges: expected {expected}, got {table_names}" + + +def test_reset_iterates_all_tenants(monkeypatch, capsys): + """Multi-tenant deployments must purge every tenant, not just the first.""" + monkeypatch.setattr(system_commands.dify_config, "EDITION", "SELF_HOSTED") + monkeypatch.setattr(system_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}") + + tenants = [MagicMock(id=f"tenant-{i}", encrypt_public_key="old") for i in range(3)] + session = MagicMock() + session.scalars.return_value.all.return_value = tenants + + fake_sessionmaker = MagicMock() + fake_sessionmaker.begin.return_value.__enter__.return_value = session + fake_sessionmaker.begin.return_value.__exit__.return_value = False + + with ( + patch.object(system_commands, "db", MagicMock()), + patch.object(system_commands, "sessionmaker", return_value=fake_sessionmaker), + ): + _invoke_reset() + + # Five purges per tenant × 3 tenants = 15 execute calls. + assert session.execute.call_count == 15 + for tenant in tenants: + assert tenant.encrypt_public_key == f"new-key-{tenant.id}" From 89cef06df2dffb3c718b2088c1033218daeb535f Mon Sep 17 00:00:00 2001 From: xr843 <137012659+xr843@users.noreply.github.com> Date: Sat, 9 May 2026 12:35:47 +0800 Subject: [PATCH 2/2] chore: re-trigger CI after title-validation network timeout