Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Select an option

  • Save colelawrence/b9b5ebc48abef6ceba8cf5fb91117db5 to your computer and use it in GitHub Desktop.

Select an option

Save colelawrence/b9b5ebc48abef6ceba8cf5fb91117db5 to your computer and use it in GitHub Desktop.
Durable sessions and reconstruction in Pi

Durable Sessions and Reconstruction in Pi

This note explains how to think about durable state in Pi when the user navigates the session tree with /tree, starts a new session, resumes another session, or forks.

The short version:

  • Pi sessions are an append-only tree of entries with a movable leaf pointer.
  • /tree does not restart the extension runtime. It changes the active leaf and rebuilds the active conversation context.
  • If your extension keeps in-memory state, you are responsible for reconstructing it from durable session entries.
  • The central design decision is whether your state should be:
    • branch-local / causally truthful (getBranch()), or
    • session-global / monotonic (getEntries()).

Reference map

The key identifiers in the pi-mono-effect reference are:

  • AgentSession.navigateTree() in packages/coding-agent/src/core/agent-session.ts

    • drives /tree navigation
    • emits session_before_tree
    • updates the leaf via SessionManager.branch(), SessionManager.resetLeaf(), or SessionManager.branchWithSummary()
    • rebuilds context via SessionManager.buildSessionContext()
    • replaces this.agent.state.messages
    • emits session_tree
  • SessionManager.getBranch(), SessionManager.getEntries(), SessionManager.buildSessionContext(), SessionManager.branch(), SessionManager.resetLeaf(), and SessionManager.branchWithSummary() in packages/coding-agent/src/core/session-manager.ts

  • TreePreparation, SessionBeforeTreeEvent, SessionTreeEvent, SessionStartEvent, SessionBeforeSwitchEvent, and SessionBeforeForkEvent in packages/coding-agent/src/core/extensions/types.ts

  • collectEntriesForBranchSummary() and generateBranchSummary() in packages/coding-agent/src/core/compaction/branch-summarization.ts

  • createBranchSummaryMessage() and convertToLlm() in packages/coding-agent/src/core/messages.ts

  • AgentSessionRuntime.switchSession(), AgentSessionRuntime.newSession(), and AgentSessionRuntime.fork() in packages/coding-agent/src/core/agent-session-runtime.ts

  • Example extensions:

    • packages/coding-agent/binaries/linux-arm64/examples/extensions/todo.ts
    • packages/coding-agent/binaries/linux-arm64/examples/extensions/tools.ts

Important note about reference drift

The binary docs/examples under packages/coding-agent/binaries/linux-arm64/ still show post-transition events like session_switch and session_fork.

The current source in packages/coding-agent/src/core/extensions/types.ts models the post side of new/resume/fork as session_start with:

  • reason: "startup" | "reload" | "new" | "resume" | "fork"

and keeps the pre side as:

  • session_before_switch
  • session_before_fork

For new code, prefer the current source identifiers when source and binary docs disagree.


Mental model: the session file is truth, runtime state is a cache

A Pi session is stored as JSONL. Each entry has an id and parentId, so entries form a tree. The current position is the session's leaf.

Your extension runtime is different:

  • module-level variables
  • in-memory maps/sets/lists
  • active UI widgets
  • cached derived state

Those runtime values are not automatically rewound when the user uses /tree.

They stay alive until:

  • the process exits, or
  • the runtime is explicitly torn down and recreated (for example by switching sessions, forking, reloading, etc.)

So the safe model is:

Treat in-memory state as a cache derived from durable session entries.

If the leaf changes, or the whole runtime is recreated, rebuild your cache from the session.


What actually happens on /tree

The relevant flow is in AgentSession.navigateTree().

