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
This commit is contained in:
xr843 2026-05-06 17:42:03 +08:00
parent f3c3534e33
commit 44562aa281
2 changed files with 122 additions and 3 deletions

View File

@ -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.",

View File

@ -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}"