@taico/adk-session-store 0.1.0

package/README.md

@@ -0,0 +1,40 @@
+# @taico/adk-session-store
+
+SQLite-backed session storage for Google ADK.
+
+`SqliteSessionService` extends ADK `BaseSessionService` and is intended as a drop-in replacement for `InMemorySessionService` when you need persistence.
+
+## Installation
+
+```bash
+npm install @taico/adk-session-store
+```
+
+## Usage
+
+```ts
+import { Runner } from '@google/adk';
+import { SqliteSessionService } from '@taico/adk-session-store';
+
+const sessionService = new SqliteSessionService({
+  filename: './adk-sessions.sqlite',
+});
+
+const runner = new Runner({
+  appName: 'my-app',
+  agent,
+  sessionService,
+});
+```
+
+Use in-memory SQLite by omitting `filename` or by setting `filename: ':memory:'`.
+
+## API
+
+- `new SqliteSessionService(options?)`
+- `options.filename?: string` (defaults to `:memory:`)
+- `close(): Promise<void>` to cleanly close the SQLite connection
+
+## Example
+
+See `examples/basic.ts` for a minimal end-to-end usage snippet.

package/dist/index.d.ts

@@ -0,0 +1 @@
+export * from './sqlite-session-service.js';

package/dist/index.js

@@ -0,0 +1 @@
+export * from './sqlite-session-service.js';

package/dist/sqlite-session-service.d.ts

@@ -0,0 +1,22 @@
+import { AppendEventRequest, BaseSessionService, CreateSessionRequest, DeleteSessionRequest, Event, GetSessionRequest, ListSessionsRequest, ListSessionsResponse, Session } from '@google/adk';
+export interface SqliteSessionServiceOptions {
+    filename?: string;
+}
+export declare class SqliteSessionService extends BaseSessionService {
+    private readonly db;
+    private readonly ready;
+    constructor(options?: SqliteSessionServiceOptions);
+    createSession({ appName, userId, state, sessionId }: CreateSessionRequest): Promise<Session>;
+    getSession({ appName, userId, sessionId, config }: GetSessionRequest): Promise<Session | undefined>;
+    listSessions({ appName, userId }: ListSessionsRequest): Promise<ListSessionsResponse>;
+    deleteSession({ appName, userId, sessionId }: DeleteSessionRequest): Promise<void>;
+    appendEvent({ session, event }: AppendEventRequest): Promise<Event>;
+    close(): Promise<void>;
+    private initialize;
+    private mergeState;
+    private persistScopedState;
+    private cloneState;
+    private run;
+    private get;
+    private all;
+}

package/dist/sqlite-session-service.js

