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