At a high level:

  1. Pi reads the current leaf (oldLeafId).
  2. Pi resolves the target entry.
  3. Pi computes the common ancestor and the abandoned segment with collectEntriesForBranchSummary().
  4. Pi emits session_before_tree.
  5. If requested, Pi generates a branch summary with generateBranchSummary().
  6. Pi changes the leaf with one of:
    • SessionManager.branch()
    • SessionManager.resetLeaf()
    • SessionManager.branchWithSummary()
  7. Pi rebuilds active context with SessionManager.buildSessionContext().
  8. Pi replaces live agent messages with the rebuilt context.
  9. Pi emits session_tree.

What does not happen on /tree:

  • no session_shutdown
  • no session_start
  • no runtime recreation

So if your extension only reconstructs at session_start, it will be wrong after /tree.


Where reconstruction belongs

Reconstruction belongs in the extension.

Pi core updates the session leaf and the agent's message state, but Pi does not know how to rebuild your custom domain state.

The usual reconstruction points are:

  • session_start
    • initial startup
    • reload
    • new session
    • resume another session
    • fork
  • session_tree
    • active leaf changed inside the same session file

If your extension maintains state that should survive across those transitions, reconstruct at least there.

A good default is:

pi.on("session_start", async (_event, ctx) => {
  reconstruct(ctx)
})

pi.on("session_tree", async (_event, ctx) => {
  reconstruct(ctx)
})

You may also update in-memory state directly in your tool/command handlers so the UI feels immediate. Reconstruction is still the correctness boundary.


Which session view should reconstruction use?

This is the most important design choice.

1. ctx.sessionManager.getBranch()

Use this when the state should mean:

"What is true on the current branch at the current leaf?"

This is the branch-local / causally truthful model.

Examples:

  • TODOs that should match the active conversation branch
  • current tool configuration for the active branch
  • branch-local plan state
  • any state where rewinding to an earlier point should rewind the state too

SessionManager.getBranch(fromId?) walks from the given entry to the root. With no argument, it walks from the current leaf.

2. ctx.sessionManager.getEntries()

Use this when the state should mean:

"What durable facts have happened anywhere in this session file?"

This is the session-global / monotonic model.

Examples:

  • facts that should remain true even if you navigate earlier with /tree
  • a ledger of completed migrations or verified steps
  • a session-global bookmark or audit trail

This is more powerful, but it changes semantics: the user can navigate to an earlier branch and still see facts learned or completed on a different later branch.

That can be correct, but only if you want it.


Rule of thumb: choose the truth boundary first

Before writing code, decide which of these you want:

Desired meaning Reconstruction source
"State should look exactly like this branch at this point in history." getBranch()
"Once a fact is learned/completed in this session file, keep it." getEntries()
"This should survive even across forks / other session files." external storage outside the session file

If you skip this choice, you usually end up with accidental behavior.


How to see the before and after sides of a /tree switch

If you need transition-aware logic, you often want:

  • the old branch state
  • the new branch state
  • the abandoned segment only

The relevant data comes from two places.

Pre-switch: session_before_tree

SessionBeforeTreeEvent exposes preparation: TreePreparation.

Important fields:

  • targetId
  • oldLeafId
  • commonAncestorId
  • entriesToSummarize
  • userWantsSummary
  • customInstructions
  • replaceInstructions
  • label

Use entriesToSummarize when you want exactly the abandoned segment, not the whole old branch.

Post-switch: session_tree

SessionTreeEvent exposes:

  • oldLeafId
  • newLeafId
  • summaryEntry?
  • fromExtension?

At session_tree time:

  • ctx.sessionManager.getBranch() is already the after branch
  • if you also want the before branch, derive it explicitly from event.oldLeafId

A small helper:

function branchAt(ctx: ExtensionContext, leafId: string | null) {
  return leafId ? ctx.sessionManager.getBranch(leafId) : []
}

pi.on("session_tree", async (event, ctx) => {
  const before = branchAt(ctx, event.oldLeafId)
  const after = branchAt(ctx, event.newLeafId)

  // Compare, diff, or reconstruct both sides here.
})

