fix: recover committed conversation leaf

README.md

@@ -44,13 +44,17 @@ themes, or ambient context.
 
 `ConversationRegistry` resolves user and chat identity to one `Conversation`.
 The conversation checks the caller's visible history, reserves one turn, streams
-through its `SessionAgent`, and commits by saving the next visible history.
-Model failure, active abort or cancellation, and history-save failure restore the
-prior Pi branch before the registry permits a new session for that identity.
-
-`FileConversationHistoryStore` atomically replaces mode-0600 history files in a
-mode-0700 directory. It stores only visible user and assistant role/content
-pairs.
+through its `SessionAgent`, and commits by saving the next visible history with
+the post-response Pi leaf. Model failure, active abort or cancellation, and
+history-save failure restore the prior Pi branch before the registry permits a
+new session for that identity.
+
+`FileConversationHistoryStore` atomically replaces mode-0600 versioned snapshots
+in a mode-0700 directory. Each snapshot stores visible user and assistant
+role/content pairs plus the committed Pi leaf. After a process restart, Stein
+restores that leaf before Pi builds model context, excluding entries abandoned
+before snapshot replacement. This is process-crash recovery, not a power-loss
+durability or multi-process safety guarantee.
 
 ## Serve OpenAI-compatible chat
 

src/conversation/conversation-registry.ts

