@taico/adk-session-store 0.3.5 → 0.4.0

package/README.md

@@ -1,8 +1,10 @@
 # @taico/adk-session-store
 
-[![Socket Badge](https://badge.socket.dev/npm/package/@taico/adk-session-store/0.3.5)](https://badge.socket.dev/npm/package/@taico/adk-session-store/0.3.5)
+[![Socket Badge](https://badge.socket.dev/npm/package/@taico/adk-session-store/0.4.0)](https://badge.socket.dev/npm/package/@taico/adk-session-store/0.4.0)
 
-SQLite-backed session storage for Google ADK.
+PostgreSQL and SQLite session storage for Google ADK. Taico production uses
+`PostgresSessionService`; the SQLite implementation remains available to package
+consumers that need an embedded store.
 
 `SqliteSessionService` extends ADK `BaseSessionService` and is intended as a drop-in replacement for `InMemorySessionService` when you need persistence.
 
@@ -14,6 +16,21 @@ npm install @taico/adk-session-store
 
 ## Usage
 
+### PostgreSQL
+
+```ts
+import { PostgresSessionService } from '@taico/adk-session-store';
+
+const sessionService = new PostgresSessionService({
+  connectionString: process.env.DATABASE_URL!,
+});
+```
+
+The application migration must create the `adk_*` tables before this service is
+used.
+
+### SQLite
+
 ```ts
 import { Runner } from '@google/adk';
 import { SqliteSessionService } from '@taico/adk-session-store';
@@ -33,6 +50,7 @@ Use in-memory SQLite by omitting `filename` or by setting `filename: ':memory:'`
 
 ## API
 
+- `new PostgresSessionService({ connectionString, ssl? })`
 - `new SqliteSessionService(options?)`
 - `options.filename?: string` (defaults to `:memory:`)
 - `close(): Promise<void>` to cleanly close the SQLite connection

package/dist/index.d.ts

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

package/dist/index.js

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

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

@@ -0,0 +1,23 @@
+import { AppendEventRequest, BaseSessionService, CreateSessionRequest, DeleteSessionRequest, Event, GetSessionRequest, ListSessionsRequest, ListSessionsResponse, Session } from "@google/adk";
+export interface PostgresSessionServiceOptions {
+    connectionString: string;
+    ssl?: boolean;
+}
+export declare class PostgresSessionService extends BaseSessionService {
+    private readonly pool;
+    private readonly ready;
+    constructor(options: PostgresSessionServiceOptions);
+    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 parseJson;
+    private cloneState;
+    private one;
+    private many;
+}

package/dist/postgres-session-service.js

@@ -0,0 +1,178 @@
+import { BaseSessionService, State, createSession, } from "@google/adk";
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+export class PostgresSessionService extends BaseSessionService {
+    pool;
+    ready;
+    constructor(options) {
+        super();
+        this.pool = new Pool({
+            connectionString: options.connectionString,
+            ssl: options.ssl ? { rejectUnauthorized: true } : false,
+        });
+        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.pool.query(`INSERT INTO adk_sessions (app_name, user_id, session_id, state_json, last_update_time)
+       VALUES ($1, $2, $3, $4::jsonb, $5)
+       ON CONFLICT (app_name, user_id, session_id) DO UPDATE SET
+         state_json = EXCLUDED.state_json,
+         last_update_time = EXCLUDED.last_update_time`, [
+            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.one(`SELECT app_name, user_id, session_id, state_json, last_update_time
+       FROM adk_sessions WHERE app_name = $1 AND user_id = $2 AND session_id = $3`, [appName, userId, sessionId]);
+        if (!sessionRow)
+            return undefined;
+        const eventRows = await this.many(`SELECT event_json FROM adk_events
+       WHERE app_name = $1 AND user_id = $2 AND session_id = $3 ORDER BY sequence ASC`, [appName, userId, sessionId]);
+        let events = eventRows.map((row) => this.parseJson(row.event_json));
+        if (config?.numRecentEvents)
+            events = events.slice(-config.numRecentEvents);
+        if (config?.afterTimestamp) {
+            let index = events.length - 1;
+            while (index >= 0 && events[index].timestamp >= config.afterTimestamp)
+                index -= 1;
+            if (index >= 0)
+                events = events.slice(index + 1);
+        }
+        return this.mergeState(createSession({
+            id: sessionRow.session_id,
+            appName: sessionRow.app_name,
+            userId: sessionRow.user_id,
+            state: this.parseJson(sessionRow.state_json),
+            events,
+            lastUpdateTime: Number(sessionRow.last_update_time),
+        }));
+    }
+    async listSessions({ appName, userId, }) {
+        await this.ready;
+        const rows = await this.many(`SELECT app_name, user_id, session_id, state_json, last_update_time
+       FROM adk_sessions WHERE app_name = $1 AND user_id = $2`, [appName, userId]);
+        return {
+            sessions: rows.map((row) => createSession({
+                id: row.session_id,
+                appName: row.app_name,
+                userId: row.user_id,
+                state: {},
+                events: [],
+                lastUpdateTime: Number(row.last_update_time),
+            })),
+        };
+    }
+    async deleteSession({ appName, userId, sessionId, }) {
+        await this.ready;
+        await this.pool.query("DELETE FROM adk_sessions WHERE app_name = $1 AND user_id = $2 AND session_id = $3", [appName, userId, sessionId]);
+    }
+    async appendEvent({ session, event }) {
+        await this.ready;
+        await super.appendEvent({ session, event });
+        session.lastUpdateTime = event.timestamp;
+        const stored = await this.one(`SELECT app_name, user_id, session_id, state_json, last_update_time
+       FROM adk_sessions WHERE app_name = $1 AND user_id = $2 AND session_id = $3`, [session.appName, session.userId, session.id]);
+        if (!stored)
+            return event;
+        const storedState = this.parseJson(stored.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.pool.query(`UPDATE adk_sessions SET state_json = $1::jsonb, last_update_time = $2
+       WHERE app_name = $3 AND user_id = $4 AND session_id = $5`, [
+            JSON.stringify(storedState),
+            event.timestamp,
+            session.appName,
+            session.userId,
+            session.id,
+        ]);
+        if (!event.partial) {
+            await this.pool.query(`INSERT INTO adk_events (app_name, user_id, session_id, event_id, timestamp, event_json)
+         VALUES ($1, $2, $3, $4, $5, $6::jsonb)`, [
+                session.appName,
+                session.userId,
+                session.id,
+                event.id,
+                event.timestamp,
+                JSON.stringify(event),
+            ]);
+        }
+        return event;
+    }
+    async close() {
+        await this.ready;
+        await this.pool.end();
+    }
+    async initialize() {
+        await this.pool.query("SELECT 1");
+    }
+    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.many("SELECT state_key, value_json FROM adk_app_state WHERE app_name = $1", [session.appName]);
+        for (const row of appRows) {
+            merged.state[`${State.APP_PREFIX}${row.state_key}`] = this.parseJson(row.value_json);
+        }
+        const userRows = await this.many("SELECT state_key, value_json FROM adk_user_state WHERE app_name = $1 AND user_id = $2", [session.appName, session.userId]);
+        for (const row of userRows) {
+            merged.state[`${State.USER_PREFIX}${row.state_key}`] = this.parseJson(row.value_json);
+        }
+        return merged;
+    }
+    async persistScopedState(appName, userId, delta) {
+        for (const [key, value] of Object.entries(delta)) {
+            if (key.startsWith(State.APP_PREFIX)) {
+                await this.pool.query(`INSERT INTO adk_app_state (app_name, state_key, value_json) VALUES ($1, $2, $3::jsonb)
+           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.pool.query(`INSERT INTO adk_user_state (app_name, user_id, state_key, value_json) VALUES ($1, $2, $3, $4::jsonb)
+           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),
+                ]);
+            }
+        }
+    }
+    parseJson(value) {
+        return (typeof value === "string" ? JSON.parse(value) : value);
+    }
+    cloneState(state) {
+        return JSON.parse(JSON.stringify(state));
+    }
+    async one(sql, params) {
+        const result = await this.pool.query(sql, params);
+        return result.rows[0];
+    }
+    async many(sql, params) {
+        const result = await this.pool.query(sql, params);
+        return result.rows;
+    }
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@taico/adk-session-store",
-  "version": "0.3.5",
+  "version": "0.4.0",
   "type": "module",
   "private": false,
   "license": "PolyForm-Small-Business-1.0.0",
@@ -28,9 +28,11 @@
   },
   "dependencies": {
     "@google/adk": "^0.2.5",
+    "pg": "^8.16.3",
     "sqlite3": "^5.1.7"
   },
   "devDependencies": {
+    "@types/pg": "^8.15.5",
     "typescript": "^5.6.3"
   }
 }