返回 DeepSeek-Reasonix
session_dag_replay.go
根目录 / internal / agent / session_dag_replay.go
1 package agent
2
3 import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "io"
9 "log/slog"
10 "os"
11 "slices"
12 "time"
13
14 "reasonix/internal/provider"
15 )
16
17 const (
18 // A half-written last line is usually another writer mid-append: re-stat a
19 // few times before calling the tail torn.
20 sessionDAGTornRetries = 3
21 sessionDAGTornRetryDelay = 50 * time.Millisecond
22 )
23
24 type sessionDAGNode struct {
25 id, parent, head, writer, turn string
26 offset int64
27 at time.Time
28 msg provider.Message
29 digest string
30 }
31
32 type sessionDAGTurn struct {
33 turn, leaf, writer string
34 preserveUser bool
35 at time.Time
36 }
37
38 type sessionDAGCompaction struct {
39 coveredLeaf string
40 coveredCount int
41 prefixHash string
42 at time.Time
43 }
44
45 type sessionDAGHead struct {
46 id, kind, name, parentHead, forkFrom, writer, leaf string
47 system *provider.Message
48 createdAt, lastActivity time.Time
49 lastOffset int64
50 retired bool
51 compaction *sessionDAGCompaction
52 openTurn *sessionDAGTurn
53 }
54
55 type sessionDAGWriter struct {
56 id, hostname string
57 pid int
58 leaseGeneration uint64
59 lastActivity time.Time
60 }
61
62 // sessionDAGState is the replayed graph: every node, every head, and the
63 // overlays (patches, redactions) that materialize applies on read.
64 type sessionDAGState struct {
65 path string
66 generation int64
67 upgradedFrom int
68 nodes map[string]*sessionDAGNode
69 heads map[string]*sessionDAGHead
70 headOrder []string
71 patches map[string]provider.Message
72 redactions map[string]provider.Message
73 writers map[string]*sessionDAGWriter
74 selected string
75 orphans []string
76 records int
77 collectionItems int
78 size int64
79 lastGoodEnd int64
80 damaged bool
81 holes int // unreadable lines skipped between good entries
82 }
83
84 func newSessionDAGState(path string) *sessionDAGState {
85 st := &sessionDAGState{
86 path: path,
87 nodes: map[string]*sessionDAGNode{},
88 heads: map[string]*sessionDAGHead{},
89 patches: map[string]provider.Message{},
90 redactions: map[string]provider.Message{},
91 writers: map[string]*sessionDAGWriter{},
92 }
93 st.heads[SessionMainHead] = &sessionDAGHead{id: SessionMainHead, kind: HeadKindMain}
94 st.headOrder = []string{SessionMainHead}
95 return st
96 }
97
98 // replaySessionDAG decodes a schema-2 log. Decoding stops at the first entry
99 // that fails to parse (damaged=true, lastGoodEnd set); an unsupported schema
100 // or entry type is a hard error because a newer writer owns the log.
101 func replaySessionDAG(ctx context.Context, path string, limits sessionReplayLimits) (*sessionDAGState, error) {
102 st := newSessionDAGState(path)
103 if err := st.replayFrom(ctx, 0, limits); err != nil {
104 return st, err
105 }
106 for attempt := 0; st.damaged && attempt < sessionDAGTornRetries; attempt++ {
107 time.Sleep(sessionDAGTornRetryDelay)
108 info, err := os.Stat(path)
109 if err != nil || info.Size() <= st.size {
110 break
111 }
112 st.damaged = false
113 if err := st.replayFrom(ctx, st.lastGoodEnd, limits); err != nil {
114 return st, err
115 }
116 }
117 return st, nil
118 }
119
120 // replayFrom applies every entry from byte offset from to the end of the log.
121 // Callers use it for the initial pass and for incremental tail reads.
122 func (st *sessionDAGState) replayFrom(ctx context.Context, from int64, limits sessionReplayLimits) error {
123 if err := ctx.Err(); err != nil {
124 return err
125 }
126 f, err := os.Open(st.path)
127 if err != nil {
128 return err
129 }
130 defer f.Close()
131 info, err := f.Stat()
132 if err != nil {
133 return err
134 }
135 st.size = info.Size()
136 if st.size > limits.maxBytes {
137 return sessionReplayLimitError(st.path, "encoded_bytes", st.size, limits.maxBytes)
138 }
139 if from > 0 {
140 if _, err := f.Seek(from, io.SeekStart); err != nil {
141 return err
142 }
143 }
144 limited := &io.LimitedReader{R: &contextReader{ctx: ctx, reader: f}, N: limits.maxBytes + 1 - from}
145 dec := json.NewDecoder(limited)
146 for {
147 if err := ctx.Err(); err != nil {
148 return err
149 }
150 var e sessionDAGEntry
151 if err := dec.Decode(&e); err != nil {
152 if ctxErr := ctx.Err(); ctxErr != nil {
153 return ctxErr
154 }
155 if limited.N == 0 {
156 return sessionReplayLimitError(st.path, "encoded_bytes", limits.maxBytes+1, limits.maxBytes)
157 }
158 if errors.Is(err, io.EOF) {
159 return nil
160 }
161 return st.resumePastTornLine(ctx, limits)
162 }
163 if e.SchemaVersion != sessionDAGSchemaVersion {
164 return fmt.Errorf("decode session event log %s: unsupported schema version %d", st.path, e.SchemaVersion)
165 }
166 if st.records >= limits.maxRecords {
167 return sessionReplayLimitError(st.path, "event_records", int64(st.records+1), int64(limits.maxRecords))
168 }
169 offset := from + dec.InputOffset()
170 ok, err := st.apply(ctx, e, offset, limits)
171 if err != nil {
172 return err
173 }
174 if !ok {
175 st.damaged = true
176 return nil
177 }
178 st.records++
179 st.lastGoodEnd = offset
180 }
181 }
182
183 // apply folds one entry into the state. ok=false marks an undecodable entry
184 // (the replay stops there); err is reserved for hard errors and budget limits.
185 func (st *sessionDAGState) apply(ctx context.Context, e sessionDAGEntry, offset int64, limits sessionReplayLimits) (bool, error) {
186 switch e.Type {
187 case sessionDAGTypeLog:
188 st.generation = e.Generation
189 st.upgradedFrom = e.UpgradedFrom
190 case sessionDAGTypeMessage:
191 return st.applyMessage(ctx, e, offset, limits)
192 case sessionDAGTypePatch, sessionDAGTypeSystem, sessionDAGTypeRedact, sessionDAGTypeWriter:
193 return st.applyOverlay(ctx, e, limits)
194 case sessionDAGTypeFork, sessionDAGTypeRewind, sessionDAGTypeSelect, sessionDAGTypeRename, sessionDAGTypeRetire,
195 sessionDAGTypeTurnBegin, sessionDAGTypeTurnEnd, sessionDAGTypeCompaction:
196 return st.applyHeadMarker(e, offset), nil
197 case sessionDAGTypeCheckpoint:
198 default:
199 return false, fmt.Errorf("decode session event log %s: unsupported entry type %q", st.path, e.Type)
200 }
201 return true, nil
202 }
203
204 func (st *sessionDAGState) applyMessage(ctx context.Context, e sessionDAGEntry, offset int64, limits sessionReplayLimits) (bool, error) {
205 if e.ID == "" {
206 return false, nil
207 }
208 m, ok, err := st.decodeOne(ctx, e.Msgs, limits)
209 if err != nil || !ok {
210 return ok, err
211 }
212 m.ID = e.ID
213 if _, dup := st.nodes[e.ID]; dup {
214 slog.Warn("session: ignoring duplicate message entry", "path", st.path, "id", e.ID, "offset", offset)
215 return true, nil
216 }
217 if e.Parent != "" {
218 if _, known := st.nodes[e.Parent]; !known {
219 st.orphans = append(st.orphans, e.ID)
220 }
221 }
222 st.nodes[e.ID] = &sessionDAGNode{
223 id: e.ID, parent: e.Parent, head: e.Head, writer: e.Writer, turn: e.Turn,
224 offset: offset, at: e.At, msg: m, digest: e.Digest,
225 }
226 h := st.headFor(e.Head, e.At)
227 h.leaf = e.ID
228 h.lastActivity, h.lastOffset = e.At, offset
229 if h.writer == "" {
230 h.writer = e.Writer
231 }
232 st.touchWriter(e.Writer, e.At)
233 return true, nil
234 }
235
236 // applyOverlay handles the entries that change how messages read without
237 // moving any head: patch replacements, system overrides, redactions, writer
238 // identity.
239 func (st *sessionDAGState) applyOverlay(ctx context.Context, e sessionDAGEntry, limits sessionReplayLimits) (bool, error) {
240 switch e.Type {
241 case sessionDAGTypePatch:
242 if _, known := st.nodes[e.Target]; !known {
243 return true, nil
244 }
245 m, ok, err := st.decodeOne(ctx, e.Msgs, limits)
246 if err != nil || !ok {
247 return ok, err
248 }
249 m.ID = e.Target
250 st.patches[e.Target] = m
251 case sessionDAGTypeSystem:
252 m, ok, err := st.decodeOne(ctx, e.Msgs, limits)
253 if err != nil || !ok {
254 return ok, err
255 }
256 st.headFor(e.Head, e.At).system = &m
257 case sessionDAGTypeRedact:
258 for id, raw := range e.Targets {
259 m, ok, err := st.decodeOne(ctx, raw, limits)
260 if err != nil || !ok {
261 return ok, err
262 }
263 m.ID = id
264 st.redactions[id] = m
265 }
266 case sessionDAGTypeWriter:
267 w := st.touchWriter(e.Writer, e.At)
268 w.pid, w.hostname, w.leaseGeneration = e.PID, e.Hostname, e.LeaseGeneration
269 }
270 return true, nil
271 }
272
273 // applyHeadMarker handles the entries that create or move heads. Only a fork
274 // without a new head id is undecodable.
275 func (st *sessionDAGState) applyHeadMarker(e sessionDAGEntry, offset int64) bool {
276 switch e.Type {
277 case sessionDAGTypeFork:
278 if e.NewHead == "" {
279 return false
280 }
281 if _, exists := st.heads[e.NewHead]; exists {
282 return true
283 }
284 parent := st.headFor(e.Head, e.At)
285 kind := e.Kind
286 if kind == "" {
287 kind = HeadKindFork
288 }
289 st.heads[e.NewHead] = &sessionDAGHead{
290 id: e.NewHead, kind: kind, name: e.Name, parentHead: parent.id, forkFrom: e.From,
291 writer: e.Writer, leaf: e.From, system: parent.system,
292 createdAt: e.At, lastActivity: e.At, lastOffset: offset,
293 }
294 st.headOrder = append(st.headOrder, e.NewHead)
295 st.touchWriter(e.Writer, e.At)
296 case sessionDAGTypeRewind:
297 h := st.headFor(e.Head, e.At)
298 h.leaf = e.To
299 h.lastActivity, h.lastOffset = e.At, offset
300 st.touchWriter(e.Writer, e.At)
301 case sessionDAGTypeSelect:
302 st.selected = e.Head
303 case sessionDAGTypeRename:
304 st.headFor(e.Head, e.At).name = e.Name
305 case sessionDAGTypeRetire:
306 st.headFor(e.Head, e.At).retired = true
307 case sessionDAGTypeTurnBegin:
308 st.headFor(e.Head, e.At).openTurn = &sessionDAGTurn{turn: e.Turn, leaf: e.Leaf, writer: e.Writer, preserveUser: e.PreserveUser, at: e.At}
309 case sessionDAGTypeTurnEnd:
310 h := st.headFor(e.Head, e.At)
311 if h.openTurn != nil && (e.Turn == "" || h.openTurn.turn == e.Turn) {
312 h.openTurn = nil
313 }
314 case sessionDAGTypeCompaction:
315 st.headFor(e.Head, e.At).compaction = &sessionDAGCompaction{
316 coveredLeaf: e.CoveredLeaf, coveredCount: e.CoveredCount, prefixHash: e.PrefixHash, at: e.At,
317 }
318 }
319 return true
320 }
321
322 // decodeOne decodes the single-message array carried by message, system, and
323 // redact entries through the same bounded decoder as schema-1 records.
324 func (st *sessionDAGState) decodeOne(ctx context.Context, raw json.RawMessage, limits sessionReplayLimits) (provider.Message, bool, error) {
325 msgs, items, err := decodeSessionEventMessages(ctx, st.path, raw, len(st.nodes), st.collectionItems, limits)
326 if err != nil {
327 if ctxErr := ctx.Err(); ctxErr != nil {
328 return provider.Message{}, false, ctxErr
329 }
330 if errors.Is(err, ErrSessionReplayLimitExceeded) {
331 return provider.Message{}, false, err
332 }
333 return provider.Message{}, false, nil
334 }
335 if len(msgs) != 1 {
336 return provider.Message{}, false, nil
337 }
338 st.collectionItems = items
339 return msgs[0], true, nil
340 }
341
342 // headFor resolves a head id, creating an undeclared head rather than dropping
343 // the entries that reference it; an empty id means the main head.
344 func (st *sessionDAGState) headFor(id string, at time.Time) *sessionDAGHead {
345 if id == "" {
346 id = SessionMainHead
347 }
348 h := st.heads[id]
349 if h == nil {
350 h = &sessionDAGHead{id: id, createdAt: at}
351 st.heads[id] = h
352 st.headOrder = append(st.headOrder, id)
353 }
354 if h.createdAt.IsZero() {
355 h.createdAt = at
356 }
357 return h
358 }
359
360 func (st *sessionDAGState) touchWriter(id string, at time.Time) *sessionDAGWriter {
361 w := st.writers[id]
362 if w == nil {
363 w = &sessionDAGWriter{id: id}
364 st.writers[id] = w
365 }
366 if at.After(w.lastActivity) {
367 w.lastActivity = at
368 }
369 return w
370 }
371
372 // chainIDs returns the message ids of a head from root to leaf. It stops at an
373 // orphan root (parent never seen) so a rotated-away prefix degrades to a
374 // shorter transcript instead of a failed load.
375 func (st *sessionDAGState) chainIDs(headID string) []string {
376 h := st.heads[headID]
377 if h == nil {
378 return nil
379 }
380 var ids []string
381 seen := map[string]struct{}{}
382 for id := h.leaf; id != ""; {
383 n := st.nodes[id]
384 if n == nil {
385 break
386 }
387 if _, cyc := seen[id]; cyc {
388 break
389 }
390 seen[id] = struct{}{}
391 ids = append(ids, id)
392 id = n.parent
393 }
394 slices.Reverse(ids)
395 return ids
396 }
397
398 // appliedMessage is one node with its latest patch and any redaction
399 // substituted; a patch is provider-equivalent to the original by contract, so
400 // the swap never touches provider-visible bytes.
401 func (st *sessionDAGState) appliedMessage(n *sessionDAGNode) provider.Message {
402 m := n.msg
403 if p, ok := st.patches[n.id]; ok {
404 m = p
405 }
406 if r, ok := st.redactions[n.id]; ok {
407 m = r
408 }
409 m.ID = n.id
410 return m
411 }
412
413 // materialize builds the transcript of one head: the parent chain with
414 // patches merged, redactions substituted, and the head's system override at
415 // position 0. times mirror msgs with each entry's append time.
416 func (st *sessionDAGState) materialize(headID string) ([]provider.Message, []time.Time) {
417 ids := st.chainIDs(headID)
418 msgs := make([]provider.Message, 0, len(ids)+1)
419 times := make([]time.Time, 0, len(ids)+1)
420 for _, id := range ids {
421 n := st.nodes[id]
422 msgs = append(msgs, st.appliedMessage(n))
423 times = append(times, n.at)
424 }
425 if h := st.heads[headID]; h != nil && h.system != nil {
426 sys := *h.system
427 if len(msgs) > 0 && msgs[0].Role == provider.RoleSystem {
428 sys.ID = msgs[0].ID
429 msgs[0] = sys
430 } else {
431 // The writer stamps the prepended system message's id on the marker
432 // so every reader materializes the same id.
433 msgs = append([]provider.Message{sys}, msgs...)
434 times = append([]time.Time{h.createdAt}, times...)
435 }
436 }
437 return msgs, times
438 }
439
439 lines GO