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, newest first.
447func ListSessions(root, cwd string) []string {
448 dir := SessionsDir(root, cwd)
449 entries, err := os.ReadDir(dir)
450 if err != nil {
451 return nil
452 }
453 var files []string
454 for _, e := range entries {
455 if !e.IsDir() && strings.HasSuffix(e.Name(), ".jsonl") {
456 files = append(files, filepath.Join(dir, e.Name()))
457 }
458 }
459 sort.Sort(sort.Reverse(sort.StringSlice(files)))
460 return files
461}
462
463// AppendMessage writes a message to the session.
464func (s *Session) AppendMessage(m provider.Message) error {
465 if s == nil {
466 return nil
467 }
468 if err := s.writeLine(sessionLine{Type: "message", Message: &m}); err != nil {
469 return err
470 }
471 s.messagesAppended++
472 return nil
473}
474
475// UpdateModel records a provider/model switch in the session file.
476// The reader keeps the most recent meta entry, so the session resumes
477// with the updated model.
478func (s *Session) UpdateModel(providerName, model string) error {
479 if s == nil {
480 return nil
481 }
482 s.Meta.Provider = providerName
483 s.Meta.Model = model
484 return s.writeLine(sessionLine{Type: "meta", Meta: &s.Meta})
485}
486
487// AppendUsage writes a usage row to the session.
488func (s *Session) AppendUsage(u, cum provider.Usage) error {
489 if s == nil {
490 return nil
491 }
492 return s.writeLine(sessionLine{Type: "usage", Usage: &u, Cumulative: &cum})
493}
494
495// Close flushes and closes the session file. If the session was
496// freshly created in this process and never had any messages
497// appended (the user opened zot, looked around, and exited without
498// prompting), the file is deleted on close so the sessions list
499// doesn't fill up with empty meta-only stubs.
500func (s *Session) Close() error {
501 if s == nil {
502 return nil
503 }
504 flushErr := s.buf.Flush()
505 closeErr := s.writer.Close()
506 if s.freshFile && s.messagesAppended == 0 {
507 // Best-effort cleanup. We deliberately don't propagate the
508 // remove error: if it fails (file already gone, perms changed)
509 // the worst case is one stale empty file in the listing.
510 _ = os.Remove(s.Path)
511 }
512 if flushErr != nil {
513 return flushErr
514 }
515 return closeErr
516}
517
518func (s *Session) writeLine(row sessionLine) error {
519 b, err := json.Marshal(row)
520 if err != nil {
521 return err
522 }
523 if _, err := s.buf.Write(b); err != nil {
524 return err
525 }
526 if err := s.buf.WriteByte('\n'); err != nil {
527 return err
528 }
529 return s.buf.Flush()
530}
531
532// ---- content (de)serialization ----
533//
534// provider.Content is an interface; encoding/json drops type information.
535// We persist messages by reading the raw "message" object back and
536// rebuilding Content from discriminated fields.
537
538func hydrateMessage(lineBytes []byte) (provider.Message, error) {
539 var row struct {
540 Message struct {
541 Role provider.Role `json:"role"`
542 Content []json.RawMessage `json:"content"`
543 Time time.Time `json:"time"`
544 } `json:"message"`
545 }
546 if err := json.Unmarshal(lineBytes, &row); err != nil {
547 return provider.Message{}, err
548 }
549 msg := provider.Message{Role: row.Message.Role, Time: row.Message.Time}
550 for _, raw := range row.Message.Content {
551 var head struct {
552 Text string `json:"text"`
553 MimeType string `json:"mime_type"`
554 Data []byte `json:"data"`
555 ID string `json:"id"`
556 Name string `json:"name"`
557 CallID string `json:"call_id"`
558 // ToolCallBlock also has Arguments, ToolResultBlock has Content + IsError
559 }
560 if err := json.Unmarshal(raw, &head); err != nil {
561 continue
562 }
563 // Discriminate by presence of fields.
564 switch {
565 case head.Name != "" && head.ID != "":
566 var tc struct {
567 ID string `json:"id"`
568 Name string `json:"name"`
569 Arguments json.RawMessage `json:"arguments"`
570 }
571 _ = json.Unmarshal(raw, &tc)
572 msg.Content = append(msg.Content, provider.ToolCallBlock{ID: tc.ID, Name: tc.Name, Arguments: tc.Arguments})
573 case head.CallID != "":
574 var tr struct {
575 CallID string `json:"call_id"`
576 Content []json.RawMessage `json:"content"`
577 IsError bool `json:"is_error"`
578 }
579 _ = json.Unmarshal(raw, &tr)
580 block := provider.ToolResultBlock{CallID: tr.CallID, IsError: tr.IsError}
581 for _, c := range tr.Content {
582 var inner struct {
583 Text string `json:"text"`
584 MimeType string `json:"mime_type"`
585 Data []byte `json:"data"`
586 }
587 _ = json.Unmarshal(c, &inner)
588 if inner.MimeType != "" {
589 block.Content = append(block.Content, provider.ImageBlock{MimeType: inner.MimeType, Data: inner.Data})
590 } else {
591 block.Content = append(block.Content, provider.TextBlock{Text: inner.Text})
592 }
593 }
594 msg.Content = append(msg.Content, block)
595 case head.MimeType != "":
596 msg.Content = append(msg.Content, provider.ImageBlock{MimeType: head.MimeType, Data: head.Data})
597 default:
598 msg.Content = append(msg.Content, provider.TextBlock{Text: head.Text})
599 }
600 }
601 return msg, nil
602}