Skip to content

feat: restore checkpoint after any AI response turn#2356

Draft
Basit-Balogun10 wants to merge 3 commits into
PostHog:mainfrom
Basit-Balogun10:claude/gracious-hawking-38aade
Draft

feat: restore checkpoint after any AI response turn#2356
Basit-Balogun10 wants to merge 3 commits into
PostHog:mainfrom
Basit-Balogun10:claude/gracious-hawking-38aade

Conversation

@Basit-Balogun10
Copy link
Copy Markdown
Contributor

@Basit-Balogun10 Basit-Balogun10 commented May 25, 2026

Note: This PR is stacked on #2321. The diff shows +2k lines because GitHub uses main as the base - cross-fork branches can't be set as PR bases (tried using Graphite to track the stack too, same limitation). The actual changes in this PR are a few commits on top of #2321. Please review with that context, or we wait for #2321 to be merged first.

Problem

After an AI turn completes, there's no way to roll back to the git state captured at the end of that turn. If the agent goes down a wrong path you have to manually reset the branch or throw away work — there's no checkpoint-based undo.

Tracked in #724 / #2328.

Changes

Core restore flow

  • New checkpoint tRPC router (routers/checkpoint.ts) with a restore procedure that:
    1. Runs RevertCheckpointSaga to reset the worktree to the saved git state
    2. Truncates the session .jsonl to the restore point so the session replays correctly
    3. Collects any orphaned checkpoint refs from the discarded lines and deletes them — no accumulation of stale refs from abandoned future turns
  • getSessionInfo(taskRunId) added to AgentService to expose { sessionId, repoPath } without leaking internal types

Per-turn restore button

  • buildConversationItems now tracks lastCheckpointId per turn — set when a _posthog/git_checkpoint notification is seen, cleared on each new turn start
  • AgentMessage gets a restore button in the top-right corner: active when the turn has a checkpoint, disabled (with tooltip) when it doesn't. Only shown on completed turns.
  • SessionUpdateView forwards the new showRestoreButton / onRestoreCheckpoint props down

Checkpoint timeline modal (mod+shift+h)

  • CheckpointTimelineModal — command-palette-style dialog listing every checkpoint in the session, newest first, with the first 120 chars of the user message that started that turn and a relative timestamp
  • Parses events: AcpMessage[] directly (Option A — no new store or API surface, scoped to the current session by definition)
  • Shortcut registered as "checkpoint-timeline" in CONFIGURABLE_SHORTCUT_IDS / DEFAULT_KEYBINDINGS / KEYBOARD_SHORTCUTS so it's user-remappable via the keybindings store — built on top of feat: configurable keyboard shortcuts #2321

Restore confirmation

  • RestoreCheckpointDialog — confirmation dialog with an amber warning before the restore runs, so accidental clicks don't immediately drop work
  • useRestoreCheckpoint hook wires the full flow: opens the confirmation dialog → calls the tRPC mutation → truncates the in-memory events via sessionStoreSetters.truncateEventsToCheckpoint → shows a success/error toast

How did you test this?

  • Ran pnpm --filter code typecheck against tsconfig.node.json and tsconfig.web.json — both exit 0 (pre-existing @posthog/platform / @posthog/agent module errors are unrelated to this change)
  • Ran pnpm lint — biome passes clean
  • Manually traced the restore flow end-to-end: checkpoint tRPC router logic, JSONL truncation returning orphaned IDs, git ref deletion
  • Verified buildConversationItems correctly sets lastCheckpointId for turns that have a _posthog/git_checkpoint notification and null for those that don't
  • Checked CheckpointTimelineModal parses a sample events array with multiple turns and renders entries newest-first with correct snippets

@Basit-Balogun10 Basit-Balogun10 force-pushed the claude/gracious-hawking-38aade branch from 7dc1907 to 021fe43 Compare May 26, 2026 15:35
- Add checkpoint tRPC router with restore procedure that reverts git state,
  truncates session JSONL to the restore point, and deletes orphaned
  checkpoint refs for abandoned future turns
- Track lastCheckpointId per turn in buildConversationItems so each
  completed agent turn knows its git ref
- Show per-turn restore button in AgentMessage (disabled with tooltip when
  no checkpoint exists for that turn)
- Add CheckpointTimelineModal (mod+shift+h) — command-palette-style list of
  all checkpoints in the session, newest first, with user message snippet
  and relative timestamp; shortcut is user-remappable via keybindings store
- Add RestoreCheckpointDialog with confirmation warning before reverting
- Add useRestoreCheckpoint hook to wire restore flow end-to-end
- Register checkpoint-timeline as a configurable shortcut

Closes PostHog#2328
The cloud path (agent-server.ts) guards checkpoint capture on posthogAPI
being configured, so local tasks never emit _posthog/git_checkpoint.

Hook into extNotification in the local AgentService: on TURN_COMPLETE, run
CaptureCheckpointSaga, then emit a synthetic _posthog/git_checkpoint ACP
message to the renderer and append it to the session JSONL so it survives
reload. The renderer's buildConversationItems already handles the
notification correctly — it just wasn't arriving.

Add console logs in buildConversationItems and structured logs in service.ts
for visibility during debugging.
extNotification is not reliably called by the ACP SDK for _posthog/
notifications in the local path. The raw stream tap (onAcpMessage)
is guaranteed to fire for every ndjson frame — move the TURN_COMPLETE
checkpoint hook there instead.
@Basit-Balogun10 Basit-Balogun10 force-pushed the claude/gracious-hawking-38aade branch from f191382 to 158a49c Compare June 1, 2026 12:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant