@pikku/kysely 0.10.0 → 0.12.0

package/dist/src/kysely-ai-storage-service.js

@@ -0,0 +1,586 @@
+import { sql } from 'kysely';
+import { parseJson } from './kysely-json.js';
+export class KyselyAIStorageService {
+    db;
+    initialized = false;
+    constructor(db) {
+        this.db = db;
+    }
+    async init() {
+        if (this.initialized) {
+            return;
+        }
+        await this.db.schema
+            .createTable('ai_threads')
+            .ifNotExists()
+            .addColumn('id', 'text', (col) => col.primaryKey())
+            .addColumn('resource_id', 'text', (col) => col.notNull())
+            .addColumn('title', 'text')
+            .addColumn('metadata', 'text')
+            .addColumn('created_at', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .addColumn('updated_at', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .execute();
+        await this.db.schema
+            .createIndex('idx_ai_threads_resource')
+            .ifNotExists()
+            .on('ai_threads')
+            .column('resource_id')
+            .execute();
+        await this.db.schema
+            .createTable('ai_message')
+            .ifNotExists()
+            .addColumn('id', 'text', (col) => col.primaryKey())
+            .addColumn('thread_id', 'text', (col) => col.notNull().references('ai_threads.id').onDelete('cascade'))
+            .addColumn('role', 'text', (col) => col.notNull())
+            .addColumn('content', 'text')
+            .addColumn('created_at', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .execute();
+        await this.db.schema
+            .createIndex('idx_ai_message_thread')
+            .ifNotExists()
+            .on('ai_message')
+            .columns(['thread_id', 'created_at'])
+            .execute();
+        await this.db.schema
+            .createTable('ai_tool_call')
+            .ifNotExists()
+            .addColumn('id', 'text', (col) => col.primaryKey())
+            .addColumn('thread_id', 'text', (col) => col.notNull().references('ai_threads.id').onDelete('cascade'))
+            .addColumn('message_id', 'text', (col) => col.notNull().references('ai_message.id').onDelete('cascade'))
+            .addColumn('run_id', 'text')
+            .addColumn('tool_name', 'text', (col) => col.notNull())
+            .addColumn('args', 'text', (col) => col.notNull().defaultTo('{}'))
+            .addColumn('result', 'text')
+            .addColumn('approval_status', 'text')
+            .addColumn('approval_type', 'text')
+            .addColumn('agent_run_id', 'text')
+            .addColumn('display_tool_name', 'text')
+            .addColumn('display_args', 'text')
+            .addColumn('created_at', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .execute();
+        await this.db.schema
+            .createIndex('idx_ai_tool_call_thread')
+            .ifNotExists()
+            .on('ai_tool_call')
+            .column('thread_id')
+            .execute();
+        await this.db.schema
+            .createIndex('idx_ai_tool_call_message')
+            .ifNotExists()
+            .on('ai_tool_call')
+            .column('message_id')
+            .execute();
+        await this.db.schema
+            .createTable('ai_working_memory')
+            .ifNotExists()
+            .addColumn('id', 'text', (col) => col.notNull())
+            .addColumn('scope', 'text', (col) => col.notNull())
+            .addColumn('data', 'text', (col) => col.notNull())
+            .addColumn('updated_at', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .addPrimaryKeyConstraint('ai_working_memory_pk', ['id', 'scope'])
+            .execute();
+        await this.db.schema
+            .createTable('ai_run')
+            .ifNotExists()
+            .addColumn('run_id', 'text', (col) => col.primaryKey())
+            .addColumn('agent_name', 'text', (col) => col.notNull())
+            .addColumn('thread_id', 'text', (col) => col.notNull().references('ai_threads.id').onDelete('cascade'))
+            .addColumn('resource_id', 'text', (col) => col.notNull())
+            .addColumn('status', 'text', (col) => col.notNull().defaultTo('running'))
+            .addColumn('suspend_reason', 'text')
+            .addColumn('missing_rpcs', 'text')
+            .addColumn('usage_input_tokens', 'integer', (col) => col.notNull().defaultTo(0))
+            .addColumn('usage_output_tokens', 'integer', (col) => col.notNull().defaultTo(0))
+            .addColumn('usage_model', 'text', (col) => col.notNull().defaultTo(''))
+            .addColumn('created_at', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .addColumn('updated_at', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .execute();
+        await this.db.schema
+            .createIndex('idx_ai_run_thread')
+            .ifNotExists()
+            .on('ai_run')
+            .columns(['thread_id', 'created_at'])
+            .execute();
+        this.initialized = true;
+    }
+    async createThread(resourceId, options) {
+        const id = options?.threadId ?? crypto.randomUUID();
+        const now = new Date();
+        await this.db
+            .insertInto('ai_threads')
+            .values({
+            id,
+            resource_id: resourceId,
+            title: options?.title ?? null,
+            metadata: JSON.stringify(options?.metadata ?? null),
+            created_at: now,
+            updated_at: now,
+        })
+            .execute();
+        return {
+            id,
+            resourceId,
+            title: options?.title,
+            metadata: options?.metadata,
+            createdAt: now,
+            updatedAt: now,
+        };
+    }
+    async getThread(threadId) {
+        const row = await this.db
+            .selectFrom('ai_threads')
+            .select([
+            'id',
+            'resource_id',
+            'title',
+            'metadata',
+            'created_at',
+            'updated_at',
+        ])
+            .where('id', '=', threadId)
+            .executeTakeFirst();
+        if (!row) {
+            throw new Error(`Thread not found: ${threadId}`);
+        }
+        return {
+            id: row.id,
+            resourceId: row.resource_id,
+            title: row.title ?? undefined,
+            metadata: parseJson(row.metadata),
+            createdAt: new Date(row.created_at),
+            updatedAt: new Date(row.updated_at),
+        };
+    }
+    async getThreads(resourceId) {
+        const result = await this.db
+            .selectFrom('ai_threads')
+            .select([
+            'id',
+            'resource_id',
+            'title',
+            'metadata',
+            'created_at',
+            'updated_at',
+        ])
+            .where('resource_id', '=', resourceId)
+            .orderBy('updated_at', 'desc')
+            .execute();
+        return result.map((row) => ({
+            id: row.id,
+            resourceId: row.resource_id,
+            title: row.title ?? undefined,
+            metadata: parseJson(row.metadata),
+            createdAt: new Date(row.created_at),
+            updatedAt: new Date(row.updated_at),
+        }));
+    }
+    async deleteThread(threadId) {
+        await this.db.deleteFrom('ai_threads').where('id', '=', threadId).execute();
+    }
+    async getMessages(threadId, options) {
+        let msgQuery = this.db
+            .selectFrom('ai_message')
+            .select(['id', 'role', 'content', 'created_at'])
+            .where('thread_id', '=', threadId);
+        if (options?.cursor) {
+            const cursorRow = await this.db
+                .selectFrom('ai_message')
+                .select('created_at')
+                .where('id', '=', options.cursor)
+                .executeTakeFirst();
+            if (cursorRow) {
+                msgQuery = msgQuery.where('created_at', '<', cursorRow.created_at);
+            }
+        }
+        if (options?.cursor || options?.lastN) {
+            const innerResult = await msgQuery
+                .orderBy('created_at', 'desc')
+                .limit(options?.lastN ?? 50)
+                .execute();
+            innerResult.reverse();
+            var msgResult = innerResult;
+        }
+        else {
+            var msgResult = await msgQuery.orderBy('created_at', 'asc').execute();
+        }
+        const tcResult = await this.db
+            .selectFrom('ai_tool_call')
+            .select(['id', 'message_id', 'tool_name', 'args', 'result'])
+            .where('thread_id', '=', threadId)
+            .orderBy('created_at', 'asc')
+            .execute();
+        const tcByMessage = new Map();
+        for (const tc of tcResult) {
+            const msgId = tc.message_id;
+            if (!tcByMessage.has(msgId))
+                tcByMessage.set(msgId, []);
+            tcByMessage.get(msgId).push(tc);
+        }
+        const messages = [];
+        for (const row of msgResult) {
+            const msg = {
+                id: row.id,
+                role: row.role,
+                content: row.content ?? undefined,
+                createdAt: new Date(row.created_at),
+            };
+            const tcs = tcByMessage.get(msg.id);
+            if (tcs?.length) {
+                msg.toolCalls = tcs.map((tc) => ({
+                    id: tc.id,
+                    name: tc.tool_name,
+                    args: parseJson(tc.args),
+                }));
+                const completed = tcs.filter((tc) => tc.result != null);
+                if (completed.length) {
+                    messages.push(msg);
+                    messages.push({
+                        id: `tool-results-${msg.id}`,
+                        role: 'tool',
+                        toolResults: completed.map((tc) => ({
+                            id: tc.id,
+                            name: tc.tool_name,
+                            result: tc.result,
+                        })),
+                        createdAt: msg.createdAt,
+                    });
+                    continue;
+                }
+            }
+            messages.push(msg);
+        }
+        return messages;
+    }
+    async saveMessages(threadId, messages) {
+        if (messages.length === 0)
+            return;
+        const toolMessages = messages.filter((m) => m.role === 'tool');
+        const nonToolMessages = messages.filter((m) => m.role !== 'tool');
+        if (nonToolMessages.length > 0) {
+            await this.db
+                .insertInto('ai_message')
+                .values(nonToolMessages.map((msg) => ({
+                id: msg.id,
+                thread_id: threadId,
+                role: msg.role,
+                content: msg.content ?? null,
+                created_at: msg.createdAt ?? new Date(),
+            })))
+                .execute();
+        }
+        const toolCalls = nonToolMessages.flatMap((msg) => msg.toolCalls?.map((tc) => ({ ...tc, messageId: msg.id })) ?? []);
+        if (toolCalls.length > 0) {
+            await this.db
+                .insertInto('ai_tool_call')
+                .values(toolCalls.map((tc) => ({
+                id: tc.id,
+                thread_id: threadId,
+                message_id: tc.messageId,
+                tool_name: tc.name,
+                args: JSON.stringify(tc.args),
+            })))
+                .execute();
+        }
+        for (const toolMsg of toolMessages) {
+            if (!toolMsg.toolResults)
+                continue;
+            for (const tr of toolMsg.toolResults) {
+                await this.db
+                    .updateTable('ai_tool_call')
+                    .set({ result: tr.result })
+                    .where('id', '=', tr.id)
+                    .execute();
+            }
+        }
+        await this.db
+            .updateTable('ai_threads')
+            .set({ updated_at: new Date() })
+            .where('id', '=', threadId)
+            .execute();
+    }
+    async getWorkingMemory(id, scope) {
+        const row = await this.db
+            .selectFrom('ai_working_memory')
+            .select('data')
+            .where('id', '=', id)
+            .where('scope', '=', scope)
+            .executeTakeFirst();
+        if (!row)
+            return null;
+        return parseJson(row.data);
+    }
+    async saveWorkingMemory(id, scope, data) {
+        await this.db
+            .insertInto('ai_working_memory')
+            .values({
+            id,
+            scope,
+            data: JSON.stringify(data),
+            updated_at: new Date(),
+        })
+            .onConflict((oc) => oc.columns(['id', 'scope']).doUpdateSet({
+            data: JSON.stringify(data),
+            updated_at: new Date(),
+        }))
+            .execute();
+    }
+    async createRun(run) {
+        const runId = crypto.randomUUID();
+        await this.db
+            .insertInto('ai_run')
+            .values({
+            run_id: runId,
+            agent_name: run.agentName,
+            thread_id: run.threadId,
+            resource_id: run.resourceId,
+            status: run.status,
+            suspend_reason: run.suspendReason ?? null,
+            missing_rpcs: run.missingRpcs ? JSON.stringify(run.missingRpcs) : null,
+            usage_input_tokens: run.usage.inputTokens,
+            usage_output_tokens: run.usage.outputTokens,
+            usage_model: run.usage.model,
+            created_at: run.createdAt,
+            updated_at: run.updatedAt,
+        })
+            .execute();
+        if (run.pendingApprovals?.length) {
+            await this.setApprovals(runId, run.pendingApprovals);
+        }
+        return runId;
+    }
+    async updateRun(runId, updates) {
+        const setValues = { updated_at: new Date() };
+        if (updates.status !== undefined) {
+            setValues.status = updates.status;
+        }
+        if (updates.suspendReason !== undefined) {
+            setValues.suspend_reason = updates.suspendReason;
+        }
+        if (updates.missingRpcs !== undefined) {
+            setValues.missing_rpcs = JSON.stringify(updates.missingRpcs);
+        }
+        if (updates.usage !== undefined) {
+            setValues.usage_input_tokens = updates.usage.inputTokens;
+            setValues.usage_output_tokens = updates.usage.outputTokens;
+            setValues.usage_model = updates.usage.model;
+        }
+        await this.db
+            .updateTable('ai_run')
+            .set(setValues)
+            .where('run_id', '=', runId)
+            .execute();
+        if (updates.pendingApprovals !== undefined) {
+            await this.db
+                .updateTable('ai_tool_call')
+                .set({
+                approval_status: null,
+                run_id: null,
+                approval_type: null,
+                agent_run_id: null,
+                display_tool_name: null,
+                display_args: null,
+            })
+                .where('run_id', '=', runId)
+                .where('approval_status', 'is not', null)
+                .execute();
+            if (updates.pendingApprovals.length) {
+                await this.setApprovals(runId, updates.pendingApprovals);
+            }
+        }
+    }
+    async getRun(runId) {
+        const row = await this.db
+            .selectFrom('ai_run')
+            .select([
+            'run_id',
+            'agent_name',
+            'thread_id',
+            'resource_id',
+            'status',
+            'suspend_reason',
+            'missing_rpcs',
+            'usage_input_tokens',
+            'usage_output_tokens',
+            'usage_model',
+            'created_at',
+            'updated_at',
+        ])
+            .where('run_id', '=', runId)
+            .executeTakeFirst();
+        if (!row)
+            return null;
+        const approvals = await this.db
+            .selectFrom('ai_tool_call')
+            .select([
+            'id',
+            'tool_name',
+            'args',
+            'approval_type',
+            'agent_run_id',
+            'display_tool_name',
+            'display_args',
+        ])
+            .where('run_id', '=', runId)
+            .where('approval_status', '=', 'pending')
+            .execute();
+        return this.mapRunRow(row, approvals);
+    }
+    async getRunsByThread(threadId) {
+        const result = await this.db
+            .selectFrom('ai_run')
+            .select([
+            'run_id',
+            'agent_name',
+            'thread_id',
+            'resource_id',
+            'status',
+            'suspend_reason',
+            'missing_rpcs',
+            'usage_input_tokens',
+            'usage_output_tokens',
+            'usage_model',
+            'created_at',
+            'updated_at',
+        ])
+            .where('thread_id', '=', threadId)
+            .orderBy('created_at', 'desc')
+            .execute();
+        const runs = [];
+        for (const row of result) {
+            const approvals = await this.db
+                .selectFrom('ai_tool_call')
+                .select([
+                'id',
+                'tool_name',
+                'args',
+                'approval_type',
+                'agent_run_id',
+                'display_tool_name',
+                'display_args',
+            ])
+                .where('run_id', '=', row.run_id)
+                .where('approval_status', '=', 'pending')
+                .execute();
+            runs.push(this.mapRunRow(row, approvals));
+        }
+        return runs;
+    }
+    async setApprovals(runId, approvals) {
+        for (const a of approvals) {
+            if (a.type === 'agent-call') {
+                await this.db
+                    .updateTable('ai_tool_call')
+                    .set({
+                    approval_status: 'pending',
+                    run_id: runId,
+                    approval_type: 'agent-call',
+                    agent_run_id: a.agentRunId,
+                    display_tool_name: a.displayToolName,
+                    display_args: JSON.stringify(a.displayArgs),
+                })
+                    .where('id', '=', a.toolCallId)
+                    .execute();
+            }
+            else {
+                await this.db
+                    .updateTable('ai_tool_call')
+                    .set({
+                    approval_status: 'pending',
+                    run_id: runId,
+                    approval_type: 'tool-call',
+                })
+                    .where('id', '=', a.toolCallId)
+                    .execute();
+            }
+        }
+    }
+    mapRunRow(row, approvalRows) {
+        const pendingApprovals = approvalRows?.length
+            ? approvalRows.map((a) => {
+                if (a.approval_type === 'agent-call') {
+                    return {
+                        type: 'agent-call',
+                        toolCallId: a.id,
+                        agentName: a.tool_name,
+                        agentRunId: a.agent_run_id,
+                        displayToolName: a.display_tool_name,
+                        displayArgs: parseJson(a.display_args),
+                    };
+                }
+                return {
+                    type: 'tool-call',
+                    toolCallId: a.id,
+                    toolName: a.tool_name,
+                    args: parseJson(a.args),
+                };
+            })
+            : undefined;
+        return {
+            runId: row.run_id,
+            agentName: row.agent_name,
+            threadId: row.thread_id,
+            resourceId: row.resource_id,
+            status: row.status,
+            suspendReason: row.suspend_reason,
+            missingRpcs: parseJson(row.missing_rpcs),
+            pendingApprovals,
+            usage: {
+                inputTokens: Number(row.usage_input_tokens),
+                outputTokens: Number(row.usage_output_tokens),
+                model: row.usage_model,
+            },
+            createdAt: new Date(row.created_at),
+            updatedAt: new Date(row.updated_at),
+        };
+    }
+    async findRunByToolCallId(toolCallId) {
+        const tc = await this.db
+            .selectFrom('ai_tool_call')
+            .select([
+            'id',
+            'tool_name',
+            'args',
+            'run_id',
+            'approval_type',
+            'agent_run_id',
+            'display_tool_name',
+            'display_args',
+        ])
+            .where('id', '=', toolCallId)
+            .where('approval_status', '=', 'pending')
+            .executeTakeFirst();
+        if (!tc || !tc.run_id)
+            return null;
+        const run = await this.getRun(tc.run_id);
+        if (!run)
+            return null;
+        let approval;
+        if (tc.approval_type === 'agent-call') {
+            approval = {
+                type: 'agent-call',
+                toolCallId: tc.id,
+                agentName: tc.tool_name,
+                agentRunId: tc.agent_run_id,
+                displayToolName: tc.display_tool_name,
+                displayArgs: parseJson(tc.display_args),
+            };
+        }
+        else {
+            approval = {
+                type: 'tool-call',
+                toolCallId: tc.id,
+                toolName: tc.tool_name,
+                args: parseJson(tc.args),
+            };
+        }
+        return { run, approval };
+    }
+    async resolveApproval(toolCallId, status) {
+        await this.db
+            .updateTable('ai_tool_call')
+            .set({ approval_status: status })
+            .where('id', '=', toolCallId)
+            .execute();
+    }
+    async close() { }
+}

package/dist/src/kysely-channel-store.d.ts

@@ -1,13 +1,17 @@
 import { CoreUserSession } from '@pikku/core';
 import { Channel, ChannelStore } from '@pikku/core/channel';
 import { Kysely } from 'kysely';
+import type { KyselyPikkuDB } from './kysely-tables.js';
 export declare class KyselyChannelStore extends ChannelStore {
-    private database;
-    constructor(database: Kysely<any>);
+    private db;
+    private initialized;
+    constructor(db: Kysely<KyselyPikkuDB>);
+    init(): Promise<void>;
     addChannel({ channelId, channelName, openingData, }: Channel): Promise<void>;
     removeChannels(channelIds: string[]): Promise<void>;
-    setUserSession(channelId: string, session: any): Promise<void>;
+    setUserSession(channelId: string, session: CoreUserSession | null): Promise<void>;
     getChannelAndSession(channelId: string): Promise<Channel & {
         session: CoreUserSession;
     }>;
+    close(): Promise<void>;
 }