1package core
2
3import (
4 "bufio"
5 "crypto/sha256"
6 "encoding/hex"
7 "encoding/json"
8 "fmt"
9 "os"
10 "path/filepath"
11 "sort"
12 "strings"
13 "time"
14
15 "github.com/google/uuid"
16 "github.com/patriceckhart/zot/internal/provider"
17)
18
19// Session is a JSONL-backed conversation transcript tied to a cwd.
20type Session struct {
21 ID string
22 Path string
23 Meta SessionMeta
24 writer *os.File
25 buf *bufio.Writer
26
27 // freshFile is true when the file was created by NewSession (this
28 // process owns it) and false when OpenSession reopened an existing
29 // transcript. Used by Close() to delete the file if the run never
30 // appended any messages — prevents a flood of empty session files
31 // from sessions the user opens then exits without prompting.
32 freshFile bool
33
34 // messagesAppended counts AppendMessage calls. Combined with
35 // freshFile it tells Close() whether the session left any content
36 // worth keeping.
37 messagesAppended int
38}
39
40// SessionMeta is written as the first line of every session file.
41type SessionMeta struct {
42 ID string `json:"id"`
43 CWD string `json:"cwd"`
44 Model string `json:"model"`
45 Provider string `json:"provider"`
46 Started time.Time `json:"started"`
47 Version string `json:"version"`
48 Title string `json:"title,omitempty"`
49
50 // Parent is the ID of the session this one was forked from, or
51 // empty for top-level sessions. The tree picker walks parents
52 // upward and sibling files (same cwd dir, same parent ID)
53 // laterally to render the branch topology.
54 Parent string `json:"parent,omitempty"`
55 // ForkPoint is the 0-indexed message position within the parent
56 // transcript where this branch diverges. Messages 0..ForkPoint-1
57 // are copied from the parent verbatim; the user's next turn on
58 // the child session continues from there.
59 ForkPoint int `json:"fork_point,omitempty"`
60}
61
62// sessionLine is the on-disk row type. Message is kept as a raw
63// JSON message on reads (because Content is an interface slice that
64// the default unmarshaler cannot reconstruct); it is written with a
65// regular provider.Message value.
66type sessionLine struct {
67 Type string `json:"type"`
68 Meta *SessionMeta `json:"meta,omitempty"`
69 Message *provider.Message `json:"message,omitempty"`
70 Usage *provider.Usage `json:"usage,omitempty"`
71 Cumulative *provider.Usage `json:"cumulative,omitempty"`
72}
73
74type sessionLineHead struct {
75 Type string `json:"type"`
76}
77
78// SessionsDir returns the per-cwd sessions directory under root.
79func SessionsDir(root, cwd string) string {
80 sum := sha256.Sum256([]byte(cwd))
81 short := hex.EncodeToString(sum[:8])
82 return filepath.Join(root, "sessions", short)
83}
84
85// NewSession creates and opens a new session file.
86func NewSession(root, cwd, providerName, model, version string) (*Session, error) {
87 dir := SessionsDir(root, cwd)
88 if err := os.MkdirAll(dir, 0o755); err != nil {
89 return nil, err
90 }
91 id := uuid.NewString()
92 name := fmt.Sprintf("%s-%s.jsonl", time.Now().UTC().Format("20060102-150405"), id[:8])
93 p := filepath.Join(dir, name)
94 f, err := os.OpenFile(p, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644)
95 if err != nil {
96 return nil, err
97 }
98 s := &Session{
99 ID: id,
100 Path: p,
101 Meta: SessionMeta{ID: id, CWD: cwd, Provider: providerName, Model: model, Started: time.Now().UTC(), Version: version},
102 writer: f,
103 buf: bufio.NewWriter(f),
104 freshFile: true,
105 }
106 if err := s.writeLine(sessionLine{Type: "meta", Meta: &s.Meta}); err != nil {
107 f.Close()
108 return nil, err
109 }
110 return s, nil
111}
112
113// SessionUsage returns the most recent cumulative usage row stored in
114// a session file. Sessions append one usage row per completed turn; the
115// latest row's cumulative field is the session total. Missing usage rows
116// are valid for old/empty sessions and return the zero value.
117func SessionUsage(path string) (provider.Usage, error) {
118 f, err := os.Open(path)
119 if err != nil {
120 return provider.Usage{}, err
121 }
122 defer f.Close()
123
124 var usage provider.Usage
125 sc := bufio.NewScanner(f)
126 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
127 for sc.Scan() {
128 var head sessionLineHead
129 if err := json.Unmarshal(sc.Bytes(), &head); err != nil || head.Type != "usage" {
130 continue
131 }
132 var row struct {
133 Cumulative provider.Usage `json:"cumulative"`
134 }
135 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
136 usage = row.Cumulative
137 }
138 }
139 if err := sc.Err(); err != nil {
140 return provider.Usage{}, err
141 }
142 return usage, nil
143}
144
145// OpenSession opens an existing session for appending.
146func OpenSession(path string) (*Session, []provider.Message, error) {
147 f, err := os.Open(path)
148 if err != nil {
149 return nil, nil, err
150 }
151 defer f.Close()
152
153 var meta SessionMeta
154 var messages []provider.Message
155 sc := bufio.NewScanner(f)
156 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
157 for sc.Scan() {
158 var head sessionLineHead
159 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
160 continue
161 }
162 switch head.Type {
163 case "meta":
164 var row struct {
165 Meta SessionMeta `json:"meta"`
166 }
167 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
168 meta = row.Meta
169 }
170 case "message":
171 if msg, err := hydrateMessage(sc.Bytes()); err == nil && len(msg.Content) > 0 {
172 messages = append(messages, msg)
173 }
174 }
175 }
176 if err := sc.Err(); err != nil {
177 return nil, nil, err
178 }
179 messages = repairToolUseResultPairs(messages)
180 out, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0o644)
181 if err != nil {
182 return nil, nil, err
183 }
184 s := &Session{ID: meta.ID, Path: path, Meta: meta, writer: out, buf: bufio.NewWriter(out)}
185 return s, messages, nil
186}
187
188// repairToolUseResultPairs walks a restored transcript and
189// synthesises stub tool_result blocks for any assistant
190// tool_use blocks that aren't paired with a matching result in
191// the next message. Anthropic (and OpenAI via the responses API)
192// reject any request whose transcript leaves a tool_use without
193// its matching tool_result immediately after, with errors like:
194//
195// messages.8: `tool_use` ids were found without `tool_result`
196// blocks immediately after
197//
198// Corruption gets into the transcript two ways we know of:
199//
200// - Older zot builds that persisted the assistant tool_use row
201// before the tool_result row, then crashed between the two.
202// - Abort paths in older builds that didn't drop the mid-turn
203// assistant message cleanly.
204//
205// Rather than change runtime semantics (which would risk hiding a
206// real bug), we scrub on load: any unmatched tool_use gets a stub
207// tool_result injected as a RoleTool message so the next
208// outbound request passes the provider's validity check. The stub
209// reads "tool call was aborted; no result recorded." so the
210// model can see what happened and decide whether to retry.
211//
212// Runs once per OpenSession call. No cost on the hot path.
213func repairToolUseResultPairs(msgs []provider.Message) []provider.Message {
214 if len(msgs) == 0 {
215 return msgs
216 }
217 out := make([]provider.Message, 0, len(msgs)+2)
218 for i, m := range msgs {
219 out = append(out, m)
220 if m.Role != provider.RoleAssistant {
221 continue
222 }
223 // Collect tool_use ids in this assistant message.
224 var ids []string
225 for _, c := range m.Content {
226 if tc, ok := c.(provider.ToolCallBlock); ok {
227 ids = append(ids, tc.ID)
228 }
229 }
230 if len(ids) == 0 {
231 continue
232 }
233 // Look at the next message (if any) and collect tool_result
234 // CallIDs it covers.
235 have := map[string]bool{}
236 if i+1 < len(msgs) && msgs[i+1].Role == provider.RoleTool {
237 for _, c := range msgs[i+1].Content {
238 if tr, ok := c.(provider.ToolResultBlock); ok {
239 have[tr.CallID] = true
240 }
241 }
242 }
243 // Build stubs for any missing id.
244 var stubs []provider.Content
245 for _, id := range ids {
246 if have[id] {
247 continue
248 }
249 stubs = append(stubs, provider.ToolResultBlock{
250 CallID: id,
251 Content: []provider.Content{provider.TextBlock{Text: "tool call was aborted; no result recorded."}},
252 IsError: true,
253 })
254 }
255 if len(stubs) == 0 {
256 continue
257 }
258 // Merge into the next tool-role message if present,
259 // otherwise insert a synthetic one right after the
260 // assistant message. Merging keeps the tool-role row
261 // count stable; inserting handles the common case where
262 // no tool message was persisted at all.
263 if i+1 < len(msgs) && msgs[i+1].Role == provider.RoleTool {
264 msgs[i+1].Content = append(msgs[i+1].Content, stubs...)
265 // We already appended m to out; the modified next
266 // message will be appended on the following iteration.
267 continue
268 }
269 out = append(out, provider.Message{
270 Role: provider.RoleTool,
271 Content: stubs,
272 Time: m.Time,
273 })
274 }
275 return out
276}
277
278// LatestSession returns the most recent session file for cwd, or "".
279func LatestSession(root, cwd string) string {
280 paths := ListSessions(root, cwd)
281 if len(paths) == 0 {
282 return ""
283 }
284 return paths[0]
285}
286
287// SessionSummary describes one on-disk session at a glance for UI pickers.
288type SessionSummary struct {
289 Path string
290 Started time.Time
291 Model string
292 Provider string
293 MessageCount int
294 FirstUserText string
295 TotalCost float64
296 Title string
297}
298
299// RenameSession updates the title field in the session's meta line.
300// It rewrites the first line of the file (the meta line) with the
301// updated title.
302// RenameSession appends a rename line to the session file. This is
303// safe even for the currently active session because it opens the
304// file independently and appends (doesn't rewrite).
305func RenameSession(path, title string) error {
306 f, err := os.OpenFile(path, os.O_WRONLY|os.O_APPEND, 0o644)
307 if err != nil {
308 return err
309 }
310 defer f.Close()
311 line, _ := json.Marshal(map[string]string{"type": "rename", "title": title})
312 line = append(line, '\n')
313 _, err = f.Write(line)
314 return err
315}
316
317// DescribeSessions returns lightweight summaries for every session in
318// cwd, newest first. Parses only the first few lines and the last usage
319// line so it's cheap to run on every dialog open.
320func DescribeSessions(root, cwd string) []SessionSummary {
321 paths := ListSessions(root, cwd)
322 summaries := make([]SessionSummary, 0, len(paths))
323 for _, p := range paths {
324 summaries = append(summaries, describeSession(p))
325 }
326 return summaries
327}
328
329func describeSession(path string) SessionSummary {
330 s := SessionSummary{Path: path}
331 f, err := os.Open(path)
332 if err != nil {
333 return s
334 }
335 defer f.Close()
336 sc := bufio.NewScanner(f)
337 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
338 for sc.Scan() {
339 var head sessionLineHead
340 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
341 continue
342 }
343 switch head.Type {
344 case "meta":
345 var row struct {
346 Meta SessionMeta `json:"meta"`
347 }
348 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
349 s.Started = row.Meta.Started
350 s.Model = row.Meta.Model
351 s.Provider = row.Meta.Provider
352 s.Title = row.Meta.Title
353 }
354 case "message":
355 s.MessageCount++
356 if s.FirstUserText == "" {
357 s.FirstUserText = firstUserText(sc.Bytes())
358 }
359 case "rename":
360 var row struct {
361 Title string `json:"title"`
362 }
363 if err := json.Unmarshal(sc.Bytes(), &row); err == nil && row.Title != "" {
364 s.Title = row.Title
365 }
366 case "usage":
367 var row struct {
368 Cumulative provider.Usage `json:"cumulative"`
369 }
370 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
371 s.TotalCost = row.Cumulative.CostUSD
372 }
373 }
374 }
375 return s
376}
377
378func firstUserText(line []byte) string {
379 var row struct {
380 Message struct {
381 Role string `json:"role"`
382 Content []struct {
383 Text string `json:"text"`
384 } `json:"content"`
385 } `json:"message"`
386 }
387 if err := json.Unmarshal(line, &row); err != nil {
388 return ""
389 }
390 if row.Message.Role != "user" {
391 return ""
392 }
393 for _, c := range row.Message.Content {
394 if c.Text != "" {
395 return c.Text
396 }
397 }
398 return ""
399}
400
401// PruneEmptySessions deletes session files in cwd's session directory
402// that contain only a meta line (no messages were ever appended).
403// Cleans up the backlog of empty stubs created by old zot versions
404// that wrote a meta line at NewSession time and never followed up.
405// Errors are swallowed; the caller treats this as best-effort.
406func PruneEmptySessions(root, cwd string) {
407 dir := SessionsDir(root, cwd)
408 entries, err := os.ReadDir(dir)
409 if err != nil {
410 return
411 }
412 for _, e := range entries {
413 if e.IsDir() || !strings.HasSuffix(e.Name(), ".jsonl") {
414 continue
415 }
416 p := filepath.Join(dir, e.Name())
417 if sessionHasNoMessages(p) {
418 _ = os.Remove(p)
419 }
420 }
421}
422
423// sessionHasNoMessages returns true when the file at path contains
424// no lines of type "message". Meta-only / usage-only files count as
425// empty. Used by PruneEmptySessions and the Describe path.
426func sessionHasNoMessages(path string) bool {
427 f, err := os.Open(path)
428 if err != nil {
429 return false
430 }
431 defer f.Close()
432 sc := bufio.NewScanner(f)
433 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
434 for sc.Scan() {
435 var head sessionLineHead
436 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
437 continue
438 }
439 if head.Type == "message" {
440 return false
441 }
442 }
443 return true
444}
445
446// ListSessions returns session file paths for cwd, most-recently-
447// modified first. Sorting on filesystem ModTime instead of the
448// timestamp embedded in the filename means a long-running session
449// the user actually returned to recently floats to the top of
450// /sessions, /continue, and the resume picker, even when it was
451// originally created days earlier than newer but idle sessions.
452// Files with identical ModTime fall back to filename desc so the
453// order stays stable across calls.
454func ListSessions(root, cwd string) []string {
455 dir := SessionsDir(root, cwd)
456 entries, err := os.ReadDir(dir)
457 if err != nil {
458 return nil
459 }
460 type rec struct {
461 path string
462 mod time.Time
463 }
464 var files []rec
465 for _, e := range entries {
466 if e.IsDir() || !strings.HasSuffix(e.Name(), ".jsonl") {
467 continue
468 }
469 p := filepath.Join(dir, e.Name())
470 info, err := e.Info()
471 if err != nil {
472 continue
473 }
474 files = append(files, rec{path: p, mod: info.ModTime()})
475 }
476 sort.Slice(files, func(i, j int) bool {
477 if !files[i].mod.Equal(files[j].mod) {
478 return files[i].mod.After(files[j].mod)
479 }
480 return files[i].path > files[j].path
481 })
482 out := make([]string, 0, len(files))
483 for _, r := range files {
484 out = append(out, r.path)
485 }
486 return out
487}
488
489// AppendMessage writes a message to the session.
490func (s *Session) AppendMessage(m provider.Message) error {
491 if s == nil {
492 return nil
493 }
494 if err := s.writeLine(sessionLine{Type: "message", Message: &m}); err != nil {
495 return err
496 }
497 s.messagesAppended++
498 return nil
499}
500
501// UpdateModel records a provider/model switch in the session file.
502// The reader keeps the most recent meta entry, so the session resumes
503// with the updated model.
504func (s *Session) UpdateModel(providerName, model string) error {
505 if s == nil {
506 return nil
507 }
508 s.Meta.Provider = providerName
509 s.Meta.Model = model
510 return s.writeLine(sessionLine{Type: "meta", Meta: &s.Meta})
511}
512
513// AppendUsage writes a usage row to the session.
514func (s *Session) AppendUsage(u, cum provider.Usage) error {
515 if s == nil {
516 return nil
517 }
518 return s.writeLine(sessionLine{Type: "usage", Usage: &u, Cumulative: &cum})
519}
520
521// Close flushes and closes the session file. If the session was
522// freshly created in this process and never had any messages
523// appended (the user opened zot, looked around, and exited without
524// prompting), the file is deleted on close so the sessions list
525// doesn't fill up with empty meta-only stubs.
526func (s *Session) Close() error {
527 if s == nil {
528 return nil
529 }
530 flushErr := s.buf.Flush()
531 closeErr := s.writer.Close()
532 if s.freshFile && s.messagesAppended == 0 {
533 // Best-effort cleanup. We deliberately don't propagate the
534 // remove error: if it fails (file already gone, perms changed)
535 // the worst case is one stale empty file in the listing.
536 _ = os.Remove(s.Path)
537 }
538 if flushErr != nil {
539 return flushErr
540 }
541 return closeErr
542}
543
544func (s *Session) writeLine(row sessionLine) error {
545 b, err := json.Marshal(row)
546 if err != nil {
547 return err
548 }
549 if _, err := s.buf.Write(b); err != nil {
550 return err
551 }
552 if err := s.buf.WriteByte('\n'); err != nil {
553 return err
554 }
555 return s.buf.Flush()
556}
557
558// ---- content (de)serialization ----
559//
560// provider.Content is an interface; encoding/json drops type information.
561// We persist messages by reading the raw "message" object back and
562// rebuilding Content from discriminated fields.
563
564func hydrateMessage(lineBytes []byte) (provider.Message, error) {
565 var row struct {
566 Message struct {
567 Role provider.Role `json:"role"`
568 Content []json.RawMessage `json:"content"`
569 Time time.Time `json:"time"`
570 } `json:"message"`
571 }
572 if err := json.Unmarshal(lineBytes, &row); err != nil {
573 return provider.Message{}, err
574 }
575 msg := provider.Message{Role: row.Message.Role, Time: row.Message.Time}
576 for _, raw := range row.Message.Content {
577 var head struct {
578 Text string `json:"text"`
579 MimeType string `json:"mime_type"`
580 Data []byte `json:"data"`
581 ID string `json:"id"`
582 Name string `json:"name"`
583 CallID string `json:"call_id"`
584 // ToolCallBlock also has Arguments, ToolResultBlock has Content + IsError
585 }
586 if err := json.Unmarshal(raw, &head); err != nil {
587 continue
588 }
589 // Discriminate by presence of fields.
590 switch {
591 case head.Name != "" && head.ID != "":
592 var tc struct {
593 ID string `json:"id"`
594 Name string `json:"name"`
595 Arguments json.RawMessage `json:"arguments"`
596 }
597 _ = json.Unmarshal(raw, &tc)
598 msg.Content = append(msg.Content, provider.ToolCallBlock{ID: tc.ID, Name: tc.Name, Arguments: tc.Arguments})
599 case head.CallID != "":
600 var tr struct {
601 CallID string `json:"call_id"`
602 Content []json.RawMessage `json:"content"`
603 IsError bool `json:"is_error"`
604 }
605 _ = json.Unmarshal(raw, &tr)
606 block := provider.ToolResultBlock{CallID: tr.CallID, IsError: tr.IsError}
607 for _, c := range tr.Content {
608 var inner struct {
609 Text string `json:"text"`
610 MimeType string `json:"mime_type"`
611 Data []byte `json:"data"`
612 }
613 _ = json.Unmarshal(c, &inner)
614 if inner.MimeType != "" {
615 block.Content = append(block.Content, provider.ImageBlock{MimeType: inner.MimeType, Data: inner.Data})
616 } else {
617 block.Content = append(block.Content, provider.TextBlock{Text: inner.Text})
618 }
619 }
620 msg.Content = append(msg.Content, block)
621 case head.MimeType != "":
622 msg.Content = append(msg.Content, provider.ImageBlock{MimeType: head.MimeType, Data: head.Data})
623 default:
624 msg.Content = append(msg.Content, provider.TextBlock{Text: head.Text})
625 }
626 }
627 return msg, nil
628}