@@ -0,0 +1,215 @@
+import { BaseSessionService, State, createSession, } from '@google/adk';
+import * as sqlite3 from 'sqlite3';
+import { randomUUID } from 'node:crypto';
+const sqlite3Module = sqlite3.default ?? sqlite3;
+export class SqliteSessionService extends BaseSessionService {
+    db;
+    ready;
+    constructor(options = {}) {
+        super();
+        this.db = new sqlite3Module.Database(options.filename ?? ':memory:');
+        this.ready = this.initialize();
+    }
+    async createSession({ appName, userId, state, sessionId }) {
+        await this.ready;
+        const session = createSession({
+            id: sessionId ?? randomUUID(),
+            appName,
+            userId,
+            state: state ?? {},
+            events: [],
+            lastUpdateTime: Date.now(),
+        });
+        await this.run('INSERT OR REPLACE INTO sessions (app_name, user_id, session_id, state_json, last_update_time) VALUES (?, ?, ?, ?, ?)', [appName, userId, session.id, JSON.stringify(session.state), session.lastUpdateTime]);
+        return this.mergeState(session);
+    }
+    async getSession({ appName, userId, sessionId, config }) {
+        await this.ready;
+        const sessionRow = await this.get('SELECT app_name, user_id, session_id, state_json, last_update_time FROM sessions WHERE app_name = ? AND user_id = ? AND session_id = ?', [appName, userId, sessionId]);
+        if (!sessionRow) {
+            return undefined;
+        }
+        const eventRows = await this.all('SELECT event_json FROM events WHERE app_name = ? AND user_id = ? AND session_id = ? ORDER BY sequence ASC', [appName, userId, sessionId]);
+        let events = eventRows.map((row) => JSON.parse(row.event_json));
+        if (config?.numRecentEvents) {
+            events = events.slice(-config.numRecentEvents);
+        }
+        if (config?.afterTimestamp) {
+            let index = events.length - 1;
+            while (index >= 0) {
+                if (events[index].timestamp < config.afterTimestamp) {
+                    break;
+                }
+                index -= 1;
+            }
+            if (index >= 0) {
+                events = events.slice(index + 1);
+            }
+        }
+        const session = createSession({
+            id: sessionRow.session_id,
+            appName: sessionRow.app_name,
+            userId: sessionRow.user_id,
+            state: JSON.parse(sessionRow.state_json),
+            events,
+            lastUpdateTime: sessionRow.last_update_time,
+        });
+        return this.mergeState(session);
+    }
+    async listSessions({ appName, userId }) {
+        await this.ready;
+        const rows = await this.all('SELECT app_name, user_id, session_id, last_update_time FROM sessions WHERE app_name = ? AND user_id = ?', [appName, userId]);
+        return {
+            sessions: rows.map((row) => createSession({
+                id: row.session_id,
+                appName: row.app_name,
+                userId: row.user_id,
+                state: {},
+                events: [],
+                lastUpdateTime: row.last_update_time,
+            })),
+        };
+    }
+    async deleteSession({ appName, userId, sessionId }) {
+        await this.ready;
+        await this.run('DELETE FROM events WHERE app_name = ? AND user_id = ? AND session_id = ?', [appName, userId, sessionId]);
+        await this.run('DELETE FROM sessions WHERE app_name = ? AND user_id = ? AND session_id = ?', [appName, userId, sessionId]);
+    }
+    async appendEvent({ session, event }) {
+        await this.ready;
+        await super.appendEvent({ session, event });
+        session.lastUpdateTime = event.timestamp;
+        const storedSession = await this.get('SELECT state_json FROM sessions WHERE app_name = ? AND user_id = ? AND session_id = ?', [session.appName, session.userId, session.id]);
+        if (!storedSession) {
+            return event;
+        }
+        let storedState = JSON.parse(storedSession.state_json);
+        if (!event.partial && event.actions?.stateDelta) {
+            for (const [key, value] of Object.entries(event.actions.stateDelta)) {
+                if (!key.startsWith(State.TEMP_PREFIX)) {
+                    storedState[key] = value;
+                }
+            }
+            await this.persistScopedState(session.appName, session.userId, event.actions.stateDelta);
+        }
+        await this.run('UPDATE sessions SET state_json = ?, last_update_time = ? WHERE app_name = ? AND user_id = ? AND session_id = ?', [JSON.stringify(storedState), event.timestamp, session.appName, session.userId, session.id]);
+        if (!event.partial) {
+            await this.run('INSERT INTO events (app_name, user_id, session_id, event_id, timestamp, event_json) VALUES (?, ?, ?, ?, ?, ?)', [session.appName, session.userId, session.id, event.id, event.timestamp, JSON.stringify(event)]);
+        }
+        return event;
+    }
+    async close() {
+        await this.ready;
+        await new Promise((resolve, reject) => {
+            this.db.close((error) => {
+                if (error) {
+                    reject(error);
+                    return;
+                }
+                resolve();
+            });
+        });
+    }
+    async initialize() {
+        await this.run('PRAGMA foreign_keys = ON');
+        await this.run(`CREATE TABLE IF NOT EXISTS sessions (
+        app_name TEXT NOT NULL,
+        user_id TEXT NOT NULL,
+        session_id TEXT NOT NULL,
+        state_json TEXT NOT NULL,
+        last_update_time INTEGER NOT NULL,
+        PRIMARY KEY (app_name, user_id, session_id)
+      )`);
+        await this.run(`CREATE TABLE IF NOT EXISTS events (
+        sequence INTEGER PRIMARY KEY AUTOINCREMENT,
+        app_name TEXT NOT NULL,
+        user_id TEXT NOT NULL,
+        session_id TEXT NOT NULL,
+        event_id TEXT NOT NULL,
+        timestamp INTEGER NOT NULL,
+        event_json TEXT NOT NULL,
+        FOREIGN KEY (app_name, user_id, session_id)
+          REFERENCES sessions (app_name, user_id, session_id)
+          ON DELETE CASCADE
+      )`);
+        await this.run(`CREATE TABLE IF NOT EXISTS app_state (
+        app_name TEXT NOT NULL,
+        state_key TEXT NOT NULL,
+        value_json TEXT NOT NULL,
+        PRIMARY KEY (app_name, state_key)
+      )`);
+        await this.run(`CREATE TABLE IF NOT EXISTS user_state (
+        app_name TEXT NOT NULL,
+        user_id TEXT NOT NULL,
+        state_key TEXT NOT NULL,
+        value_json TEXT NOT NULL,
+        PRIMARY KEY (app_name, user_id, state_key)
+      )`);
+    }
+    async mergeState(session) {
+        const merged = createSession({
+            id: session.id,
+            appName: session.appName,
+            userId: session.userId,
+            state: this.cloneState(session.state),
+            events: session.events,
+            lastUpdateTime: session.lastUpdateTime,
+        });
+        const appRows = await this.all('SELECT state_key, value_json FROM app_state WHERE app_name = ?', [session.appName]);
+        for (const row of appRows) {
+            merged.state[`${State.APP_PREFIX}${row.state_key}`] = JSON.parse(row.value_json);
+        }
+        const userRows = await this.all('SELECT state_key, value_json FROM user_state WHERE app_name = ? AND user_id = ?', [session.appName, session.userId]);
+        for (const row of userRows) {
+            merged.state[`${State.USER_PREFIX}${row.state_key}`] = JSON.parse(row.value_json);
+        }
+        return merged;
+    }
+    async persistScopedState(appName, userId, delta) {
+        const entries = Object.entries(delta);
+        for (const [key, value] of entries) {
+            if (key.startsWith(State.APP_PREFIX)) {
+                await this.run('INSERT INTO app_state (app_name, state_key, value_json) VALUES (?, ?, ?) ON CONFLICT(app_name, state_key) DO UPDATE SET value_json = excluded.value_json', [appName, key.replace(State.APP_PREFIX, ''), JSON.stringify(value)]);
+            }
+            if (key.startsWith(State.USER_PREFIX)) {
+                await this.run('INSERT INTO user_state (app_name, user_id, state_key, value_json) VALUES (?, ?, ?, ?) ON CONFLICT(app_name, user_id, state_key) DO UPDATE SET value_json = excluded.value_json', [appName, userId, key.replace(State.USER_PREFIX, ''), JSON.stringify(value)]);
+            }
+        }
+    }
+    cloneState(state) {
+        return JSON.parse(JSON.stringify(state));
+    }
+    run(sql, params = []) {
+        return new Promise((resolve, reject) => {
+            this.db.run(sql, params, (error) => {
+                if (error) {
+                    reject(error);
+                    return;
+                }
+                resolve();
+            });
+        });
+    }
+    get(sql, params = []) {
+        return new Promise((resolve, reject) => {
+            this.db.get(sql, params, (error, row) => {
+                if (error) {
+                    reject(error);
+                    return;
+                }
+                resolve(row ?? undefined);
+            });
+        });
+    }
+    all(sql, params = []) {
+        return new Promise((resolve, reject) => {
+            this.db.all(sql, params, (error, rows) => {
+                if (error) {
+                    reject(error);
+                    return;
+                }
+                resolve(rows);
+            });
+        });
+    }
+}

package/package.json

@@ -0,0 +1,35 @@
+{
+  "name": "@taico/adk-session-store",
+  "version": "0.1.0",
+  "type": "module",
+  "private": false,
+  "license": "MIT",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "files": [
+    "dist",
+    "README.md",
+    "package.json"
+  ],
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js",
+      "default": "./dist/index.js"
+    },
+    "./package.json": "./package.json"
+  },
+  "scripts": {
+    "build": "tsc -p tsconfig.build.json",
+    "typecheck": "tsc -p tsconfig.build.json --noEmit",
+    "pack": "npm run build && npm pack",
+    "release": "npm run build && npm publish --access public"
+  },
+  "dependencies": {
+    "@google/adk": "^0.2.5",
+    "sqlite3": "^5.1.7"
+  },
+  "devDependencies": {
+    "typescript": "^5.6.3"
+  }
+}