feat: Implement database migration system and update agent/session handling

2025-12-24 10:40:07 +08:00 · 2025-09-17 14:59:08 +08:00 · 2025-09-17 14:59:08 +08:00 · c196a02c95
commit c196a02c95
parent d1ff8591a6
8 changed files with 192 additions and 147 deletions
--- a/package.json
+++ b/package.json
@ -44,10 +44,10 @@
    "release": "node scripts/version.js",
    "publish": "yarn build:check && yarn release patch push",
    "pulish:artifacts": "cd packages/artifacts && npm publish && cd -",
-    "agents:generate": "drizzle-kit generate --config src/main/services/agents/drizzle.config.ts",
-    "agents:push": "drizzle-kit push --config src/main/services/agents/drizzle.config.ts",
-    "agents:studio": "drizzle-kit studio --config src/main/services/agents/drizzle.config.ts",
-    "agents:drop": "drizzle-kit drop --config src/main/services/agents/drizzle.config.ts",
+    "agents:generate": "NODE_ENV='development' drizzle-kit generate --config src/main/services/agents/drizzle.config.ts",
+    "agents:push": "NODE_ENV='development' drizzle-kit push --config src/main/services/agents/drizzle.config.ts",
+    "agents:studio": "NODE_ENV='development' drizzle-kit studio --config src/main/services/agents/drizzle.config.ts",
+    "agents:drop": "NODE_ENV='development' drizzle-kit drop --config src/main/services/agents/drizzle.config.ts",
    "generate:icons": "electron-icon-builder --input=./build/logo.png --output=build",
    "analyze:renderer": "VISUALIZER_RENDERER=true yarn build",
    "analyze:main": "VISUALIZER_MAIN=true yarn build",
