From e32210301a0d69e29f5289b4b24d462e3ea2c15c Mon Sep 17 00:00:00 2001 From: fdc310 <2213070223@qq.com> Date: Fri, 24 Jul 2026 20:50:25 +0800 Subject: [PATCH] fix(agent): migrate official runner ids --- .../0015_migrate_official_runner_ids.py | 225 ++++++++++++++++++ .../persistence/test_migrations.py | 163 ++++++++++++- 2 files changed, 387 insertions(+), 1 deletion(-) create mode 100644 src/langbot/pkg/persistence/alembic/versions/0015_migrate_official_runner_ids.py diff --git a/src/langbot/pkg/persistence/alembic/versions/0015_migrate_official_runner_ids.py b/src/langbot/pkg/persistence/alembic/versions/0015_migrate_official_runner_ids.py new file mode 100644 index 000000000..dc0837a9c --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0015_migrate_official_runner_ids.py @@ -0,0 +1,225 @@ +"""Migrate official AgentRunner IDs to their marketplace identities. + +Revision ID: 0015_official_runner_ids +Revises: 0014_interaction_delivery +""" + +from __future__ import annotations + +import hashlib +import json +import typing + +import sqlalchemy as sa +from alembic import op + + +revision = '0015_official_runner_ids' +down_revision = '0014_interaction_delivery' +branch_labels = None +depends_on = None + + +_RUNNER_ID_RENAMES = { + 'plugin:langbot/acp-agent-runner/default': 'plugin:langbot-team/ACPAgentRunner/default', + 'plugin:langbot/claude-code-agent/default': 'plugin:langbot-team/ClaudeCodeAgent/default', + 'plugin:langbot/codex-agent/default': 'plugin:langbot-team/CodexAgent/default', + 'plugin:langbot/coze-agent/default': 'plugin:langbot-team/CozeAgent/default', + 'plugin:langbot/dashscope-agent/default': 'plugin:langbot-team/DashScopeAgent/default', + 'plugin:langbot/deerflow-agent/default': 'plugin:langbot-team/DeerFlowAgent/default', + 'plugin:langbot/dify-agent/default': 'plugin:langbot-team/DifyAgent/default', + 'plugin:langbot/langflow-agent/default': 'plugin:langbot-team/LangflowAgent/default', + 'plugin:langbot/n8n-agent/default': 'plugin:langbot-team/N8nAgent/default', + 'plugin:langbot/tbox-agent/default': 'plugin:langbot-team/TboxAgent/default', + 'plugin:langbot/weknora-agent/default': 'plugin:langbot-team/WeKnoraAgent/default', +} + +_JSON_COLUMNS = { + 'legacy_pipelines': ('uuid', ('config',)), + 'agents': ('uuid', ('config',)), + 'workflows': ('uuid', ('definition', 'global_config')), +} + +_TEXT_COLUMNS = { + 'agents': ('uuid', ('component_ref',)), + 'agent_run': ('id', ('runner_id', 'binding_id')), + 'agent_interaction': ('id', ('runner_id', 'binding_id')), + 'event_log': ('id', ('runner_id',)), + 'transcript': ('id', ('runner_id',)), +} + + +def _table_columns(table_name: str) -> set[str]: + inspector = sa.inspect(op.get_bind()) + if table_name not in inspector.get_table_names(): + return set() + return {column['name'] for column in inspector.get_columns(table_name)} + + +def _rewrite_text(value: typing.Any, renames: dict[str, str]) -> typing.Any: + if not isinstance(value, str): + return value + for old_id, new_id in renames.items(): + value = value.replace(old_id, new_id) + return value + + +def _rewrite_json_value(value: typing.Any, renames: dict[str, str]) -> typing.Any: + if isinstance(value, dict): + result: dict[str, typing.Any] = {} + renamed_items: list[tuple[str, typing.Any]] = [] + for key, item in value.items(): + rewritten_key = _rewrite_text(key, renames) + if rewritten_key != key: + renamed_items.append((rewritten_key, item)) + else: + result[key] = _rewrite_json_value(item, renames) + for rewritten_key, item in renamed_items: + result.setdefault(rewritten_key, _rewrite_json_value(item, renames)) + return result + if isinstance(value, list): + return [_rewrite_json_value(item, renames) for item in value] + return _rewrite_text(value, renames) + + +def _rewrite_json_columns(table_name: str, primary_key: str, columns: tuple[str, ...], renames: dict[str, str]) -> None: + available_columns = _table_columns(table_name) + if primary_key not in available_columns or not set(columns) <= available_columns: + return + + bind = op.get_bind() + selected_columns = ', '.join((primary_key, *columns)) + rows = bind.execute(sa.text(f'SELECT {selected_columns} FROM {table_name}')).mappings().all() + for row in rows: + updates: dict[str, str] = {} + for column in columns: + raw_value = row[column] + if raw_value is None: + continue + try: + decoded = json.loads(raw_value) if isinstance(raw_value, str) else raw_value + except (TypeError, ValueError): + continue + rewritten = _rewrite_json_value(decoded, renames) + if rewritten != decoded: + updates[column] = json.dumps(rewritten, ensure_ascii=False, separators=(',', ':')) + if not updates: + continue + assignments = ', '.join(f'{column} = :{column}' for column in updates) + bind.execute( + sa.text(f'UPDATE {table_name} SET {assignments} WHERE {primary_key} = :_pk'), + {**updates, '_pk': row[primary_key]}, + ) + + +def _rewrite_text_columns(table_name: str, primary_key: str, columns: tuple[str, ...], renames: dict[str, str]) -> None: + available_columns = _table_columns(table_name) + if primary_key not in available_columns or not set(columns) <= available_columns: + return + + bind = op.get_bind() + selected_columns = ', '.join((primary_key, *columns)) + rows = bind.execute(sa.text(f'SELECT {selected_columns} FROM {table_name}')).mappings().all() + for row in rows: + updates = { + column: rewritten for column in columns if (rewritten := _rewrite_text(row[column], renames)) != row[column] + } + if not updates: + continue + assignments = ', '.join(f'{column} = :{column}' for column in updates) + bind.execute( + sa.text(f'UPDATE {table_name} SET {assignments} WHERE {primary_key} = :_pk'), + {**updates, '_pk': row[primary_key]}, + ) + + +def _state_scope_key(row: typing.Mapping[str, typing.Any], runner_id: str, binding_identity: str) -> str | None: + scope = row['scope'] + parts = { + 'runner_id': runner_id, + 'binding_identity': binding_identity, + 'bot_id': row['bot_id'], + 'workspace_id': row['workspace_id'], + } + if scope == 'conversation': + parts.update(conversation_id=row['conversation_id'], thread_id=row['thread_id']) + elif scope == 'actor': + parts.update(actor_type=row['actor_type'] or 'user', actor_id=row['actor_id']) + elif scope == 'subject': + parts.update(subject_type=row['subject_type'] or 'unknown', subject_id=row['subject_id']) + elif scope != 'runner': + return None + + payload = {'version': 2, 'scope': scope, **parts} + raw = json.dumps(payload, sort_keys=True, separators=(',', ':'), ensure_ascii=False) + return f'{scope}:v2:{hashlib.sha256(raw.encode("utf-8")).hexdigest()}' + + +def _rewrite_runner_state(renames: dict[str, str]) -> None: + table_name = 'agent_runner_state' + required_columns = { + 'id', + 'runner_id', + 'binding_identity', + 'scope', + 'scope_key', + 'state_key', + 'bot_id', + 'workspace_id', + 'conversation_id', + 'thread_id', + 'actor_type', + 'actor_id', + 'subject_type', + 'subject_id', + } + if not required_columns <= _table_columns(table_name): + return + + bind = op.get_bind() + rows = bind.execute(sa.text(f'SELECT {", ".join(sorted(required_columns))} FROM {table_name}')).mappings().all() + for row in rows: + runner_id = _rewrite_text(row['runner_id'], renames) + binding_identity = _rewrite_text(row['binding_identity'], renames) + if runner_id == row['runner_id'] and binding_identity == row['binding_identity']: + continue + scope_key = _state_scope_key(row, runner_id, binding_identity) or row['scope_key'] + collision = bind.execute( + sa.text( + 'SELECT id FROM agent_runner_state ' + 'WHERE scope_key = :scope_key AND state_key = :state_key AND id != :id' + ), + {'scope_key': scope_key, 'state_key': row['state_key'], 'id': row['id']}, + ).scalar_one_or_none() + if collision is not None: + bind.execute(sa.text('DELETE FROM agent_runner_state WHERE id = :id'), {'id': row['id']}) + continue + bind.execute( + sa.text( + 'UPDATE agent_runner_state ' + 'SET runner_id = :runner_id, binding_identity = :binding_identity, scope_key = :scope_key ' + 'WHERE id = :id' + ), + { + 'runner_id': runner_id, + 'binding_identity': binding_identity, + 'scope_key': scope_key, + 'id': row['id'], + }, + ) + + +def _migrate(renames: dict[str, str]) -> None: + for table_name, (primary_key, columns) in _JSON_COLUMNS.items(): + _rewrite_json_columns(table_name, primary_key, columns, renames) + for table_name, (primary_key, columns) in _TEXT_COLUMNS.items(): + _rewrite_text_columns(table_name, primary_key, columns, renames) + _rewrite_runner_state(renames) + + +def upgrade() -> None: + _migrate(_RUNNER_ID_RENAMES) + + +def downgrade() -> None: + _migrate({new_id: old_id for old_id, new_id in _RUNNER_ID_RENAMES.items()}) diff --git a/tests/integration/persistence/test_migrations.py b/tests/integration/persistence/test_migrations.py index 53468e39a..a3e5f6055 100644 --- a/tests/integration/persistence/test_migrations.py +++ b/tests/integration/persistence/test_migrations.py @@ -10,6 +10,7 @@ Run: uv run pytest tests/integration/persistence/test_migrations.py -q from __future__ import annotations import json +import hashlib import pytest import sqlalchemy as sa @@ -58,7 +59,7 @@ class TestAlembicRevisionGraph: revisions = list(script.walk_revisions()) assert script.get_bases() == ['0001_baseline'] - assert script.get_heads() == ['0014_interaction_delivery'] + assert script.get_heads() == ['0015_official_runner_ids'] assert all(len(item.revision) <= 32 for item in revisions), { item.revision: len(item.revision) for item in revisions if len(item.revision) > 32 } @@ -345,6 +346,166 @@ class TestSQLiteMigrationUpgrade: 'pipeline-default': [], } + @pytest.mark.asyncio + async def test_official_runner_ids_migrate_without_losing_config_or_state(self, sqlite_engine): + """0015 rewrites persisted official Runner identities and state scope keys.""" + old_dify = 'plugin:langbot/dify-agent/default' + new_dify = 'plugin:langbot-team/DifyAgent/default' + old_codex = 'plugin:langbot/codex-agent/default' + new_codex = 'plugin:langbot-team/CodexAgent/default' + binding_id = f'pipeline_pipeline-1_{old_dify}' + scope_payload = { + 'version': 2, + 'scope': 'conversation', + 'runner_id': old_dify, + 'binding_identity': binding_id, + 'bot_id': 'bot-1', + 'workspace_id': None, + 'conversation_id': 'conversation-1', + 'thread_id': None, + } + old_scope_key = ( + 'conversation:v2:' + + hashlib.sha256(json.dumps(scope_payload, sort_keys=True, separators=(',', ':')).encode()).hexdigest() + ) + + async with sqlite_engine.begin() as conn: + await conn.execute(sa.text('CREATE TABLE legacy_pipelines (uuid TEXT PRIMARY KEY, config JSON NOT NULL)')) + await conn.execute( + sa.text('CREATE TABLE agents (uuid TEXT PRIMARY KEY, component_ref TEXT, config JSON NOT NULL)') + ) + await conn.execute( + sa.text( + 'CREATE TABLE workflows (' + 'uuid TEXT PRIMARY KEY, definition JSON NOT NULL, global_config JSON NOT NULL' + ')' + ) + ) + await conn.execute( + sa.text( + 'CREATE TABLE agent_runner_state (' + 'id INTEGER PRIMARY KEY, runner_id TEXT NOT NULL, binding_identity TEXT NOT NULL, ' + 'scope TEXT NOT NULL, scope_key TEXT NOT NULL, state_key TEXT NOT NULL, ' + 'bot_id TEXT, workspace_id TEXT, conversation_id TEXT, thread_id TEXT, ' + 'actor_type TEXT, actor_id TEXT, subject_type TEXT, subject_id TEXT, ' + 'UNIQUE(scope_key, state_key)' + ')' + ) + ) + for table_name, columns in { + 'agent_run': 'id INTEGER PRIMARY KEY, runner_id TEXT NOT NULL, binding_id TEXT', + 'agent_interaction': 'id INTEGER PRIMARY KEY, runner_id TEXT NOT NULL, binding_id TEXT', + 'event_log': 'id INTEGER PRIMARY KEY, runner_id TEXT', + 'transcript': 'id INTEGER PRIMARY KEY, runner_id TEXT', + }.items(): + await conn.execute(sa.text(f'CREATE TABLE {table_name} ({columns})')) + + pipeline_config = { + 'ai': { + 'runner': {'id': old_dify}, + 'runner_config': { + old_dify: {'api-key': 'preserved-secret'}, + old_codex: {'workspace': 'K:/workspace'}, + }, + } + } + await conn.execute( + sa.text('INSERT INTO legacy_pipelines (uuid, config) VALUES (:uuid, :config)'), + {'uuid': 'pipeline-1', 'config': json.dumps(pipeline_config)}, + ) + await conn.execute( + sa.text('INSERT INTO agents (uuid, component_ref, config) VALUES (:uuid, :ref, :config)'), + {'uuid': 'agent-1', 'ref': old_codex, 'config': json.dumps({'runner': {'id': old_codex}})}, + ) + await conn.execute( + sa.text( + 'INSERT INTO workflows (uuid, definition, global_config) ' + 'VALUES (:uuid, :definition, :global_config)' + ), + { + 'uuid': 'workflow-1', + 'definition': json.dumps({'runner_id': old_dify}), + 'global_config': json.dumps({'runner_id': old_codex}), + }, + ) + await conn.execute( + sa.text( + 'INSERT INTO agent_runner_state (' + 'id, runner_id, binding_identity, scope, scope_key, state_key, bot_id, ' + 'workspace_id, conversation_id, thread_id, actor_type, actor_id, subject_type, subject_id' + ') VALUES (' + '1, :runner_id, :binding_identity, :scope, :scope_key, :state_key, :bot_id, ' + 'NULL, :conversation_id, NULL, NULL, NULL, NULL, NULL' + ')' + ), + { + 'runner_id': old_dify, + 'binding_identity': binding_id, + 'scope': 'conversation', + 'scope_key': old_scope_key, + 'state_key': 'external.conversation_id', + 'bot_id': 'bot-1', + 'conversation_id': 'conversation-1', + }, + ) + for table_name in ('agent_run', 'agent_interaction'): + await conn.execute( + sa.text( + f'INSERT INTO {table_name} (id, runner_id, binding_id) VALUES (1, :runner_id, :binding_id)' + ), + {'runner_id': old_dify, 'binding_id': binding_id}, + ) + for table_name in ('event_log', 'transcript'): + await conn.execute( + sa.text(f'INSERT INTO {table_name} (id, runner_id) VALUES (1, :runner_id)'), + {'runner_id': old_codex}, + ) + + await run_alembic_stamp(sqlite_engine, '0014_interaction_delivery') + await run_alembic_upgrade(sqlite_engine, 'head') + + async with sqlite_engine.connect() as conn: + pipeline_raw = ( + await conn.execute(sa.text("SELECT config FROM legacy_pipelines WHERE uuid = 'pipeline-1'")) + ).scalar_one() + agent_row = ( + await conn.execute(sa.text("SELECT component_ref, config FROM agents WHERE uuid = 'agent-1'")) + ).one() + workflow_row = ( + await conn.execute(sa.text("SELECT definition, global_config FROM workflows WHERE uuid = 'workflow-1'")) + ).one() + state_row = ( + await conn.execute( + sa.text('SELECT runner_id, binding_identity, scope_key FROM agent_runner_state WHERE id = 1') + ) + ).one() + direct_runner_ids = { + table_name: ( + await conn.execute(sa.text(f'SELECT runner_id FROM {table_name} WHERE id = 1')) + ).scalar_one() + for table_name in ('agent_run', 'agent_interaction', 'event_log', 'transcript') + } + + pipeline = json.loads(pipeline_raw) + assert pipeline['ai']['runner']['id'] == new_dify + assert pipeline['ai']['runner_config'][new_dify] == {'api-key': 'preserved-secret'} + assert pipeline['ai']['runner_config'][new_codex] == {'workspace': 'K:/workspace'} + assert old_dify not in pipeline['ai']['runner_config'] + assert agent_row.component_ref == new_codex + assert json.loads(agent_row.config)['runner']['id'] == new_codex + assert json.loads(workflow_row.definition)['runner_id'] == new_dify + assert json.loads(workflow_row.global_config)['runner_id'] == new_codex + assert state_row.runner_id == new_dify + assert state_row.binding_identity == f'pipeline_pipeline-1_{new_dify}' + assert state_row.scope_key != old_scope_key + assert direct_runner_ids == { + 'agent_run': new_dify, + 'agent_interaction': new_dify, + 'event_log': new_codex, + 'transcript': new_codex, + } + assert await get_alembic_current(sqlite_engine) == '0015_official_runner_ids' + class TestSQLiteMigrationFreshDatabase: """Tests for fresh database workflow."""