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 out, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0o644)
147 if err != nil {
148 return nil, nil, err
149 }
150 s := &Session{ID: meta.ID, Path: path, Meta: meta, writer: out, buf: bufio.NewWriter(out)}
151 return s, messages, nil
152}
153
154// LatestSession returns the most recent session file for cwd, or "".
155func LatestSession(root, cwd string) string {
156 paths := ListSessions(root, cwd)
157 if len(paths) == 0 {
158 return ""
159 }
160 return paths[0]
161}
162
163// SessionSummary describes one on-disk session at a glance for UI pickers.
164type SessionSummary struct {
165 Path string
166 Started time.Time
167 Model string
168 Provider string
169 MessageCount int
170 FirstUserText string
171 TotalCost float64
172}
173
174// DescribeSessions returns lightweight summaries for every session in
175// cwd, newest first. Parses only the first few lines and the last usage
176// line so it's cheap to run on every dialog open.
177func DescribeSessions(root, cwd string) []SessionSummary {
178 paths := ListSessions(root, cwd)
179 summaries := make([]SessionSummary, 0, len(paths))
180 for _, p := range paths {
181 summaries = append(summaries, describeSession(p))
182 }
183 return summaries
184}
185
186func describeSession(path string) SessionSummary {
187 s := SessionSummary{Path: path}
188 f, err := os.Open(path)
189 if err != nil {
190 return s
191 }
192 defer f.Close()
193 sc := bufio.NewScanner(f)
194 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
195 for sc.Scan() {
196 var head sessionLineHead
197 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
198 continue
199 }
200 switch head.Type {
201 case "meta":
202 var row struct {
203 Meta SessionMeta `json:"meta"`
204 }
205 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
206 s.Started = row.Meta.Started
207 s.Model = row.Meta.Model
208 s.Provider = row.Meta.Provider
209 }
210 case "message":
211 s.MessageCount++
212 if s.FirstUserText == "" {
213 s.FirstUserText = firstUserText(sc.Bytes())
214 }
215 case "usage":
216 var row struct {
217 Cumulative provider.Usage `json:"cumulative"`
218 }
219 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
220 s.TotalCost = row.Cumulative.CostUSD
221 }
222 }
223 }
224 return s
225}
226
227func firstUserText(line []byte) string {
228 var row struct {
229 Message struct {
230 Role string `json:"role"`
231 Content []struct {
232 Text string `json:"text"`
233 } `json:"content"`
234 } `json:"message"`
235 }
236 if err := json.Unmarshal(line, &row); err != nil {
237 return ""
238 }
239 if row.Message.Role != "user" {
240 return ""
241 }
242 for _, c := range row.Message.Content {
243 if c.Text != "" {
244 return c.Text
245 }
246 }
247 return ""
248}
249
250// PruneEmptySessions deletes session files in cwd's session directory
251// that contain only a meta line (no messages were ever appended).
252// Cleans up the backlog of empty stubs created by old zot versions
253// that wrote a meta line at NewSession time and never followed up.
254// Errors are swallowed; the caller treats this as best-effort.
255func PruneEmptySessions(root, cwd string) {
256 dir := SessionsDir(root, cwd)
257 entries, err := os.ReadDir(dir)
258 if err != nil {
259 return
260 }
261 for _, e := range entries {
262 if e.IsDir() || !strings.HasSuffix(e.Name(), ".jsonl") {
263 continue
264 }
265 p := filepath.Join(dir, e.Name())
266 if sessionHasNoMessages(p) {
267 _ = os.Remove(p)
268 }
269 }
270}
271
272// sessionHasNoMessages returns true when the file at path contains
273// no lines of type "message". Meta-only / usage-only files count as
274// empty. Used by PruneEmptySessions and the Describe path.
275func sessionHasNoMessages(path string) bool {
276 f, err := os.Open(path)
277 if err != nil {
278 return false
279 }
280 defer f.Close()
281 sc := bufio.NewScanner(f)
282 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
283 for sc.Scan() {
284 var head sessionLineHead
285 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
286 continue
287 }
288 if head.Type == "message" {
289 return false
290 }
291 }
292 return true
293}
294
295// ListSessions returns session file paths for cwd, newest first.
296func ListSessions(root, cwd string) []string {
297 dir := SessionsDir(root, cwd)
298 entries, err := os.ReadDir(dir)
299 if err != nil {
300 return nil
301 }
302 var files []string
303 for _, e := range entries {
304 if !e.IsDir() && strings.HasSuffix(e.Name(), ".jsonl") {
305 files = append(files, filepath.Join(dir, e.Name()))
306 }
307 }
308 sort.Sort(sort.Reverse(sort.StringSlice(files)))
309 return files
310}
311
312// AppendMessage writes a message to the session.
313func (s *Session) AppendMessage(m provider.Message) error {
314 if s == nil {
315 return nil
316 }
317 if err := s.writeLine(sessionLine{Type: "message", Message: &m}); err != nil {
318 return err
319 }
320 s.messagesAppended++
321 return nil
322}
323
324// UpdateModel records a provider/model switch in the session file.
325// The reader keeps the most recent meta entry, so the session resumes
326// with the updated model.
327func (s *Session) UpdateModel(providerName, model string) error {
328 if s == nil {
329 return nil
330 }
331 s.Meta.Provider = providerName
332 s.Meta.Model = model
333 return s.writeLine(sessionLine{Type: "meta", Meta: &s.Meta})
334}
335
336// AppendUsage writes a usage row to the session.
337func (s *Session) AppendUsage(u, cum provider.Usage) error {
338 if s == nil {
339 return nil
340 }
341 return s.writeLine(sessionLine{Type: "usage", Usage: &u, Cumulative: &cum})
342}
343
344// Close flushes and closes the session file. If the session was
345// freshly created in this process and never had any messages
346// appended (the user opened zot, looked around, and exited without
347// prompting), the file is deleted on close so the sessions list
348// doesn't fill up with empty meta-only stubs.
349func (s *Session) Close() error {
350 if s == nil {
351 return nil
352 }
353 flushErr := s.buf.Flush()
354 closeErr := s.writer.Close()
355 if s.freshFile && s.messagesAppended == 0 {
356 // Best-effort cleanup. We deliberately don't propagate the
357 // remove error: if it fails (file already gone, perms changed)
358 // the worst case is one stale empty file in the listing.
359 _ = os.Remove(s.Path)
360 }
361 if flushErr != nil {
362 return flushErr
363 }
364 return closeErr
365}
366
367func (s *Session) writeLine(row sessionLine) error {
368 b, err := json.Marshal(row)
369 if err != nil {
370 return err
371 }
372 if _, err := s.buf.Write(b); err != nil {
373 return err
374 }
375 if err := s.buf.WriteByte('\n'); err != nil {
376 return err
377 }
378 return s.buf.Flush()
379}
380
381// ---- content (de)serialization ----
382//
383// provider.Content is an interface; encoding/json drops type information.
384// We persist messages by reading the raw "message" object back and
385// rebuilding Content from discriminated fields.
386
387func hydrateMessage(lineBytes []byte) (provider.Message, error) {
388 var row struct {
389 Message struct {
390 Role provider.Role `json:"role"`
391 Content []json.RawMessage `json:"content"`
392 Time time.Time `json:"time"`
393 } `json:"message"`
394 }
395 if err := json.Unmarshal(lineBytes, &row); err != nil {
396 return provider.Message{}, err
397 }
398 msg := provider.Message{Role: row.Message.Role, Time: row.Message.Time}
399 for _, raw := range row.Message.Content {
400 var head struct {
401 Text string `json:"text"`
402 MimeType string `json:"mime_type"`
403 Data []byte `json:"data"`
404 ID string `json:"id"`
405 Name string `json:"name"`
406 CallID string `json:"call_id"`
407 // ToolCallBlock also has Arguments, ToolResultBlock has Content + IsError
408 }
409 if err := json.Unmarshal(raw, &head); err != nil {
410 continue
411 }
412 // Discriminate by presence of fields.
413 switch {
414 case head.Name != "" && head.ID != "":
415 var tc struct {
416 ID string `json:"id"`
417 Name string `json:"name"`
418 Arguments json.RawMessage `json:"arguments"`
419 }
420 _ = json.Unmarshal(raw, &tc)
421 msg.Content = append(msg.Content, provider.ToolCallBlock{ID: tc.ID, Name: tc.Name, Arguments: tc.Arguments})
422 case head.CallID != "":
423 var tr struct {
424 CallID string `json:"call_id"`
425 Content []json.RawMessage `json:"content"`
426 IsError bool `json:"is_error"`
427 }
428 _ = json.Unmarshal(raw, &tr)
429 block := provider.ToolResultBlock{CallID: tr.CallID, IsError: tr.IsError}
430 for _, c := range tr.Content {
431 var inner struct {
432 Text string `json:"text"`
433 MimeType string `json:"mime_type"`
434 Data []byte `json:"data"`
435 }
436 _ = json.Unmarshal(c, &inner)
437 if inner.MimeType != "" {
438 block.Content = append(block.Content, provider.ImageBlock{MimeType: inner.MimeType, Data: inner.Data})
439 } else {
440 block.Content = append(block.Content, provider.TextBlock{Text: inner.Text})
441 }
442 }
443 msg.Content = append(msg.Content, block)
444 case head.MimeType != "":
445 msg.Content = append(msg.Content, provider.ImageBlock{MimeType: head.MimeType, Data: head.Data})
446 default:
447 msg.Content = append(msg.Content, provider.TextBlock{Text: head.Text})
448 }
449 }
450 return msg, nil
451}