| 1 | package control |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "crypto/sha256" |
| 6 | "encoding/json" |
| 7 | "errors" |
| 8 | "fmt" |
| 9 | "reflect" |
| 10 | "sort" |
| 11 | "time" |
| 12 | |
| 13 | "reasonix/internal/agent" |
| 14 | "reasonix/internal/event" |
| 15 | "reasonix/internal/provider" |
| 16 | "reasonix/internal/session" |
| 17 | ) |
| 18 | |
| 19 | const terminationFlushTimeout = 15 * time.Second |
| 20 | |
| 21 | var errTerminationDurability = errors.New("turn termination was not made durable") |
| 22 | |
| 23 | type executionTokenKey struct{} |
| 24 | |
| 25 | func (c *Controller) failTermination(err error) { |
| 26 | if err == nil { |
| 27 | return |
| 28 | } |
| 29 | c.mu.Lock() |
| 30 | c.enterRecoveryLocked("termination_commit_failed") |
| 31 | c.mu.Unlock() |
| 32 | c.disarmGoalLifecycle("persistence-error") |
| 33 | c.failTurnEventLedger(fmt.Errorf("%w: %w", errTerminationDurability, err)) |
| 34 | } |
| 35 | |
| 36 | // messageCommitAllowedLocked is checked after acquiring commitMu, so a |
| 37 | // recorder waiting behind a terminal cannot revive its closed turn. |
| 38 | func (c *Controller) messageCommitAllowedLocked(ctx context.Context, store *session.Session) bool { |
| 39 | if !c.sessionEventCommitAllowed() { |
| 40 | return false |
| 41 | } |
| 42 | token, turnID, _ := c.currentTurnToken() |
| 43 | if ctx != nil { |
| 44 | if origin, ok := ctx.Value(executionTokenKey{}).(uint64); ok && origin != token { |
| 45 | return false |
| 46 | } |
| 47 | } |
| 48 | if turnID != "" && c.turnEvents.finalizedTurn == turnID { |
| 49 | return false |
| 50 | } |
| 51 | projection := store.StateSnapshot().Projection |
| 52 | return projection.Recovery == nil || projection.Recovery.State != "recovery_required" |
| 53 | } |
| 54 | |
| 55 | type terminationBoundary struct { |
| 56 | token uint64 |
| 57 | prefix map[string]bool |
| 58 | fallback provider.Message |
| 59 | preserveUser bool |
| 60 | } |
| 61 | |
| 62 | func (c *Controller) noteTerminationBoundary(fallback provider.Message, preserveUser bool) { |
| 63 | if c.executor == nil { |
| 64 | return |
| 65 | } |
| 66 | token, _, _ := c.currentTurnToken() |
| 67 | b := &terminationBoundary{token: token, prefix: map[string]bool{}, fallback: fallback, preserveUser: preserveUser} |
| 68 | for _, message := range c.executor.Session().Snapshot() { |
| 69 | b.prefix[message.ID] = true |
| 70 | } |
| 71 | c.turnEvents.commitMu.Lock() |
| 72 | c.turnEvents.terminationBoundary = b |
| 73 | c.turnEvents.commitMu.Unlock() |
| 74 | } |
| 75 | |
| 76 | // terminationMessages is for terminal metadata only. UI and execution readers |
| 77 | // continue to observe the accepted projection until the terminal batch lands. |
| 78 | func (c *Controller) terminationMessages() []provider.Message { |
| 79 | c.turnEvents.commitMu.Lock() |
| 80 | defer c.turnEvents.commitMu.Unlock() |
| 81 | if p := c.turnEvents.pendingTermination; p != nil { |
| 82 | return append([]provider.Message(nil), p.Messages...) |
| 83 | } |
| 84 | return c.executor.Session().Snapshot() |
| 85 | } |
| 86 | |
| 87 | // TerminationPlan is an owned snapshot of one explicit cleanup. Retractions |
| 88 | // come only from its input workset, never from a diff against durable history. |
| 89 | type TerminationPlan struct { |
| 90 | SessionID string |
| 91 | TurnID string |
| 92 | Generation uint64 |
| 93 | Token uint64 |
| 94 | Messages []provider.Message |
| 95 | Events []session.Event |
| 96 | } |
| 97 | |
| 98 | func buildTerminationPlan(before, after []provider.Message) (*TerminationPlan, error) { |
| 99 | // Detach nested tool/reasoning records from the worker's mutable cache. |
| 100 | data, err := json.Marshal(after) |
| 101 | if err != nil { |
| 102 | return nil, err |
| 103 | } |
| 104 | if err := json.Unmarshal(data, &after); err != nil { |
| 105 | return nil, err |
| 106 | } |
| 107 | p := &TerminationPlan{Messages: after} |
| 108 | old := make(map[string]provider.Message, len(before)) |
| 109 | for _, message := range before { |
| 110 | old[message.ID] = message |
| 111 | } |
| 112 | kept := make(map[string]bool, len(after)) |
| 113 | for _, message := range after { |
| 114 | if message.ID == "" { |
| 115 | return nil, fmt.Errorf("termination message has no stable identity") |
| 116 | } |
| 117 | kept[message.ID] = true |
| 118 | if previous, ok := old[message.ID]; ok && reflect.DeepEqual(previous, message) { |
| 119 | continue |
| 120 | } |
| 121 | payload, err := json.Marshal(map[string]any{"message": message}) |
| 122 | if err != nil { |
| 123 | return nil, err |
| 124 | } |
| 125 | p.Events = append(p.Events, session.Event{Kind: "message/upsert", Payload: payload}) |
| 126 | } |
| 127 | var removed []string |
| 128 | for id := range old { |
| 129 | if id != "" && !kept[id] { |
| 130 | removed = append(removed, id) |
| 131 | } |
| 132 | } |
| 133 | if len(removed) > 0 { |
| 134 | sort.Strings(removed) |
| 135 | payload, _ := json.Marshal(map[string]any{"messageIds": removed, "reason": "interrupted-turn-cleanup"}) |
| 136 | p.Events = append(p.Events, session.Event{Kind: "message/retract", Payload: payload, Required: true}) |
| 137 | } |
| 138 | model := provider.ModelMessages(after) |
| 139 | if model == nil { |
| 140 | model = []provider.Message{} |
| 141 | } |
| 142 | payload, err := json.Marshal(map[string]any{"messages": model, "reason": "interrupted-turn-cleanup"}) |
| 143 | if err != nil { |
| 144 | return nil, err |
| 145 | } |
| 146 | p.Events = append(p.Events, session.Event{Kind: "model/context-replace", Payload: payload}) |
| 147 | return p, nil |
| 148 | } |
| 149 | |
| 150 | func (c *Controller) replaceSessionAfterCancelFrom(before, after []provider.Message) { |
| 151 | c.replaceSessionAfterCancelFromScoped(before, after, false) |
| 152 | } |
| 153 | |
| 154 | func (c *Controller) noteCommittedMessagesLocked(events []session.Event) { |
| 155 | for _, e := range events { |
| 156 | if e.Kind == "turn/start" { |
| 157 | c.turnEvents.turnMessageIDs = map[string]bool{} |
| 158 | } |
| 159 | if e.Kind == "message/complete" { |
| 160 | var payload struct { |
| 161 | Message provider.Message `json:"message"` |
| 162 | } |
| 163 | if json.Unmarshal(e.Payload, &payload) == nil && payload.Message.ID != "" { |
| 164 | if c.turnEvents.turnMessageIDs == nil { |
| 165 | c.turnEvents.turnMessageIDs = map[string]bool{} |
| 166 | } |
| 167 | c.turnEvents.turnMessageIDs[payload.Message.ID] = true |
| 168 | } |
| 169 | } |
| 170 | } |
| 171 | } |
| 172 | |
| 173 | func (c *Controller) retractMissingTurnMessagesLocked(p *TerminationPlan) { |
| 174 | kept := map[string]bool{} |
| 175 | for _, m := range p.Messages { |
| 176 | kept[m.ID] = true |
| 177 | } |
| 178 | var ids []string |
| 179 | for id := range c.turnEvents.turnMessageIDs { |
| 180 | if !kept[id] { |
| 181 | ids = append(ids, id) |
| 182 | } |
| 183 | } |
| 184 | if len(ids) == 0 { |
| 185 | return |
| 186 | } |
| 187 | sort.Strings(ids) |
| 188 | payload, _ := json.Marshal(map[string]any{"messageIds": ids, "reason": "synthetic-turn-interrupted"}) |
| 189 | p.Events = append([]session.Event{{Kind: "message/retract", Required: true, Payload: payload}}, p.Events...) |
| 190 | } |
| 191 | |
| 192 | func (c *Controller) replaceSessionAfterCancelFromScoped(before, after []provider.Message, retractTurn bool) { |
| 193 | if c.executor == nil { |
| 194 | return |
| 195 | } |
| 196 | after = agent.PlanRetainedToolRecords(before, after) |
| 197 | for i := range after { |
| 198 | if after[i].ID == "" { |
| 199 | after[i].ID = agent.NewMessageID() |
| 200 | } |
| 201 | } |
| 202 | p, err := buildTerminationPlan(before, after) |
| 203 | if err != nil { |
| 204 | c.failTermination(err) |
| 205 | return |
| 206 | } |
| 207 | c.snapshotMu.Lock() |
| 208 | store := c.sessionEventStore() |
| 209 | if store == nil { |
| 210 | c.replaceLegacySessionAfterCancelLocked(after) |
| 211 | c.snapshotMu.Unlock() |
| 212 | return |
| 213 | } |
| 214 | p.SessionID = store.ID() |
| 215 | p.Token, p.TurnID, _ = c.currentTurnToken() |
| 216 | p.Generation = c.ExecutionGeneration() |
| 217 | c.turnEvents.commitMu.Lock() |
| 218 | if !c.sessionEventCommitAllowed() || (p.TurnID != "" && c.turnEvents.finalizedTurn == p.TurnID) { |
| 219 | c.turnEvents.commitMu.Unlock() |
| 220 | c.snapshotMu.Unlock() |
| 221 | return |
| 222 | } |
| 223 | projection := store.StateSnapshot().Projection |
| 224 | if projection.Recovery != nil && projection.Recovery.State == "recovery_required" { |
| 225 | c.turnEvents.commitMu.Unlock() |
| 226 | c.snapshotMu.Unlock() |
| 227 | return |
| 228 | } |
| 229 | if p.TurnID != "" && projection.TurnID == p.TurnID { |
| 230 | if retractTurn { |
| 231 | c.retractMissingTurnMessagesLocked(p) |
| 232 | } |
| 233 | // Compatibility runners may update the executor before their recorder |
| 234 | // commits. Preserve these unchanged messages explicitly as well. |
| 235 | known := map[string]bool{} |
| 236 | for _, m := range store.ExecutionSnapshot().Projection.ModelMessages { |
| 237 | known[m.ID] = true |
| 238 | } |
| 239 | var unrecorded []session.Event |
| 240 | for _, m := range before { |
| 241 | if known[m.ID] { |
| 242 | continue |
| 243 | } |
| 244 | for _, retained := range after { |
| 245 | if retained.ID == m.ID && reflect.DeepEqual(m, retained) { |
| 246 | payload, _ := json.Marshal(map[string]any{"message": retained}) |
| 247 | unrecorded = append(unrecorded, session.Event{Kind: "message/upsert", Payload: payload}) |
| 248 | } |
| 249 | } |
| 250 | } |
| 251 | p.Events = append(unrecorded, p.Events...) |
| 252 | c.turnEvents.pendingTermination = p |
| 253 | c.turnEvents.commitMu.Unlock() |
| 254 | c.snapshotMu.Unlock() |
| 255 | return |
| 256 | } |
| 257 | // Resume cleanup has no live turn to close. Its deterministic repair |
| 258 | // operation never adds a second terminal record. |
| 259 | data, _ := json.Marshal(p.Events) |
| 260 | digest := sha256.Sum256(data) |
| 261 | ctx, cancel := context.WithTimeout(context.Background(), terminationFlushTimeout) |
| 262 | _, err = c.appendSessionBatch(ctx, store, session.Batch{OperationID: fmt.Sprintf("turn-repair:%x", digest), Events: p.Events}) |
| 263 | if err == nil { |
| 264 | c.executor.Session().Replace(p.Messages) |
| 265 | _, err = store.Flush(ctx) |
| 266 | } |
| 267 | c.turnEvents.commitMu.Unlock() |
| 268 | if err == nil && !c.sessionEngineEnabled() { |
| 269 | c.replaceLegacySessionAfterCancelLocked(p.Messages) |
| 270 | } |
| 271 | c.snapshotMu.Unlock() |
| 272 | cancel() |
| 273 | if err != nil { |
| 274 | c.failTermination(err) |
| 275 | } |
| 276 | } |
| 277 | |
| 278 | // appendTerminationLocked runs under snapshotMu -> commitMu. The runtime gate |
| 279 | // revalidates ownership and turn identity at acceptance, after preparation. |
| 280 | func (c *Controller) appendTerminationLocked(ctx context.Context, e event.Event, store *session.Session, terminal []session.Event) error { |
| 281 | if c.turnEvents.finalizedTurn == e.TurnID { |
| 282 | _, err := store.Flush(ctx) |
| 283 | if err != nil { |
| 284 | return fmt.Errorf("%w: %w", errTerminationDurability, err) |
| 285 | } |
| 286 | return nil |
| 287 | } |
| 288 | snapshot := store.StateSnapshot() |
| 289 | if e.TurnID == "" || snapshot.Projection.TurnID != e.TurnID { |
| 290 | return session.ErrStaleExecution |
| 291 | } |
| 292 | p := c.turnEvents.pendingTermination |
| 293 | if e.Recovery != nil && e.Recovery.State == "recovery_required" { |
| 294 | var err error |
| 295 | p, err = c.watchdogTerminationPlanLocked(store, e.TurnID) |
| 296 | if err != nil { |
| 297 | return err |
| 298 | } |
| 299 | } |
| 300 | if p != nil { |
| 301 | token, turnID, _ := c.currentTurnToken() |
| 302 | if p.SessionID != store.ID() || p.Generation != c.ExecutionGeneration() || p.Token != token || p.TurnID != turnID { |
| 303 | return session.ErrStaleExecution |
| 304 | } |
| 305 | terminal = append(append([]session.Event(nil), p.Events...), terminal...) |
| 306 | } |
| 307 | batch := session.Batch{OperationID: "turn-finalize:" + e.TurnID, TurnID: e.TurnID, Events: terminal} |
| 308 | prepared, err := store.PrepareBatchContext(ctx, batch.OperationID, batch) |
| 309 | if err != nil { |
| 310 | if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) || errors.Is(err, session.ErrOperationConflict) { |
| 311 | return fmt.Errorf("%w: %w", errTerminationDurability, err) |
| 312 | } |
| 313 | return err |
| 314 | } |
| 315 | _, runtime, exclusive := c.v3Binding() |
| 316 | if exclusive && runtime != nil { |
| 317 | _, err = runtime.CommitPreparedForTurn(c.ExecutionGeneration(), e.TurnID, prepared) |
| 318 | } else { |
| 319 | _, err = store.CommitPrepared(prepared) |
| 320 | } |
| 321 | if err != nil { |
| 322 | return err |
| 323 | } |
| 324 | c.turnEvents.finalizedTurn = e.TurnID |
| 325 | c.turnEvents.pendingTermination = nil |
| 326 | if p != nil && c.executor != nil { |
| 327 | c.executor.Session().Replace(p.Messages) |
| 328 | } |
| 329 | _, err = store.Flush(ctx) |
| 330 | if err != nil { |
| 331 | return fmt.Errorf("%w: %w", errTerminationDurability, err) |
| 332 | } |
| 333 | return err |
| 334 | } |
| 335 | |
| 336 | func (c *Controller) watchdogTerminationPlanLocked(store *session.Session, turnID string) (*TerminationPlan, error) { |
| 337 | // Seal accepted work only: an uncooperative worker may mutate its cache. |
| 338 | before := store.ExecutionSnapshot().Projection.ModelMessages |
| 339 | var next []provider.Message |
| 340 | b := c.turnEvents.terminationBoundary |
| 341 | token, _, _ := c.currentTurnToken() |
| 342 | if b != nil && b.token == token { |
| 343 | for _, message := range before { |
| 344 | if b.prefix[message.ID] || agent.IsCompactionSummary(message) { |
| 345 | next = append(next, message) |
| 346 | } |
| 347 | } |
| 348 | if b.preserveUser && b.fallback.ID != "" { |
| 349 | user := b.fallback |
| 350 | for _, message := range before { |
| 351 | if message.ID == user.ID { |
| 352 | user = message |
| 353 | break |
| 354 | } |
| 355 | } |
| 356 | user.Content = StripComposePrefixes(user.Content) |
| 357 | next = append(next, user) |
| 358 | } |
| 359 | } else { |
| 360 | next = append([]provider.Message(nil), before...) |
| 361 | } |
| 362 | p, err := buildTerminationPlan(before, next) |
| 363 | if err != nil { |
| 364 | return nil, err |
| 365 | } |
| 366 | p.SessionID, p.TurnID, p.Generation, p.Token = store.ID(), turnID, c.ExecutionGeneration(), token |
| 367 | if b != nil && !b.preserveUser { |
| 368 | c.retractMissingTurnMessagesLocked(p) |
| 369 | } |
| 370 | return p, nil |
| 371 | } |
| 372 |