@geminixiang/mikan 1.2.2 → 1.3.0-beta.0

package/dist/sessions/session-store.js

@@ -3,10 +3,11 @@ import { randomUUID } from "node:crypto";
 import { dirname } from "node:path";
 import { BACKGROUND_CONTEXT } from "@earendil-works/chord/context";
 import { calculateContextTokens } from "@earendil-works/pi-ai/utils/estimate";
-import { AssistantEntry, CompactionEntry, Harness, LiveDoc, MemoryStorage, ProviderDoc, ROOT_CONVERSATION_ID, ToolResultEntry, UserEntry, createRegistry, defineDoc, defineEntry, defineExtension, defineTask, } from "@earendil-works/pi-durable";
+import { AssistantEntry, CompactionEntry, Harness, LiveDoc, MemoryStorage, ProviderDoc, ROOT_CONVERSATION_ID, ToolResultEntry, UsageDoc, UserEntry, createRegistry, defineDoc, defineEntry, defineExtension, defineTask, } from "@earendil-works/pi-durable";
 import { openNodeSqliteStorage } from "@earendil-works/pi-durable/storage/sqlite/node";
 import { RUN_CAUSE_CUSTOM_TYPE } from "./types.js";
 import { loadMcpTools } from "../harness/mcp.js";
+import { addUsage, createEmptyUsage, sumUsageState } from "../harness/session.js";
 import * as log from "../log.js";
 import { compactionSummaryOf, wrapCompactionSummary } from "./compaction-summary.js";
 import { isThreadSessionKey } from "./session-key.js";
@@ -31,6 +32,14 @@ const SessionDoc = defineDoc({
     fork: "initial",
     initial: () => ({}),
 });
