| 1 | package control |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "fmt" |
| 6 | "sync" |
| 7 | "sync/atomic" |
| 8 | "time" |
| 9 | |
| 10 | "reasonix/internal/agent" |
| 11 | "reasonix/internal/checkpoint" |
| 12 | "reasonix/internal/config" |
| 13 | "reasonix/internal/diff" |
| 14 | "reasonix/internal/provider" |
| 15 | ) |
| 16 | |
| 17 | // checkpointManager owns the snapshot-based rewind bookkeeping: the per-session |
| 18 | // checkpoint store, the monotonic turn counter, and the conversation-rewind |
| 19 | // boundary map. Like approvalManager it holds only the bookkeeping behind its own |
| 20 | // lock, off the controller's c.mu — the Controller keeps the rewind/fork |
| 21 | // orchestration (truncating the session, restoring code, emitting events) that |
| 22 | // needs its other collaborators. |
| 23 | // |
| 24 | // turn is decoupled from the store so it remains monotonic across session work; |
| 25 | // bound[turn] records len(Session.Messages) at that turn's start — the truncation |
| 26 | // boundary for a conversation rewind/fork. Boundaries are persisted in each |
| 27 | // checkpoint and rebuilt from the store on resume (so a reopened session can still |
| 28 | // rewind conversation / fork). Context compression never changes the transcript, |
| 29 | // so it leaves these boundaries intact. Every store call does its disk I/O off mu — |
| 30 | // mu is taken only to read/swap the store pointer and mutate turn/bound. |
| 31 | type checkpointManager struct { |
| 32 | // mu guards store, turn, and bound; every critical section under it is short |
| 33 | // and non-blocking (no disk I/O). |
| 34 | mu sync.Mutex |
| 35 | store *checkpoint.Store |
| 36 | turn int |
| 37 | bound map[int]int |
| 38 | } |
| 39 | |
| 40 | // rebind points the store at the (possibly new) session, loading any checkpoints |
| 41 | // already on disk, and resets the turn counter and boundaries from them. root is |
| 42 | // the workspace root used to guard restore writes. opts carry the configured |
| 43 | // retention. Called on construction and whenever the session path changes |
| 44 | // (NewSession/Resume/SetSessionPath/fork). |
| 45 | func (m *checkpointManager) rebind(dir, root string, opts ...checkpoint.Option) { |
| 46 | store := checkpoint.New(dir, root, opts...) |
| 47 | next := store.NextTurn() // continue numbering past any checkpoints on disk |
| 48 | bound := store.Bounds() // rebuilt from persisted checkpoints so a resumed |
| 49 | if bound == nil { // session can still rewind conversation / fork |
| 50 | bound = map[int]int{} |
| 51 | } |
| 52 | m.mu.Lock() |
| 53 | m.store = store |
| 54 | m.turn = next |
| 55 | m.bound = bound |
| 56 | m.mu.Unlock() |
| 57 | } |
| 58 | |
| 59 | // checkpointOptions turns the workspace [checkpoints] config into store |
| 60 | // options. A read failure leaves the built-in retention defaults in place — |
| 61 | // checkpoints are a safety net, so a malformed config must not block a session. |
| 62 | func (c *Controller) checkpointOptions() []checkpoint.Option { |
| 63 | cfg, err := config.LoadForRootReadOnly(c.workspaceRoot) |
| 64 | if err != nil { |
| 65 | return nil |
| 66 | } |
| 67 | var opts []checkpoint.Option |
| 68 | if turns := cfg.Checkpoints.RetainTurns; turns > 0 { |
| 69 | opts = append(opts, checkpoint.WithRetainCheckpoints(turns)) |
| 70 | } |
| 71 | if quota := cfg.Checkpoints.BlobQuotaBytes; quota > 0 { |
| 72 | opts = append(opts, checkpoint.WithBlobQuota(quota)) |
| 73 | } |
| 74 | return opts |
| 75 | } |
| 76 | |
| 77 | // enabled reports whether a checkpoint store is bound. |
| 78 | func (m *checkpointManager) enabled() bool { |
| 79 | m.mu.Lock() |
| 80 | defer m.mu.Unlock() |
| 81 | return m.store != nil |
| 82 | } |
| 83 | |
| 84 | // beginWithObserver opens a checkpoint and updates the mutation observer's |
| 85 | // ownership turn for subsequent captures. |
| 86 | func (m *checkpointManager) beginWithObserver(input string, msgIndex int, obs *checkpoint.MutationObserver) (int, *checkpoint.Store, bool) { |
| 87 | m.mu.Lock() |
| 88 | store := m.store |
| 89 | if store == nil { |
| 90 | m.mu.Unlock() |
| 91 | return 0, nil, false |
| 92 | } |
| 93 | turn := m.turn |
| 94 | m.turn++ |
| 95 | m.bound[turn] = msgIndex |
| 96 | m.mu.Unlock() |
| 97 | if obs != nil { |
| 98 | obs.NoteCrossTurnBackgroundWriter(turn) |
| 99 | obs.SetOwnershipTurn(turn) |
| 100 | } |
| 101 | store.Begin(turn, input, msgIndex) |
| 102 | return turn, store, true |
| 103 | } |
| 104 | |
| 105 | type guardedTurnCheckpoint struct { |
| 106 | session *agent.Session |
| 107 | store *checkpoint.Store |
| 108 | turn int |
| 109 | messageIndex int |
| 110 | openedAt int64 |
| 111 | } |
| 112 | |
| 113 | type guardedTurnCompletion struct { |
| 114 | checkpoint *guardedTurnCheckpoint |
| 115 | } |
| 116 | |
| 117 | type guardedTurnCompletionKey struct{} |
| 118 | |
| 119 | func withGuardedTurnCompletion(ctx context.Context) (context.Context, *guardedTurnCompletion) { |
| 120 | completion := &guardedTurnCompletion{} |
| 121 | return context.WithValue(ctx, guardedTurnCompletionKey{}, completion), completion |
| 122 | } |
| 123 | |
| 124 | // beginCheckpoint opens a rewind checkpoint before the visible user message is |
| 125 | // appended. Guarded turns retain the exact boundary so TurnDone can identify |
| 126 | // the corresponding optimistic frontend item without positional guessing. |
| 127 | func (c *Controller) beginCheckpoint(ctx context.Context, input string) { |
| 128 | if c.executor == nil || c.executor.Session() == nil { |
| 129 | return |
| 130 | } |
| 131 | session := c.executor.Session() |
| 132 | messageIndex := session.Len() |
| 133 | openedAt := time.Now().UnixMilli() |
| 134 | atomic.AddInt64(&c.sessionRevision, 1) |
| 135 | turn, store, ok := c.checkpoints.beginWithObserver(input, messageIndex, c.mutationObserver) |
| 136 | if ok { |
| 137 | if completion, _ := ctx.Value(guardedTurnCompletionKey{}).(*guardedTurnCompletion); completion != nil { |
| 138 | completion.checkpoint = &guardedTurnCheckpoint{ |
| 139 | session: session, store: store, turn: turn, messageIndex: messageIndex, openedAt: openedAt, |
| 140 | } |
| 141 | } |
| 142 | } |
| 143 | // User-visible turn start records an irreversible message-send receipt so |
| 144 | // recovery never claims a clean rollback of an already-committed prompt. |
| 145 | // Keep this owner bookkeeping even when checkpoints are disabled. |
| 146 | gen := c.RuntimeGeneration() |
| 147 | if gen == 0 { |
| 148 | gen = c.RuntimeOwner().Gate.Published() |
| 149 | } |
| 150 | msgID := fmt.Sprintf("turn-%d-%d", gen, atomic.LoadInt64(&c.sessionRevision)) |
| 151 | // Dedup: a retried turn with the same revision must not double-record. |
| 152 | owner := c.RuntimeOwner() |
| 153 | owner.RecordMessageSentOnce(gen, msgID, "control") |
| 154 | d := owner.DecideResume(gen) |
| 155 | c.mu.Lock() |
| 156 | c.lastResumeDecision = d |
| 157 | c.mu.Unlock() |
| 158 | } |
| 159 | |
| 160 | // validatedCheckpointTurn returns the checkpoint only while its original |
| 161 | // boundary still names the real user message committed by this guarded turn. |
| 162 | // Stale or synthetic candidates fail closed rather than being relocated. |
| 163 | func (c *Controller) validatedCheckpointTurn(completion *guardedTurnCompletion) *int { |
| 164 | if completion == nil || completion.checkpoint == nil || c.executor == nil { |
| 165 | return nil |
| 166 | } |
| 167 | candidate := completion.checkpoint |
| 168 | if c.executor.Session() != candidate.session { |
| 169 | return nil |
| 170 | } |
| 171 | if !c.checkpoints.matchesBoundary(candidate.store, candidate.turn, candidate.messageIndex) { |
| 172 | return nil |
| 173 | } |
| 174 | messages := c.terminationMessages() |
| 175 | if candidate.messageIndex < 0 || candidate.messageIndex >= len(messages) { |
| 176 | return nil |
| 177 | } |
| 178 | message := messages[candidate.messageIndex] |
| 179 | if message.Role != provider.RoleUser || message.LocalOnly || |
| 180 | !agent.IsUserAuthoredTurnMessage(message) || |
| 181 | (message.CreatedAt > 0 && candidate.openedAt > 0 && message.CreatedAt < candidate.openedAt) { |
| 182 | return nil |
| 183 | } |
| 184 | turn := candidate.turn |
| 185 | return &turn |
| 186 | } |
| 187 | |
| 188 | func (m *checkpointManager) matchesBoundary(store *checkpoint.Store, turn, messageIndex int) bool { |
| 189 | m.mu.Lock() |
| 190 | defer m.mu.Unlock() |
| 191 | boundary, ok := m.bound[turn] |
| 192 | return ok && m.store == store && boundary == messageIndex |
| 193 | } |
| 194 | |
| 195 | // turnsByMessageIndex returns message-log index -> checkpoint turn over live |
| 196 | // boundaries. The desktop transcript uses this authoritative map instead of |
| 197 | // recounting visible user bubbles, which can diverge when synthetic user-role |
| 198 | // messages are hidden from the UI. |
| 199 | func (m *checkpointManager) turnsByMessageIndex() map[int]int { |
| 200 | m.mu.Lock() |
| 201 | defer m.mu.Unlock() |
| 202 | out := make(map[int]int, len(m.bound)) |
| 203 | for turn, index := range m.bound { |
| 204 | if existing, ok := out[index]; ok && existing < turn { |
| 205 | continue |
| 206 | } |
| 207 | out[index] = turn |
| 208 | } |
| 209 | return out |
| 210 | } |
| 211 | |
| 212 | // boundary returns the recorded turn-start message index, if any. |
| 213 | func (m *checkpointManager) boundary(turn int) (int, bool) { |
| 214 | m.mu.Lock() |
| 215 | defer m.mu.Unlock() |
| 216 | b, ok := m.bound[turn] |
| 217 | return b, ok |
| 218 | } |
| 219 | |
| 220 | // list returns the checkpoint metadata (nil when disabled). |
| 221 | func (m *checkpointManager) list() []checkpoint.Meta { |
| 222 | m.mu.Lock() |
| 223 | store := m.store |
| 224 | m.mu.Unlock() |
| 225 | if store == nil { |
| 226 | return nil |
| 227 | } |
| 228 | return store.List() |
| 229 | } |
| 230 | |
| 231 | func (m *checkpointManager) fileState(path string) (checkpoint.FileState, bool) { |
| 232 | m.mu.Lock() |
| 233 | store := m.store |
| 234 | m.mu.Unlock() |
| 235 | if store == nil { |
| 236 | return checkpoint.FileState{}, false |
| 237 | } |
| 238 | return store.FileState(path) |
| 239 | } |
| 240 | |
| 241 | // CheckpointTurnChanges is read-only and never computes from the current tree. |
| 242 | func (c *Controller) CheckpointTurnChanges(turn int) *checkpoint.TurnChanges { |
| 243 | return c.checkpoints.storeRef().TurnChanges(turn) |
| 244 | } |
| 245 | |
| 246 | // snapshot records a pre-edit file change into the open checkpoint — the |
| 247 | // executor's pre-edit hook. No-op when disabled. |
| 248 | func (m *checkpointManager) snapshot(ch diff.Change) { |
| 249 | m.mu.Lock() |
| 250 | store := m.store |
| 251 | m.mu.Unlock() |
| 252 | if store != nil { |
| 253 | store.Snapshot(ch) |
| 254 | } |
| 255 | } |
| 256 | |
| 257 | // truncateFrom renumbers future turns from `turn` and drops every boundary at or |
| 258 | // after it — the conversation-rewind renumber after the message log is cut back. |
| 259 | func (m *checkpointManager) truncateFrom(turn int) error { |
| 260 | m.mu.Lock() |
| 261 | store := m.store |
| 262 | m.mu.Unlock() |
| 263 | if store != nil { |
| 264 | if err := store.TruncateFrom(turn); err != nil { |
| 265 | return err |
| 266 | } |
| 267 | } |
| 268 | m.mu.Lock() |
| 269 | m.turn = turn |
| 270 | for k := range m.bound { |
| 271 | if k >= turn { |
| 272 | delete(m.bound, k) |
| 273 | } |
| 274 | } |
| 275 | m.mu.Unlock() |
| 276 | return nil |
| 277 | } |
| 278 | |
| 279 | // storeRef returns the live store pointer without holding mu across caller work. |
| 280 | func (m *checkpointManager) storeRef() *checkpoint.Store { |
| 281 | m.mu.Lock() |
| 282 | defer m.mu.Unlock() |
| 283 | return m.store |
| 284 | } |
| 285 |