@davideasden/pi-undo 0.1.2 → 0.2.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (11) hide show
  1. package/README.md +147 -85
  2. package/extensions/pi-undo.ts +172 -16
  3. package/package.json +1 -1
  4. package/src/atomic-fs.ts +7 -2
  5. package/src/controller.ts +97 -38
  6. package/src/mutation-journal.ts +180 -40
  7. package/src/pi-runtime.ts +2 -2
  8. package/src/quarantine.ts +357 -17
  9. package/src/restore-engine.ts +269 -49
  10. package/src/snapshot-store.ts +316 -63
  11. package/src/status-reporter.ts +15 -2
package/README.md CHANGED
@@ -1,45 +1,45 @@
1
1
  # pi-undo
2
2
 
3
- Persistent undo and redo for [Pi](https://github.com/badlogic/pi-mono) that restores both the agent session and workspace files without requiring users to manage Git.
3
+ **Persistent workspace undo and redo for [Pi](https://github.com/badlogic/pi-mono).**
4
4
 
5
- Each history entry represents one completed agent run. When you undo or redo, `pi-undo` moves Pi's session tree to the matching point and restores the corresponding workspace state.
5
+ Each completed agent run creates a checkpoint that captures both the Pi session boundary and a content-addressed workspace snapshot. `/undo` moves Pi's session tree to a prior run and restores the matching workspace. `/redo` reverses the undo with a safety snapshot taken before the operation. The system is crash-safe: an interrupted restore is rolled forward or back on next startup using durable journals, write-ahead logging, and quarantine artifacts.
6
6
 
7
7
  ## Features
8
8
 
9
- - Persistent `/undo` and `/redo` across Pi restarts.
10
- - Restores the Pi session and workspace files together.
11
- - Preserves redo state with a safety snapshot taken immediately before undo.
12
- - Integrates with Pi's native `/tree` navigation.
13
- - Supports ordinary directories, Git repositories, nested repositories, and initialized submodules.
14
- - Uses durable journals and restart recovery for interrupted operations.
15
- - Detects conflicting external file changes and fails closed instead of overwriting unknown content.
16
- - Does not require `git commit`, `git stash`, `git reset`, hidden branches, or any other user-managed Git workflow.
9
+ - **Persistent undo / redo** across Pi restarts. History is stored per-session and survives crashes.
10
+ - **Diff viewer** — `/diff` and `/diff N` show before-and-after file changes with colored, scrollable output in TUI mode and a summary in non-TUI modes.
11
+ - **Session + workspace** — Pi's session tree and workspace files are restored together as a unit.
12
+ - **Safety snapshots** — `undo` captures the current workspace as a redo safety snapshot; `tree navigation` records a rescue snapshot before switching branches.
13
+ - **Deferred prompts** — text, images, `@file` references, and paste content entered during an undo/redo are queued and processed in FIFO order once the operation completes. Prompts are never lost.
14
+ - **Crash recovery** — a write-ahead log (WAL), mutation journal, and quarantine artifacts allow in-progress operations to be rolled forward or back on restart.
15
+ - **External concurrency detection** — file fingerprint and inode checks detect external modification. Conflicting changes are never silently overwritten; the system fails closed or enters `recovery required`.
16
+ - **No Git workflow** — snapshots use a private object database. No `git commit`, `git stash`, `git reset`, branches, or forges are required.
17
+ - **Nested repositories** and **initialized submodules** are handled as independent roots. Their `.git` metadata is never modified.
18
+ - **Performance** — batch WAL operations (up to 128 files per batch), parallel target artifact I/O (32 concurrent), concurrent request preparation (32 concurrent scoped blob reads plus fingerprint computation), scoped safety snapshots, and batch multi-file deletes keep restore fast at scale.
17
19
 
18
20
  ## Requirements
19
21
 
20
22
  - Pi `0.80.10` or a compatible release.
21
23
  - Node.js `22.19.0` or later.
22
- - Git available on `PATH`.
23
-
24
- Git is used internally to create private, content-addressed snapshots. You do not need to initialize a repository or run Git commands yourself.
24
+ - Git available on `PATH` (used internally for content-addressed snapshots).
25
25
 
26
26
  ## Installation
27
27
 
28
- Install the published Pi package from npm:
28
+ Install the published package from npm:
29
29
 
30
30
  ```bash
31
31
  pi install npm:@davideasden/pi-undo
32
32
  ```
33
33
 
34
- Restart Pi after installation if the extension is not loaded in the current session.
34
+ Restart Pi if the extension is not already loaded.
35
35
 
36
- To install a local checkout instead:
36
+ To install from a local checkout:
37
37
 
38
38
  ```bash
39
- pi install ./path/to/pi-undo
39
+ pi install /path/to/pi-undo
40
40
  ```
41
41
 
42
- During development, you can load the extension directly:
42
+ During development, load the extension directly:
43
43
 
44
44
  ```bash
45
45
  pi -e /absolute/path/to/pi-undo/extensions/pi-undo.ts
@@ -47,7 +47,7 @@ pi -e /absolute/path/to/pi-undo/extensions/pi-undo.ts
47
47
 
48
48
  ## Usage
49
49
 
50
- Work with Pi normally. `pi-undo` records a boundary after each completed agent run.
50
+ Work with Pi normally. `pi-undo` automatically records a boundary after each completed agent run.
51
51
 
52
52
  ### Diff
53
53
 
@@ -55,15 +55,15 @@ Work with Pi normally. `pi-undo` records a boundary after each completed agent r
55
55
  /diff
56
56
  ```
57
57
 
58
- `/diff` reviews the files changed by the most recent completed agent run. In TUI mode, select a file to open a colored, scrollable before-and-after comparison. Binary files are listed but do not receive a line-by-line diff.
58
+ Shows files changed by the most recent completed run. In TUI mode, select a file to open a colored, scrollable before-and-after comparison. Binary files are listed without line-by-line diff.
59
59
 
60
- Use a one-based history position to inspect an earlier applied run, where `1` is the most recent:
60
+ Use a one-based position to inspect an earlier run (`1` is the most recent):
61
61
 
62
62
  ```text
63
- /diff 2
63
+ /diff 3
64
64
  ```
65
65
 
66
- Print, JSON, and RPC modes report a one-line file and line-count summary instead of opening the interactive viewer.
66
+ Print, JSON, and RPC modes report a one-line file and line-count summary.
67
67
 
68
68
  ### Undo
69
69
 
@@ -71,9 +71,9 @@ Print, JSON, and RPC modes report a one-line file and line-count summary instead
71
71
  /undo
72
72
  ```
73
73
 
74
- `/undo` returns to the previous completed agent run on the current branch. Before restoring that checkpoint, it captures the current workspace as a redo safety snapshot.
74
+ Returns to the previous completed run on the current branch. Before restoring, it captures the current workspace as a redo safety snapshot.
75
75
 
76
- After a successful undo, `pi-undo` tries to put the original prompt back into an empty TUI editor. It never overwrites text already present in the editor. RPC mode only reports the refill request, while print and JSON modes do not promise editor refill behavior.
76
+ After a successful undo, text entered during the operation is replayed into the editor. RPC mode reports the refill request; print and JSON modes do not replay prompts.
77
77
 
78
78
  ### Redo
79
79
 
@@ -81,7 +81,7 @@ After a successful undo, `pi-undo` tries to put the original prompt back into an
81
81
  /redo
82
82
  ```
83
83
 
84
- `/redo` restores the session and workspace captured immediately before the corresponding undo. It does not simply reapply the checkpoint's original after-snapshot, so edits preserved by the redo safety snapshot are restored correctly.
84
+ Restores the session and workspace captured immediately before the corresponding undo. The redo safety snapshot preserves edits made between the undo and the redo.
85
85
 
86
86
  ### Tree Navigation
87
87
 
@@ -91,16 +91,45 @@ Continue to use Pi's native command:
91
91
  /tree
92
92
  ```
93
93
 
94
- Before Pi moves to another session-tree boundary, `pi-undo` validates the target and records a rescue snapshot. After navigation, it restores the workspace associated with the selected boundary.
94
+ Before Pi moves to another session-tree boundary, `pi-undo` validates the target and records a rescue snapshot. After navigation it restores the matching workspace.
95
95
 
96
- If the agent is streaming, `/undo` and `/redo` request an abort and wait for the agent to become idle. If the wait times out or tools are still running, no files are restored. Native `/tree` navigation is cancelled while streaming and can be retried when Pi is idle.
96
+ While the agent is streaming, `/undo` and `/redo` request an abort and wait for idle. If the wait times out, no files are restored. Native `/tree` is cancelled while streaming and can be retried when idle.
97
97
 
98
- The footer shows the available history, for example:
98
+ The footer shows available history, for example:
99
99
 
100
100
  ```text
101
101
  undo:2 redo:1
102
102
  ```
103
103
 
104
+ When crash recovery is pending, the footer shows:
105
+
106
+ ```text
107
+ recovery_required
108
+ ```
109
+
110
+ ### Performance Notes
111
+
112
+ Real-world measurements from a 104-file undo operation before optimizations:
113
+
114
+ ```text
115
+ ok files:104 total:8797ms apply:8591ms capture:89ms journal:64ms commit:27ms plan:11ms
116
+ ```
117
+
118
+ After batch WAL, scoped safety snapshots, parallel restore I/O, and batch file deletes:
119
+
120
+ ```text
121
+ ok files:104 total:~1050ms apply:~850ms capture:~90ms journal:~65ms
122
+ ```
123
+
124
+ Performance gains come from:
125
+
126
+ - **Batch WAL durability** — write-once, fsync-once per batch of up to 128 entries instead of per-file.
127
+ - **Scoped safety snapshots** — only paths recorded in a checkpoint's `changedPaths` are snapshotted for undo/redo, not the entire workspace.
128
+ - **Parallel target I/O** — up to 32 target artifact files are written and synced concurrently.
129
+ - **Concurrent request preparation** — blob reads, live-state checks, and fingerprint computation run at 32-wide concurrency.
130
+ - **Batch file deletes** — delete operations share the same WAL batch and directory fsync merging.
131
+ - **Directory fsync merging** — barriers are applied once per unique directory per phase, not once per file.
132
+
104
133
  ## How It Works
105
134
 
106
135
  Private state is stored alongside the Pi session:
@@ -109,48 +138,78 @@ Private state is stored alongside the Pi session:
109
138
  <sessionDir>/.pi-undo/
110
139
  ```
111
140
 
112
- For every completed agent run, `pi-undo` captures the session boundary and a workspace manifest. Snapshots use a private Git object database and temporary index. They do not use the repository's normal history.
141
+ For every completed agent run, `pi-undo` captures:
142
+
143
+ 1. A **session boundary** — the logical point in Pi's session tree.
144
+ 2. A **workspace manifest** — a content-addressed snapshot of the workspace root(s).
145
+
146
+ Snapshots use a private Git object database and a temporary Git index. They never touch the repository's normal history, `HEAD`, reflog, or stash.
113
147
 
114
- Undo, redo, and tree restoration follow the same high-level flow:
148
+ The restore flow is:
115
149
 
116
- 1. Validate the current session, workspace topology, and target manifest.
117
- 2. Write a durable write-ahead log (WAL) before changing files.
150
+ 1. Validate current session, workspace topology, and target manifest.
151
+ 2. Write a durable WAL (`INTENT`) before changing any files.
118
152
  3. Capture a safety or rescue snapshot when required.
119
- 4. Move the Pi session to the target logical boundary.
120
- 5. Restore workspace paths with fingerprint checks and no-clobber installation.
121
- 6. Verify the resulting session and workspace before committing the journal.
153
+ 4. Move the Pi session cursor to the target boundary.
154
+ 5. Restore workspace paths with fingerprint and inode checks and no-clobber installation.
155
+ 6. Verify the resulting session and workspace before committing the journal (`TARGET_VERIFIED` → `CLEANED`).
122
156
 
123
- Ordinary files and symbolic links are restored through same-directory, same-filesystem quarantine artifacts. For regular files, source capture uses a hard link and checks the path fingerprint and inode identity again immediately before removing the original path. This narrows the external-concurrency window, but cannot make the final check and unlink operation atomic in pure Node.js.
157
+ ### Mutation Journal (WAL)
124
158
 
125
- ## Safety and Recovery
159
+ Each file change is tracked through a six-state hash chain:
160
+
161
+ ```
162
+ INTENT → SOURCE_QUARANTINED → SOURCE_VERIFIED → TARGET_INSTALLED → TARGET_VERIFIED → CLEANED
163
+ ```
164
+
165
+ - **INTENT** — durable intent recorded before any mutation.
166
+ - **SOURCE_QUARANTINED** — original file hard-linked to a random artifact name.
167
+ - **SOURCE_VERIFIED** — original removed, artifact content verified.
168
+ - **TARGET_INSTALLED** — new file hard-linked into place (no-clobber).
169
+ - **TARGET_VERIFIED** — workspace confirmed matching target state.
170
+ - **CLEANED** — source and target artifacts removed.
126
171
 
127
- `pi-undo` does not modify the user's Git `HEAD`, index, refs, reflogs, stash, configuration, or other repository metadata. Nested repositories and initialized submodules are captured as independent roots, but only their working files are restored. The extension does not switch their real `HEAD`, run `git submodule update`, or recreate deleted `.git` metadata.
172
+ Batch operations (`beginMany`, `advanceBatch`) maintain the same per-file six-state contract while reducing physical fsync calls by grouping entries.
128
173
 
129
- If Pi or the process stops during a restore, the next session startup reads the transaction journal:
174
+ Every record is checksum-linked to its predecessor. The journal is append-only and never mutated in place.
130
175
 
131
- - Without a trusted cursor marker, it restores the rollback manifest and marks the transaction aborted.
132
- - With a trusted cursor marker, it restores the target manifest, completes cursor durability, and commits the transaction.
133
- - If the journal, session identity, logical leaf, manifest, cursor, or workspace contents conflict, it enters `recovery required` and blocks further history mutation.
176
+ ### Quarantine and External Concurrency
134
177
 
135
- When the footer reports `recovery_required`, first back up the workspace and Pi session JSONL. Transaction diagnostics are stored in:
178
+ Ordinary files and symlinks are restored through same-directory, same-filesystem artifacts:
179
+
180
+ - **Source capture**: `link(original, sourceArtifact)` — fails with `EEXIST` if artifact already exists.
181
+ - **Target installation**: `link(targetArtifact, original)` — fails with `EEXIST` if original was externally recreated.
182
+ - **Fingerprint + inode**: before each state transition, both content fingerprint and `(dev, ino)` identity are verified. A matching fingerprint with a different inode is treated as external replacement and triggers a safe fail-closed state.
183
+ - **Artifact cleanup**: only artifacts registered in the journal with matching paths, names, fingerprints, and ownership are removed. Unclaimed files are never deleted.
184
+ - **Rollback**: from any state, `restoreMutation` converges to the pre-operation state. If a previous target was externally replaced with identical content but a different inode, the external file is preserved and the journal remains active.
185
+
186
+ ## Safety and Recovery
187
+
188
+ - Pi `HEAD`, index, refs, reflog, stash, config, and other repository metadata are never modified.
189
+ - Nested repositories and submodules are restored as independent roots. Their `.git` directories are not created or deleted.
190
+ - If the process stops during a restore, startup recovery reads the transaction journal:
191
+ - **Without a trusted cursor marker**: restores the rollback manifest and marks the transaction aborted.
192
+ - **With a trusted cursor marker**: completes cursor durability and commits the transaction.
193
+ - **On conflict** (identity, leaf, manifest, cursor, or workspace mismatch): enters `recovery required` and blocks further history mutation.
194
+
195
+ When `recovery_required` appears, first back up the workspace and Pi session JSONL. Transaction diagnostics are stored in:
136
196
 
137
197
  ```text
138
198
  <sessionDir>/.pi-undo/transactions/
139
199
  ```
140
200
 
141
- The relevant transaction may contain `descriptor.json`, `restore-plan.json`, `state.json`, and mutation journal data. Do not delete the `.pi-undo` directory without a backup: unresolved quarantine artifacts may contain the only preserved copy of a file version.
201
+ A transaction directory may contain `descriptor.json`, `restore-plan.json`, `state.json`, and the mutation journal. Do not delete `.pi-undo` without a backup: unresolved quarantine artifacts may be the only surviving copy of a file version.
142
202
 
143
203
  ## Limitations
144
204
 
145
205
  - Git-ignored files are not included in snapshots and are not created or deleted during restore.
146
206
  - Empty directories are not represented in snapshots.
147
- - Real `.git` metadata is never restored.
148
- - The workspace lock coordinates `pi-undo` instances, but cannot prevent editors, watchers, or other processes from writing files concurrently.
149
- - Quarantine and repeated fingerprint checks reduce, but cannot eliminate, the final check-to-unlink race with arbitrary external processes.
150
- - When ownership or contents cannot be proven, `pi-undo` preserves the available versions and enters `recovery required` instead of guessing, deleting, or overwriting.
207
+ - Real `.git` metadata is never created, modified, or deleted.
208
+ - The workspace lock prevents concurrent `pi-undo` instances, but cannot prevent editors, watchers, or other processes from writing files concurrently.
209
+ - Fingerprint and inode checks reduce the external-concurrency race window, but cannot make the final check-and-unlink atomic in pure Node.js.
151
210
  - `--no-session` mode has no durable session cursor. Undo and redo can work in-process, but crash persistence is not guaranteed.
152
211
  - Uninitialized gitlinks are not initialized automatically. A broken nested repository causes capture to fail instead of being silently skipped.
153
- - Workspace files and Pi JSONL are not an operating-system-level ACID transaction. Snapshots, WAL records, cursor markers, verification, and idempotent recovery are used to converge to the old or new state.
212
+ - Workspace files and Pi session JSONL are not an operating-system-level ACID transaction. Snapshots, WAL, cursor markers, verification, and idempotent recovery are used to converge to the old or new state.
154
213
 
155
214
  ## Development
156
215
 
@@ -162,53 +221,56 @@ cd pi-undo
162
221
  npm install
163
222
  ```
164
223
 
165
- The development dependency for Pi points to:
224
+ ### Project Layout
166
225
 
167
226
  ```text
168
- resources/pi-0.80.10/packages/coding-agent
169
- ```
170
-
171
- For type checking and the real `AgentSession` integration tests, place the Pi `0.80.10` source tree at that path, install its workspace dependencies, and build the required workspace packages. The published package does not include `resources/`.
172
-
173
- Useful project paths:
174
-
175
- ```text
176
- extensions/pi-undo.ts Pi extension entry point
177
- src/ Snapshot, journal, restore, and runtime implementation
178
- test/ Unit, integration, recovery, and fault-injection tests
179
- ```
180
-
181
- ## Testing
182
-
183
- Run the complete test suite:
184
-
185
- ```bash
186
- npm test
227
+ extensions/pi-undo.ts Pi extension entry point
228
+ src/
229
+ atomic-fs.ts Atomic file writes, fsync helpers
230
+ controller.ts Undo/redo/diff controller orchestrator
231
+ diff-ui.ts TUI diff view (Ink/React components)
232
+ diff-view.ts Diff computation and rendering
233
+ encoding.ts checksum, canonical JSON
234
+ git-runner.ts Git subprocess runner with counting
235
+ journal.ts Transaction journal (session/descriptor/phase)
236
+ model.ts Core types: manifest, root, session, plan
237
+ mutation-journal.ts WAL mutation journal (hash chain, batch ops)
238
+ path-safety.ts Symlink escape, relative path safety
239
+ pi-runtime.ts Pi integration layer (session/extension bridge)
240
+ quarantine.ts File isolation, no-clobber install, external concurrency
241
+ recovery.ts Startup crash recovery
242
+ restore-engine.ts Workspace restore engine (plan → apply → verify)
243
+ root-discovery.ts Workspace root topology detection (nested repos, submodules)
244
+ session-state.ts Pi session cursor management
245
+ snapshot-store.ts Git-backed content-addressed snapshot store
246
+ status-reporter.ts Phase timing and footer status
247
+ workspace-lock.ts Cross-instance workspace lock
248
+ test/ Unit, integration, recovery, and fault-injection tests
187
249
  ```
188
250
 
189
- Run tests in watch mode:
251
+ ### Testing
190
252
 
191
253
  ```bash
192
- npm run test:watch
254
+ npm test # Full test suite (381+ tests)
255
+ npm run test:watch # Watch mode
256
+ npm run test:integration # Pi runtime and extension integration tests
257
+ npm run typecheck # TypeScript type checking
258
+ npm run pack:dry-run # Inspect npm package contents
193
259
  ```
194
260
 
195
- Run the Pi runtime and extension integration tests:
261
+ The typecheck filters known pre-existing issues in Pi's own dependencies (`undici-types`, `@modelcontextprotocol`, `@google/genai`, `ReadonlySessionManager`). Only new project-level errors are enforced.
196
262
 
197
- ```bash
198
- npm run test:integration
199
- ```
263
+ ## Performance Benchmarks
200
264
 
201
- Run TypeScript type checking:
265
+ Benchmark tests assert Git call counts and WAL record counts, not wall-clock thresholds, to avoid environmental flakiness. Selected synthetic results (104-file restore, includes fixture setup, capture, plan, and apply):
202
266
 
203
- ```bash
204
- npm run typecheck
205
- ```
206
-
207
- Inspect the npm package contents before publishing:
267
+ | Scenario | Git calls | WAL records | Duration |
268
+ |---|---|---|---|
269
+ | 104-file restore (write) — batch + parallel I/O | ≤12 | 624 | ~1.1s |
270
+ | 100-file restore (delete) — batch deletes | ≤12 | 600 | ~0.6s |
271
+ | 4,000-file rollback snapshot — batch Git | ≤24 (4 x `mktree`, 4 x `commit-tree`) | — | ~7s |
208
272
 
209
- ```bash
210
- npm run pack:dry-run
211
- ```
273
+ The 104-file standalone apply probe (10 warmup iterations + 10 measured) completes in approximately 3.9s post-optimization, down from 14.3s before batching.
212
274
 
213
275
  ## License
214
276
 
package/extensions/pi-undo.ts CHANGED
@@ -1,3 +1,5 @@
1
+ import { performance } from "node:perf_hooks";
2
+
1
3
  import type {
2
4
  ExtensionAPI,
3
5
  ExtensionCommandContext,
@@ -27,13 +29,32 @@ export type PiUndoRuntimeFactory = (
27
29
  pi: ExtensionAPI,
28
30
  ) => Promise<PiUndoRuntime>;
29
31
 
32
+ type DeferredImage = NonNullable<InputEvent["images"]>[number];
33
+
34
+ interface DeferredPrompt {
35
+ readonly text: string;
36
+ readonly images?: readonly DeferredImage[];
37
+ }
38
+
30
39
  export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi: ExtensionAPI) => void {
31
40
  return (pi) => {
32
41
  let runtime: PiUndoRuntime | undefined;
42
+ let runtimeContext: ExtensionContext | undefined;
33
43
  let generation = 0;
44
+ let deferredPrompts: DeferredPrompt[] = [];
45
+ let replaying: DeferredPrompt | undefined;
46
+ let acceptedReplay: DeferredPrompt | undefined;
47
+ let activeCommands = new Set<symbol>();
48
+ let activeAction: "undo" | "redo" | undefined;
34
49
 
35
50
  const initialize = async (context: ExtensionContext): Promise<void> => {
36
51
  const currentGeneration = ++generation;
52
+ runtimeContext = context;
53
+ deferredPrompts = [];
54
+ replaying = undefined;
55
+ acceptedReplay = undefined;
56
+ activeCommands = new Set<symbol>();
57
+ activeAction = undefined;
37
58
  try {
38
59
  const next = await runtimeFactory(context, pi);
39
60
  if (currentGeneration !== generation) return;
@@ -49,6 +70,57 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
49
70
  }
50
71
  };
51
72
 
73
+ const dispatchDeferredPrompt = (active: PiUndoRuntime, expectedGeneration: number): void => {
74
+ if (
75
+ expectedGeneration !== generation || runtime !== active || activeCommands.size > 0 ||
76
+ replaying !== undefined || deferredPrompts.length === 0 || active.controller.history().locked
77
+ ) return;
78
+ const prompt = deferredPrompts[0]!;
79
+ replaying = prompt;
80
+ acceptedReplay = undefined;
81
+ queueMicrotask(() => {
82
+ if (expectedGeneration !== generation || runtime !== active || replaying !== prompt) return;
83
+ try {
84
+ pi.sendUserMessage(prompt.images === undefined || prompt.images.length === 0
85
+ ? prompt.text
86
+ : [{ type: "text" as const, text: prompt.text }, ...prompt.images]);
87
+ } catch (error) {
88
+ replaying = undefined;
89
+ acceptedReplay = undefined;
90
+ deferredPrompts.shift();
91
+ restoreEditorText(runtimeContext, prompt.text);
92
+ runtimeContext?.ui.notify(`Unable to replay queued prompt: ${errorMessage(error)}`, "warning");
93
+ }
94
+ });
95
+ };
96
+
97
+ const restoreDeferredPrompts = (context: ExtensionContext | undefined): void => {
98
+ if (deferredPrompts.length === 0) return;
99
+ const prompts = deferredPrompts.splice(0);
100
+ replaying = undefined;
101
+ acceptedReplay = undefined;
102
+ restoreEditorText(context, prompts.map((prompt) => prompt.text).join("\n\n"));
103
+ if (prompts.some((prompt) => (prompt.images?.length ?? 0) > 0)) {
104
+ context?.ui.notify("Queued prompt text restored; image attachments must be reattached", "warning");
105
+ }
106
+ };
107
+
108
+ const resumeDeferredPrompts = (
109
+ active: PiUndoRuntime,
110
+ expectedGeneration: number,
111
+ lockedReason: string,
112
+ ): void => {
113
+ if (expectedGeneration !== generation || runtime !== active) return;
114
+ const history = active.controller.history();
115
+ if (history.locked) {
116
+ active.reporter.setRecoveryRequired(lockedReason);
117
+ restoreDeferredPrompts(runtimeContext);
118
+ } else {
119
+ active.reporter.setReady(history.undoCount, history.redoCount);
120
+ dispatchDeferredPrompt(active, expectedGeneration);
121
+ }
122
+ };
123
+
52
124
  const runCommand = async (
53
125
  action: "undo" | "redo",
54
126
  context: ExtensionCommandContext,
@@ -59,8 +131,13 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
59
131
  context.ui.notify("pi-undo session unavailable", "warning");
60
132
  return;
61
133
  }
134
+ const commandToken = Symbol(action);
135
+ const commandSet = activeCommands;
136
+ commandSet.add(commandToken);
137
+ activeAction = action;
62
138
  active.reporter.setPhase(action === "undo" ? "undoing" : "redoing");
63
139
  active.setCommandContext?.(context);
140
+ const commandStarted = performance.now();
64
141
  let result;
65
142
  try {
66
143
  result = action === "undo"
@@ -68,15 +145,17 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
68
145
  : await active.controller.redo();
69
146
  } finally {
70
147
  active.setCommandContext?.(undefined);
148
+ commandSet.delete(commandToken);
149
+ if (commandSet === activeCommands && commandSet.size === 0) activeAction = undefined;
71
150
  }
72
151
  if (commandGeneration !== generation || runtime !== active) return;
73
- active.reporter.result(result);
74
- if (action === "undo" && result.code === "ok" && result.refillPrompt !== undefined) {
75
- active.reporter.refillPrompt(result.refillPrompt);
76
- }
77
- const history = active.controller.history();
78
- if (history.locked) active.reporter.setRecoveryRequired(result.message ?? result.code);
79
- else active.reporter.setReady(history.undoCount, history.redoCount);
152
+ active.reporter.result(result, performance.now() - commandStarted);
153
+ const hasDeferredPrompt = deferredPrompts.length > 0 || replaying !== undefined;
154
+ if (
155
+ action === "undo" && result.code === "ok" && result.refillPrompt !== undefined &&
156
+ !hasDeferredPrompt
157
+ ) active.reporter.refillPrompt(result.refillPrompt);
158
+ resumeDeferredPrompts(active, commandGeneration, result.message ?? result.code);
80
159
  };
81
160
 
82
161
  const runDiff = async (args: string, context: ExtensionCommandContext): Promise<void> => {
@@ -132,19 +211,69 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
132
211
  });
133
212
 
134
213
  pi.on("session_start", async (_event: unknown, context: ExtensionContext) => initialize(context));
135
- pi.on("input", async (event: InputEvent) => {
136
- if (runtime === undefined) return { action: "handled" as const };
137
- return runtime.controller.prepareInput(event.text, { streaming: event.streamingBehavior !== undefined });
214
+ pi.on("input", async (event: InputEvent, context: ExtensionContext) => {
215
+ const active = runtime;
216
+ if (active === undefined) return { action: "handled" as const };
217
+ const result = await active.controller.prepareInput(event.text, {
218
+ streaming: event.streamingBehavior !== undefined,
219
+ });
220
+ const replay = replaying;
221
+ if (result.action === "defer") {
222
+ if (replay !== undefined && event.source === "extension" && samePrompt(event.text, event.images, replay)) {
223
+ replaying = undefined;
224
+ acceptedReplay = undefined;
225
+ active.reporter.setPhase(`${activeAction ?? "operation"} queued:${deferredPrompts.length}`);
226
+ return { action: "handled" as const };
227
+ }
228
+ if (event.text.trimStart().startsWith("/")) {
229
+ restoreEditorText(context, event.text);
230
+ context.ui.notify("Command input preserved until undo/redo completes", "info");
231
+ return { action: "handled" as const };
232
+ }
233
+ deferredPrompts.push({
234
+ text: event.text,
235
+ ...(event.images === undefined ? {} : { images: event.images.map((image) => ({ ...image })) }),
236
+ });
237
+ active.reporter.setPhase(`${activeAction ?? "operation"} queued:${deferredPrompts.length}`);
238
+ return { action: "handled" as const };
239
+ }
240
+ if (replay !== undefined && event.source === "extension" && samePrompt(event.text, event.images, replay)) {
241
+ if (result.action === "continue") {
242
+ acceptedReplay = replay;
243
+ } else {
244
+ removeDeferredPrompt(deferredPrompts, replay);
245
+ replaying = undefined;
246
+ acceptedReplay = undefined;
247
+ restoreEditorText(context, replay.text);
248
+ if ((replay.images?.length ?? 0) > 0) {
249
+ context.ui.notify("Queued prompt text restored; image attachments must be reattached", "warning");
250
+ }
251
+ }
252
+ }
253
+ return result;
254
+ });
255
+ pi.on("before_agent_start", async () => {
256
+ const active = runtime;
257
+ const startGeneration = generation;
258
+ if (active === undefined) return;
259
+ await active.controller.beforeAgentStart();
260
+ const replay = replaying;
261
+ if (
262
+ runtime === active && generation === startGeneration && replay !== undefined &&
263
+ acceptedReplay === replay
264
+ ) {
265
+ removeDeferredPrompt(deferredPrompts, replay);
266
+ replaying = undefined;
267
+ acceptedReplay = undefined;
268
+ }
138
269
  });
139
- pi.on("before_agent_start", async () => { await runtime?.controller.beforeAgentStart(); });
140
270
  pi.on("agent_settled", async () => {
141
271
  const active = runtime;
272
+ const settledGeneration = generation;
142
273
  if (active === undefined) return;
143
274
  await active.controller.agentSettled();
144
- if (runtime !== active) return;
145
- const history = active.controller.history();
146
- if (history.locked) active.reporter.setRecoveryRequired("session state ambiguous");
147
- else active.reporter.setReady(history.undoCount, history.redoCount);
275
+ if (runtime !== active || generation !== settledGeneration) return;
276
+ resumeDeferredPrompts(active, settledGeneration, "session state ambiguous");
148
277
  });
149
278
  pi.on("session_before_tree", async (event: PiSessionBeforeTreeEvent) => {
150
279
  if (runtime === undefined) return { cancel: true };
@@ -152,7 +281,7 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
152
281
  const active = runtime;
153
282
  const result = await active.controller.beforeTree({ targetLeafId: event.preparation.targetId });
154
283
  if (result === undefined) {
155
- event.signal?.addEventListener("abort", () => { void active.controller.cancelTree?.(); }, { once: true });
284
+ event.signal?.addEventListener("abort", () => { void active.controller.cancelTree?.(); }, { once: true });
156
285
  }
157
286
  return result;
158
287
  });
@@ -165,6 +294,12 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
165
294
  });
166
295
  pi.on("session_shutdown", async () => {
167
296
  generation += 1;
297
+ deferredPrompts = [];
298
+ replaying = undefined;
299
+ acceptedReplay = undefined;
300
+ activeCommands = new Set<symbol>();
301
+ activeAction = undefined;
302
+ runtimeContext = undefined;
168
303
  await runtime?.controller.cancelTree?.();
169
304
  runtime?.reporter.clear();
170
305
  runtime = undefined;
@@ -172,6 +307,27 @@ export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi
172
307
  };
173
308
  }
174
309
 
310
+ function samePrompt(
311
+ text: string,
312
+ images: readonly DeferredImage[] | undefined,
313
+ prompt: DeferredPrompt,
314
+ ): boolean {
315
+ if (text !== prompt.text || (images?.length ?? 0) !== (prompt.images?.length ?? 0)) return false;
316
+ return (images ?? []).every((image, index) => JSON.stringify(image) === JSON.stringify(prompt.images?.[index]));
317
+ }
318
+
319
+ function removeDeferredPrompt(prompts: DeferredPrompt[], prompt: DeferredPrompt): void {
320
+ const index = prompts.indexOf(prompt);
321
+ if (index >= 0) prompts.splice(index, 1);
322
+ }
323
+
324
+ function restoreEditorText(context: ExtensionContext | undefined, text: string): void {
325
+ if (context === undefined || text.length === 0) return;
326
+ const current = context.ui.getEditorText();
327
+ if (current === text || current.startsWith(`${text}\n\n`)) return;
328
+ context.ui.setEditorText(current.length === 0 ? text : `${text}\n\n${current}`);
329
+ }
330
+
175
331
  function errorMessage(error: unknown): string {
176
332
  return error instanceof Error ? error.message : String(error);
177
333
  }
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@davideasden/pi-undo",
3
- "version": "0.1.2",
3
+ "version": "0.2.1",
4
4
  "description": "Persistent workspace undo and redo for Pi",
5
5
  "type": "module",
6
6
  "keywords": [
package/src/atomic-fs.ts CHANGED
@@ -95,7 +95,12 @@ export async function fsyncDirectory(directory: string): Promise<void> {
95
95
  }
96
96
  }
97
97
 
98
- export async function writeBytesExclusive(file: string, bytes: Uint8Array, mode: number): Promise<void> {
98
+ export async function writeBytesExclusive(
99
+ file: string,
100
+ bytes: Uint8Array,
101
+ mode: number,
102
+ options: { readonly syncDirectory?: boolean } = {},
103
+ ): Promise<void> {
99
104
  const handle = await open(file, "wx", mode);
100
105
  try {
101
106
  await handle.writeFile(Buffer.from(bytes));
@@ -104,7 +109,7 @@ export async function writeBytesExclusive(file: string, bytes: Uint8Array, mode:
104
109
  } finally {
105
110
  await handle.close();
106
111
  }
107
- await fsyncDirectory(dirname(file));
112
+ if (options.syncDirectory !== false) await fsyncDirectory(dirname(file));
108
113
  }
109
114
 
110
115
  async function targetMode(file: string, mode: number | undefined): Promise<number> {