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
28// SessionMeta is written as the first line of every session file.
29type SessionMeta struct {
30 ID string `json:"id"`
31 CWD string `json:"cwd"`
32 Model string `json:"model"`
33 Provider string `json:"provider"`
34 Started time.Time `json:"started"`
35 Version string `json:"version"`
36}
37
38// sessionLine is the on-disk row type. Message is kept as a raw
39// JSON message on reads (because Content is an interface slice that
40// the default unmarshaler cannot reconstruct); it is written with a
41// regular provider.Message value.
42type sessionLine struct {
43 Type string `json:"type"`
44 Meta *SessionMeta `json:"meta,omitempty"`
45 Message *provider.Message `json:"message,omitempty"`
46 Usage *provider.Usage `json:"usage,omitempty"`
47 Cumulative *provider.Usage `json:"cumulative,omitempty"`
48}
49
50type sessionLineHead struct {
51 Type string `json:"type"`
52}
53
54// SessionsDir returns the per-cwd sessions directory under root.
55func SessionsDir(root, cwd string) string {
56 sum := sha256.Sum256([]byte(cwd))
57 short := hex.EncodeToString(sum[:8])
58 return filepath.Join(root, "sessions", short)
59}
60
61// NewSession creates and opens a new session file.
62func NewSession(root, cwd, providerName, model, version string) (*Session, error) {
63 dir := SessionsDir(root, cwd)
64 if err := os.MkdirAll(dir, 0o755); err != nil {
65 return nil, err
66 }
67 id := uuid.NewString()
68 name := fmt.Sprintf("%s-%s.jsonl", time.Now().UTC().Format("20060102-150405"), id[:8])
69 p := filepath.Join(dir, name)
70 f, err := os.OpenFile(p, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644)
71 if err != nil {
72 return nil, err
73 }
74 s := &Session{
75 ID: id,
76 Path: p,
77 Meta: SessionMeta{ID: id, CWD: cwd, Provider: providerName, Model: model, Started: time.Now().UTC(), Version: version},
78 writer: f,
79 buf: bufio.NewWriter(f),
80 }
81 if err := s.writeLine(sessionLine{Type: "meta", Meta: &s.Meta}); err != nil {
82 f.Close()
83 return nil, err
84 }
85 return s, nil
86}
87
88// OpenSession opens an existing session for appending.
89func OpenSession(path string) (*Session, []provider.Message, error) {
90 f, err := os.Open(path)
91 if err != nil {
92 return nil, nil, err
93 }
94 defer f.Close()
95
96 var meta SessionMeta
97 var messages []provider.Message
98 sc := bufio.NewScanner(f)
99 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
100 for sc.Scan() {
101 var head sessionLineHead
102 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
103 continue
104 }
105 switch head.Type {
106 case "meta":
107 var row struct {
108 Meta SessionMeta `json:"meta"`
109 }
110 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
111 meta = row.Meta
112 }
113 case "message":
114 if msg, err := hydrateMessage(sc.Bytes()); err == nil && len(msg.Content) > 0 {
115 messages = append(messages, msg)
116 }
117 }
118 }
119 if err := sc.Err(); err != nil {
120 return nil, nil, err
121 }
122 out, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0o644)
123 if err != nil {
124 return nil, nil, err
125 }
126 s := &Session{ID: meta.ID, Path: path, Meta: meta, writer: out, buf: bufio.NewWriter(out)}
127 return s, messages, nil
128}
129
130// LatestSession returns the most recent session file for cwd, or "".
131func LatestSession(root, cwd string) string {
132 paths := ListSessions(root, cwd)
133 if len(paths) == 0 {
134 return ""
135 }
136 return paths[0]
137}
138
139// SessionSummary describes one on-disk session at a glance for UI pickers.
140type SessionSummary struct {
141 Path string
142 Started time.Time
143 Model string
144 Provider string
145 MessageCount int
146 FirstUserText string
147 TotalCost float64
148}
149
150// DescribeSessions returns lightweight summaries for every session in
151// cwd, newest first. Parses only the first few lines and the last usage
152// line so it's cheap to run on every dialog open.
153func DescribeSessions(root, cwd string) []SessionSummary {
154 paths := ListSessions(root, cwd)
155 summaries := make([]SessionSummary, 0, len(paths))
156 for _, p := range paths {
157 summaries = append(summaries, describeSession(p))
158 }
159 return summaries
160}
161
162func describeSession(path string) SessionSummary {
163 s := SessionSummary{Path: path}
164 f, err := os.Open(path)
165 if err != nil {
166 return s
167 }
168 defer f.Close()
169 sc := bufio.NewScanner(f)
170 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
171 for sc.Scan() {
172 var head sessionLineHead
173 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
174 continue
175 }
176 switch head.Type {
177 case "meta":
178 var row struct {
179 Meta SessionMeta `json:"meta"`
180 }
181 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
182 s.Started = row.Meta.Started
183 s.Model = row.Meta.Model
184 s.Provider = row.Meta.Provider
185 }
186 case "message":
187 s.MessageCount++
188 if s.FirstUserText == "" {
189 s.FirstUserText = firstUserText(sc.Bytes())
190 }
191 case "usage":
192 var row struct {
193 Cumulative provider.Usage `json:"cumulative"`
194 }
195 if err := json.Unmarshal(sc.Bytes(), &row); err == nil {
196 s.TotalCost = row.Cumulative.CostUSD
197 }
198 }
199 }
200 return s
201}
202
203func firstUserText(line []byte) string {
204 var row struct {
205 Message struct {
206 Role string `json:"role"`
207 Content []struct {
208 Text string `json:"text"`
209 } `json:"content"`
210 } `json:"message"`
211 }
212 if err := json.Unmarshal(line, &row); err != nil {
213 return ""
214 }
215 if row.Message.Role != "user" {
216 return ""
217 }
218 for _, c := range row.Message.Content {
219 if c.Text != "" {
220 return c.Text
221 }
222 }
223 return ""
224}
225
226// ListSessions returns session file paths for cwd, newest first.
227func ListSessions(root, cwd string) []string {
228 dir := SessionsDir(root, cwd)
229 entries, err := os.ReadDir(dir)
230 if err != nil {
231 return nil
232 }
233 var files []string
234 for _, e := range entries {
235 if !e.IsDir() && strings.HasSuffix(e.Name(), ".jsonl") {
236 files = append(files, filepath.Join(dir, e.Name()))
237 }
238 }
239 sort.Sort(sort.Reverse(sort.StringSlice(files)))
240 return files
241}
242
243// AppendMessage writes a message to the session.
244func (s *Session) AppendMessage(m provider.Message) error {
245 if s == nil {
246 return nil
247 }
248 return s.writeLine(sessionLine{Type: "message", Message: &m})
249}
250
251// UpdateModel records a provider/model switch in the session file.
252// The reader keeps the most recent meta entry, so the session resumes
253// with the updated model.
254func (s *Session) UpdateModel(providerName, model string) error {
255 if s == nil {
256 return nil
257 }
258 s.Meta.Provider = providerName
259 s.Meta.Model = model
260 return s.writeLine(sessionLine{Type: "meta", Meta: &s.Meta})
261}
262
263// AppendUsage writes a usage row to the session.
264func (s *Session) AppendUsage(u, cum provider.Usage) error {
265 if s == nil {
266 return nil
267 }
268 return s.writeLine(sessionLine{Type: "usage", Usage: &u, Cumulative: &cum})
269}
270
271// Close flushes and closes the session file.
272func (s *Session) Close() error {
273 if s == nil {
274 return nil
275 }
276 if err := s.buf.Flush(); err != nil {
277 s.writer.Close()
278 return err
279 }
280 return s.writer.Close()
281}
282
283func (s *Session) writeLine(row sessionLine) error {
284 b, err := json.Marshal(row)
285 if err != nil {
286 return err
287 }
288 if _, err := s.buf.Write(b); err != nil {
289 return err
290 }
291 if err := s.buf.WriteByte('\n'); err != nil {
292 return err
293 }
294 return s.buf.Flush()
295}
296
297// ---- content (de)serialization ----
298//
299// provider.Content is an interface; encoding/json drops type information.
300// We persist messages by reading the raw "message" object back and
301// rebuilding Content from discriminated fields.
302
303func hydrateMessage(lineBytes []byte) (provider.Message, error) {
304 var row struct {
305 Message struct {
306 Role provider.Role `json:"role"`
307 Content []json.RawMessage `json:"content"`
308 Time time.Time `json:"time"`
309 } `json:"message"`
310 }
311 if err := json.Unmarshal(lineBytes, &row); err != nil {
312 return provider.Message{}, err
313 }
314 msg := provider.Message{Role: row.Message.Role, Time: row.Message.Time}
315 for _, raw := range row.Message.Content {
316 var head struct {
317 Text string `json:"text"`
318 MimeType string `json:"mime_type"`
319 Data []byte `json:"data"`
320 ID string `json:"id"`
321 Name string `json:"name"`
322 CallID string `json:"call_id"`
323 // ToolCallBlock also has Arguments, ToolResultBlock has Content + IsError
324 }
325 if err := json.Unmarshal(raw, &head); err != nil {
326 continue
327 }
328 // Discriminate by presence of fields.
329 switch {
330 case head.Name != "" && head.ID != "":
331 var tc struct {
332 ID string `json:"id"`
333 Name string `json:"name"`
334 Arguments json.RawMessage `json:"arguments"`
335 }
336 _ = json.Unmarshal(raw, &tc)
337 msg.Content = append(msg.Content, provider.ToolCallBlock{ID: tc.ID, Name: tc.Name, Arguments: tc.Arguments})
338 case head.CallID != "":
339 var tr struct {
340 CallID string `json:"call_id"`
341 Content []json.RawMessage `json:"content"`
342 IsError bool `json:"is_error"`
343 }
344 _ = json.Unmarshal(raw, &tr)
345 block := provider.ToolResultBlock{CallID: tr.CallID, IsError: tr.IsError}
346 for _, c := range tr.Content {
347 var inner struct {
348 Text string `json:"text"`
349 MimeType string `json:"mime_type"`
350 Data []byte `json:"data"`
351 }
352 _ = json.Unmarshal(c, &inner)
353 if inner.MimeType != "" {
354 block.Content = append(block.Content, provider.ImageBlock{MimeType: inner.MimeType, Data: inner.Data})
355 } else {
356 block.Content = append(block.Content, provider.TextBlock{Text: inner.Text})
357 }
358 }
359 msg.Content = append(msg.Content, block)
360 case head.MimeType != "":
361 msg.Content = append(msg.Content, provider.ImageBlock{MimeType: head.MimeType, Data: head.Data})
362 default:
363 msg.Content = append(msg.Content, provider.TextBlock{Text: head.Text})
364 }
365 }
366 return msg, nil
367}