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