Skip to content

Commit ff90057

Browse files
ndemiancclaude
andcommitted
Atom++ AI: workspace checkpoints (per-turn restore) + composer history recall
Workspace checkpoints — M5's third pillar. Each agent turn that changes files is a checkpoint; an always-visible "↩ Restore checkpoint · N files" chip under each user message rolls the workspace back to BEFORE that turn ran, undoing that turn and every later one. - reviewSession.js: applyEdit gains a first-touch hook that hands the SAME from-disk pre-image to the checkpoint (the two layers can't diverge); new shared restoreOne(key, before, created) factors undoFile's write-back (replace+save / recreate-if-deleted / deleteFile) with the expected re-sync. - extension.js: a per-session checkpoints stack {turnId, goalMsg, files} opened at agentFlow start, filled on first touch, closed in finally; restoreCheckpoint unions files across the restored suffix (earliest pre-image wins, unit-tested), reverts in reverse order, finalizes pending review, truncates the transcript by message IDENTITY (trimAgentMemory-proof), and refuses while the agent runs. Reset on New Chat; never routed through reapCommands (pure data, not process). - chat.html: the Restore chip (confirm → restore → prune log + card maps), plus shell-style ↑/↓ command history recall in the composer (session-scoped, multiline-safe, draft preserved). Restore-union + history-recall logic unit-tested; loaded cleanly in dev. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 1983231 commit ff90057

3 files changed

Lines changed: 154 additions & 5 deletions

File tree