This works because SessionManager.getBranch(fromId?) can walk from any entry id, not just the current leaf.


What branch summaries change, and what they do not

If the user asks Pi to summarize the branch they are leaving:

  • Pi may append a branch_summary entry with SessionManager.branchWithSummary().
  • Later, createBranchSummaryMessage() and convertToLlm() turn that into a synthetic user message for the model.

That affects LLM context on later turns.

It does not reconstruct your extension state for you.

Your extension still has to choose whether to replay:

  • only the current branch (getBranch()), or
  • the whole session file (getEntries())

What should be persisted?

Two common patterns are supported directly by the reference material.

Pattern A: tool result details

The docs in docs/extensions.md explicitly recommend storing branch-sensitive state in tool result details.

Why it works:

  • tool results live on branches
  • replaying branch entries in order recreates the right state for that branch

This is the pattern used by the binary example examples/extensions/todo.ts.

Pattern B: custom entries via pi.appendEntry(customType, data?)

This is good for extension-owned durable state that should not be shown to the LLM.

Why it works:

  • custom entries are session entries like everything else
  • you can replay them from getBranch() or getEntries() depending on your semantics

This is the pattern used by the binary example examples/extensions/tools.ts.


Example 1: branch-local TODOs

This is the ordinary, causally truthful design.

Semantics:

  • the TODO list should match the active branch
  • jumping to an earlier point with /tree should restore the TODO list for that point
  • fork/new/resume should reconstruct from the new runtime's session_start

This is the same core idea as examples/extensions/todo.ts, but adapted to the current source model where session_start.reason covers new, resume, and fork.

import { StringEnum } from "@mariozechner/pi-ai"
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent"
import { Type } from "@sinclair/typebox"

interface Todo {
  id: number
  text: string
  done: boolean
}

interface TodoSnapshot {
  todos: Todo[]
  nextId: number
}

const TodoParams = Type.Object({
  action: StringEnum(["list", "add", "toggle", "clear"] as const),
  text: Type.Optional(Type.String()),
  id: Type.Optional(Type.Number()),
})

function cloneTodos(todos: Todo[]): Todo[] {
  return todos.map((todo) => ({ ...todo }))
}

export default function todoExtension(pi: ExtensionAPI) {
  let todos: Todo[] = []
  let nextId = 1

  function reconstructFromCurrentBranch(ctx: ExtensionContext) {
    todos = []
    nextId = 1

    for (const entry of ctx.sessionManager.getBranch()) {
      if (entry.type !== "message") continue
      const message = entry.message
      if (message.role !== "toolResult" || message.toolName !== "todo") continue

      const snapshot = message.details as TodoSnapshot | undefined
      if (!snapshot) continue

      todos = cloneTodos(snapshot.todos)
      nextId = snapshot.nextId
    }
  }

  pi.on("session_start", async (_event, ctx) => {
    reconstructFromCurrentBranch(ctx)
  })

  pi.on("session_tree", async (_event, ctx) => {
    reconstructFromCurrentBranch(ctx)
  })

  pi.registerTool({
    name: "todo",
    description: "Manage branch-local TODOs",
    parameters: TodoParams,
    async execute(_toolCallId, params) {
      switch (params.action) {
        case "list": {
          return {
            content: [{
              type: "text",
              text: todos.length === 0
                ? "No todos"
                : todos.map((t) => `[${t.done ? "x" : " "}] #${t.id}: ${t.text}`).join("\n"),
            }],
            details: {
              todos: cloneTodos(todos),
              nextId,
            } satisfies TodoSnapshot,
          }
        }
        case "add": {
          if (!params.text) {
            return {
              content: [{ type: "text", text: "Missing text" }],
              details: { todos: cloneTodos(todos), nextId } satisfies TodoSnapshot,
            }
          }

          todos.push({ id: nextId++, text: params.text, done: false })
          return {
            content: [{ type: "text", text: "Added TODO" }],
            details: { todos: cloneTodos(todos), nextId } satisfies TodoSnapshot,
          }
        }
        case "toggle": {
          const todo = todos.find((t) => t.id === params.id)
          if (!todo) {
            return {
              content: [{ type: "text", text: "TODO not found" }],
              details: { todos: cloneTodos(todos), nextId } satisfies TodoSnapshot,
            }
          }

          todo.done = !todo.done
          return {
            content: [{ type: "text", text: `TODO #${todo.id} is now ${todo.done ? "done" : "not done"}` }],
            details: { todos: cloneTodos(todos), nextId } satisfies TodoSnapshot,
          }
        }
        case "clear": {
          todos = []
          nextId = 1
          return {
            content: [{ type: "text", text: "Cleared TODOs" }],
            details: { todos: [], nextId } satisfies TodoSnapshot,
          }
        }
      }
    },
  })
}

