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// OpenSession opens an existing session for appending.
114func OpenSession(path string) (*Session, []provider.Message, error) {
115 f, err := os.Open(path)
116 if err != nil {
117 return nil, nil, err
118 }
119 defer f.Close()
120
121 var meta SessionMeta
122 var messages []provider.Message
123 sc := bufio.NewScanner(f)
124 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
125 for sc.Scan() {
126 var head sessionLineHead
127 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
128 continue
129 }
130 switch head.Type {
131 case "meta":
132 var row struct {
133 Meta SessionMeta `json:"meta"`
134 }
135 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
136 meta = row.Meta
137 }
138 case "message":
139 if msg, err := hydrateMessage(sc.Bytes()); err == nil && len(msg.Content) > 0 {
140 messages = append(messages, msg)
141 }
142 }
143 }
144 if err := sc.Err(); err != nil {
145 return nil, nil, err
146 }
147 messages = repairToolUseResultPairs(messages)
148 out, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0o644)
149 if err != nil {
150 return nil, nil, err
151 }
152 s := &Session{ID: meta.ID, Path: path, Meta: meta, writer: out, buf: bufio.NewWriter(out)}
153 return s, messages, nil
154}
155
156// repairToolUseResultPairs walks a restored transcript and
157// synthesises stub tool_result blocks for any assistant
158// tool_use blocks that aren't paired with a matching result in
159// the next message. Anthropic (and OpenAI via the responses API)
160// reject any request whose transcript leaves a tool_use without
161// its matching tool_result immediately after, with errors like:
162//
163// messages.8: `tool_use` ids were found without `tool_result`
164// blocks immediately after
165//
166// Corruption gets into the transcript two ways we know of:
167//
168// - Older zot builds that persisted the assistant tool_use row
169// before the tool_result row, then crashed between the two.
170// - Abort paths in older builds that didn't drop the mid-turn
171// assistant message cleanly.
172//
173// Rather than change runtime semantics (which would risk hiding a
174// real bug), we scrub on load: any unmatched tool_use gets a stub
175// tool_result injected as a RoleTool message so the next
176// outbound request passes the provider's validity check. The stub
177// reads "tool call was aborted; no result recorded." so the
178// model can see what happened and decide whether to retry.
179//
180// Runs once per OpenSession call. No cost on the hot path.
181func repairToolUseResultPairs(msgs []provider.Message) []provider.Message {
182 if len(msgs) == 0 {
183 return msgs
184 }
185 out := make([]provider.Message, 0, len(msgs)+2)
186 for i, m := range msgs {
187 out = append(out, m)
188 if m.Role != provider.RoleAssistant {
189 continue
190 }
191 // Collect tool_use ids in this assistant message.
192 var ids []string
193 for _, c := range m.Content {
194 if tc, ok := c.(provider.ToolCallBlock); ok {
195 ids = append(ids, tc.ID)
196 }
197 }
198 if len(ids) == 0 {
199 continue
200 }
201 // Look at the next message (if any) and collect tool_result
202 // CallIDs it covers.
203 have := map[string]bool{}
204 if i+1 < len(msgs) && msgs[i+1].Role == provider.RoleTool {
205 for _, c := range msgs[i+1].Content {
206 if tr, ok := c.(provider.ToolResultBlock); ok {
207 have[tr.CallID] = true
208 }
209 }
210 }
211 // Build stubs for any missing id.
212 var stubs []provider.Content
213 for _, id := range ids {
214 if have[id] {
215 continue
216 }
217 stubs = append(stubs, provider.ToolResultBlock{
218 CallID: id,
219 Content: []provider.Content{provider.TextBlock{Text: "tool call was aborted; no result recorded."}},
220 IsError: true,
221 })
222 }
223 if len(stubs) == 0 {
224 continue
225 }
226 // Merge into the next tool-role message if present,
227 // otherwise insert a synthetic one right after the
228 // assistant message. Merging keeps the tool-role row
229 // count stable; inserting handles the common case where
230 // no tool message was persisted at all.
231 if i+1 < len(msgs) && msgs[i+1].Role == provider.RoleTool {
232 msgs[i+1].Content = append(msgs[i+1].Content, stubs...)
233 // We already appended m to out; the modified next
234 // message will be appended on the following iteration.
235 continue
236 }
237 out = append(out, provider.Message{
238 Role: provider.RoleTool,
239 Content: stubs,
240 Time: m.Time,
241 })
242 }
243 return out
244}
245
246// LatestSession returns the most recent session file for cwd, or "".
247func LatestSession(root, cwd string) string {
248 paths := ListSessions(root, cwd)
249 if len(paths) == 0 {
250 return ""
251 }
252 return paths[0]
253}
254
255// SessionSummary describes one on-disk session at a glance for UI pickers.
256type SessionSummary struct {
257 Path string
258 Started time.Time
259 Model string
260 Provider string
261 MessageCount int
262 FirstUserText string
263 TotalCost float64
264 Title string
265}
266
267// RenameSession updates the title field in the session's meta line.
268// It rewrites the first line of the file (the meta line) with the
269// updated title.
270// RenameSession appends a rename line to the session file. This is
271// safe even for the currently active session because it opens the
272// file independently and appends (doesn't rewrite).
273func RenameSession(path, title string) error {
274 f, err := os.OpenFile(path, os.O_WRONLY|os.O_APPEND, 0o644)
275 if err != nil {
276 return err
277 }
278 defer f.Close()
279 line, _ := json.Marshal(map[string]string{"type": "rename", "title": title})
280 line = append(line, '\n')
281 _, err = f.Write(line)
282 return err
283}
284
285// DescribeSessions returns lightweight summaries for every session in
286// cwd, newest first. Parses only the first few lines and the last usage
287// line so it's cheap to run on every dialog open.
288func DescribeSessions(root, cwd string) []SessionSummary {
289 paths := ListSessions(root, cwd)
290 summaries := make([]SessionSummary, 0, len(paths))
291 for _, p := range paths {
292 summaries = append(summaries, describeSession(p))
293 }
294 return summaries
295}
296
297func describeSession(path string) SessionSummary {
298 s := SessionSummary{Path: path}
299 f, err := os.Open(path)
300 if err != nil {
301 return s
302 }
303 defer f.Close()
304 sc := bufio.NewScanner(f)
305 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
306 for sc.Scan() {
307 var head sessionLineHead
308 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
309 continue
310 }
311 switch head.Type {
312 case "meta":
313 var row struct {
314 Meta SessionMeta `json:"meta"`
315 }
316 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
317 s.Started = row.Meta.Started
318 s.Model = row.Meta.Model
319 s.Provider = row.Meta.Provider
320 s.Title = row.Meta.Title
321 }
322 case "message":
323 s.MessageCount++
324 if s.FirstUserText == "" {
325 s.FirstUserText = firstUserText(sc.Bytes())
326 }
327 case "rename":
328 var row struct {
329 Title string `json:"title"`
330 }
331 if err := json.Unmarshal(sc.Bytes(), &row); err == nil && row.Title != "" {
332 s.Title = row.Title
333 }
334 case "usage":
335 var row struct {
336 Cumulative provider.Usage `json:"cumulative"`
337 }
338 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
339 s.TotalCost = row.Cumulative.CostUSD
340 }
341 }
342 }
343 return s
344}
345
346func firstUserText(line []byte) string {
347 var row struct {
348 Message struct {
349 Role string `json:"role"`
350 Content []struct {
351 Text string `json:"text"`
352 } `json:"content"`
353 } `json:"message"`
354 }
355 if err := json.Unmarshal(line, &row); err != nil {
356 return ""
357 }
358 if row.Message.Role != "user" {
359 return ""
360 }
361 for _, c := range row.Message.Content {
362 if c.Text != "" {
363 return c.Text
364 }
365 }
366 return ""
367}
368
369// PruneEmptySessions deletes session files in cwd's session directory
370// that contain only a meta line (no messages were ever appended).
371// Cleans up the backlog of empty stubs created by old zot versions
372// that wrote a meta line at NewSession time and never followed up.
373// Errors are swallowed; the caller treats this as best-effort.
374func PruneEmptySessions(root, cwd string) {
375 dir := SessionsDir(root, cwd)
376 entries, err := os.ReadDir(dir)
377 if err != nil {
378 return
379 }
380 for _, e := range entries {
381 if e.IsDir() || !strings.HasSuffix(e.Name(), ".jsonl") {
382 continue
383 }
384 p := filepath.Join(dir, e.Name())
385 if sessionHasNoMessages(p) {
386 _ = os.Remove(p)
387 }
388 }
389}
390
391// sessionHasNoMessages returns true when the file at path contains
392// no lines of type "message". Meta-only / usage-only files count as
393// empty. Used by PruneEmptySessions and the Describe path.
394func sessionHasNoMessages(path string) bool {
395 f, err := os.Open(path)
396 if err != nil {
397 return false
398 }
399 defer f.Close()
400 sc := bufio.NewScanner(f)
401 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
402 for sc.Scan() {
403 var head sessionLineHead
404 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
405 continue
406 }
407 if head.Type == "message" {
408 return false
409 }
410 }
411 return true
412}
413
414// ListSessions returns session file paths for cwd, newest first.
415func ListSessions(root, cwd string) []string {
416 dir := SessionsDir(root, cwd)
417 entries, err := os.ReadDir(dir)
418 if err != nil {
419 return nil
420 }
421 var files []string
422 for _, e := range entries {
423 if !e.IsDir() && strings.HasSuffix(e.Name(), ".jsonl") {
424 files = append(files, filepath.Join(dir, e.Name()))
425 }
426 }
427 sort.Sort(sort.Reverse(sort.StringSlice(files)))
428 return files
429}
430
431// AppendMessage writes a message to the session.
432func (s *Session) AppendMessage(m provider.Message) error {
433 if s == nil {
434 return nil
435 }
436 if err := s.writeLine(sessionLine{Type: "message", Message: &m}); err != nil {
437 return err
438 }
439 s.messagesAppended++
440 return nil
441}
442
443// UpdateModel records a provider/model switch in the session file.
444// The reader keeps the most recent meta entry, so the session resumes
445// with the updated model.
446func (s *Session) UpdateModel(providerName, model string) error {
447 if s == nil {
448 return nil
449 }
450 s.Meta.Provider = providerName
451 s.Meta.Model = model
452 return s.writeLine(sessionLine{Type: "meta", Meta: &s.Meta})
453}
454
455// AppendUsage writes a usage row to the session.
456func (s *Session) AppendUsage(u, cum provider.Usage) error {
457 if s == nil {
458 return nil
459 }
460 return s.writeLine(sessionLine{Type: "usage", Usage: &u, Cumulative: &cum})
461}
462
463// Close flushes and closes the session file. If the session was
464// freshly created in this process and never had any messages
465// appended (the user opened zot, looked around, and exited without
466// prompting), the file is deleted on close so the sessions list
467// doesn't fill up with empty meta-only stubs.
468func (s *Session) Close() error {
469 if s == nil {
470 return nil
471 }
472 flushErr := s.buf.Flush()
473 closeErr := s.writer.Close()
474 if s.freshFile && s.messagesAppended == 0 {
475 // Best-effort cleanup. We deliberately don't propagate the
476 // remove error: if it fails (file already gone, perms changed)
477 // the worst case is one stale empty file in the listing.
478 _ = os.Remove(s.Path)
479 }
480 if flushErr != nil {
481 return flushErr
482 }
483 return closeErr
484}
485
486func (s *Session) writeLine(row sessionLine) error {
487 b, err := json.Marshal(row)
488 if err != nil {
489 return err
490 }
491 if _, err := s.buf.Write(b); err != nil {
492 return err
493 }
494 if err := s.buf.WriteByte('\n'); err != nil {
495 return err
496 }
497 return s.buf.Flush()
498}
499
500// ---- content (de)serialization ----
501//
502// provider.Content is an interface; encoding/json drops type information.
503// We persist messages by reading the raw "message" object back and
504// rebuilding Content from discriminated fields.
505
506func hydrateMessage(lineBytes []byte) (provider.Message, error) {
507 var row struct {
508 Message struct {
509 Role provider.Role `json:"role"`
510 Content []json.RawMessage `json:"content"`
511 Time time.Time `json:"time"`
512 } `json:"message"`
513 }
514 if err := json.Unmarshal(lineBytes, &row); err != nil {
515 return provider.Message{}, err
516 }
517 msg := provider.Message{Role: row.Message.Role, Time: row.Message.Time}
518 for _, raw := range row.Message.Content {
519 var head struct {
520 Text string `json:"text"`
521 MimeType string `json:"mime_type"`
522 Data []byte `json:"data"`
523 ID string `json:"id"`
524 Name string `json:"name"`
525 CallID string `json:"call_id"`
526 // ToolCallBlock also has Arguments, ToolResultBlock has Content + IsError
527 }
528 if err := json.Unmarshal(raw, &head); err != nil {
529 continue
530 }
531 // Discriminate by presence of fields.
532 switch {
533 case head.Name != "" && head.ID != "":
534 var tc struct {
535 ID string `json:"id"`
536 Name string `json:"name"`
537 Arguments json.RawMessage `json:"arguments"`
538 }
539 _ = json.Unmarshal(raw, &tc)
540 msg.Content = append(msg.Content, provider.ToolCallBlock{ID: tc.ID, Name: tc.Name, Arguments: tc.Arguments})
541 case head.CallID != "":
542 var tr struct {
543 CallID string `json:"call_id"`
544 Content []json.RawMessage `json:"content"`
545 IsError bool `json:"is_error"`
546 }
547 _ = json.Unmarshal(raw, &tr)
548 block := provider.ToolResultBlock{CallID: tr.CallID, IsError: tr.IsError}
549 for _, c := range tr.Content {
550 var inner struct {
551 Text string `json:"text"`
552 MimeType string `json:"mime_type"`
553 Data []byte `json:"data"`
554 }
555 _ = json.Unmarshal(c, &inner)
556 if inner.MimeType != "" {
557 block.Content = append(block.Content, provider.ImageBlock{MimeType: inner.MimeType, Data: inner.Data})
558 } else {
559 block.Content = append(block.Content, provider.TextBlock{Text: inner.Text})
560 }
561 }
562 msg.Content = append(msg.Content, block)
563 case head.MimeType != "":
564 msg.Content = append(msg.Content, provider.ImageBlock{MimeType: head.MimeType, Data: head.Data})
565 default:
566 msg.Content = append(msg.Content, provider.TextBlock{Text: head.Text})
567 }
568 }
569 return msg, nil
570}