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