Why this works

  • Every tool result stores a full snapshot in details.
  • Reconstruction replays ctx.sessionManager.getBranch() from root to current leaf.
  • When /tree moves the leaf, replay yields the branch-correct list.

What this means semantically

If a TODO was checked off only on a later abandoned branch, then jumping to an earlier point can make it appear unchecked again.

That is not a bug. It is the intended meaning of branch-local truth.


Example 2: TODOs that stay checked off when you jump earlier

This is a different design.

Semantics:

  • once a TODO is completed anywhere in this session file, it should remain completed even if the user navigates to an earlier point with /tree
  • completion is monotonic: done means done
  • this is intentionally not branch-local truth

Use this model only if the TODOs represent real-world facts that should not be rewound by conversation navigation.

A good approach is to store a ledger of custom entries and reconstruct from getEntries().

import { StringEnum } from "@mariozechner/pi-ai"
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent"
import { Type } from "@sinclair/typebox"

const LEDGER_TYPE = "durable-todo-ledger"

interface Todo {
  id: number
  text: string
  done: boolean
}

type TodoLedgerEvent =
  | { type: "todo_added"; id: number; text: string }
  | { type: "todo_completed"; id: number }

const TodoParams = Type.Object({
  action: StringEnum(["list", "add", "complete"] as const),
  text: Type.Optional(Type.String()),
  id: Type.Optional(Type.Number()),
})

export default function durableTodoExtension(pi: ExtensionAPI) {
  let todos: Todo[] = []
  let nextId = 1

  function reconstructFromWholeSession(ctx: ExtensionContext) {
    const byId = new Map<number, Todo>()

    for (const entry of ctx.sessionManager.getEntries()) {
      if (entry.type !== "custom" || entry.customType !== LEDGER_TYPE) continue
      const event = entry.data as TodoLedgerEvent | undefined
      if (!event) continue

      if (event.type === "todo_added") {
        byId.set(event.id, { id: event.id, text: event.text, done: false })
        nextId = Math.max(nextId, event.id + 1)
      }

      if (event.type === "todo_completed") {
        const todo = byId.get(event.id)
        if (todo) {
          todo.done = true
        }
      }
    }

    todos = Array.from(byId.values()).sort((a, b) => a.id - b.id)
  }

  pi.on("session_start", async (_event, ctx) => {
    nextId = 1
    reconstructFromWholeSession(ctx)
  })

  pi.on("session_tree", async (_event, ctx) => {
    nextId = 1
    reconstructFromWholeSession(ctx)
  })

  pi.registerTool({
    name: "durable_todo",
    description: "Manage session-global monotonic TODOs",
    parameters: TodoParams,
    async execute(_toolCallId, params) {
      switch (params.action) {
        case "list": {
          return {
            content: [{
              type: "text",
              text: todos.length === 0
                ? "No todos"
                : todos.map((t) => `[${t.done ? "x" : " "}] #${t.id}: ${t.text}`).join("\n"),
            }],
          }
        }

        case "add": {
          if (!params.text) {
            return { content: [{ type: "text", text: "Missing text" }] }
          }

          const id = nextId++
          pi.appendEntry<TodoLedgerEvent>(LEDGER_TYPE, {
            type: "todo_added",
            id,
            text: params.text,
          })

          todos.push({ id, text: params.text, done: false })
          return { content: [{ type: "text", text: `Added durable TODO #${id}` }] }
        }

        case "complete": {
          if (params.id === undefined) {
            return { content: [{ type: "text", text: "Missing id" }] }
          }

          const todo = todos.find((t) => t.id === params.id)
          if (!todo) {
            return { content: [{ type: "text", text: "TODO not found" }] }
          }

          if (!todo.done) {
            pi.appendEntry<TodoLedgerEvent>(LEDGER_TYPE, {
              type: "todo_completed",
              id: todo.id,
            })
            todo.done = true
          }

          return { content: [{ type: "text", text: `Completed TODO #${todo.id}` }] }
        }
      }
    },
  })
}