‎extensions/atom-ai/extension.js‎

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,7 @@ function workspaceMapBlock(allFiles) {
331331
function newChat() {
332332
conversation = [];
333333
agentMessages = [];
334+
checkpoints.length = 0; currentCheckpoint = null; // drop the per-turn restore stack
334335
pendingContext = null;
335336
contextFiles = [];
336337
if (abort) { abort.abort(); }
@@ -381,6 +382,40 @@ function reapCommands() {
381382
bgRuns.clear();
382383
}
383384

385+
// Workspace checkpoints: a per-user-turn stack of file pre-images so the user can roll the workspace back
386+
// to before any turn ran. In-memory per chat session (reset on New Chat); NOT routed through reapCommands
387+
// (these are pure data, not process state — restoring files must not kill running servers).
388+
const checkpoints = []; // oldest-first: { turnId, label, ts, goalMsg, files: Map<key,{before,created}> }
389+
let currentCheckpoint = null;
390+
let checkpointSeq = 0;
391+
/** First-touch hook handed to reviewSession.applyEdit — record a file's pre-image into the open turn. */
392+
function recordCheckpointTouch(key, before, created) {
393+
if (currentCheckpoint && !currentCheckpoint.files.has(key)) { currentCheckpoint.files.set(key, { before: before, created: created }); }
394+
}
395+
/** Restore the workspace to before `turnId` ran — reverting that turn AND every turn after it. */
396+
async function restoreCheckpoint(turnId) {
397+
if (abort) { post({ type: 'checkpointBusy', turnId: turnId }); return; } // never restore under a live turn
398+
const idx = checkpoints.findIndex((c) => c.turnId === turnId);
399+
if (idx < 0) { post({ type: 'checkpointBusy', turnId: turnId, reason: 'gone' }); return; }
400+
const cp = checkpoints[idx];
401+
const cutIdx = agentMessages.indexOf(cp.goalMsg); // by identity → survives trimAgentMemory
402+
// Union of files across this checkpoint + every later one; EARLIEST pre-image wins (true pre-turn state).
403+
const union = new Map();
404+
for (let i = idx; i < checkpoints.length; i++) {
405+
for (const [k, v] of checkpoints[i].files) { if (!union.has(k)) { union.set(k, v); } }
406+
}
407+
let restored = 0;
408+
for (const k of [...union.keys()].reverse()) { // reverse for interdependency safety
409+
const v = union.get(k);
410+
try { if (await review.restoreOne(k, v.before, v.created)) { restored++; } } catch (e) { dbg('checkpoint.restoreErr', { key: k, msg: String((e && e.message) || e) }); }
411+
}
412+
review.finalizeAll(); // pending Keep/Undo cards now reference stale snapshots
413+
if (cutIdx >= 0) { agentMessages.length = cutIdx; } // truncate transcript at the goal boundary
414+
checkpoints.length = idx; // drop restored + all later checkpoints (no redo)
415+
dbg('checkpoint.restore', { turnId: turnId, files: restored, truncatedAt: cutIdx });
416+
post({ type: 'checkpointRestored', turnId: turnId, filesRestored: restored });
417+
}
418+
384419
/** Pending in-chat approval requests, keyed by id, resolved by the webview. */
385420
const pendingApprovals = new Map();
386421
let approvalSeq = 0;
@@ -513,7 +548,12 @@ async function agentFlow(text) {
513548
post({ type: 'agentStart' });
514549
abort = new AbortController();
515550
repairAgentMemory();
516-
agentMessages.push({ role: 'user', content: text });
551+
// Open a workspace checkpoint for this turn (before the goal is pushed) so the user can roll back here.
552+
const goalMsg = { role: 'user', content: text };
553+
currentCheckpoint = { turnId: ++checkpointSeq, label: (text || '').slice(0, 60), ts: Date.now(), goalMsg: goalMsg, files: new Map() };
554+
checkpoints.push(currentCheckpoint);
555+
post({ type: 'checkpointOpened', turnId: currentCheckpoint.turnId });
556+
agentMessages.push(goalMsg);
517557
trimAgentMemory();
518558
dbg('agent.start', { model: cfg.get('claude.model', 'claude-sonnet-4-6'), maxSteps: cfg.get('agent.maxSteps', 25), transcriptMsgs: agentMessages.length, goalChars: text.length });
519559
// Auto-verify: capture a pre-run diagnostics baseline (so we only flag NEW problems) and a fresh
@@ -551,6 +591,7 @@ async function agentFlow(text) {
551591
clearApprovals();
552592
clearQuestions();
553593
abort = null;
594+
if (currentCheckpoint) { post({ type: 'checkpointClosed', turnId: currentCheckpoint.turnId, fileCount: currentCheckpoint.files.size }); currentCheckpoint = null; }
554595
}
555596
}
556597

@@ -692,6 +733,7 @@ class ChatViewProvider {
692733
case 'copy': try { await vscode.env.clipboard.writeText(String(msg.text || '')); } catch (e) { /* clipboard unavailable */ } break;
693734
case 'retry': if (lastAgentGoal && !abort) { dbg('retry', { goalChars: lastAgentGoal.length }); await agentFlow(lastAgentGoal); } break;
694735
case 'continueAgent': if (!abort) { const g = lastAgentGoal; dbg('continue', { transcriptMsgs: agentMessages.length }); await agentFlow('Continue from where you left off and finish the task. Pick up exactly where you stopped — do not restart or repeat work that is already done.'); lastAgentGoal = g; } break;
736+
case 'restoreCheckpoint': dbg('restoreCheckpoint', { turnId: msg.turnId, running: !!abort }); await restoreCheckpoint(msg.turnId); break;
695737
case 'feedback': dbg('feedback', { value: msg.value }); break;
696738
case 'openFile': await openWorkspaceFile(msg.path); break;
697739
case 'reviewKeepFile': dbg('review.keep', { id: msg.id }); review.keepFile(msg.id, 'kept'); break;
@@ -865,7 +907,7 @@ function activate(context) {
865907
});
866908

867909
// Apply-then-review: agent edits land immediately and are reviewed with Keep/Undo (editor + panel).
868-
review = registerReview(context, post, dbg);
910+
review = registerReview(context, post, dbg, recordCheckpointTouch);
869911

870912
// Inline (ghost-text) tab-completion. Silent key lookup — it must never prompt mid-typing.
871913
registerInlineComplete(context, {

‎extensions/atom-ai/media/chat.html‎

Lines changed: 81 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,20 @@
4141
.msg:hover .msgcopy, .msgcopy:focus-visible, .msgcopy.done { opacity: .8; }
4242
.msgcopy:hover { opacity: 1; color: var(--vscode-foreground); background: var(--vscode-toolbar-hoverBackground, rgba(127,127,127,.18)); }
4343
.msgcopy.done { color: var(--vscode-gitDecoration-addedResourceForeground, #73c991); border-color: var(--vscode-gitDecoration-addedResourceForeground, #73c991); }
44+
/* per-turn workspace-checkpoint Restore control (hover-revealed on the user message) */
45+
.ckwrap { margin-top: 6px; }
46+
.ckrestore { opacity: .62; cursor: pointer; border: 1px solid var(--border); border-radius: 6px; padding: 3px 9px; font-size: 11px; display: inline-flex; align-items: center; gap: 5px; background: transparent; color: var(--muted); transition: opacity .12s, color .12s, background .12s; }
47+
.msg.user:hover .ckrestore, .ckrestore:focus-visible { opacity: .9; }
48+
.ckrestore:hover { opacity: 1; color: var(--vscode-foreground); background: var(--vscode-toolbar-hoverBackground, rgba(127,127,127,.18)); }
49+
.ckrestore .ckn { color: var(--muted); }
50+
.ckrestore svg, .ckmuted svg { width: 1em; height: 1em; }
51+
.ckconfirm { font-size: 11.5px; color: var(--muted); display: inline-flex; align-items: center; gap: 6px; flex-wrap: wrap; }
52+
.ckconfirm button { cursor: pointer; border: 1px solid var(--border); border-radius: 5px; padding: 2px 9px; font-size: 11px; color: var(--vscode-foreground); background: transparent; }
53+
.ckconfirm .ckyes { background: var(--vscode-button-background); color: var(--vscode-button-foreground); border-color: transparent; }
54+
.ckconfirm .ckyes:hover { background: var(--vscode-button-hoverBackground); }
55+
.ckconfirm .ckno:hover { background: var(--vscode-toolbar-hoverBackground, rgba(127,127,127,.18)); }
56+
.ckmuted { font-size: 11.5px; color: var(--muted); display: inline-flex; align-items: center; gap: 5px; }
57+
.ckdivider { text-align: center; font-size: 11px; color: var(--muted); margin: 8px 0; opacity: .85; }
4458
/* ---- jump-to-latest affordance (shown only when auto-scroll is paused) ---- */
4559
#jumpLatest {
4660
position: fixed; right: 16px; bottom: 96px; z-index: 20;
@@ -1428,9 +1442,17 @@
14281442
sendBtn.innerHTML = on ? '&#9632;' : '&#8629;'; // ■ while streaming, ↵ otherwise
14291443
sendBtn.title = on ? 'Stop' : 'Send (Enter)';
14301444
}
1445+
// Shell-style history recall in the composer: ↑ = previous sent message, ↓ = newer (back to your draft).
1446+
// Session-only. In multiline text, ↑/↓ recall only when the caret is on the first/last line (arrows still move otherwise).
1447+
const cmdHistory = []; let histIdx = -1; let histDraft = '';
1448+
function caretAtFirstLine(el){ return el.value.lastIndexOf('\n', el.selectionStart - 1) === -1; }
1449+
function caretAtLastLine(el){ return el.value.indexOf('\n', el.selectionEnd) === -1; }
1450+
function histSet(text){ input.value = text; auto(); const n = text.length; try { input.setSelectionRange(n, n); } catch (e) {} }
14311451
function doSend(){
14321452
if (streaming) { flushAll = true; ensurePump(); vscode.postMessage({ type: 'stop' }); return; }
14331453
const t = input.value.trim(); if (!t) return;
1454+
cmdHistory.push(t); if (cmdHistory.length > 200) { cmdHistory.shift(); } // record for ↑/↓ recall
1455+
histIdx = -1; histDraft = '';
14341456
add('user', render(t)); forceStick(); input.value = ''; auto();
14351457
vscode.postMessage({ type: 'send', text: t });
14361458
}
@@ -1497,12 +1519,66 @@
14971519
});
14981520
renderAccount({ signedIn: false });
14991521

1500-
input.addEventListener('keydown', e => { if (e.key === 'Enter' && !e.shiftKey){ e.preventDefault(); doSend(); } });
1522+
input.addEventListener('keydown', e => {
1523+
if (e.key === 'Enter' && !e.shiftKey){ e.preventDefault(); doSend(); return; }
1524+
const plain = !e.shiftKey && !e.altKey && !e.metaKey && !e.ctrlKey;
1525+
if (e.key === 'ArrowUp' && plain){
1526+
if (!cmdHistory.length){ return; }
1527+
if (histIdx === -1){ if (!caretAtFirstLine(input)){ return; } histDraft = input.value; histIdx = cmdHistory.length; }
1528+
if (histIdx > 0){ histIdx--; histSet(cmdHistory[histIdx]); }
1529+
e.preventDefault(); return;
1530+
}
1531+
if (e.key === 'ArrowDown' && plain){
1532+
if (histIdx === -1){ return; }
1533+
if (!caretAtLastLine(input)){ return; }
1534+
histIdx++;
1535+
if (histIdx >= cmdHistory.length){ histIdx = -1; histSet(histDraft); }
1536+
else { histSet(cmdHistory[histIdx]); }
1537+
e.preventDefault(); return;
1538+
}
1539+
});
15011540
function auto(){ input.style.height = 'auto'; input.style.height = Math.min(input.scrollHeight, 200) + 'px'; }
15021541
input.addEventListener('input', auto);
15031542
input.addEventListener('focus', () => composer.classList.add('focused'));
15041543
input.addEventListener('blur', () => composer.classList.remove('focused'));
15051544

1545+
// Workspace checkpoints: a hover-revealed "Restore" control on each user turn that changed files.
1546+
const REVERT = '<svg viewBox="0 0 16 16" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><path d="M3.8 8a4.2 4.2 0 1 0 1.3-3.05"/><path d="M2.4 3v2.5h2.5"/></svg>';
1547+
function ckBubble(turnId){ return log.querySelector('.msg.user[data-turn-id="' + turnId + '"]'); }
1548+
function stampCheckpoint(turnId){
1549+
const us = log.querySelectorAll('.msg.user'); const u = us[us.length - 1];
1550+
if (u && !u.dataset.turnId){ u.dataset.turnId = turnId; }
1551+
}
1552+
function renderRestoreIdle(wrap, turnId, fileCount){
1553+
wrap.innerHTML = '<button class="ckrestore" title="Restore the workspace to before this message ran">' + REVERT
1554+
+ ' Restore checkpoint <span class="ckn">· ' + fileCount + ' file' + (fileCount === 1 ? '' : 's') + '</span></button>';
1555+
wrap.querySelector('.ckrestore').onclick = () => {
1556+
wrap.innerHTML = '<span class="ckconfirm">Revert the changes from this message onward? <button class="ckyes">Restore</button> <button class="ckno">Cancel</button></span>';
1557+
wrap.querySelector('.ckno').onclick = () => renderRestoreIdle(wrap, turnId, fileCount);
1558+
wrap.querySelector('.ckyes').onclick = () => { wrap.innerHTML = '<span class="ckmuted">' + REVERT + ' Restoring…</span>'; vscode.postMessage({ type: 'restoreCheckpoint', turnId: Number(turnId) }); };
1559+
};
1560+
}
1561+
function addRestoreControl(turnId, fileCount){
1562+
if (!fileCount){ return; } // turn changed no files → nothing to restore
1563+
const u = ckBubble(turnId); if (!u){ return; }
1564+
u.dataset.ckFiles = fileCount;
1565+
let wrap = u.querySelector('.ckwrap'); if (!wrap){ wrap = document.createElement('div'); wrap.className = 'ckwrap'; u.appendChild(wrap); }
1566+
renderRestoreIdle(wrap, turnId, fileCount);
1567+
}
1568+
function pruneDetached(map){ for (const k in map){ const el = map[k]; if (el && !log.contains(el)){ delete map[k]; } } }
1569+
function applyCheckpointRestored(m){
1570+
const u = ckBubble(m.turnId); if (!u){ return; }
1571+
let n = u.nextSibling; while (n){ const next = n.nextSibling; n.remove(); n = next; } // drop this turn's output + every later turn
1572+
const wrap = u.querySelector('.ckwrap'); if (wrap){ wrap.innerHTML = '<span class="ckmuted">' + REVERT + ' Restored to here · ' + (m.filesRestored || 0) + ' file' + (m.filesRestored === 1 ? '' : 's') + ' reverted</span>'; }
1573+
pruneDetached(editCards); pruneDetached(termCards); pruneDetached(verifyCards);
1574+
agentBubble = null; agentRaw = ''; setStreaming(false); forceStick();
1575+
}
1576+
function checkpointBusyNotice(m){
1577+
if (m.turnId != null){ const u = ckBubble(m.turnId); const wrap = u && u.querySelector('.ckwrap'); if (wrap){ renderRestoreIdle(wrap, m.turnId, Number(u.dataset.ckFiles || 1)); } }
1578+
const note = document.createElement('div'); note.className = 'ckdivider';
1579+
note.textContent = (m.reason === 'gone') ? 'That checkpoint is no longer available.' : 'Stop the running agent before restoring a checkpoint.';
1580+
log.appendChild(note); scrollIfStuck(); setTimeout(() => { try { note.remove(); } catch(e){} }, 3500);
1581+
}
15061582
window.addEventListener('message', ev => {
15071583
const m = ev.data;
15081584
if (m.type === 'config'){
@@ -1535,6 +1611,10 @@
15351611
else if (m.type === 'termOutput'){ termOutputAppend(m); }
15361612
else if (m.type === 'termExit'){ termExitFinish(m); }
15371613
else if (m.type === 'bgTask'){ applyBgTask(m); }
1614+
else if (m.type === 'checkpointOpened'){ stampCheckpoint(m.turnId); }
1615+
else if (m.type === 'checkpointClosed'){ addRestoreControl(m.turnId, m.fileCount); }
1616+
else if (m.type === 'checkpointRestored'){ applyCheckpointRestored(m); }
1617+
else if (m.type === 'checkpointBusy'){ checkpointBusyNotice(m); }
15381618
else if (m.type === 'verifyRun'){ addVerifyRun(m); }
15391619
else if (m.type === 'verifyOutput'){ verifyOutputAppend(m); }
15401620
else if (m.type === 'verifyDone'){ verifyDoneFinish(m); }

‎extensions/atom-ai/reviewSession.js‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ function countDiff(diff) {
5050
* @param {vscode.ExtensionContext} context
5151
* @param {(m:any)=>void} post Send a message to the chat webview.
5252
*/
53-
function registerReview(context, post, dbg) {
53+
function registerReview(context, post, dbg, recordTouch) {
5454
dbg = dbg || (() => {});
5555
/** @type {Map<string,{uri:vscode.Uri, rel:string, exists:boolean, snapshot:string, range:vscode.Range, add:number, del:number}>} */
5656
const pending = new Map();
@@ -122,6 +122,10 @@ function registerReview(context, post, dbg) {
122122
catch { return false; } // can't read the file we're about to edit → abort (never a destructive empty-snapshot edit)
123123
} else { snapshot = ''; }
124124

125+
// Workspace-checkpoint hook: hand the SAME from-disk pre-image to the per-turn checkpoint (snapshot-
126+
// on-first-touch). Shares this exact read so the two layers can never diverge. Never blocks an edit.
127+
if (recordTouch) { try { recordTouch(key, existsNow ? snapshot : null, !existsNow); } catch (e) { /* checkpoint capture must never break editing */ } }
128+
125129
const edit = new vscode.WorkspaceEdit();
126130
if (existsNow) {
127131
const doc = await vscode.workspace.openTextDocument(uri);
@@ -218,6 +222,29 @@ function registerReview(context, post, dbg) {
218222
postState();
219223
}
220224

225+
/** Workspace-checkpoint restore primitive: write a file back to `before` (replace+save, recreating it
226+
* if it was deleted since), or delete it if the agent had `created` it. Reuses undoFile's write-back +
227+
* the `expected` re-sync so a restored file isn't later mis-read as a user edit. Returns true on success. */
228+
async function restoreOne(key, before, created) {
229+
let uri;
230+
try { uri = vscode.Uri.parse(key); } catch { return false; }
231+
const edit = new vscode.WorkspaceEdit();
232+
if (!created && before != null) {
233+
let doc = null;
234+
try { doc = await vscode.workspace.openTextDocument(uri); } catch { doc = null; }
235+
if (doc) { edit.replace(uri, doc.validateRange(new vscode.Range(0, 0, doc.lineCount + 1, 0)), before); }
236+
else { edit.createFile(uri, { ignoreIfExists: true }); edit.insert(uri, new vscode.Position(0, 0), before); } // deleted since → recreate
237+
expected.set(key, before);
238+
} else {
239+
edit.deleteFile(uri, { ignoreIfNotExists: true });
240+
expected.delete(key);
241+
}
242+
const ok = await vscode.workspace.applyEdit(edit);
243+
if (ok && !created && before != null) { try { const doc = await vscode.workspace.openTextDocument(uri); if (doc.isDirty) { await doc.save(); } } catch { /* */ } }
244+
undecorate(uri);
245+
return ok;
246+
}
247+
221248
async function openDiff(key) {
222249
const fr = pending.get(key);
223250
if (!fr) { return; }
@@ -310,7 +337,7 @@ function registerReview(context, post, dbg) {
310337
postState();
311338
}
312339

313-
return { applyEdit, keepFile, undoFile, keepAll, undoAll, finalizeAll, openDiff, resync, pendingCount: () => pending.size };
340+
return { applyEdit, keepFile, undoFile, keepAll, undoAll, finalizeAll, restoreOne, openDiff, resync, pendingCount: () => pending.size };
314341
}
315342

316343
module.exports = { registerReview };

0 commit comments

Comments
 (0)