@geminixiang/mikan 1.0.0-beta.9 → 1.0.0-beta.91

package/dist/sessions/session-store.js

@@ -0,0 +1,721 @@
+import { chmodSync, existsSync, mkdirSync } from "node:fs";
+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, ROOT_CONVERSATION_ID, ToolResultEntry, UserEntry, createRegistry, defineDoc, defineEntry, } 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 * as log from "../log.js";
+import { compactionSummaryOf, wrapCompactionSummary } from "./compaction-summary.js";
+import { isThreadSessionKey } from "./session-key.js";
+import { errorMessage, isRecord } from "../unknown-values.js";
+const context = BACKGROUND_CONTEXT;
+const ENTRY_PAGE_SIZE = 500;
+const SessionDoc = defineDoc({
+    kind: "mikan.session",
+    version: 1,
+    scope: "conversation",
+    history: "latest",
+    fork: "initial",
+    initial: () => ({}),
+});
+const SessionIndexDoc = defineDoc({
+    kind: "mikan.sessions",
+    version: 1,
+    scope: "session",
+    initial: () => ({ sessions: {} }),
+});
+const CustomSessionEntry = defineEntry("mikan.custom");
+function toJson(input) {
+    if (input === undefined)
+        return null;
+    const copy = JSON.parse(JSON.stringify(input));
+    return copy;
+}
+function durableMessage(message) {
+    const copy = JSON.parse(JSON.stringify(message));
+    return copy;
+}
+function messageTimestamp(message) {
+    return message && "timestamp" in message && typeof message.timestamp === "number"
+        ? message.timestamp
+        : 0;
+}
+function contentText(message) {
+    if (!message || !("content" in message))
+        return "";
+    if (typeof message.content === "string")
+        return message.content;
+    return message.content
+        .map((part) => (part.type === "text" ? part.text : ""))
+        .filter(Boolean)
+        .join("\n");
+}
+function toSessionEntry(record, sessionId) {
+    const id = `${sessionId}:${record.id}`;
+    const message = record.model?.[0];
+    if (UserEntry.is(record) || AssistantEntry.is(record) || ToolResultEntry.is(record)) {
+        if (!message)
+            return undefined;
+        return { type: "message", id, timestamp: messageTimestamp(message), message };
+    }
+    if (CompactionEntry.is(record)) {
+        return {
+            type: "compaction",
+            id,
+            timestamp: messageTimestamp(message),
+            summary: (message && compactionSummaryOf(message)) ?? contentText(message),
+        };
+    }
+    if (CustomSessionEntry.is(record)) {
+        return {
+            type: "custom",
+            id,
+            timestamp: record.data.timestamp,
+            customType: record.data.customType,
+            data: record.data.data ?? undefined,
+        };
+    }
+    return undefined;
+}
+function toSessionEntries(records, sessionId) {
+    return records.flatMap((record) => toSessionEntry(record, sessionId) ?? []);
+}
+async function readRecords(conversation) {
+    const records = [];
+    let cursor;
+    do {
+        const page = await conversation.entries({}, ENTRY_PAGE_SIZE, cursor, context);
+        records.push(...page.items);
+        cursor = page.next;
+    } while (cursor !== undefined);
+    return records.toReversed();
+}
+async function readContext(conversation) {
+    const view = await conversation.context(context);
+    const messages = view.messages
+        .filter((message) => message.role !== "system")
+        .map(durableMessage);
+    return { messages };
+}
+function bindModels(models, property) {
+    const value = Reflect.get(models, property, models);
+    return typeof value === "function" ? value.bind(models) : value;
+}
+function requestOwnerKey(request) {
+    if (typeof request !== "object" || request === null || !("messages" in request))
+        return undefined;
+    const { messages } = request;
+    if (!Array.isArray(messages))
+        return undefined;
+    const last = messages.at(-1);
+    return typeof last === "object" && last !== null ? last : undefined;
+}
+function protectStorageFiles(path) {
+    for (const file of [path, `${path}-wal`, `${path}-shm`]) {
+        if (existsSync(file))
+            chmodSync(file, 0o600);
+    }
+}
+const REQUEST_METHODS = new Set(["streamSimple", "completeSimple"]);
+class OfficeStorage {
+    static { this.pool = new Map(); }
+    constructor(path) {
+        this.path = path;
+        this.registry = createRegistry();
+        this.refs = 0;
+        this.bindings = new Map();
+        this.owners = new WeakMap();
+        this.writers = new Set();
+    }
+    get harness() {
+        if (!this.opened)
+            throw new Error("Office session storage is not open");
+        return this.opened;
+    }
+    static async acquire(path) {
+        let pending = OfficeStorage.pool.get(path);
+        if (!pending) {
+            pending = OfficeStorage.openShared(path);
+            OfficeStorage.pool.set(path, pending);
+            pending.catch(() => OfficeStorage.pool.delete(path));
+        }
+        const storage = await pending;
+        storage.refs += 1;
+        return storage;
+    }
+    static async inMemory() {
+        const storage = await OfficeStorage.open(null, new MemoryStorage());
+        storage.refs = 1;
+        return storage;
+    }
+    static async openShared(path) {
+        mkdirSync(dirname(path), { recursive: true, mode: 0o700 });
+        const storage = await OfficeStorage.open(path, await openNodeSqliteStorage(path));
+        protectStorageFiles(path);
+        await storage.abortUnfinishedWork();
+        return storage;
+    }
+    static async open(path, backend) {
+        const storage = new OfficeStorage(path);
+        storage.opened = await Harness.open(backend, storage.harnessOptions(), context);
+        return storage;
+    }
+    harnessOptions() {
+        const latest = () => [...this.bindings.values()].at(-1);
+        const fallback = () => {
+            const models = latest()?.models;
+            if (!models)
+                throw new Error("Session has no model binding; start a run to use the model");
+            return models;
+        };
+        const models = new Proxy({}, {
+            get: (_target, property) => {
+                if (!REQUEST_METHODS.has(property))
+                    return bindModels(fallback(), property);
+                return (...args) => {
+                    const key = requestOwnerKey(args[1]);
+                    const owner = key ? this.owners.get(key) : undefined;
+                    const binding = owner === undefined ? undefined : this.bindings.get(owner);
+                    const target = binding?.requestModels ?? binding?.models ?? fallback();
+                    const method = Reflect.get(target, property, target);
+                    return method.apply(target, args);
+                };
+            },
+        });
+        const settings = () => latest()?.settings;
+        return {
+            models,
+            registry: this.registry,
+            settings: {
+                extensions: [],
+                get stream() {
+                    return settings()?.stream;
+                },
+                get retry() {
+                    return settings()?.retry;
+                },
+                get compaction() {
+                    return settings()?.compaction;
+                },
+                get toolExecution() {
+                    return settings()?.toolExecution;
+                },
+                get steeringMode() {
+                    return settings()?.steeringMode;
+                },
+                get followUpMode() {
+                    return settings()?.followUpMode;
+                },
+            },
+            env: (target, envContext) => this.bindings.get(target.conversationId)?.env?.(target, envContext),
+            onReport: (error) => {
+                const report = latest()?.onReport;
+                if (report)
+                    report(error);
+                else
+                    log.logWarning("Durable harness report", errorMessage(error));
+            },
+        };
+    }
+    async abortUnfinishedWork() {
+        const { tasks } = await this.harness.inspect(context);
+        const unowned = tasks.filter((task) => !task.record.owner);
+        if (unowned.length === 0)
+            return;
+        log.logWarning(`Aborting ${unowned.length} durable tasks left unfinished by a previous process`);
+        for (const task of unowned)
+            await this.harness.abortTask(task.record.id, context);
+    }
+    async index() {
+        return (await this.harness.snapshot(SessionIndexDoc, context))?.sessions ?? {};
+    }
+    claimWriter(conversationId, key) {
+        if (this.writers.has(conversationId)) {
+            throw new Error(`Session already has an active writer: ${key}`);
+        }
+        this.writers.add(conversationId);
+    }
+    releaseWriter(conversationId) {
+        this.writers.delete(conversationId);
+    }
+    bind(conversationId, binding) {
+        if (this.bindings.has(conversationId)) {
+            throw new Error("Session already has an attached harness");
+        }
+        this.bindings.set(conversationId, binding);
+    }
+    unbind(conversationId) {
+        this.bindings.delete(conversationId);
+    }
+    tag(conversationId, messages) {
+        const last = messages.at(-1);
+        if (!last)
+            return messages;
+        const tagged = { ...last };
+        this.owners.set(tagged, conversationId);
+        return [...messages.slice(0, -1), tagged];
+    }
+    async release() {
+        this.refs -= 1;
+        if (this.refs > 0)
+            return;
+        if (this.path !== null)
+            OfficeStorage.pool.delete(this.path);
+        await this.harness.close(context);
+    }
+}
+function writeSessionIndex(tx, key, record) {
+    return tx.doc(SessionIndexDoc).then((doc) => {
+        doc.sessions[key] = record;
+    });
+}
+async function lookupSession(storage, key) {
+    const record = (await storage.index())[key];
+    if (!record)
+        return undefined;
+    const conversation = await storage.harness.conversation(record.conversationId, context);
+    return conversation ? { conversation, record } : undefined;
+}
+async function createSession(storage, key) {
+    let record;
+    const init = async (tx, conversationId) => {
+        record = { conversationId, id: randomUUID(), createdAt: Date.now() };
+        await writeSessionIndex(tx, key, record);
+    };
+    const conversation = isThreadSessionKey(key)
+        ? await storage.harness.createConversation({ ownership: { kind: "ownerless" }, init }, context)
+        : await storage.harness.root(context, { init });
+    if (!record) {
+        record = { conversationId: conversation.id, id: randomUUID(), createdAt: Date.now() };
+        const created = record;
+        await conversation.commit((tx) => writeSessionIndex(tx, key, created), context);
+    }
+    return { conversation, record };
+}
+async function parentSessionId(storage, conversationId) {
+    const parent = await storage.harness.commit(async (tx) => (await tx.conversation(conversationId))?.parent, context);
+    if (!parent)
+        return undefined;
+    const records = Object.values(await storage.index());
+    return records.find((record) => record.conversationId === parent.conversationId)?.id;
+}
+function isRunCause(record, messageId) {
+    if (!CustomSessionEntry.is(record) || record.data.customType !== RUN_CAUSE_CUSTOM_TYPE) {
+        return false;
+    }
+    const { data } = record.data;
+    return messageId === undefined || (isRecord(data) && data.messageId === messageId);
+}
+function runEnd(records, messageId) {
+    const start = records.findLastIndex((record) => isRunCause(record, messageId));
+    if (start < 0)
+        return undefined;
+    const next = records.findIndex((record, index) => index > start && isRunCause(record));
+    const run = records.slice(start, next < 0 ? undefined : next);
+    return (run.findLast((record) => AssistantEntry.is(record)) ?? run.at(-1))?.id;
+}
+const EARLIER_SESSION_PREFIX = "earlier:";
+export function earlierSessionKey(sessionId) {
+    return `${EARLIER_SESSION_PREFIX}${sessionId}`;
+}
+export function isEarlierSessionKey(key) {
+    return key.startsWith(EARLIER_SESSION_PREFIX);
+}
+async function appendImportedEntry(tx, conversationId, entry) {
+    switch (entry.type) {
+        case "message": {
+            const model = [durableMessage(entry.message)];
+            if (entry.message.role === "user")
+                await tx.appendEntry(UserEntry, conversationId, { model });
+            else if (entry.message.role === "assistant") {
+                await tx.appendEntry(AssistantEntry, conversationId, { model });
+            }
+            else {
+                await tx.appendEntry(ToolResultEntry, conversationId, {
+                    model,
+                    data: { diagnostics: [] },
+                });
+            }
+            return;
+        }
+        case "compaction":
+            await tx.appendEntry(CompactionEntry, conversationId, {
+                model: [
+                    {
+                        role: "user",
+                        content: [{ type: "text", text: wrapCompactionSummary(entry.summary) }],
+                        timestamp: entry.timestamp,
+                    },
+                ],
+                data: { reason: "manual" },
+                head: "self",
+            });
+            return;
+        case "custom":
+            await tx.appendEntry(CustomSessionEntry, conversationId, {
+                data: {
+                    customType: entry.customType,
+                    data: toJson(entry.data),
+                    timestamp: entry.timestamp,
+                },
+            });
+            return;
+        default:
+            entry;
+    }
+}
+export async function importOfficeSessions(path, sessions) {
+    mkdirSync(dirname(path), { recursive: true, mode: 0o700 });
+    const harness = await Harness.open(await openNodeSqliteStorage(path), { models: {}, registry: createRegistry(), settings: { extensions: [] } }, context);
+    try {
+        for (const session of sessions) {
+            const init = async (tx, conversationId) => {
+                for (const entry of session.entries)
+                    await appendImportedEntry(tx, conversationId, entry);
+                if (session.name)
+                    (await tx.doc(SessionDoc, conversationId)).name = session.name;
+                await writeSessionIndex(tx, session.key, {
+                    conversationId,
+                    id: session.id,
+                    createdAt: session.createdAt,
+                });
+            };
+            if (session.root)
+                await harness.root(context, { init });
+            else
+                await harness.createConversation({ ownership: { kind: "ownerless" }, init }, context);
+        }
+    }
+    finally {
+        await harness.close(context);
+    }
+    protectStorageFiles(path);
+}
+export async function readImportedContexts(path) {
+    const harness = await Harness.open(await openNodeSqliteStorage(path), { models: {}, registry: createRegistry(), settings: { extensions: [] } }, context);
+    try {
+        const index = (await harness.snapshot(SessionIndexDoc, context))?.sessions ?? {};
+        const contexts = new Map();
+        for (const [key, record] of Object.entries(index)) {
+            const conversation = await harness.conversation(record.conversationId, context);
+            if (conversation)
+                contexts.set(key, (await readContext(conversation)).messages);
+        }
+        return contexts;
+    }
+    finally {
+        await harness.close(context);
+    }
+}
+class CachedSessionInspection {
+    constructor(header, entries, name, sessionContext) {
+        this.header = header;
+        this.entries = entries;
+        this.name = name;
+        this.sessionContext = sessionContext;
+    }
+    getHeader() {
+        return structuredClone(this.header);
+    }
+    async getEntries() {
+        return structuredClone(this.entries);
+    }
+    async getSessionName() {
+        return this.name;
+    }
+    async buildSessionContext() {
+        return structuredClone(this.sessionContext);
+    }
+}
+async function withOfficeStorage(office, read) {
+    const storage = await OfficeStorage.acquire(office.sessionsPath);
+    try {
+        return await read(storage);
+    }
+    finally {
+        await storage.release();
+    }
+}
+export class SessionStore {
+    constructor(storage, conversation, header) {
+        this.storage = storage;
+        this.conversation = conversation;
+        this.header = header;
+        this.mutationTail = Promise.resolve();
+        this.closed = false;
+        this.bound = false;
+    }
+    static async open(office, key) {
+        const storage = await OfficeStorage.acquire(office.sessionsPath);
+        try {
+            const session = (await lookupSession(storage, key)) ?? (await createSession(storage, key));
+            storage.claimWriter(session.conversation.id, key);
+            return new SessionStore(storage, session.conversation, {
+                id: session.record.id,
+                createdAt: session.record.createdAt,
+                parentSessionId: await parentSessionId(storage, session.conversation.id),
+            });
+        }
+        catch (error) {
+            await storage.release();
+            throw error;
+        }
+    }
+    static async forkRun(office, key, cause, entries) {
+        return withOfficeStorage(office, async (storage) => {
+            if ((await storage.index())[key])
+                return false;
+            const parent = await lookupSession(storage, cause.sessionKey);
+            if (!parent)
+                return false;
+            const at = runEnd(await readRecords(parent.conversation), cause.messageId);
+            if (at === undefined)
+                return false;
+            await parent.conversation.fork(at, {
+                ownership: { kind: "ownerless" },
+                init: async (tx, conversationId) => {
+                    for (const entry of entries)
+                        await appendImportedEntry(tx, conversationId, entry);
+                    await writeSessionIndex(tx, key, {
+                        conversationId,
+                        id: randomUUID(),
+                        createdAt: Date.now(),
+                    });
+                },
+            }, context);
+            return true;
+        });
+    }
+    static async inMemory() {
+        const storage = await OfficeStorage.inMemory();
+        const conversation = await storage.harness.root(context);
+        return new SessionStore(storage, conversation, { id: randomUUID(), createdAt: Date.now() });
+    }
+    static async exists(office, key) {
+        if (!existsSync(office.sessionsPath))
+            return false;
+        return withOfficeStorage(office, async (storage) => (await storage.index())[key] !== undefined);
+    }
+    static async list(office) {
+        if (!existsSync(office.sessionsPath))
+            return [];
+        return withOfficeStorage(office, async (storage) => {
+            const index = await storage.index();
+            const byConversation = new Map(Object.values(index).map((record) => [record.conversationId, record.id]));
+            return Promise.all(Object.entries(index).map(async ([key, record]) => {
+                const parent = await storage.harness.commit(async (tx) => (await tx.conversation(record.conversationId))?.parent, context);
+                const parentId = parent ? byConversation.get(parent.conversationId) : undefined;
+                return {
+                    key,
+                    id: record.id,
+                    createdAt: record.createdAt,
+                    root: record.conversationId === ROOT_CONVERSATION_ID,
+                    parentSessionId: parentId,
+                    forkEntryId: parent && parentId ? `${parentId}:${parent.at}` : undefined,
+                };
+            }));
+        });
+    }
+    static async inspect(office, key) {
+        if (!existsSync(office.sessionsPath))
+            return undefined;
+        return withOfficeStorage(office, async (storage) => {
+            const session = await lookupSession(storage, key);
+            if (!session)
+                return undefined;
+            const { conversation, record } = session;
+            const name = (await storage.harness.snapshot(SessionDoc, conversation.id, context))?.name;
+            return new CachedSessionInspection({
+                id: record.id,
+                createdAt: record.createdAt,
+                parentSessionId: await parentSessionId(storage, conversation.id),
+            }, toSessionEntries(await readRecords(conversation), record.id), name, await readContext(conversation));
+        });
+    }
+    static async inspectExecution(office, key) {
+        if (!existsSync(office.sessionsPath))
+            return { open: false, started: false };
+        return withOfficeStorage(office, async (storage) => {
+            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 {
+                open,
+                started: open || run !== undefined,
+                result: run?.endedAt !== undefined && run.status !== undefined
+                    ? { status: run.status, endedAt: run.endedAt }
+                    : undefined,
+            };
+        });
+    }
+    getHeader() {
+        this.assertOpen();
+        return structuredClone(this.header);
+    }
+    getSessionId() {
+        this.assertOpen();
+        return this.header.id;
+    }
+    async getEntries() {
+        this.assertOpen();
+        return toSessionEntries(await readRecords(this.conversation), this.header.id);
+    }
+    async getContextEntries() {
+        this.assertOpen();
+        const view = await this.conversation.context(context);
+        return toSessionEntries(view.entries, this.header.id);
+    }
+    async getContextTokens() {
+        this.assertOpen();
+        const view = await this.conversation.context(context);
+        const after = view.head?.id ?? Number.NEGATIVE_INFINITY;
+        for (const record of view.entries.toReversed()) {
+            const message = record.model?.[0];
+            if (record.id <= after || !AssistantEntry.is(record) || message?.role !== "assistant")
+                continue;
+            if (message.stopReason === "aborted" || message.stopReason === "error")
+                continue;
+            return calculateContextTokens(message.usage);
+        }
+        return undefined;
+    }
+    async getSessionName() {
+        this.assertOpen();
+        return (await this.storage.harness.snapshot(SessionDoc, this.conversation.id, context))?.name;
+    }
+    async buildSessionContext() {
+        this.assertOpen();
+        return readContext(this.conversation);
+    }
+    async connectMcp(servers, signal) {
+        return this.mutate(async () => {
+            if (this.mcp)
+                throw new Error("SessionStore already has MCP connections");
+            this.mcp = await loadMcpTools(servers, signal);
+            for (const error of this.mcp.errors) {
+                log.logWarning(`MCP server unavailable: ${error.server}`, error.error);
+            }
+            signal?.throwIfAborted();
+            return this.mcp.tools;
+        });
+    }
+    mcpServers() {
+        return this.mcp?.servers ?? [];
+    }
+    async bindHarness(binding) {
+        return this.mutate(async () => {
+            this.storage.bind(this.conversation.id, binding);
+            this.bound = true;
+            const conversationId = this.conversation.id;
+            return {
+                harness: this.storage.harness,
+                conversation: this.conversation,
+                registry: this.storage.registry,
+                extensionName: `mikan.${conversationId}`,
+                tagRequest: (messages) => this.storage.tag(conversationId, messages),
+            };
+        });
+    }
+    async appendMessage(message) {
+        return this.mutate(async () => {
+            const conversation = this.conversation;
+            const model = [durableMessage(message)];
+            const entry = await conversation.commit((tx) => {
+                switch (message.role) {
+                    case "user":
+                        return tx.appendEntry(UserEntry, conversation.id, { model });
+                    case "assistant":
+                        return tx.appendEntry(AssistantEntry, conversation.id, { model });
+                    case "toolResult":
+                        return tx.appendEntry(ToolResultEntry, conversation.id, {
+                            model,
+                            data: { diagnostics: [] },
+                        });
+                    default:
+                        throw new Error("Only user, assistant, and tool result messages enter a session");
+                }
+            }, context);
+            return `${this.header.id}:${entry.id}`;
+        });
+    }
+    async appendCustomEntry(customType, data, timestamp = Date.now()) {
+        return this.mutate(async () => {
+            const conversation = this.conversation;
+            const entry = await conversation.commit((tx) => tx.appendEntry(CustomSessionEntry, conversation.id, {
+                data: { customType, data: toJson(data), timestamp },
+            }), context);
+            return `${this.header.id}:${entry.id}`;
+        });
+    }
+    async appendCompactionSummary(summary, timestamp) {
+        const text = wrapCompactionSummary(summary);
+        return this.mutate(async () => {
+            const conversation = this.conversation;
+            const entry = await conversation.commit((tx) => tx.appendEntry(CompactionEntry, conversation.id, {
+                model: [{ role: "user", content: [{ type: "text", text }], timestamp }],
+                data: { reason: "manual" },
+                head: "self",
+            }), context);
+            return `${this.header.id}:${entry.id}`;
+        });
+    }
+    async reset() {
+        await this.mutate(() => this.conversation.reset(undefined, context));
+    }
+    async setSessionName(name) {
+        await this.updateSessionDoc((doc) => {
+            doc.name = name.trim() || undefined;
+        });
+    }
+    async recordRun(run) {
+        await this.updateSessionDoc((doc) => {
+            doc.run = { ...run };
+        });
+    }
+    async close() {
+        if (this.closePromise)
+            return this.closePromise;
+        this.closed = true;
+        const conversationId = this.conversation.id;
+        this.closePromise = this.mutationTail.then(async () => {
+            try {
+                await this.mcp?.dispose();
+            }
+            finally {
+                if (this.bound) {
+                    this.storage.registry.uninstall({ name: `mikan.${conversationId}` });
+                    this.storage.unbind(conversationId);
+                }
+                this.storage.releaseWriter(conversationId);
+                await this.storage.release();
+            }
+        });
+        return this.closePromise;
+    }
+    async updateSessionDoc(change) {
+        await this.mutate(async () => {
+            const conversation = this.conversation;
+            await conversation.commit(async (tx) => {
+                change(await tx.doc(SessionDoc, conversation.id));
+            }, context);
+        });
+    }
+    assertOpen() {
+        if (this.closed)
+            throw new Error("SessionStore is closed");
+    }
+    mutate(operation) {
+        this.assertOpen();
+        const result = this.mutationTail.then(operation);
+        this.mutationTail = result.then(() => undefined, () => undefined);
+        return result;
+    }
+}
+//# sourceMappingURL=session-store.js.map