--- a/src/main/services/agents/database/drizzle/0000_dry_luke_cage.sql
+++ b/src/main/services/agents/database/drizzle/0000_dry_luke_cage.sql
@ -1,3 +1,10 @@
+--> statement-breakpoint
+CREATE TABLE `migrations` (
+	`version` integer PRIMARY KEY NOT NULL,
+	`tag` text NOT NULL,
+	`executed_at` integer NOT NULL
+);
+
 CREATE TABLE `agents` (
 	`id` text PRIMARY KEY NOT NULL,
 	`type` text NOT NULL,
@ -14,9 +21,11 @@ CREATE TABLE `agents` (
 	`created_at` text NOT NULL,
 	`updated_at` text NOT NULL
 );
+
 --> statement-breakpoint
 CREATE TABLE `sessions` (
 	`id` text PRIMARY KEY NOT NULL,
+	`agent_type` text NOT NULL,
 	`agent_id` text NOT NULL,
 	`name` text NOT NULL,
 	`description` text,
@ -31,6 +40,7 @@ CREATE TABLE `sessions` (
 	`created_at` text NOT NULL,
 	`updated_at` text NOT NULL
 );
+
 --> statement-breakpoint
 CREATE TABLE `session_messages` (
 	`id` integer PRIMARY KEY AUTOINCREMENT NOT NULL,
--- a/src/main/services/agents/database/drizzle/meta/0000_snapshot.json
+++ b/src/main/services/agents/database/drizzle/meta/0000_snapshot.json
@ -1,7 +1,7 @@
 {
  "version": "6",
  "dialect": "sqlite",
-  "id": "7ef99575-0fcf-471c-9da7-77e5cf8de6a2",
+  "id": "35efb412-0230-4767-9c76-7b7c4d40369f",
  "prevId": "00000000-0000-0000-0000-000000000000",
  "tables": {
    "agents": {
@ -112,114 +112,6 @@
      "uniqueConstraints": {},
      "checkConstraints": {}
    },
-    "sessions": {
-      "name": "sessions",
-      "columns": {
-        "id": {
-          "name": "id",
-          "type": "text",
-          "primaryKey": true,
-          "notNull": true,
-          "autoincrement": false
-        },
-        "agent_id": {
-          "name": "agent_id",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": true,
-          "autoincrement": false
-        },
-        "name": {
-          "name": "name",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": true,
-          "autoincrement": false
-        },
-        "description": {
-          "name": "description",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": false,
-          "autoincrement": false
-        },
-        "accessible_paths": {
-          "name": "accessible_paths",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": false,
-          "autoincrement": false
-        },
-        "instructions": {
-          "name": "instructions",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": false,
-          "autoincrement": false
-        },
-        "model": {
-          "name": "model",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": true,
-          "autoincrement": false
-        },
-        "plan_model": {
-          "name": "plan_model",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": false,
-          "autoincrement": false
-        },
-        "small_model": {
-          "name": "small_model",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": false,
-          "autoincrement": false
-        },
-        "mcps": {
-          "name": "mcps",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": false,
-          "autoincrement": false
-        },
-        "allowed_tools": {
-          "name": "allowed_tools",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": false,
-          "autoincrement": false
-        },
-        "configuration": {
-          "name": "configuration",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": false,
-          "autoincrement": false
-        },
-        "created_at": {
-          "name": "created_at",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": true,
-          "autoincrement": false
-        },
-        "updated_at": {
-          "name": "updated_at",
-          "type": "text",
-          "primaryKey": false,
-          "notNull": true,
-          "autoincrement": false
-        }
-      },
-      "indexes": {},
-      "foreignKeys": {},
-      "compositePrimaryKeys": {},
-      "uniqueConstraints": {},
-      "checkConstraints": {}
-    },
    "session_messages": {
      "name": "session_messages",
      "columns": {
@ -278,6 +170,152 @@
      "compositePrimaryKeys": {},
      "uniqueConstraints": {},
      "checkConstraints": {}
+    },
+    "migrations": {
+      "name": "migrations",
+      "columns": {
+        "version": {
+          "name": "version",
+          "type": "integer",
+          "primaryKey": true,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "tag": {
+          "name": "tag",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "executed_at": {
+          "name": "executed_at",
+          "type": "integer",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {},
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {},
+      "checkConstraints": {}
+    },
+    "sessions": {
+      "name": "sessions",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "text",
+          "primaryKey": true,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "agent_type": {
+          "name": "agent_type",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "agent_id": {
+          "name": "agent_id",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "name": {
+          "name": "name",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "description": {
+          "name": "description",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "accessible_paths": {
+          "name": "accessible_paths",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "instructions": {
+          "name": "instructions",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "model": {
+          "name": "model",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "plan_model": {
+          "name": "plan_model",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "small_model": {
+          "name": "small_model",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "mcps": {
+          "name": "mcps",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "allowed_tools": {
+          "name": "allowed_tools",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "configuration": {
+          "name": "configuration",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "updated_at": {
+          "name": "updated_at",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {},
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {},
+      "checkConstraints": {}
    }
  },
  "views": {},
--- a/src/main/services/agents/database/drizzle/meta/_journal.json
+++ b/src/main/services/agents/database/drizzle/meta/_journal.json
@ -5,8 +5,8 @@
    {
      "idx": 0,
      "version": "6",
-      "when": 1758035192486,
-      "tag": "0000_dry_luke_cage",
+      "when": 1758091173882,
+      "tag": "0000_confused_wendigo",
      "breakpoints": true
    }
  ]
--- a/src/main/services/agents/database/MigrationService.ts
+++ b/src/main/services/agents/database/MigrationService.ts
@ -1,5 +1,6 @@
 import { type Client } from '@libsql/client'
 import { loggerService } from '@logger'
+import { getResourcePath } from '@main/utils'
 import { type LibSQLDatabase } from 'drizzle-orm/libsql'
 import fs from 'fs'
 import path from 'path'
@ -29,15 +30,18 @@ export class MigrationService {
  constructor(db: LibSQLDatabase<typeof schema>, client: Client) {
    this.db = db
    this.client = client
-    this.migrationDir = path.join(__dirname, 'drizzle')
+    this.migrationDir = path.join(getResourcePath(), 'database', 'drizzle')
  }

  async runMigrations(): Promise<void> {
    try {
      logger.info('Starting migration check...')

-      // Ensure migrations table exists
-      await this.ensureMigrationsTable()
+      const hasMigrationsTable = await this.migrationsTableExists()
+
+      if (!hasMigrationsTable) {
+        logger.info('Migrations table not found; assuming fresh database state')
+      }

      // Read migration journal
      const journal = await this.readMigrationJournal()
@ -47,7 +51,9 @@ export class MigrationService {
      }

      // Get applied migrations
-      const appliedMigrations = await this.getAppliedMigrations()
+      const appliedMigrations = hasMigrationsTable
+        ? await this.getAppliedMigrations()
+        : []
      const appliedVersions = new Set(appliedMigrations.map((m) => Number(m.version)))

      const latestAppliedVersion = appliedMigrations.reduce(
@ -82,27 +88,14 @@ export class MigrationService {
    }
  }

-  private async ensureMigrationsTable(): Promise<void> {
+  private async migrationsTableExists(): Promise<boolean> {
    try {
-      const tableExists = await this.client.execute(`SELECT name FROM sqlite_master WHERE type='table' AND name='migrations'`)
-
-      if (tableExists.rows.length === 0) {
-        logger.info('Migrations table missing, creating...')
-
-        await this.client.execute(`
-          CREATE TABLE IF NOT EXISTS migrations (
-            version INTEGER PRIMARY KEY,
-            tag TEXT NOT NULL,
-            executed_at INTEGER NOT NULL
-          )
-        `)
-
-        logger.info('Migrations table created successfully')
-      } else {
-        logger.debug('Migrations table already exists')
-      }
+      const table = await this.client.execute(
+        `SELECT name FROM sqlite_master WHERE type='table' AND name='migrations'`
+      )
+      return table.rows.length > 0
    } catch (error) {
-      logger.error('Failed to ensure migrations table exists:', { error })
+      logger.error('Failed to check migrations table status:', { error })
      throw error
    }
  }
@ -147,7 +140,7 @@ export class MigrationService {

      // Read and execute SQL
      const sqlContent = fs.readFileSync(sqlFilePath, 'utf-8')
-      await this.client.execute(sqlContent)
+      await this.client.executeMultiple(sqlContent)

      // Record migration as applied (store journal idx as version for tracking)
      const newMigration: NewMigration = {
@ -156,6 +149,10 @@ export class MigrationService {
        executedAt: Date.now()
      }

+      if (!(await this.migrationsTableExists())) {
+        throw new Error('Migrations table missing after executing migration; cannot record progress')
+      }
+
      await this.db.insert(migrations).values(newMigration)

      const executionTime = Date.now() - startTime
--- a/src/main/services/agents/drizzle.config.ts
+++ b/src/main/services/agents/drizzle.config.ts
@ -5,12 +5,11 @@
 import os from 'node:os'
 import path from 'node:path'

-import { isDev } from '@main/constant'
 import { defineConfig } from 'drizzle-kit'
 import { app } from 'electron'

 function getDbPath() {
-  if (isDev) {
+  if (process.env.NODE_ENV === 'development') {
    return path.join(os.homedir(), '.cherrystudio', 'data', 'agents.db')
  }
  return path.join(app.getPath('userData'), 'agents.db')
@ -23,7 +22,7 @@ export const dbPath = resolvedDbPath
 export default defineConfig({
  dialect: 'sqlite',
  schema: './src/main/services/agents/database/schema/index.ts',
-  out: './src/main/services/agents/database/drizzle',
+  out: './resources/database/drizzle',
  dbCredentials: {
    url: `file:${resolvedDbPath}`
  },
--- a/tests/apis/agents/agents.http
+++ b/tests/apis/agents/agents.http
@ -1,5 +1,6 @@
@host=http://localhost:23333
@token=cs-sk-af798ed4-7cf5-4fd7-ae4b-df203b164194
+@agent_id=agent_1758092281575_tn9dxio9k

 ### List Agents
 GET {{host}}/v1/agents
@ -27,17 +28,17 @@ Content-Type: application/json
 }

 ### Get Agent Details
-GET {{host}}/v1/agents/agent_1757663884173_4tyeh3vqq
+GET {{host}}/v1/agents/{{agent_id}}
 Authorization: Bearer {{token}}
 Content-Type: application/json

 ### Delete Agent
-DELETE {{host}}/v1/agents/agent_1757663884173
+DELETE {{host}}/v1/agents/{{agent_id}}
 Authorization: Bearer {{token}}
 Content-Type: application/json

 ### Update Agent
-PATCH {{host}}/v1/agents/agent_1757663884173
+PATCH {{host}}/v1/agents/{{agent_id}}
 Authorization: Bearer {{token}}
 Content-Type: application/json

--- a/tests/apis/agents/sessions.http
+++ b/tests/apis/agents/sessions.http
@ -1,8 +1,8 @@

@host=http://localhost:23333
@token=cs-sk-af798ed4-7cf5-4fd7-ae4b-df203b164194
-@agent_id=agent_1758084953648_my4oxnbm3
-@session_id=session_1758087299124_2etavjo2x
+@agent_id=agent_1758092281575_tn9dxio9k
+@session_id=session_1758092305477_b0g0cmnkp

 ### List Sessions
 GET {{host}}/v1/agents/{{agent_id}}/sessions