@@ -2,9 +2,11 @@ import { createHash } from "node:crypto";
 import { SessionAgent } from "../session/session-agent.ts";
 import {
   Conversation,
+  type ConversationHistorySnapshot,
   type ConversationHistoryStore,
   type ConversationMessage,
   type ConversationTurn,
+  validateConversationHistorySnapshot,
 } from "./conversation.ts";
 
 export type ConversationIdentity = Readonly<{
@@ -15,6 +17,7 @@ export type ConversationIdentity = Readonly<{
 
 export type ConversationAgentFactory = (
   identity: ConversationIdentity,
+  snapshot: ConversationHistorySnapshot,
 ) => Promise<SessionAgent>;
 
 export class ConversationRegistry {
@@ -42,19 +45,20 @@ export class ConversationRegistry {
     if (existing) return existing;
 
     let created: Promise<Conversation>;
-    created = this.#historyStore.load(identity.conversationId).then((history) =>
-      new Conversation({
+    created = this.#historyStore.load(identity.conversationId).then((snapshot) => {
+      validateConversationHistorySnapshot(snapshot);
+      return new Conversation({
         conversationId: identity.conversationId,
-        history,
+        history: snapshot.messages,
         historyStore: this.#historyStore,
-        createAgent: () => this.#createAgent(identity),
+        createAgent: () => this.#createAgent(identity, snapshot),
         evict: () => {
           if (this.#conversations.get(identity.conversationId) === created) {
             this.#conversations.delete(identity.conversationId);
           }
         },
-      })
-    );
+      });
+    });
     this.#conversations.set(identity.conversationId, created);
     created.catch(() => {
       if (this.#conversations.get(identity.conversationId) === created) {

src/conversation/conversation.ts

@@ -12,16 +12,51 @@ export type VisibleConversationMessage = Readonly<{
   content: string;
 }>;
 
+export type ConversationHistorySnapshot = Readonly<{
+  version: 1;
+  messages: readonly VisibleConversationMessage[];
+  committedLeafId: string | null;
+}>;
+
 export interface ConversationHistoryStore {
-  load(conversationId: string): Promise<readonly VisibleConversationMessage[]>;
+  load(conversationId: string): Promise<ConversationHistorySnapshot>;
   // Success is the turn commit point; abort before replacement must reject.
   save(
     conversationId: string,
-    history: readonly VisibleConversationMessage[],
+    snapshot: ConversationHistorySnapshot,
     signal: AbortSignal,
   ): Promise<void>;
 }
 
+export function emptyConversationHistorySnapshot(): ConversationHistorySnapshot {
+  return { version: 1, messages: [], committedLeafId: null };
+}
+
+export function validateConversationHistorySnapshot(
+  value: ConversationHistorySnapshot,
+): void {
+  if (value.version !== 1 || !Array.isArray(value.messages)) {
+    throw new Error("Invalid visible conversation history snapshot");
+  }
+  if (!value.messages.every(isVisibleConversationMessage)) {
+    throw new Error("Invalid visible conversation history messages");
+  }
+  if (value.committedLeafId !== null && !value.committedLeafId.trim()) {
+    throw new Error("Invalid committed Pi session leaf");
+  }
+  if ((value.messages.length === 0) !== (value.committedLeafId === null)) {
+    throw new Error("Visible conversation history and committed Pi leaf disagree");
+  }
+  if (
+    value.messages.length % 2 !== 0 ||
+    value.messages.some((message, index) =>
+      message.role !== (index % 2 === 0 ? "user" : "assistant")
+    )
+  ) {
+    throw new Error("Visible conversation history is not a sequence of committed turns");
+  }
+}
+
 export type ConversationTurn = Readonly<{
   conversationId: string;
   deltas: AsyncIterable<string>;
@@ -95,15 +130,22 @@ export class Conversation {
   async commit(
     userMessage: ConversationMessage,
     assistantContent: string,
+    committedLeafId: string | null,
     signal: AbortSignal,
   ): Promise<void> {
-    const history: VisibleConversationMessage[] = [
+    const messages: VisibleConversationMessage[] = [
       ...this.#history,
       { role: "user", content: userMessage.content },
       { role: "assistant", content: assistantContent },
     ];
-    await this.#historyStore.save(this.conversationId, history, signal);
-    this.#history = history;
+    const snapshot: ConversationHistorySnapshot = {
+      version: 1,
+      messages,
+      committedLeafId,
+    };
+    validateConversationHistorySnapshot(snapshot);
+    await this.#historyStore.save(this.conversationId, snapshot, signal);
+    this.#history = messages;
     this.#active = false;
   }
 
@@ -177,9 +219,11 @@ class ActiveConversationTurn implements ConversationTurn {
 
       if (this.#phase !== "responding") return;
       this.#phase = "saving";
+      const committedLeafId = this.#agent.checkpoint();
       this.#savePromise = this.#conversation.commit(
         this.#userMessage,
         output,
+        committedLeafId,
         this.#abortController.signal,
       );
       await this.#savePromise;
@@ -209,6 +253,17 @@ class ActiveConversationTurn implements ConversationTurn {
   }
 }
 
+function isVisibleConversationMessage(
+  value: unknown,
+): value is VisibleConversationMessage {
+  if (typeof value !== "object" || value === null) return false;
+  const record = value as Record<string, unknown>;
+  return (
+    (record.role === "user" || record.role === "assistant") &&
+    typeof record.content === "string"
+  );
+}
+
 function snapshotUserMessage(message: ConversationMessage): ConversationMessage {
   return {
     role: "user",

src/conversation/file-history-store.ts

@@ -2,9 +2,11 @@ import { randomUUID } from "node:crypto";
 import { renameSync } from "node:fs";
 import { chmod, mkdir, readFile, rm, writeFile } from "node:fs/promises";
 import { isAbsolute, join } from "node:path";
-import type {
-  ConversationHistoryStore,
-  VisibleConversationMessage,
+import {
+  emptyConversationHistorySnapshot,
+  type ConversationHistorySnapshot,
+  type ConversationHistoryStore,
+  validateConversationHistorySnapshot,
 } from "./conversation.ts";
 
 export class FileConversationHistoryStore implements ConversationHistoryStore {
@@ -17,31 +19,44 @@ export class FileConversationHistoryStore implements ConversationHistoryStore {
     this.#directory = directory;
   }
 
-  async load(conversationId: string): Promise<readonly VisibleConversationMessage[]> {
+  async load(conversationId: string): Promise<ConversationHistorySnapshot> {
     let contents: string;
     try {
       contents = await readFile(this.#path(conversationId), "utf8");
     } catch (error) {
-      if (isMissing(error)) return [];
+      if (isMissing(error)) return emptyConversationHistorySnapshot();
       throw error;
     }
 
-    const value: unknown = JSON.parse(contents);
-    if (!Array.isArray(value) || !value.every(isVisibleMessage)) {
-      throw new Error(`Invalid visible conversation history for ${conversationId}`);
+    let value: unknown;
+    try {
+      value = JSON.parse(contents);
+    } catch (error) {
+      throw new Error("Invalid visible conversation history JSON", { cause: error });
+    }
+    if (Array.isArray(value)) {
+      throw new Error(
+        "Legacy array-only visible history requires an explicit migration or reset",
+      );
     }
-    return value.map(({ role, content }) => ({ role, content }));
+    if (!isConversationHistorySnapshot(value)) {
+      throw new Error(`Invalid visible conversation history snapshot for ${conversationId}`);
+    }
+    validateConversationHistorySnapshot(value);
+    return {
+      version: 1,
+      messages: value.messages.map(({ role, content }) => ({ role, content })),
+      committedLeafId: value.committedLeafId,
+    };
   }
 
   async save(
     conversationId: string,
-    history: readonly VisibleConversationMessage[],
+    snapshot: ConversationHistorySnapshot,
     signal: AbortSignal,
   ): Promise<void> {
     const destination = this.#path(conversationId);
-    if (!history.every(isVisibleMessage)) {
-      throw new Error(`Invalid visible conversation history for ${conversationId}`);
-    }
+    validateConversationHistorySnapshot(snapshot);
 
     signal.throwIfAborted();
     await mkdir(this.#directory, { recursive: true, mode: 0o700 });
@@ -50,7 +65,7 @@ export class FileConversationHistoryStore implements ConversationHistoryStore {
     signal.throwIfAborted();
     const temporary = `${destination}.${process.pid}.${randomUUID()}.tmp`;
     try {
-      await writeFile(temporary, `${JSON.stringify(history, null, 2)}\n`, {
+      await writeFile(temporary, `${JSON.stringify(snapshot, null, 2)}\n`, {
         encoding: "utf8",
         flag: "wx",
         mode: 0o600,
@@ -71,12 +86,15 @@ export class FileConversationHistoryStore implements ConversationHistoryStore {
   }
 }
 
-function isVisibleMessage(value: unknown): value is VisibleConversationMessage {
+function isConversationHistorySnapshot(
+  value: unknown,
+): value is ConversationHistorySnapshot {
   if (typeof value !== "object" || value === null) return false;
   const record = value as Record<string, unknown>;
   return (
-    (record.role === "user" || record.role === "assistant") &&
-    typeof record.content === "string"
+    record.version === 1 &&
+    Array.isArray(record.messages) &&
+    (record.committedLeafId === null || typeof record.committedLeafId === "string")
   );
 }
 

src/server/server-runner.ts

@@ -26,7 +26,8 @@ export async function runServer(config: ServerConfig): Promise<never> {
     bearerToken: bearerToken.trim(),
     modelId,
     historyStore: new FileConversationHistoryStore(config.sessionDirectory),
-    createAgent: ({ conversationId }) => createSession(conversationId),
+    createAgent: ({ conversationId }, snapshot) =>
+      createSession(conversationId, { committedLeafId: snapshot.committedLeafId }),
   });
   const server = Bun.serve({
     hostname: config.hostname,

src/session/pi-session.ts

@@ -1,3 +1,4 @@
+import { existsSync } from "node:fs";
 import { isAbsolute, join } from "node:path";
 import {
   createAgentSession,
@@ -10,6 +11,10 @@ import { SessionAgent } from "./session-agent.ts";
 
 export type ModelDescription = Readonly<{ provider: string; id: string }>;
 
+export type PiSessionRecovery = Readonly<{
+  committedLeafId: string | null;
+}>;
+
 export type PiSessionFactoryConfig = Readonly<{
   model: ModelDescription;
   systemPrompt: string;
@@ -20,7 +25,10 @@ export type PiSessionFactoryConfig = Readonly<{
 
 export function createPiSessionFactory(config: PiSessionFactoryConfig) {
   validateConfig(config);
-  return async (conversationId: string): Promise<SessionAgent> => {
+  return async (
+    conversationId: string,
+    recovery?: PiSessionRecovery,
+  ): Promise<SessionAgent> => {
     validateConversationId(conversationId);
     const modelRuntime = await ModelRuntime.create({
       authPath: join(config.agentDirectory, "auth.json"),
@@ -48,18 +56,20 @@ export function createPiSessionFactory(config: PiSessionFactoryConfig) {
     });
     await resourceLoader.reload();
     const sessionFile = join(config.sessionDirectory, `${conversationId}.jsonl`);
+    const sessionManager = openSessionManagerForRecovery(
+      sessionFile,
+      config.sessionDirectory,
+      config.workspaceDirectory,
+      recovery,
+    );
     const { session, extensionsResult, modelFallbackMessage } = await createAgentSession({
       cwd: config.workspaceDirectory,
       agentDir: config.agentDirectory,
       modelRuntime,
       model,
       noTools: "all",
       resourceLoader,
-      sessionManager: SessionManager.open(
-        sessionFile,
-        config.sessionDirectory,
-        config.workspaceDirectory,
-      ),
+      sessionManager,
       settingsManager,
     });
     if (extensionsResult.errors.length > 0 || modelFallbackMessage) {
@@ -71,6 +81,39 @@ export function createPiSessionFactory(config: PiSessionFactoryConfig) {
   };
 }
 
+export function openSessionManagerForRecovery(
+  sessionFile: string,
+  sessionDirectory: string,
+  workspaceDirectory: string,
+  recovery?: PiSessionRecovery,
+): SessionManager {
+  const sessionExists = existsSync(sessionFile);
+  if (recovery !== undefined && recovery.committedLeafId !== null && !sessionExists) {
+    throw new Error("Committed visible history has no Pi session file");
+  }
+
+  const sessionManager = SessionManager.open(
+    sessionFile,
+    sessionDirectory,
+    workspaceDirectory,
+  );
+  if (recovery === undefined) return sessionManager;
+
+  if (recovery.committedLeafId === null) {
+    sessionManager.resetLeaf();
+    return sessionManager;
+  }
+  try {
+    sessionManager.branch(recovery.committedLeafId);
+  } catch (error) {
+    throw new Error(
+      `Committed Pi session leaf is unavailable: ${recovery.committedLeafId}`,
+      { cause: error },
+    );
+  }
+  return sessionManager;
+}
+
 function validateConfig(config: PiSessionFactoryConfig): void {
   requireText("model.provider", config.model.provider);
   requireText("model.id", config.model.id);