1#!/usr/bin/env python3

2"""MCP tools for schema migration status and control."""

3

4 from __future__ import annotations

5

6 import typing as t

7 from typing import TYPE_CHECKING

8

9 from session_buddy.memory.migration import (

10 create_backup,

11 get_migration_status,

12 migrate_v1_to_v2,

13 needs_migration,

14)

15

16 if TYPE_CHECKING:

17 from fastmcp import FastMCP

18

19

20 def register_migration_tools(mcp: FastMCP) -> None:

21 @mcp.tool() # type: ignore[no-untyped-call]

22 async def migration_status() -> dict[str, t.Any]:

23 """Check migration status and progress."""

24 return get_migration_status()

25

26 @mcp.tool() # type: ignore[no-untyped-call]

27 async def trigger_migration(

28 create_backup_first: bool = True, dry_run: bool = False

29 ) -> dict[str, t.Any]:

30 """Manually trigger migration (with preview)."""

31 if create_backup_first and not dry_run:

32 backup_path = create_backup()

33 backup = str(backup_path)

34 else:

35 backup = None

36

37 result = migrate_v1_to_v2(dry_run=dry_run)

38 return {

39 "success": result.success,

40 "error": result.error,

41 "stats": result.stats or {},

42 "duration_seconds": result.duration_seconds,

43 "backup": backup,

44 "migration_needed": needs_migration() if not result.success else False,

45 }

46

47 @mcp.tool() # type: ignore[no-untyped-call]

48 async def rollback_migration(backup_path: str) -> dict[str, t.Any]:

49 """Restore database from a previous backup file path."""

50 from pathlib import Path

51

52 from session_buddy.memory.migration import get_schema_version, restore_backup

53

54 restore_backup(Path(backup_path))

55 version = get_schema_version()

56 status = get_migration_status()

57 return {"restored_version": version, "status": status}