@markusylisiurunen/tau 0.3.21 → 0.3.22

package/dist/host/local_session_host.js

@@ -1,4 +1,5 @@
 import { randomUUID } from "node:crypto";
+import { isDeepStrictEqual } from "node:util";
 import { resolvePromptTemplateWithBackend } from "../core/config/index.js";
 import { ChatRuntime } from "../core/runtime/chat_runtime.js";
 import { filterProjectPathAutocompleteEntries, loadProjectPathAutocompleteEntriesWithBackend, } from "../core/utils/project_files.js";
@@ -55,15 +56,27 @@ export class LocalSessionHost {
         return hostedSession;
     }
     async createRecoveredSession(snapshot, executionEnvironment) {
-        const recovered = normalizeRecoveredSnapshot(snapshot);
-        const hostedSession = await this.createLocalSessionHandle(executionEnvironment, recovered.snapshot, undefined, recovered.changed);
-        hostedSession.session.restoreState({
-            sessionId: recovered.snapshot.sessionId,
-            historyEntries: recovered.snapshot.messages.flatMap((entry) => entry.modelVisible && isCoreMessage(entry.message)
-                ? [{ id: entry.id, message: entry.message }]
-                : []),
-        });
-        return hostedSession;
+        let hostedSession;
+        try {
+            const recovered = normalizeRecoveredSnapshot(snapshot);
+            hostedSession = await this.createLocalSessionHandle(executionEnvironment, recovered.snapshot, undefined, recovered.changed);
+            hostedSession.session.restoreState({
+                sessionId: recovered.snapshot.sessionId,
+                historyEntries: recovered.snapshot.messages.flatMap((entry) => entry.modelVisible && isCoreMessage(entry.message)
+                    ? [{ id: entry.id, message: entry.message }]
+                    : []),
+            });
+            return hostedSession;
+        }
+        catch (error) {
+            if (hostedSession) {
+                await hostedSession.dispose();
+            }
+            else {
+                await executionEnvironment.dispose();
+            }
+            throw error;
+        }
     }
     async createLocalSessionHandle(executionEnvironment, committedSnapshot, createParams, forceNextSnapshotRevision = false) {
         const bootstrap = await this.resolveNewSessionBootstrap(executionEnvironment);
@@ -235,6 +248,16 @@ export class LocalSessionHost {
                 errors.push(result.reason);
             }
         }