Why this works

  • pi.appendEntry() creates durable custom entries.
  • Reconstruction scans ctx.sessionManager.getEntries(), not just the active branch.
  • Completion is monotonic, so it survives /tree jumps to earlier points in the same session file.

What this means semantically

If you complete a TODO on branch B and then /tree back to an earlier point on branch A, the TODO still appears completed.

That is intentional.

When this is the right model

Use this if the TODO is really representing something like:

  • a migration already applied in the workspace
  • a review already completed
  • a fact discovered during the session that should remain true

Do not use this if the TODO is supposed to track branch-local conversation progress.


Very important limitation of example 2

getEntries() sees the whole current session file.

It does not magically survive crossing into a different session file.

In the current source, AgentSessionRuntime.fork() eventually uses SessionManager.createBranchedSession() to extract a path into a new session file. That means facts that exist only on other branches are not copied into the forked file.

So the monotonic ledger design above guarantees:

  • persistence across /tree rewinds within the same session file

It does not guarantee:

  • persistence across /fork into a new session file that does not contain the relevant ledger entries

If you need that, your truth source should live outside the session JSONL file.


A practical design checklist

Before adding durable state to an extension, answer these questions:

  1. What is the truth boundary?

    • current branch?
    • entire session file?
    • external workspace or database?
  2. Is the state rewindable?

    • if yes, prefer getBranch()
    • if no, prefer a monotonic ledger replayed from getEntries()
  3. What is the replay unit?

    • full snapshots in toolResult.details
    • event ledger in custom entries
  4. Where does reconstruction run?

    • at least session_start and session_tree
  5. Do you need before/after inspection on /tree?

    • use session_before_tree.preparation
    • or use session_tree.oldLeafId plus getBranch(oldLeafId)
  6. Should forks inherit this truth?

    • if yes, session-local storage may not be enough

Recommended default

If you are unsure, start here:

  • store branch-local snapshots in toolResult.details
  • reconstruct from ctx.sessionManager.getBranch()
  • rebuild on session_start and session_tree

That matches Pi's tree semantics most naturally.

Only move to getEntries() when you have explicitly decided that your state should remain true even when the user navigates to an earlier point.


One final principle

The safest reconstruction code is a pure replay:

function reconstruct(entries: SessionEntry[]): DerivedState {
  // no hidden dependencies
  // no reliance on old in-memory values
  // same input => same output
}

Then your event handlers are thin wrappers:

pi.on("session_start", async (_event, ctx) => {
  state = reconstruct(ctx.sessionManager.getBranch())
})

pi.on("session_tree", async (_event, ctx) => {
  state = reconstruct(ctx.sessionManager.getBranch())
})

That model stays correct when Pi:

  • rewinds with /tree
  • resumes a different session
  • forks
  • reloads the runtime

and it matches the design implied by SessionManager, AgentSession.navigateTree(), and the reference extension patterns.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment