refactor: better db connection handling

agent/src/index.ts

@@ -303,6 +303,7 @@ function intializeDbCache(character: Character, db: IDatabaseCacheAdapter) {
 }
 
 async function startAgent(character: Character, directClient) {
+    let db: IDatabaseAdapter & IDatabaseCacheAdapter;
     try {
         character.id ??= stringToUuid(character.name);
         character.username ??= character.name;
@@ -314,7 +315,7 @@ async function startAgent(character: Character, directClient) {
             fs.mkdirSync(dataDir, { recursive: true });
         }
 
-        const db = initializeDatabase(dataDir);
+        db = initializeDatabase(dataDir);
 
         await db.init();
 
@@ -334,6 +335,9 @@ async function startAgent(character: Character, directClient) {
             error
         );
         console.error(error);
+        if (db) {
+            await db.close();
+        }
         throw error;
     }
 }

packages/adapter-postgres/src/index.ts

@@ -113,6 +113,10 @@ export class PostgresDatabaseAdapter
         await this.query(schema);
     }
 
+    async close() {
+        await this.pool.end();
+    }
+
     async testConnection(): Promise<boolean> {
         let client;
         try {

packages/adapter-sqlite/src/index.ts

@@ -80,6 +80,10 @@ export class SqliteDatabaseAdapter
         this.db.exec(sqliteTables);
     }
 
+    async close() {
+        this.db.close();
+    }
+
     async getAccountById(userId: UUID): Promise<Account | null> {
         const sql = "SELECT * FROM accounts WHERE id = ?";
         const account = this.db.prepare(sql).get(userId) as Account;

packages/adapter-sqljs/src/index.ts

@@ -18,7 +18,8 @@ import { Database } from "./types.ts";
 
 export class SqlJsDatabaseAdapter
     extends DatabaseAdapter<Database>
-    implements IDatabaseCacheAdapter {
+    implements IDatabaseCacheAdapter
+{
     constructor(db: Database) {
         super();
         this.db = db;
@@ -28,6 +29,10 @@ export class SqlJsDatabaseAdapter
         this.db.exec(sqliteTables);
     }
 
+    async close() {
+        this.db.close();
+    }
+
     async getRoom(roomId: UUID): Promise<UUID | null> {
         const sql = "SELECT id FROM rooms WHERE id = ?";
         const stmt = this.db.prepare(sql);
@@ -77,10 +82,14 @@ export class SqlJsDatabaseAdapter
         const placeholders = params.roomIds.map(() => "?").join(", ");
         const sql = `SELECT * FROM memories WHERE 'type' = ? AND agentId = ? AND roomId IN (${placeholders})`;
         const stmt = this.db.prepare(sql);
-        const queryParams = [params.tableName, params.agentId, ...params.roomIds];
-        console.log({ queryParams })
+        const queryParams = [
+            params.tableName,
+            params.agentId,
+            ...params.roomIds,
+        ];
+        console.log({ queryParams });
         stmt.bind(queryParams);
-        console.log({ queryParams })
+        console.log({ queryParams });
 
         const memories: Memory[] = [];
         while (stmt.step()) {

packages/adapter-supabase/src/index.ts

@@ -100,6 +100,14 @@ export class SupabaseDatabaseAdapter extends DatabaseAdapter {
         this.supabase = createClient(supabaseUrl, supabaseKey);
     }
 
+    async init() {
+        // noop
+    }
+
+    async close() {
+        // noop
+    }
+
     async getMemoriesByRoomIds(params: {
         roomIds: UUID[];
         agentId?: UUID;

packages/core/src/database.ts

@@ -19,6 +19,19 @@ export abstract class DatabaseAdapter<DB = any> implements IDatabaseAdapter {
      * The database instance.
      */
     db: DB;
+
+    /**
+     * Optional initialization method for the database adapter.
+     * @returns A Promise that resolves when initialization is complete.
+     */
+    abstract init(): Promise<void>;
+
+    /**
+     * Optional close method for the database adapter.
+     * @returns A Promise that resolves when closing is complete.
+     */
+    abstract close(): Promise<void>;
+
     /**
      * Retrieves an account by its ID.
      * @param userId The UUID of the user account to retrieve.

packages/core/src/types.ts

@@ -297,9 +297,9 @@ export interface State {
     formattedConversation?: string;
 
     /** Optional formatted knowledge */
-    knowledge?: string,
+    knowledge?: string;
     /** Optional knowledge data */
-    knowledgeData?: KnowledgeItem[],
+    knowledgeData?: KnowledgeItem[];
 
     /** Additional dynamic properties */
     [key: string]: unknown;
@@ -714,7 +714,10 @@ export interface IDatabaseAdapter {
     db: any;
 
     /** Optional initialization */
-    init?(): Promise<void>;
+    init(): Promise<void>;
+
+    /** Close database connection */
+    close(): Promise<void>;
 
     /** Get account by ID */
     getAccountById(userId: UUID): Promise<Account | null>;