@markusylisiurunen/tau 0.3.44 → 0.3.45

package/dist/core/telegram/session_manager.js

@@ -1,4 +1,4 @@
-import { randomBytes } from "node:crypto";
+import { randomBytes, randomUUID } from "node:crypto";
 import { mkdir, readFile, rename, stat, writeFile } from "node:fs/promises";
 import { dirname, resolve } from "node:path";
 import { z } from "zod";
@@ -24,10 +24,49 @@ const persistedTelegramSessionRecordV1Schema = z
     tauSessionId: z.string().min(1).optional(),
 })
     .strip();
-const persistedTelegramSessionRecordSchema = persistedTelegramSessionRecordV1Schema.extend({
+const persistedTelegramSessionRecordV2Schema = persistedTelegramSessionRecordV1Schema.extend({
     state: telegramSessionStateValueSchema,
     error: z.string().min(1).optional(),
 });
+const persistedTelegramSessionRecordV3Schema = persistedTelegramSessionRecordV2Schema.extend({
+    activeTurnIds: z.array(z.string().min(1)),
+});
+const persistedTelegramTurnFailureSchema = z.discriminatedUnion("status", [
+    z
+        .object({
+        status: z.literal("failed"),
+        stopReason: z.literal("error"),
+        errorMessage: z.string().optional(),
+    })
+        .strip(),
+    z
+        .object({
+        status: z.literal("blocked"),
+        reason: z.literal("auto-compaction-failed"),
+        message: z.string(),
+    })
+        .strip(),
+]);
+const persistedTelegramTurnNotificationSchema = z.discriminatedUnion("kind", [
+    z
+        .object({
+        kind: z.literal("failed"),
+        historyEntryId: z.string().min(1),
+        timestamp: z.string().datetime(),
+        failure: persistedTelegramTurnFailureSchema,
+    })
+        .strip(),
+    z
+        .object({
+        kind: z.literal("rejected"),
+        historyEntryId: z.string().min(1),
+        timestamp: z.string().datetime(),
+    })
+        .strip(),
+]);
+const persistedTelegramSessionRecordSchema = persistedTelegramSessionRecordV3Schema.extend({
+    pendingTurnNotifications: z.array(persistedTelegramTurnNotificationSchema),
+});
 const telegramSessionStateSchema = z.discriminatedUnion("version", [
     z
         .object({
@@ -38,17 +77,41 @@ const telegramSessionStateSchema = z.discriminatedUnion("version", [
     z
         .object({
         version: z.literal(2),
+        sessions: z.array(persistedTelegramSessionRecordV2Schema),
+    })
+        .strip(),
+    z
+        .object({
+        version: z.literal(3),
+        sessions: z.array(persistedTelegramSessionRecordV3Schema),
+    })
+        .strip(),
+    z
+        .object({
+        version: z.literal(4),
         sessions: z.array(persistedTelegramSessionRecordSchema),
     })
         .strip(),
 ]);
 function normalizePersistedTelegramSessions(state) {
-    if (state.version === 2) {
+    if (state.version === 4) {
         return state.sessions;
     }
+    if (state.version === 3) {
+        return state.sessions.map((record) => ({ ...record, pendingTurnNotifications: [] }));
+    }
+    if (state.version === 2) {
+        return state.sessions.map((record) => ({
+            ...record,
+            activeTurnIds: [],
+            pendingTurnNotifications: [],
+        }));
+    }
     return state.sessions.map((record) => ({
         ...record,
         state: record.tauSessionId ? "waiting-input" : "queued",
+        activeTurnIds: [],
+        pendingTurnNotifications: [],
     }));
 }
 export function resolveTelegramSessionStatePath(workspaceRoot) {
@@ -59,6 +122,7 @@ const PUBLIC_SESSION_ID_LENGTH = 8;
 const NANOSECONDS_PER_MILLISECOND = 1000000n;
 const PROVISION_SCRIPT_PATH = ".tau/scripts/provision";
 const MAX_PROVISION_DIAGNOSTIC_CHARS = 2_000;
+const TURN_RECOVERY_POLL_INTERVAL_MS = 1_000;
 const ACTIVE_STATES = new Set([
     "queued",
     "preparing-workspace",
@@ -215,7 +279,11 @@ class TelegramSessionManagerImpl {
             consumedFacetEventCounts: new Map(),
             emittedAssistantMessageIds: new Set(),
             emittedNoticeIds: new Set(),
-            emittedTurnFailureIds: new Set(),
+            activeTurnIds: new Set(),
+            settledTurnIds: new Set(),
+            pendingTurnNotifications: new Map(),
+            emittedTurnNotificationIds: new Set(),
+            turnChangeResolvers: new Map(),
         };
         this.sessions.set(id, entry);
         try {
@@ -259,6 +327,21 @@ class TelegramSessionManagerImpl {
         }
         return (this.provisionFailures.get(sessionId) ?? []).map((failure) => ({ ...failure }));
     }
+    getPendingTurnNotifications(sessionId) {
+        const entry = this.getEntryBySessionId(sessionId);
+        if (!entry) {
+            return [];
+        }
+        return Array.from(entry.pendingTurnNotifications.values(), (notification) => this.toTurnNotificationEvent(entry, notification));
+    }
+    async acknowledgeTurnNotification(sessionId, historyEntryId) {
+        const entry = this.requireSession(sessionId);
+        if (!entry.pendingTurnNotifications.delete(historyEntryId)) {
+            return;
+        }
+        entry.emittedTurnNotificationIds.delete(historyEntryId);
+        await this.persistSessions();
+    }
     async getSessionSnapshot(sessionId) {
         const entry = this.getEntryBySessionId(sessionId);
         if (!entry?.tauSession) {
@@ -284,8 +367,10 @@ class TelegramSessionManagerImpl {
         if (!entry.tauSession) {
             throw new TelegramSessionManagerError("not_ready", "session is still preparing");
         }
-        const mode = options?.mode ?? "submit";
-        if (mode !== "steer" && (entry.record.state === "running" || entry.activeSubmit)) {
+        const requestedMode = options?.mode ?? "submit";
+        const running = entry.record.state === "running" || Boolean(entry.activeSubmit);
+        const mode = requestedMode === "auto" ? (running ? "steer" : "submit") : requestedMode;
+        if (mode !== "steer" && running) {
             throw new TelegramSessionManagerError("busy", "session is running");
         }
         void this.submitText(entry, trimmed, "user-message", options?.additionalSystemMessage, mode);
@@ -293,7 +378,8 @@ class TelegramSessionManagerImpl {
     }
     async interruptSession(sessionId) {
         const entry = this.requireSession(sessionId);
-        if (entry.record.state !== "running" || !entry.activeSubmit) {
+        if (entry.record.state !== "running" ||
+            (!entry.activeSubmit && entry.activeTurnIds.size === 0)) {
             return {
                 session: this.toRecord(entry),
                 interrupted: false,
@@ -381,8 +467,9 @@ class TelegramSessionManagerImpl {
             if (!project || this.sessions.has(record.id)) {
                 continue;
             }
+            const { activeTurnIds, pendingTurnNotifications, ...sessionRecord } = record;
             const entry = {
-                record: { ...record },
+                record: sessionRecord,
                 logs: [],
                 project,
                 abortController: new AbortController(),
@@ -390,13 +477,20 @@ class TelegramSessionManagerImpl {
                 consumedFacetEventCounts: new Map(),
                 emittedAssistantMessageIds: new Set(),
                 emittedNoticeIds: new Set(),
-                emittedTurnFailureIds: new Set(),
+                activeTurnIds: new Set(activeTurnIds),
+                settledTurnIds: new Set(),
+                pendingTurnNotifications: new Map(pendingTurnNotifications.map((notification) => [
+                    notification.historyEntryId,
+                    notification,
+                ])),
+                emittedTurnNotificationIds: new Set(),
+                turnChangeResolvers: new Map(),
             };
             this.sessions.set(record.id, entry);
         }
         await this.cleanupOrphanedWorkspaces();
         await Promise.all(Array.from(this.sessions.values(), async (entry) => {
-            if (entry.record.state === "failed") {
+            if (entry.record.state === "failed" && entry.activeTurnIds.size === 0) {
                 return;
             }
             try {
@@ -415,6 +509,9 @@ class TelegramSessionManagerImpl {
             }
         }));
         await this.persistSessions();
+        for (const entry of this.sessions.values()) {
+            this.emitPendingTurnNotifications(entry);
+        }
     }
     async cleanupOrphanedWorkspaces() {
         const workspaceRoots = new Set([this.workspaceRoot]);
@@ -498,7 +595,13 @@ class TelegramSessionManagerImpl {
         });
         this.initializeRecoveredSnapshotDeliveryState(entry, await tauSession.snapshot());
         entry.record.error = undefined;
-        this.setState(entry, "waiting-input");
+        if (entry.activeTurnIds.size > 0) {
+            this.setState(entry, "running");
+            this.trackRecoveredTurns(entry, tauSession);
+        }
+        else {
+            this.setState(entry, "waiting-input");
+        }
         this.log(entry, "info", "session recovered", { tauSessionId, workspacePath });
         if (shouldPrepareWorkspace) {
             this.startProvision(entry, tauSession, provisionTargets);
@@ -513,7 +616,7 @@ class TelegramSessionManagerImpl {
             .catch(() => undefined)
             .then(async () => {
             const state = {
-                version: 2,
+                version: 4,
                 sessions: Array.from(this.sessions.values(), (entry) => ({
                     id: entry.record.id,
                     projectId: entry.record.projectId,
@@ -523,6 +626,8 @@ class TelegramSessionManagerImpl {
                     updatedAt: entry.record.updatedAt,
                     ...(entry.record.tauSessionId ? { tauSessionId: entry.record.tauSessionId } : {}),
                     ...(entry.record.error ? { error: entry.record.error } : {}),
+                    activeTurnIds: [...entry.activeTurnIds],
+                    pendingTurnNotifications: [...entry.pendingTurnNotifications.values()],
                 })),
             };
             await mkdir(dirname(persistencePath), { recursive: true });
@@ -591,6 +696,7 @@ class TelegramSessionManagerImpl {
                 this.handleClientEvent(entry, event);
             });
             const snapshot = await tauSession.snapshot();
+            this.reconcileSnapshotTurns(entry, snapshot, { rejectMissingActiveTurns: false });
             this.initializeSnapshotMessageDeliveryState(entry, snapshot);
             this.handleSnapshotNotices(entry, snapshot);
             this.log(entry, "info", "session preparation complete", {
@@ -698,35 +804,57 @@ class TelegramSessionManagerImpl {
         const previousSubmit = entry.activeSubmit;
         const tauSession = entry.tauSession;
         const payload = this.buildSubmitPayload(text, additionalSystemMessage);
+        const requestedHistoryEntryId = mode === "submit" ? `telegram-turn-${randomUUID()}` : undefined;
         let submitPromise;
         submitPromise = (async () => {
+            let requestStarted = false;
             try {
-                const result = mode === "steer" ? await tauSession.steer(payload) : await tauSession.submit(payload);
-                const failure = result.turn.status === "failed" || result.turn.status === "blocked"
-                    ? result.turn
-                    : undefined;
-                this.log(entry, failure ? "error" : "info", "message finished", {
+                if (requestedHistoryEntryId) {
+                    entry.activeTurnIds.add(requestedHistoryEntryId);
+                    await this.persistSessions();
+                }
+                requestStarted = true;
+                const result = mode === "steer"
+                    ? await tauSession.steer(payload)
+                    : await tauSession.submit(payload, { historyEntryId: requestedHistoryEntryId });
+                if (requestedHistoryEntryId && result.userHistoryEntryId !== requestedHistoryEntryId) {
+                    throw new Error(`Tau accepted turn '${result.userHistoryEntryId}' instead of requested turn '${requestedHistoryEntryId}'`);
+                }
+                this.reconcileTurnRecord(entry, {
+                    userHistoryEntryId: result.userHistoryEntryId,
+                    state: "settled",
+                    outcome: result.turn,
+                }, { owned: true });
+                await this.persistSessions();
+                this.emitPendingTurnNotifications(entry);
+                this.log(entry, result.turn.status === "failed" || result.turn.status === "blocked" ? "error" : "info", "message finished", {
                     source,
                     turn: result.turn,
                     userHistoryEntryId: result.userHistoryEntryId,
                 });
-                if (!entry.cancelRequested) {
-                    if (failure && !entry.emittedTurnFailureIds.has(result.userHistoryEntryId)) {
-                        entry.emittedTurnFailureIds.add(result.userHistoryEntryId);
-                        this.emit({
-                            type: "session-turn-failed",
-                            sessionId: entry.record.id,
-                            projectId: entry.record.projectId,
-                            timestamp: this.now().toISOString(),
-                            failure,
-                        });
-                    }
-                    else if (!failure && (!entry.activeSubmit || entry.activeSubmit === submitPromise)) {
-                        this.setState(entry, "waiting-input");
-                    }
+                if (!entry.cancelRequested &&
+                    (!entry.activeSubmit || entry.activeSubmit === submitPromise) &&
+                    entry.activeTurnIds.size === 0) {
+                    this.setState(entry, "waiting-input");
                 }
             }
             catch (error) {
+                if (requestedHistoryEntryId && !requestStarted) {
+                    entry.activeTurnIds.delete(requestedHistoryEntryId);
+                }
+                if (!entry.cancelRequested && requestedHistoryEntryId && requestStarted) {
+                    try {
+                        if (await this.recoverRejectedSubmit(entry, tauSession, requestedHistoryEntryId, error)) {
+                            return;
+                        }
+                    }
+                    catch (recoveryError) {
+                        this.log(entry, "warn", "failed to inspect rejected submit", {
+                            userHistoryEntryId: requestedHistoryEntryId,
+                            cause: formatErrorDiagnostic(recoveryError),
+                        });
+                    }
+                }
                 if (!entry.cancelRequested) {
                     const diagnostic = formatErrorDiagnostic(error);
                     entry.record.error = diagnostic;
@@ -765,13 +893,234 @@ class TelegramSessionManagerImpl {
             .finally(() => {
             if (entry.activeSubmit === trackedSubmit) {
                 entry.activeSubmit = undefined;
-                if (!entry.cancelRequested && entry.record.state === "running") {
+                if (!entry.cancelRequested &&
+                    entry.record.state === "running" &&
+                    entry.activeTurnIds.size === 0) {
                     this.setState(entry, "waiting-input");
                 }
             }
         });
         return submitPromise;
     }
+    async recoverRejectedSubmit(entry, tauSession, historyEntryId, error) {
+        const snapshot = await tauSession.snapshot();
+        const turn = Object.hasOwn(snapshot.turns, historyEntryId)
+            ? snapshot.turns[historyEntryId]
+            : undefined;
+        if (!turn) {
+            this.recordTurnRejection(entry, historyEntryId);
+            await this.persistSessions();
+            this.emitPendingTurnNotifications(entry);
+            this.log(entry, "warn", "submit rejected before turn acceptance", {
+                userHistoryEntryId: historyEntryId,
+                cause: formatErrorDiagnostic(error),
+            });
+            return true;
+        }
+        this.reconcileTurnRecord(entry, turn);
+        await this.persistSessions();
+        this.emitPendingTurnNotifications(entry);
+        this.log(entry, "warn", "submit response recovered from turn ledger", {
+            userHistoryEntryId: historyEntryId,
+            turn,
+            cause: formatErrorDiagnostic(error),
+        });
+        if (turn.state === "running") {
+            await this.followRunningTurn(entry, tauSession, historyEntryId);
+        }
+        else if (entry.activeTurnIds.size === 0) {
+            this.setState(entry, "waiting-input");
+        }
+        return true;
+    }
+    trackRecoveredTurns(entry, tauSession) {
+        const trackedTurns = Promise.all([...entry.activeTurnIds].map(async (historyEntryId) => await this.followRunningTurn(entry, tauSession, historyEntryId))).then(() => undefined);
+        entry.activeSubmit = trackedTurns;
+        void trackedTurns
+            .catch((error) => {
+            this.log(entry, "error", "recovered turn tracking failed", {
+                cause: formatErrorDiagnostic(error),
+            });
+        })
+            .finally(() => {
+            if (entry.activeSubmit === trackedTurns) {
+                entry.activeSubmit = undefined;
+                if (!entry.cancelRequested && entry.record.state === "running") {
+                    this.setState(entry, "waiting-input");
+                }
+            }
+        });
+    }
+    async followRunningTurn(entry, tauSession, historyEntryId) {
+        while (entry.activeTurnIds.has(historyEntryId) && !entry.cancelRequested) {
+            await this.waitForTurnChange(entry, historyEntryId);
+            if (!entry.activeTurnIds.has(historyEntryId) || entry.cancelRequested) {
+                return;
+            }
+            let snapshot;
+            try {
+                snapshot = await tauSession.snapshot();
+            }
+            catch (error) {
+                this.log(entry, "warn", "failed to refresh recovered turn", {
+                    userHistoryEntryId: historyEntryId,
+                    cause: formatErrorDiagnostic(error),
+                });
+                continue;
+            }
+            const turn = Object.hasOwn(snapshot.turns, historyEntryId)
+                ? snapshot.turns[historyEntryId]
+                : undefined;
+            if (!turn) {
+                this.recordTurnRejection(entry, historyEntryId);
+            }
+            else {
+                this.reconcileTurnRecord(entry, turn);
+            }
+            await this.persistSessions();
+            this.emitPendingTurnNotifications(entry);
+        }
+    }
+    waitForTurnChange(entry, historyEntryId) {
+        return new Promise((resolve) => {
+            let settled = false;
+            const finish = () => {
+                if (settled) {
+                    return;
+                }
+                settled = true;
+                clearTimeout(timeout);
+                entry.abortController.signal.removeEventListener("abort", finish);
+                if (entry.turnChangeResolvers.get(historyEntryId) === finish) {
+                    entry.turnChangeResolvers.delete(historyEntryId);
+                }
+                resolve();
+            };
+            const timeout = setTimeout(finish, TURN_RECOVERY_POLL_INTERVAL_MS);
+            entry.turnChangeResolvers.set(historyEntryId, finish);
+            entry.abortController.signal.addEventListener("abort", finish, { once: true });
+        });
+    }
+    reconcileTurnRecord(entry, turn, options = {}) {
+        const historyEntryId = turn.userHistoryEntryId;
+        if (turn.state === "running") {
+            if (entry.settledTurnIds.has(historyEntryId) || entry.activeTurnIds.has(historyEntryId)) {
+                return false;
+            }
+            entry.activeTurnIds.add(historyEntryId);
+            return true;
+        }
+        if (entry.settledTurnIds.has(historyEntryId)) {
+            return false;
+        }
+        entry.settledTurnIds.add(historyEntryId);
+        const wasActive = entry.activeTurnIds.delete(historyEntryId);
+        if (!wasActive && !options.owned) {
+            return false;
+        }
+        entry.turnChangeResolvers.get(historyEntryId)?.();
+        if (turn.outcome.status === "failed" || turn.outcome.status === "blocked") {
+            this.recordTurnFailure(entry, historyEntryId, turn.outcome);
+        }
+        if (entry.activeTurnIds.size === 0 &&
+            !entry.activeSubmit &&
+            !entry.cancelRequested &&
+            entry.record.state === "running") {
+            this.setState(entry, "waiting-input");
+        }
+        return true;
+    }
+    recordTurnFailure(entry, historyEntryId, failure) {
+        if (entry.pendingTurnNotifications.has(historyEntryId)) {
+            return;
+        }
+        entry.pendingTurnNotifications.set(historyEntryId, {
+            kind: "failed",
+            historyEntryId,
+            timestamp: this.now().toISOString(),
+            failure: structuredClone(failure),
+        });
+    }
+    recordTurnRejection(entry, historyEntryId) {
+        entry.activeTurnIds.delete(historyEntryId);
+        entry.settledTurnIds.add(historyEntryId);
+        entry.turnChangeResolvers.get(historyEntryId)?.();
+        if (!entry.pendingTurnNotifications.has(historyEntryId)) {
+            entry.pendingTurnNotifications.set(historyEntryId, {
+                kind: "rejected",
+                historyEntryId,
+                timestamp: this.now().toISOString(),
+            });
+        }
+        if (entry.activeTurnIds.size === 0 &&
+            !entry.activeSubmit &&
+            !entry.cancelRequested &&
+            entry.record.state === "running") {
+            this.setState(entry, "waiting-input");
+        }
+    }
+    toTurnNotificationEvent(entry, notification) {
+        if (notification.kind === "failed") {
+            return {
+                type: "session-turn-failed",
+                sessionId: entry.record.id,
+                projectId: entry.record.projectId,
+                timestamp: notification.timestamp,
+                historyEntryId: notification.historyEntryId,
+                failure: structuredClone(notification.failure),
+            };
+        }
+        return {
+            type: "session-turn-rejected",
+            sessionId: entry.record.id,
+            projectId: entry.record.projectId,
+            timestamp: notification.timestamp,
+            historyEntryId: notification.historyEntryId,
+        };
+    }
+    emitPendingTurnNotifications(entry) {
+        for (const notification of entry.pendingTurnNotifications.values()) {
+            if (entry.emittedTurnNotificationIds.has(notification.historyEntryId)) {
+                continue;
+            }
+            entry.emittedTurnNotificationIds.add(notification.historyEntryId);
+            this.emit(this.toTurnNotificationEvent(entry, notification));
+        }
+    }
+    reconcileSnapshotTurns(entry, snapshot, options) {
+        let changed = false;
+        const snapshotTurnIds = new Set(Object.keys(snapshot.turns));
+        if (options.rejectMissingActiveTurns) {
+            for (const historyEntryId of entry.activeTurnIds) {
+                if (!snapshotTurnIds.has(historyEntryId)) {
+                    this.recordTurnRejection(entry, historyEntryId);
+                    changed = true;
+                }
+            }
+        }
+        for (const turn of Object.values(snapshot.turns)) {
+            if (turn.state === "running" || entry.activeTurnIds.has(turn.userHistoryEntryId)) {
+                changed = this.reconcileTurnRecord(entry, turn) || changed;
+            }
+            else {
+                entry.settledTurnIds.add(turn.userHistoryEntryId);
+            }
+        }
+        return changed;
+    }
+    persistTurnState(entry) {
+        void this.persistSessions()
+            .then(() => {
+            this.emitPendingTurnNotifications(entry);
+        })
+            .catch((error) => {
+            this.onLog?.({
+                level: "error",
+                message: "turn state persistence failed",
+                data: { cause: formatErrorDiagnostic(error) },
+            });
+        });
+    }
     buildSessionAttributes(entry) {
         const configuredRepositories = "repo" in entry.project
             ? [entry.project.repo]
@@ -1031,6 +1380,11 @@ class TelegramSessionManagerImpl {
     }
     handleClientEvent(entry, clientEvent) {
         if (clientEvent.delta.type === "snapshot.reset") {
+            if (this.reconcileSnapshotTurns(entry, clientEvent.delta.snapshot, {
+                rejectMissingActiveTurns: false,
+            })) {
+                this.persistTurnState(entry);
+            }
             this.handleSnapshotNotices(entry, clientEvent.delta.snapshot);
             if (clientEvent.cause.type !== "assistant-message") {
                 return;
@@ -1046,6 +1400,12 @@ class TelegramSessionManagerImpl {
             return;
         }
         for (const change of clientEvent.delta.changes) {
+            if (change.type === "turn.set") {
+                if (this.reconcileTurnRecord(entry, change.turn)) {
+                    this.persistTurnState(entry);
+                }
+                continue;
+            }
             if (change.type === "facet.set") {
                 this.handleFacetProgress(entry, change.facet);
                 continue;
@@ -1063,6 +1423,7 @@ class TelegramSessionManagerImpl {
         }
     }
     initializeRecoveredSnapshotDeliveryState(entry, snapshot) {
+        this.reconcileSnapshotTurns(entry, snapshot, { rejectMissingActiveTurns: true });
         this.initializeSnapshotMessageDeliveryState(entry, snapshot);
         for (const item of snapshot.timeline.items) {
             if (item.type === "notice") {
@@ -1232,6 +1593,17 @@ class ScopedTelegramSessionManager {
         }
         return this.sessionManager.getProvisionFailures(sessionId);
     }
+    getPendingTurnNotifications(sessionId) {
+        const session = this.sessionManager.getSession(sessionId);
+        if (!session || !this.isVisibleSession(session)) {
+            return [];
+        }
+        return this.sessionManager.getPendingTurnNotifications(sessionId);
+    }
+    async acknowledgeTurnNotification(sessionId, historyEntryId) {
+        this.requireSession(sessionId);
+        await this.sessionManager.acknowledgeTurnNotification(sessionId, historyEntryId);
+    }
     async getSessionSnapshot(sessionId) {
         this.requireSession(sessionId);
         return await this.sessionManager.getSessionSnapshot(sessionId);