+export const ScriptStoreDoc = defineDoc({
+    kind: "mikan.script-store",
+    version: 1,
+    scope: "conversation",
+    history: "rewindable",
+    fork: "asOf",
+    initial: () => ({ values: {} }),
+});
 const SessionIndexDoc = defineDoc({
     kind: "mikan.sessions",
     version: 1,
@@ -276,6 +285,44 @@ class OfficeStorage {
         await this.harness.close(context);
     }
 }
+const OWNED_CONVERSATION_PAGE_SIZE = 100;
+async function ownedConversations(storage, ownerConversationId) {
+    return storage.harness.commit(async (tx) => {
+        const owned = [];
+        let cursor;
+        do {
+            const page = await tx.scanConversations({ ownerConversationId: ownerConversationId, order: "descending" }, OWNED_CONVERSATION_PAGE_SIZE, cursor);
+            owned.push(...page.items);
+            cursor = page.next;
+        } while (cursor !== undefined);
+        return owned;
+    }, context);
+}
+async function readTaskState(storage, sessionKey, conversationId, anchorId) {
+    const anchor = await storage.harness.commit((tx) => tx.task(anchorId), context);
+    return {
+        sessionKey,
+        acknowledgement: briefAcknowledgement(anchor?.input),
+        ...(await readExecution(storage, conversationId)),
+    };
+}
+async function readExecution(storage, conversationId) {
+    const live = await storage.harness.snapshot(LiveDoc, conversationId, context);
+    const run = (await storage.harness.snapshot(SessionDoc, conversationId, context))?.run;
+    const open = live?.run !== undefined;
+    return {
+        open,
+        started: open || run !== undefined,
+        result: run?.endedAt !== undefined && run.status !== undefined
+            ? { status: run.status, endedAt: run.endedAt }
+            : undefined,
+    };
+}
+function briefAcknowledgement(input) {
+    return isRecord(input) && typeof input.acknowledgement === "string"
+        ? input.acknowledgement
+        : undefined;
+}
 function writeSessionIndex(tx, key, record) {
     return tx.doc(SessionIndexDoc).then((doc) => {
         doc.sessions[key] = record;
@@ -304,12 +351,12 @@ async function createSession(storage, key) {
     }
     return { conversation, record };
 }
-async function createTaskSession(storage, key, requesterKey) {
+async function createTaskSession(storage, key, requesterKey, brief) {
     const requester = await lookupSession(storage, requesterKey);
     if (!requester)
         throw new Error(`Task requester session not found: ${requesterKey}`);
     const record = await requester.conversation.commit(async (tx) => {
-        const anchor = await tx.createTask(TaskAnchor, null, {
+        const anchor = await tx.createTask(TaskAnchor, brief, {
             ownership: { kind: "conversation" },
             background: true,
         });
@@ -429,6 +476,70 @@ export async function importOfficeSessions(path, sessions) {
     }
     protectStorageFiles(path);
 }
+export async function backfillImportedUsage(path, dryRun) {
+    const harness = await Harness.open(await openNodeSqliteStorage(path), { models: {}, registry: createRegistry(), settings: { extensions: [] } }, context);
+    try {
+        let backfilled = 0;
+        let cursor;
+        do {
+            const page = await harness.commit((tx) => tx.scanConversations({}, 100, cursor), context);
+            for (const conversation of page.items) {
+                const imported = await importedUsage(harness, conversation.id);
+                if (imported.size === 0)
+                    continue;
+                backfilled += 1;
+                if (!dryRun)
+                    await recordImportedUsage(harness, conversation.id, imported);
+            }
+            cursor = page.next;
+        } while (cursor !== undefined);
+        return backfilled;
+    }
+    finally {
+        await harness.close(context);
+    }
+}
+async function importedUsage(harness, conversationId) {
+    const byModel = new Map();
+    if ((await harness.snapshot(SessionDoc, conversationId, context))?.importedUsage)
+        return byModel;
+    let cursor;
+    do {
+        const page = await harness.commit((tx) => tx.scanEntries({ conversationId }, ENTRY_PAGE_SIZE, cursor), context);
+        for (const record of page.items) {
+            const message = record.model?.[0];
+            if (record.conversationId !== conversationId ||
+                record.byTaskId !== undefined ||
+                !AssistantEntry.is(record) ||
+                message?.role !== "assistant") {
+                continue;
+            }
+            const key = `${message.provider}/${message.model}`;
+            const total = byModel.get(key) ?? createEmptyUsage();
+            addUsage(total, message.usage);
+            byModel.set(key, total);
+        }
+        cursor = page.next;
+    } while (cursor !== undefined);
+    return byModel;
+}
+async function recordImportedUsage(harness, conversationId, imported) {
+    await harness.commit(async (tx) => {
+        const doc = await tx.doc(SessionDoc, conversationId);
+        if (doc.importedUsage)
+            return;
+        const ledger = await tx.doc(UsageDoc, conversationId);
+        for (const [key, usage] of imported) {
+            const merged = createEmptyUsage();
+            const recorded = ledger.models[key];
+            if (recorded)
+                addUsage(merged, recorded);
+            addUsage(merged, usage);
+            ledger.models[key] = merged;
+        }
+        doc.importedUsage = true;
+    }, context);
+}
 export async function readImportedContexts(path) {
     const harness = await Harness.open(await openNodeSqliteStorage(path), { models: {}, registry: createRegistry(), settings: { extensions: [] } }, context);
     try {
@@ -486,8 +597,8 @@ export class SessionStore {
     static open(office, key) {
         return SessionStore.openOrCreate(office, key, (storage) => createSession(storage, key));
     }
-    static openTask(office, key, requesterKey) {
-        return SessionStore.openOrCreate(office, key, (storage) => createTaskSession(storage, key, requesterKey));
+    static openTask(office, key, requesterKey, brief) {
+        return SessionStore.openOrCreate(office, key, (storage) => createTaskSession(storage, key, requesterKey, brief));
     }
     static async openOrCreate(office, key, create) {
         const storage = await OfficeStorage.acquire(office.sessionsPath);
@@ -637,6 +748,25 @@ export class SessionStore {
             return interrupted;
         });
     }
+    static async listTasks(office, requesterKey) {
+        if (!existsSync(office.sessionsPath))
+            return [];
+        return withOfficeStorage(office, async (storage) => {
+            const index = await storage.index();
+            const requester = index[requesterKey];
+            if (!requester)
+                return [];
+            const keyOf = new Map(Object.entries(index).map(([key, record]) => [record.conversationId, key]));
+            const tasks = [];
+            for (const child of await ownedConversations(storage, requester.conversationId)) {
+                const sessionKey = keyOf.get(child.id);
+                if (sessionKey === undefined || child.owner === undefined)
+                    continue;
+                tasks.push(await readTaskState(storage, sessionKey, child.id, child.owner.taskId));
+            }
+            return tasks;
+        });
+    }
     static async inspectExecution(office, key) {
         if (!existsSync(office.sessionsPath))
             return { open: false, started: false };
@@ -644,19 +774,52 @@ export class SessionStore {
             const session = await lookupSession(storage, key);
             if (!session)
                 return { open: false, started: false };
-            const id = session.conversation.id;
-            const live = await storage.harness.snapshot(LiveDoc, id, context);
-            const run = (await storage.harness.snapshot(SessionDoc, id, context))?.run;
-            const open = live?.run !== undefined;
+            return readExecution(storage, session.conversation.id);
+        });
+    }
+    static async spend(office, key) {
+        if (!existsSync(office.sessionsPath))
+            return undefined;
+        return withOfficeStorage(office, async (storage) => {
+            const session = await lookupSession(storage, key);
+            if (!session)
+                return undefined;
+            const ledger = await storage.harness.snapshot(UsageDoc, session.conversation.id, context);
+            const newest = (await session.conversation.entries({}, 1, undefined, context)).items[0];
+            const timestamp = newest && toSessionEntry(newest, session.record.id)?.timestamp;
             return {
-                open,
-                started: open || run !== undefined,
-                result: run?.endedAt !== undefined && run.status !== undefined
-                    ? { status: run.status, endedAt: run.endedAt }
-                    : undefined,
+                usage: sumUsageState(ledger ?? { models: {}, tools: {} }),
+                updatedAt: timestamp || session.record.createdAt,
             };
         });
     }
+    static async ownResponseUsage(office, key) {
+        if (!existsSync(office.sessionsPath))
+            return [];
+        return withOfficeStorage(office, async (storage) => {
+            const session = await lookupSession(storage, key);
+            if (!session)
+                return [];
+            const id = session.conversation.id;
+            return (await readRecords(session.conversation)).flatMap((record) => {
+                const message = record.model?.[0];
+                return record.conversationId === id && message?.role === "assistant"
+                    ? [{ timestamp: message.timestamp, usage: message.usage }]
+                    : [];
+            });
+        });
+    }
+    static async isTask(office, key) {
+        if (!existsSync(office.sessionsPath))
+            return false;
+        return withOfficeStorage(office, async (storage) => {
+            const record = (await storage.index())[key];
+            if (!record)
+                return false;
+            const conversation = await storage.harness.commit((tx) => tx.conversation(record.conversationId), context);
+            return conversation?.owner !== undefined;
+        });
+    }
     getHeader() {
         this.assertOpen();
         return structuredClone(this.header);
@@ -774,7 +937,12 @@ export class SessionStore {
         });
     }
     async reset() {
-        await this.mutate(() => this.conversation.reset(undefined, context));
+        await this.mutate(async () => {
+            await this.conversation.reset(undefined, context);
+            await this.conversation.commit(async (tx) => {
+                (await tx.doc(ScriptStoreDoc, this.conversation.id)).values = {};
+            }, context);
+        });
     }
     async setSessionName(name) {
         await this.updateSessionDoc((doc) => {