+        for (const session of this.sessions) {
+            session.interruptTurn();
+            session.interruptMaintenance();
+        }
+        const settlementResults = await Promise.allSettled([...this.sessions].map((session) => session.waitForActiveWork()));
+        for (const result of settlementResults) {
+            if (result.status === "rejected") {
+                errors.push(result.reason);
+            }
+        }
         for (const session of this.sessions) {
             try {
                 await session.snapshot();
@@ -285,7 +308,7 @@ class LocalHostedSessionHandle {
     session;
     committedSessionId;
     committedSnapshot;
-    persistedSnapshotRevision;
+    persistedSnapshot;
     draftAssistantMessage;
     messageStates = new Map();
     restoredMessageIds;
@@ -293,6 +316,7 @@ class LocalHostedSessionHandle {
     timelineExtras = [];
     tools = new Map();
     agents = new Map();
+    agentCostTotals = new Map();
     facets = new Map();
     deltaListeners = new Set();
     ephemeralListeners = new Set();
@@ -302,6 +326,12 @@ class LocalHostedSessionHandle {
     unsubscribeSubagentEvent;
     runtimeEventQueue = Promise.resolve();
     snapshotQueue = Promise.resolve();
+    snapshotGeneration = 0;
+    maintenanceAbortControllers = new Set();
+    maintenancePromises = new Set();
+    activeTurnPromise;
+    disposePromise;
+    disposing = false;
     costTotal = 0;
     forceNextSnapshotRevision;
     disposed = false;
@@ -319,7 +349,9 @@ class LocalHostedSessionHandle {
         this.committedSnapshot = committedSnapshot
             ? cloneSessionProtocolSnapshot(committedSnapshot)
             : undefined;
-        this.persistedSnapshotRevision = committedSnapshot?.revision;
+        this.persistedSnapshot = committedSnapshot
+            ? cloneSessionProtocolSnapshot(committedSnapshot)
+            : undefined;
         this.forceNextSnapshotRevision = forceNextSnapshotRevision;
         this.restoreProtocolState(committedSnapshot);
         this.unsubscribeSubagentEvent = this.session.onSubagentEvent((event) => {
@@ -362,6 +394,18 @@ class LocalHostedSessionHandle {
     }
     async runTurn() {
         this.assertActive();
+        const run = this.runTurnNow();
+        this.activeTurnPromise = run;
+        try {
+            return await run;
+        }
+        finally {
+            if (this.activeTurnPromise === run) {
+                this.activeTurnPromise = undefined;
+            }
+        }
+    }
+    async runTurnNow() {
         try {
             const result = await this.runtime.runTurn({
                 onEvent: (event) => this.recordTurnRuntimeEvent(event),
@@ -386,6 +430,23 @@ class LocalHostedSessionHandle {
         this.assertActive();
         return this.runtime.interruptTurn();
     }
+    interruptMaintenance() {
+        let interrupted = false;
+        for (const abortController of this.maintenanceAbortControllers) {
+            if (!abortController.signal.aborted) {
+                abortController.abort();
+                interrupted = true;
+            }
+        }
+        return interrupted;
+    }
+    async waitForActiveWork() {
+        await Promise.allSettled([
+            ...(this.activeTurnPromise ? [this.activeTurnPromise] : []),
+            ...this.maintenancePromises,
+        ]);
+        await this.runtimeEventQueue;
+    }
     requestTurnBoundaryStop() {
         this.assertActive();
         return this.runtime.requestTurnBoundaryStop();
@@ -415,7 +476,7 @@ class LocalHostedSessionHandle {
         else if (snapshot.revision !== fromRevision) {
             this.emitDelta(createSessionProtocolDeltaMessage({
                 sessionId: snapshot.sessionId,
-                fromRevision,
+                fromRevision: snapshot.revision - 1,
                 toRevision: snapshot.revision,
                 reason: "configuration",
                 delta: {
@@ -547,37 +608,60 @@ class LocalHostedSessionHandle {
         return await this.pathAutocompleteLoad;
     }
     async compact(options) {
-        this.assertActive();
-        const result = await this.session.compact({
-            mode: options.mode === "summary-only" ? "only-summary" : "with-last-assistant",
-            ...(options.guidance !== undefined ? { guidance: options.guidance } : {}),
+        return await this.runMaintenance(async (signal) => {
+            const result = await this.session.compact({
+                mode: options.mode === "summary-only" ? "only-summary" : "with-last-assistant",
+                ...(options.guidance !== undefined ? { guidance: options.guidance } : {}),
+                signal,
+            });
+            signal.throwIfAborted();
+            await this.runtimeEventQueue;
+            this.reconcileProjections();
+            const snapshot = await this.commitSnapshot();
+            this.emitSnapshotReset("maintenance", snapshot);
+            return {
+                snapshot,
+                compactionMessage: result.compactionMessage,
+                includedLastAssistant: result.includedLastAssistant,
+            };
         });
-        const snapshot = await this.commitSnapshot();
-        this.emitSnapshotReset("maintenance", snapshot);
-        return {
-            snapshot,
-            compactionMessage: result.compactionMessage,
-            includedLastAssistant: result.includedLastAssistant,
-        };
     }
     async pruneToolResults(options) {
-        this.assertActive();
-        const result = await this.session.pruneToolResults({
-            strategy: options.strategy,
-            fraction: options.fraction,
-            ...(options.guidance !== undefined ? { guidance: options.guidance } : {}),
+        return await this.runMaintenance(async (signal) => {
+            const result = await this.session.pruneToolResults({
+                strategy: options.strategy,
+                fraction: options.fraction,
+                ...(options.guidance !== undefined ? { guidance: options.guidance } : {}),
+                signal,
+            });
+            signal.throwIfAborted();
+            this.reconcileProjections({ prunedToolResults: result.prunedToolResults });
+            const snapshot = await this.commitSnapshot();
+            this.emitSnapshotReset("maintenance", snapshot);
+            return {
+                snapshot,
+                message: result.message,
+                noop: result.noop,
+                bashResultsPruned: result.bashResultsPruned,
+                editCallsPruned: result.editCallsPruned,
+                editResultsPruned: result.editResultsPruned,
+                bytesPruned: result.bytesPruned,
+            };
         });
-        const snapshot = await this.commitSnapshot();
-        this.emitSnapshotReset("maintenance", snapshot);
-        return {
-            snapshot,
-            message: result.message,
-            noop: result.noop,
-            bashResultsPruned: result.bashResultsPruned,
-            editCallsPruned: result.editCallsPruned,
-            editResultsPruned: result.editResultsPruned,
-            bytesPruned: result.bytesPruned,
-        };
+    }
+    async runMaintenance(operation) {
+        this.assertActive();
+        const abortController = new AbortController();
+        this.maintenanceAbortControllers.add(abortController);
+        const promise = operation(abortController.signal);
+        this.maintenancePromises.add(promise);
+        try {
+            return await promise;
+        }
+        finally {
+            this.maintenanceAbortControllers.delete(abortController);
+            this.maintenancePromises.delete(promise);
+        }
     }
     async rewindToHistoryEntryId(historyEntryId) {
         this.assertActive();
@@ -585,6 +669,8 @@ class LocalHostedSessionHandle {
         if (!result) {
             throw new Error("rewind failed");
         }
+        await this.runtimeEventQueue;
+        this.reconcileProjections({ removeMissingAgents: true });
         const snapshot = await this.commitSnapshot();
         this.emitSnapshotReset("maintenance", snapshot);
         return {
@@ -646,95 +732,172 @@ class LocalHostedSessionHandle {
         return await this.commitSnapshot();
     }
     async dispose() {
+        if (!this.disposePromise) {
+            this.disposing = true;
+            this.disposePromise = this.disposeNow();
+        }
+        return await this.disposePromise;
+    }
+    async disposeNow() {
         if (this.disposed) {
             return;
         }
+        const errors = [];
+        this.runtime.interruptTurn();
+        this.interruptMaintenance();
+        try {
+            await this.waitForActiveWork();
+        }
+        catch (error) {
+            errors.push(error);
+        }
         this.disposed = true;
         try {
             this.unsubscribeSubagentEvent();
-            for (const session of this.ephemeralAgentSessions.values()) {
+        }
+        catch (error) {
+            errors.push(error);
+        }
+        for (const session of this.ephemeralAgentSessions.values()) {
+            try {
                 session.dispose();
             }
-            this.ephemeralAgentSessions.clear();
+            catch (error) {
+                errors.push(error);
+            }
+        }
+        this.ephemeralAgentSessions.clear();
+        try {
             this.session.dispose();
+        }
+        catch (error) {
+            errors.push(error);
+        }
+        try {
             await this.executionEnvironment.dispose();
         }
+        catch (error) {
+            errors.push(error);
+        }
         finally {
             this.removeFromHost(this);
         }
+        if (errors.length > 0) {
+            throw new AggregateError(errors, "failed to dispose hosted session");
+        }
     }
     async commitSnapshot() {
         const write = this.snapshotQueue.catch(() => undefined).then(() => this.writeSnapshot());
         this.snapshotQueue = write.catch(() => undefined);
         return await write;
     }
-    async commitSnapshotWithRevision(revision, options = {}) {
+    async commitSnapshotWithRevision(revision) {
         const write = this.snapshotQueue
             .catch(() => undefined)
-            .then(() => this.writeSnapshotWithRevision(revision, options));
+            .then(() => this.writeSnapshotWithRevision(revision));
         this.snapshotQueue = write.catch(() => undefined);
         return await write;
     }
     async writeSnapshot() {
-        this.assertActive();
+        this.assertNotDisposed();
+        const generation = this.snapshotGeneration;
         const draft = this.buildSnapshotDraft();
-        if (this.committedSessionId !== draft.sessionId) {
-            await this.store.deleteSession(this.committedSessionId, {
-                ...(this.persistedSnapshotRevision !== undefined
-                    ? { expectedRevision: this.persistedSnapshotRevision }
-                    : {}),
-            });
-            this.committedSessionId = draft.sessionId;
-            this.committedSnapshot = undefined;
-            this.persistedSnapshotRevision = undefined;
-        }
+        await this.switchSnapshotSession(draft.sessionId);
         const snapshot = {
             ...draft,
             revision: this.nextSnapshotRevision(draft),
         };
-        if (this.committedSnapshot &&
-            snapshot.revision === this.committedSnapshot.revision &&
-            this.persistedSnapshotRevision === snapshot.revision) {
-            return cloneSessionProtocolSnapshot(this.committedSnapshot);
+        if (this.persistedSnapshot && isDeepStrictEqual(this.persistedSnapshot, snapshot)) {
+            return cloneSessionProtocolSnapshot(this.committedSnapshot ?? snapshot);
         }
         await this.store.commitSessionSnapshot(snapshot, {
-            expectedRevision: this.persistedSnapshotRevision ?? 0,
+            expectedRevision: this.persistedSnapshot?.revision ?? 0,
         });
-        this.persistedSnapshotRevision = snapshot.revision;
+        this.persistedSnapshot = cloneSessionProtocolSnapshot(snapshot);
+        if (generation !== this.snapshotGeneration) {
+            return await this.writeSnapshot();
+        }
         return cloneSessionProtocolSnapshot(this.updateCommittedSnapshotAfterWrite(snapshot));
     }
-    async writeSnapshotWithRevision(revision, options = {}) {
-        this.assertActive();
+    async writeSnapshotWithRevision(revision) {
+        this.assertNotDisposed();
+        const generation = this.snapshotGeneration;
         const draft = this.buildSnapshotDraft();
-        const persist = options.persist ?? true;
-        if (this.committedSessionId !== draft.sessionId) {
-            await this.store.deleteSession(this.committedSessionId, {
-                ...(this.persistedSnapshotRevision !== undefined
-                    ? { expectedRevision: this.persistedSnapshotRevision }
-                    : {}),
-            });
-            this.committedSessionId = draft.sessionId;
-            this.committedSnapshot = undefined;
-            this.persistedSnapshotRevision = undefined;
-        }
+        await this.switchSnapshotSession(draft.sessionId);
         const snapshot = {
             ...draft,
             revision,
         };
-        if (persist) {
-            await this.store.commitSessionSnapshot(snapshot, {
-                expectedRevision: this.persistedSnapshotRevision ?? 0,
-            });
-            this.persistedSnapshotRevision = snapshot.revision;
+        await this.store.commitSessionSnapshot(snapshot, {
+            expectedRevision: this.persistedSnapshot?.revision ?? 0,
+        });
+        this.persistedSnapshot = cloneSessionProtocolSnapshot(snapshot);
+        if (generation !== this.snapshotGeneration) {
+            return await this.writeSnapshot();
         }
         return cloneSessionProtocolSnapshot(this.updateCommittedSnapshotAfterWrite(snapshot));
     }
+    async switchSnapshotSession(sessionId) {
+        if (this.committedSessionId === sessionId) {
+            return;
+        }
+        await this.store.deleteSession(this.committedSessionId, {
+            ...(this.persistedSnapshot ? { expectedRevision: this.persistedSnapshot.revision } : {}),
+        });
+        this.committedSessionId = sessionId;
+        this.committedSnapshot = undefined;
+        this.persistedSnapshot = undefined;
+    }
     updateCommittedSnapshotAfterWrite(snapshot) {
-        if (!this.committedSnapshot || this.committedSnapshot.revision <= snapshot.revision) {
+        if (!this.committedSnapshot || this.committedSnapshot.revision < snapshot.revision) {
             this.committedSnapshot = cloneSessionProtocolSnapshot(snapshot);
+            this.snapshotGeneration += 1;
         }
         return this.committedSnapshot;
     }
+    reconcileProjections(options = {}) {
+        const messageIds = new Set(this.session.rawHistoryEntries.map((entry) => entry.id));
+        messageIds.add("system");
+        for (const [id, tool] of this.tools) {
+            if (!messageIds.has(tool.call.messageId) ||
+                (tool.resultMessageId !== undefined && !messageIds.has(tool.resultMessageId))) {
+                this.tools.delete(id);
+            }
+        }
+        if (options.removeMissingAgents) {
+            for (const id of this.agents.keys()) {
+                if (!this.session.hasSubagent(id)) {
+                    this.agents.delete(id);
+                    this.agentCostTotals.delete(id);
+                }
+            }
+        }
+        const operationIds = new Set(this.timelineExtras.filter((item) => item.type === "operation").map((item) => item.id));
+        const prunedByToolId = new Map((options.prunedToolResults ?? []).map((result) => [result.toolCallId, result.content]));
+        for (const [id, facet] of this.facets) {
+            const subjectExists = facet.subject.type === "session" ||
+                (facet.subject.type === "message" && messageIds.has(facet.subject.id)) ||
+                (facet.subject.type === "tool" && this.tools.has(facet.subject.id)) ||
+                (facet.subject.type === "agent" && this.agents.has(facet.subject.id)) ||
+                (facet.subject.type === "operation" && operationIds.has(facet.subject.id));
+            if (!subjectExists) {
+                this.facets.delete(id);
+                continue;
+            }
+            if (facet.subject.type === "tool") {
+                const toolCallId = facet.subject.id;
+                const content = prunedByToolId.get(toolCallId);
+                if (content !== undefined && facet.kind === "tau.tool-ui-events") {
+                    this.facets.set(id, {
+                        ...facet,
+                        data: {
+                            events: [{ type: "tool_pruned", toolCallId, content }],
+                        },
+                    });
+                }
+            }
+        }
+    }
     buildSnapshotDraft() {
         const messages = this.buildProtocolMessages();
         return {
@@ -832,8 +995,10 @@ class LocalHostedSessionHandle {
             this.tools.set(id, structuredClone(tool));
         }
         this.agents.clear();
+        this.agentCostTotals.clear();
         for (const [id, agent] of Object.entries(snapshot.agents)) {
             this.agents.set(id, structuredClone(agent));
+            this.agentCostTotals.set(id, agent.costTotal);
         }
         this.costTotal = snapshot.costTotal;
         this.facets.clear();
@@ -849,6 +1014,11 @@ class LocalHostedSessionHandle {
         }
     }
     assertActive() {
+        if (this.disposing || this.disposed) {
+            throw new Error(`session is shut down: ${this.committedSessionId}`);
+        }
+    }
+    assertNotDisposed() {
         if (this.disposed) {
             throw new Error(`session is shut down: ${this.committedSessionId}`);
         }
@@ -1047,13 +1217,14 @@ class LocalHostedSessionHandle {
             }
             case "compaction_end":
                 this.clearRunningAutoCompactionOperations();
+                this.reconcileProjections();
                 await this.emitSnapshotReset("maintenance", await this.commitSnapshot());
                 return;
             case "tool_ui":
                 await this.recordToolUiEvent(event.uiEvent);
                 return;
             case "subagent_ui":
-                await this.recordSubagentUiEvent(event.event, event.originHistoryEntryId);
+                await this.recordSubagentUiEvent(event.event);
                 return;
         }
     }
@@ -1062,6 +1233,10 @@ class LocalHostedSessionHandle {
     }
     async recordToolUiEvent(event) {
         const toolCallId = event.toolCallId;
+        const existingTool = this.tools.get(toolCallId);
+        if (!existingTool) {
+            return;
+        }
         const facetId = `tool-ui-${toolCallId}`;
         const existing = this.facets.get(facetId);
         const events = Array.isArray(existing?.data.events) ? existing.data.events : [];
@@ -1073,18 +1248,10 @@ class LocalHostedSessionHandle {
             data: { events: [...events, structuredClone(event)] },
         };
         this.facets.set(facet.id, facet);
-        const existingTool = this.tools.get(toolCallId);
-        const toolChange = existingTool
-            ? {
-                type: "tool.set",
-                tool: this.updateToolRunFromUiEvent(existingTool, event),
-            }
-            : undefined;
-        if (toolChange) {
-            this.tools.set(toolChange.tool.id, toolChange.tool);
-        }
+        const tool = this.updateToolRunFromUiEvent(existingTool, event);
+        this.tools.set(tool.id, tool);
         await this.emitPatch("tool-run", [
-            ...(toolChange ? [toolChange] : []),
+            { type: "tool.set", tool },
             { type: "facet.set", facet: structuredClone(facet) },
         ]);
     }
@@ -1113,13 +1280,18 @@ class LocalHostedSessionHandle {
         }
         return next;
     }
-    async recordSubagentUiEvent(event, originHistoryEntryId) {
+    async recordSubagentUiEvent(event) {
         const existing = "id" in event ? this.agents.get(event.id) : this.agents.get(event.state.id);
-        const agent = agentRunFromSubagentEvent(event, originHistoryEntryId, existing);
+        const agent = agentRunFromSubagentEvent(event, existing);
         if (!agent) {
             return;
         }
-        this.costTotal += Math.max(0, agent.costTotal - (existing?.costTotal ?? 0));
+        const previousCost = this.agentCostTotals.get(agent.id) ?? 0;
+        this.costTotal += Math.max(0, agent.costTotal - previousCost);
+        this.agentCostTotals.set(agent.id, agent.costTotal);
+        if (!this.session.hasSubagent(agent.id)) {
+            return;
+        }
         this.agents.set(agent.id, agent);
         await this.emitPatch("agent-run", [
             { type: "cost.set", costTotal: this.costTotal },
@@ -1165,10 +1337,11 @@ class LocalHostedSessionHandle {
         });
         if (options.persist === false) {
             this.committedSnapshot = applySessionProtocolDelta(this.committedSnapshot, delta);
+            this.snapshotGeneration += 1;
             this.emitDelta(delta);
             return;
         }
-        const snapshot = await this.commitSnapshotWithRevision(toRevision, options);
+        const snapshot = await this.commitSnapshotWithRevision(toRevision);
         this.emitDelta(createSessionProtocolDeltaMessage({
             sessionId: delta.sessionId,
             fromRevision,
@@ -1433,7 +1606,7 @@ function createInterruptedAssistantMessageFromModelSnapshot(draft, model) {
         timestamp: draft.message.timestamp,
     };
 }
-function agentRunFromSubagentEvent(event, originHistoryEntryId, existing) {
+function agentRunFromSubagentEvent(event, existing) {
     const state = event.type === "subagent_spawned" || event.type === "subagent_finished"
         ? event.state
         : undefined;
@@ -1467,7 +1640,6 @@ function agentRunFromSubagentEvent(event, originHistoryEntryId, existing) {
                 : state.status === "aborted"
                     ? "cancelled"
                     : "running",
-        originMessageId: originHistoryEntryId,
         ...(state.modelLabel !== undefined ? { modelLabel: state.modelLabel } : {}),
         costTotal: state.costTotal,
         turns: state.turns,