1package core
2
3import (
4 "bufio"
5 "bytes"
6 "encoding/json"
7 "errors"
8 "fmt"
9 "io"
10 "os"
11 "path/filepath"
12 "strings"
13 "time"
14
15 "github.com/google/uuid"
16 "github.com/patriceckhart/zot/packages/provider"
17)
18
19// PortableExt is the filesystem extension used for exported sessions.
20// A ".zotsession" is just a zot JSONL session file with the meta
21// header rewritten so the importing user gets fresh ownership.
22const PortableExt = ".zotsession"
23
24// ExportSession writes the session at srcPath to dstPath as a
25// portable .zotsession file. If dstPath is an existing directory the
26// file is created inside it with a name derived from the session's
27// meta ("YYYYMMDD-HHMMSS-<first-prompt-excerpt>.zotsession"). The
28// destination's directory is created if needed. Returns the final
29// resolved path so the caller can tell the user where it landed.
30//
31// The on-disk format is unchanged from a live session; only the
32// meta.cwd is stripped of its per-machine prefix (the importing
33// user doesn't care what directory it came from). Everything else
34// round-trips byte-for-byte.
35func ExportSession(srcPath, dstPath string) (string, error) {
36 if srcPath == "" {
37 return "", errors.New("export: source path is empty")
38 }
39 if dstPath == "" {
40 return "", errors.New("export: destination path is empty")
41 }
42
43 // Read the source meta up-front so we can name the output sensibly
44 // when dstPath is a directory, and so we can validate it's a real
45 // session before starting to write.
46 src, err := os.Open(srcPath)
47 if err != nil {
48 return "", fmt.Errorf("export: open source: %w", err)
49 }
50 defer src.Close()
51
52 sc := bufio.NewScanner(src)
53 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
54 if !sc.Scan() {
55 return "", errors.New("export: session file is empty")
56 }
57 var head sessionLine
58 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
59 return "", fmt.Errorf("export: parse meta: %w", err)
60 }
61 if head.Type != "meta" || head.Meta == nil {
62 return "", errors.New("export: first line is not a meta row")
63 }
64
65 // Scan the rest of the file for the first user message so we can
66 // build a humane filename. Only reads if dstPath doesn't already
67 // end in .zotsession.
68 firstPrompt := ""
69 if !strings.HasSuffix(strings.ToLower(dstPath), PortableExt) {
70 if fi, _ := os.Stat(dstPath); fi == nil || fi.IsDir() {
71 p, err := firstUserPrompt(src)
72 if err != nil {
73 return "", fmt.Errorf("export: read first prompt: %w", err)
74 }
75 firstPrompt = p
76 }
77 }
78
79 // Resolve dstPath: if it's a directory, build a name inside it.
80 outPath := dstPath
81 if fi, err := os.Stat(dstPath); err == nil && fi.IsDir() {
82 name := filenameFor(head.Meta.Started, head.Meta.ID, firstPrompt)
83 outPath = filepath.Join(dstPath, name)
84 } else if !strings.HasSuffix(strings.ToLower(outPath), PortableExt) {
85 outPath += PortableExt
86 }
87
88 // Re-open the source from the top since we advanced the scanner.
89 if _, err := src.Seek(0, io.SeekStart); err != nil {
90 return "", fmt.Errorf("export: rewind: %w", err)
91 }
92
93 if err := os.MkdirAll(filepath.Dir(outPath), 0o755); err != nil {
94 return "", fmt.Errorf("export: mkdir dst: %w", err)
95 }
96 dst, err := os.OpenFile(outPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
97 if err != nil {
98 return "", fmt.Errorf("export: create dst: %w", err)
99 }
100 defer dst.Close()
101 bw := bufio.NewWriter(dst)
102
103 // Rewrite the meta row: strip the cwd (the importing user has
104 // their own) and keep everything else identical. ID stays so the
105 // export is traceable; the importer will rotate to a fresh ID.
106 exportMeta := *head.Meta
107 exportMeta.CWD = ""
108 metaLine, err := json.Marshal(sessionLine{Type: "meta", Meta: &exportMeta})
109 if err != nil {
110 return "", fmt.Errorf("export: marshal meta: %w", err)
111 }
112 if _, err := bw.Write(metaLine); err != nil {
113 return "", err
114 }
115 if err := bw.WriteByte('\n'); err != nil {
116 return "", err
117 }
118
119 // Stream every non-meta row verbatim. Use ReadBytes instead of
120 // bufio.Scanner: large sessions can contain very long JSONL rows
121 // (image blocks, big tool outputs, compacted history) that exceed
122 // Scanner's token limit and fail with "token too long".
123 r := bufio.NewReader(src)
124 for {
125 line, err := r.ReadBytes('\n')
126 if len(line) > 0 {
127 line = bytes.TrimRight(line, "\r\n")
128 var h sessionLineHead
129 if err := json.Unmarshal(line, &h); err == nil && h.Type != "meta" {
130 if _, werr := bw.Write(line); werr != nil {
131 return "", werr
132 }
133 if werr := bw.WriteByte('\n'); werr != nil {
134 return "", werr
135 }
136 }
137 }
138 if err == io.EOF {
139 break
140 }
141 if err != nil {
142 return "", fmt.Errorf("export: read source: %w", err)
143 }
144 }
145 if err := bw.Flush(); err != nil {
146 return "", err
147 }
148 return outPath, nil
149}
150
151// ImportSession copies the .zotsession file at srcPath into the
152// running user's session store under the given root+cwd, rewriting
153// the meta's id / cwd / started fields so the imported session is
154// owned by the current user / directory / clock. Returns the path
155// of the created session file, ready to pass to OpenSession.
156//
157// The imported session is a first-class zot session: it'll show up
158// in /sessions, /jump, and on-disk summaries just like any other.
159// Messages and usage rows are preserved verbatim.
160func ImportSession(srcPath, root, cwd, version string) (string, error) {
161 if srcPath == "" {
162 return "", errors.New("import: source path is empty")
163 }
164 src, err := os.Open(srcPath)
165 if err != nil {
166 return "", fmt.Errorf("import: open source: %w", err)
167 }
168 defer src.Close()
169
170 // Validate the file header before committing to a destination.
171 sc := bufio.NewScanner(src)
172 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
173 if !sc.Scan() {
174 return "", errors.New("import: session file is empty")
175 }
176 var head sessionLine
177 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
178 return "", fmt.Errorf("import: parse meta: %w", err)
179 }
180 if head.Type != "meta" || head.Meta == nil {
181 return "", errors.New("import: first line is not a meta row")
182 }
183
184 // Build the destination inside the current cwd's session dir
185 // with a fresh timestamped name.
186 dir := SessionsDir(root, cwd)
187 if err := os.MkdirAll(dir, 0o755); err != nil {
188 return "", err
189 }
190 newID := uuid.NewString()
191 name := fmt.Sprintf("%s-%s.jsonl", time.Now().UTC().Format("20060102-150405"), newID[:8])
192 outPath := filepath.Join(dir, name)
193 dst, err := os.OpenFile(outPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644)
194 if err != nil {
195 return "", fmt.Errorf("import: create dst: %w", err)
196 }
197 defer dst.Close()
198 bw := bufio.NewWriter(dst)
199
200 // Write a fresh meta row claiming ownership.
201 importMeta := SessionMeta{
202 ID: newID,
203 CWD: cwd,
204 Model: head.Meta.Model,
205 Provider: head.Meta.Provider,
206 Started: time.Now().UTC(),
207 Version: version,
208 }
209 metaLine, err := json.Marshal(sessionLine{Type: "meta", Meta: &importMeta})
210 if err != nil {
211 return "", fmt.Errorf("import: marshal meta: %w", err)
212 }
213 if _, err := bw.Write(metaLine); err != nil {
214 return "", err
215 }
216 if err := bw.WriteByte('\n'); err != nil {
217 return "", err
218 }
219
220 // Rewind the source and stream every non-meta row. Avoid
221 // bufio.Scanner so exported sessions with huge JSONL rows import
222 // cleanly.
223 if _, err := src.Seek(0, io.SeekStart); err != nil {
224 return "", fmt.Errorf("import: rewind: %w", err)
225 }
226 if err := forEachJSONLLine(src, func(line []byte) error {
227 var h sessionLineHead
228 if err := json.Unmarshal(line, &h); err != nil || h.Type == "meta" {
229 return nil
230 }
231 if _, err := bw.Write(line); err != nil {
232 return err
233 }
234 return bw.WriteByte('\n')
235 }); err != nil {
236 return "", fmt.Errorf("import: read source: %w", err)
237 }
238 if err := bw.Flush(); err != nil {
239 return "", err
240 }
241 return outPath, nil
242}
243
244// BranchSession creates a new session in root/cwd that contains the
245// parent's messages 0..upToMessageIdx-1 (i.e. the first N user+
246// assistant+tool rows). The new meta records Parent=<parent id> and
247// ForkPoint=N so /session tree can rebuild the branch topology
248// later. All non-message rows (usage) are preserved up to the cut
249// point so the running cost tracker stays accurate.
250//
251// upToMessageIdx is a count over the flat message stream as
252// returned by OpenSession. To "branch at user turn 3" the caller
253// passes the index of that user message in msgs + 1 (so the
254// message itself is included). The caller figures that out; this
255// helper just copies the first N message rows.
256//
257// Returns the path of the new session file, ready for OpenSession.
258func BranchSession(parentPath, root, cwd, version string, upToMessageIdx int) (string, error) {
259 if parentPath == "" {
260 return "", errors.New("branch: parent path is empty")
261 }
262 if upToMessageIdx < 0 {
263 return "", errors.New("branch: upToMessageIdx must be >= 0")
264 }
265
266 src, err := os.Open(parentPath)
267 if err != nil {
268 return "", fmt.Errorf("branch: open parent: %w", err)
269 }
270 defer src.Close()
271
272 // Read the parent meta so we can copy model/provider and record
273 // the parent id on the child.
274 sc := bufio.NewScanner(src)
275 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
276 if !sc.Scan() {
277 return "", errors.New("branch: parent session is empty")
278 }
279 var head sessionLine
280 if err := json.Unmarshal(sc.Bytes(), &head); err != nil {
281 return "", fmt.Errorf("branch: parse parent meta: %w", err)
282 }
283 if head.Type != "meta" || head.Meta == nil {
284 return "", errors.New("branch: parent first line is not a meta row")
285 }
286 parentMeta := *head.Meta
287
288 // Build the destination file.
289 dir := SessionsDir(root, cwd)
290 if err := os.MkdirAll(dir, 0o755); err != nil {
291 return "", err
292 }
293 newID := uuid.NewString()
294 name := fmt.Sprintf("%s-%s.jsonl", time.Now().UTC().Format("20060102-150405"), newID[:8])
295 outPath := filepath.Join(dir, name)
296 dst, err := os.OpenFile(outPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644)
297 if err != nil {
298 return "", fmt.Errorf("branch: create dst: %w", err)
299 }
300 defer dst.Close()
301 bw := bufio.NewWriter(dst)
302
303 // Write the branch meta.
304 branchMeta := SessionMeta{
305 ID: newID,
306 CWD: cwd,
307 Model: parentMeta.Model,
308 Provider: parentMeta.Provider,
309 Started: time.Now().UTC(),
310 Version: version,
311 Parent: parentMeta.ID,
312 ForkPoint: upToMessageIdx,
313 }
314 metaLine, err := json.Marshal(sessionLine{Type: "meta", Meta: &branchMeta})
315 if err != nil {
316 return "", fmt.Errorf("branch: marshal meta: %w", err)
317 }
318 if _, err := bw.Write(metaLine); err != nil {
319 return "", err
320 }
321 if err := bw.WriteByte('\n'); err != nil {
322 return "", err
323 }
324
325 // Reconstruct the effective transcript the same way OpenSession
326 // does: message rows append, and compaction rows replace everything
327 // before them. The fork index is defined over that effective stream,
328 // not over the raw audit rows kept on disk before a compaction.
329 if _, err := src.Seek(0, io.SeekStart); err != nil {
330 return "", fmt.Errorf("branch: rewind parent: %w", err)
331 }
332 var effective []provider.Message
333 var nonCompactedRows [][]byte
334 effectiveCount := 0
335 sawCompaction := false
336 if err := forEachJSONLLine(src, func(line []byte) error {
337 var h sessionLineHead
338 if err := json.Unmarshal(line, &h); err != nil {
339 return nil
340 }
341 switch h.Type {
342 case "message":
343 if msg, err := hydrateMessage(line); err == nil && len(msg.Content) > 0 {
344 effective = append(effective, msg)
345 if !sawCompaction && effectiveCount < upToMessageIdx {
346 raw := append([]byte(nil), line...)
347 nonCompactedRows = append(nonCompactedRows, raw)
348 }
349 effectiveCount++
350 }
351 case "compaction":
352 if compacted, err := hydrateCompaction(line); err == nil {
353 effective = compacted
354 effectiveCount = len(effective)
355 sawCompaction = true
356 }
357 case "usage":
358 if !sawCompaction && effectiveCount < upToMessageIdx {
359 raw := append([]byte(nil), line...)
360 nonCompactedRows = append(nonCompactedRows, raw)
361 }
362 }
363 return nil
364 }); err != nil && err != io.EOF {
365 return "", fmt.Errorf("branch: read parent: %w", err)
366 }
367 if sawCompaction {
368 limit := upToMessageIdx
369 if limit > len(effective) {
370 limit = len(effective)
371 }
372 for i := 0; i < limit; i++ {
373 msg := effective[i]
374 line, err := json.Marshal(sessionLine{Type: "message", Message: &msg})
375 if err != nil {
376 return "", fmt.Errorf("branch: marshal message: %w", err)
377 }
378 if _, err := bw.Write(line); err != nil {
379 return "", err
380 }
381 if err := bw.WriteByte('\n'); err != nil {
382 return "", err
383 }
384 }
385 } else {
386 for _, row := range nonCompactedRows {
387 if _, err := bw.Write(row); err != nil {
388 return "", err
389 }
390 if err := bw.WriteByte('\n'); err != nil {
391 return "", err
392 }
393 }
394 }
395 if err := bw.Flush(); err != nil {
396 return "", err
397 }
398 return outPath, nil
399}
400
401// TreeNode is one entry in the branch tree returned by
402// BuildSessionTree. Children are populated by linking on Parent ID.
403type TreeNode struct {
404 Summary SessionSummary
405 Meta SessionMeta
406 Children []*TreeNode
407}
408
409// BuildSessionTree loads every session in the cwd dir and returns
410// the forest rooted at parentless sessions, with each non-root
411// session placed under its parent. Used by /session tree to render
412// the branch hierarchy.
413func BuildSessionTree(root, cwd string) []*TreeNode {
414 dir := SessionsDir(root, cwd)
415 entries, err := os.ReadDir(dir)
416 if err != nil {
417 return nil
418 }
419 nodes := make(map[string]*TreeNode)
420 order := []string{}
421 for _, e := range entries {
422 if e.IsDir() || !strings.HasSuffix(e.Name(), ".jsonl") {
423 continue
424 }
425 path := filepath.Join(dir, e.Name())
426 summary := describeSession(path)
427 meta, _ := readSessionMeta(path)
428 if meta.ID == "" {
429 continue
430 }
431 nodes[meta.ID] = &TreeNode{Summary: summary, Meta: meta}
432 order = append(order, meta.ID)
433 }
434 var roots []*TreeNode
435 for _, id := range order {
436 n := nodes[id]
437 if n.Meta.Parent == "" {
438 roots = append(roots, n)
439 continue
440 }
441 if parent, ok := nodes[n.Meta.Parent]; ok {
442 parent.Children = append(parent.Children, n)
443 } else {
444 // Parent file missing (was manually deleted). Treat as
445 // a root so it still shows up in the tree.
446 roots = append(roots, n)
447 }
448 }
449 return roots
450}
451
452// readSessionMeta opens path, reads the meta row, and returns it.
453// Empty SessionMeta when the file is missing or not a valid session.
454func readSessionMeta(path string) (SessionMeta, error) {
455 f, err := os.Open(path)
456 if err != nil {
457 return SessionMeta{}, err
458 }
459 defer f.Close()
460 sc := bufio.NewScanner(f)
461 sc.Buffer(make([]byte, 0, 64*1024), 20*1024*1024)
462 if !sc.Scan() {
463 return SessionMeta{}, errors.New("empty file")
464 }
465 var line sessionLine
466 if err := json.Unmarshal(sc.Bytes(), &line); err != nil {
467 return SessionMeta{}, err
468 }
469 if line.Type != "meta" || line.Meta == nil {
470 return SessionMeta{}, errors.New("first line is not meta")
471 }
472 return *line.Meta, nil
473}
474
475// FindSessionByID looks up a session file in root/cwd whose meta id
476// matches. Used by /session tree when the user picks an entry. O(n)
477// over the files in the dir; the list is small in practice.
478func FindSessionByID(root, cwd, id string) string {
479 dir := SessionsDir(root, cwd)
480 entries, err := os.ReadDir(dir)
481 if err != nil {
482 return ""
483 }
484 for _, e := range entries {
485 if e.IsDir() || !strings.HasSuffix(e.Name(), ".jsonl") {
486 continue
487 }
488 path := filepath.Join(dir, e.Name())
489 meta, err := readSessionMeta(path)
490 if err != nil {
491 continue
492 }
493 if meta.ID == id {
494 return path
495 }
496 }
497 return ""
498}
499
500// firstUserPrompt scans forward from the current source position
501// looking for the first user-role message and returns its text.
502// Used to build a humane export filename. Uses Reader instead of
503// Scanner so a very large JSONL row before the first user prompt
504// cannot trip Scanner's token limit.
505func firstUserPrompt(src io.Reader) (string, error) {
506 r := bufio.NewReader(src)
507 for {
508 lineBytes, err := r.ReadBytes('\n')
509 if len(lineBytes) > 0 {
510 lineBytes = bytes.TrimRight(lineBytes, "\r\n")
511 var line sessionLine
512 if err := json.Unmarshal(lineBytes, &line); err == nil {
513 if line.Type == "message" && line.Message != nil && line.Message.Role == "user" {
514 b, _ := json.Marshal(line.Message)
515 var m struct {
516 Content []struct {
517 Text string `json:"text"`
518 } `json:"content"`
519 }
520 _ = json.Unmarshal(b, &m)
521 for _, c := range m.Content {
522 if c.Text != "" {
523 return c.Text, nil
524 }
525 }
526 }
527 }
528 }
529 if err == io.EOF {
530 return "", nil
531 }
532 if err != nil {
533 return "", err
534 }
535 }
536}
537
538// filenameFor builds a descriptive .zotsession filename from the
539// session's start time and, when available, an excerpt of the
540// first user prompt.
541func filenameFor(started time.Time, id, firstPrompt string) string {
542 base := started.UTC().Format("20060102-150405")
543 if id != "" && len(id) >= 8 {
544 base += "-" + id[:8]
545 }
546 slug := slugify(firstPrompt, 40)
547 if slug != "" {
548 base += "-" + slug
549 }
550 return base + PortableExt
551}
552
553// slugify lowercases, strips punctuation, collapses whitespace to
554// hyphens, and truncates to max runes so it's safe as a filename.
555func slugify(s string, max int) string {
556 s = strings.TrimSpace(strings.ToLower(s))
557 if s == "" {
558 return ""
559 }
560 var out strings.Builder
561 prevDash := false
562 for _, r := range s {
563 switch {
564 case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
565 out.WriteRune(r)
566 prevDash = false
567 default:
568 if !prevDash && out.Len() > 0 {
569 out.WriteByte('-')
570 prevDash = true
571 }
572 }
573 if out.Len() >= max {
574 break
575 }
576 }
577 return strings.TrimRight(out.String(), "-")
578}