Yliken/BlackBean
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first commit

Browse Source
This commit is contained in:
Yliken
2026-08-14 23:41:57 +08:00
commit 086803a8dd
471 changed files with 91938 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pigo/internal/session/export.go
+112
View File
@@ -0,0 +1,112 @@
// This file implements session export/import (US-008, #124): a session can be
// exported to a self-contained JSONL or HTML file, and a JSONL export can be
// imported back as a fresh, resumable session. The JSONL form is the same
// role-discriminated schema the store persists, so an export → import round-trip
// is lossless (message contents and the id/parentId tree survive verbatim); the
// HTML form is a read-only, self-contained transcript with inline styles and no
// external network resources, suitable for sharing.
package session
import (
"fmt"
"io"
"os"
"path/filepath"
"strings"
"time"
)
// WriteJSONL writes header + entries as JSONL in the store's on-disk schema
// (header line first, then one entry line each, ids/parentIds preserved). It is
// the export counterpart to writeSessionEntries and the exact input ReadJSONL
// expects, so a WriteJSONL → ReadJSONL round-trip is lossless.
func WriteJSONL(w io.Writer, header SessionHeader, entries []Entry) error {
header.Version = SchemaVersion
return writeSessionEntries(w, header, entries)
}
// ReadJSONL decodes a JSONL export (as produced by WriteJSONL or a raw session
// file) into a header and entries, migrating v1/v2 bare-message files the same
// way LoadEntries does. It is the import counterpart to WriteJSONL.
func ReadJSONL(r io.Reader) (SessionHeader, []Entry, error) {
return readSession(r)
}
// Export writes the session identified by id to outPath. The format is chosen
// by outPath's extension: ".html"/".htm" produces a self-contained HTML
// transcript; anything else (including ".jsonl") produces JSONL. The parent
// directory of outPath must already exist. It returns the number of entries
// written so a caller can report progress.
func (s *Store) Export(id, outPath string) (int, error) {
header, entries, err := s.LoadEntries(id)
if err != nil {
return 0, err
}
f, err := os.Create(outPath)
if err != nil {
return 0, fmt.Errorf("session: create export %s: %w", outPath, err)
}
defer f.Close()
ext := strings.ToLower(filepath.Ext(outPath))
if ext == ".html" || ext == ".htm" {
if err := WriteHTML(f, header, entries); err != nil {
return 0, err
}
} else {
if err := WriteJSONL(f, header, entries); err != nil {
return 0, err
}
}
if err := f.Close(); err != nil {
return 0, fmt.Errorf("session: finalize export %s: %w", outPath, err)
}
return len(entries), nil
}
// Import reads a JSONL export at inPath and materializes it as a fresh session
// in the store: a new id (derived from now) is assigned, the original id is
// recorded as ParentSession for lineage, and the entries are written verbatim
// (ids/parentIds preserved) so the tree — and thus PathToLeaf/resume — behaves
// exactly as in the source. It returns the new header and the imported entries.
// An HTML file (or any non-JSONL input) fails to parse and returns an error
// rather than importing garbage.
func (s *Store) Import(inPath string, now time.Time) (SessionHeader, []Entry, error) {
f, err := os.Open(inPath)
if err != nil {
return SessionHeader{}, nil, fmt.Errorf("session: open import %s: %w", inPath, err)
}
defer f.Close()
srcHeader, entries, err := ReadJSONL(f)
if err != nil {
return SessionHeader{}, nil, err
}
newHeader := SessionHeader{
ID: NewID(now),
CreatedAt: now,
UpdatedAt: now,
Model: srcHeader.Model,
Provider: srcHeader.Provider,
SystemPrompt: srcHeader.SystemPrompt,
ParentSession: srcHeader.ID,
}
if err := s.SaveEntries(newHeader, entries); err != nil {
return SessionHeader{}, nil, err
}
return newHeader, entries, nil
}
// WriteHTML writes a self-contained HTML transcript of the session: inline CSS
// only (no external stylesheets, fonts, scripts, or network resources), role
// color-coding, and tool-call/result blocks. All message text is HTML-escaped
// so a transcript containing markup or a crafted "</script>" cannot break out of
// its container or inject active content (defensive against a hostile session).
func WriteHTML(w io.Writer, header SessionHeader, entries []Entry) error {
var b strings.Builder
b.WriteString(htmlHead(header))
for _, e := range entries {
b.WriteString(renderEntryHTML(e))
}
b.WriteString(htmlFoot())
_, err := io.WriteString(w, b.String())
return err
}
pigo/internal/session/export_html.go
+125
View File
@@ -0,0 +1,125 @@
// This file holds the HTML rendering helpers for session export (US-008, #124).
// The output is a single self-contained document: all CSS is inlined in a
// <style> block, there are no <script> tags, no external fonts, and no network
// requests, so the transcript renders identically offline and cannot phone home.
// Every piece of session-derived text is escaped via html.EscapeString before it
// reaches the document.
package session
import (
"fmt"
"html"
"strings"
"time"
"github.com/smallnest/pigo/internal/agentcore"
)
// htmlHead emits the document head + opening container, embedding the session
// metadata (all escaped). The CSS is deliberately inline and minimal so the file
// is self-contained and depends on no external resource.
func htmlHead(header SessionHeader) string {
title := header.ID
if title == "" {
title = "session"
}
meta := []string{}
if header.Model != "" {
meta = append(meta, "model: "+html.EscapeString(header.Model))
}
if header.Provider != "" {
meta = append(meta, "provider: "+html.EscapeString(header.Provider))
}
if !header.UpdatedAt.IsZero() {
meta = append(meta, "updated: "+html.EscapeString(header.UpdatedAt.Format(time.RFC3339)))
}
return fmt.Sprintf(`<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>pigo session — %s</title>
<style>
* { margin: 0; padding: 0; box-sizing: border-box; }
body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; background: #FAF9F6; color: #1a1a1a; padding: 2rem 1.5rem; line-height: 1.6; }
.container { max-width: 820px; margin: 0 auto; }
h1 { font-size: 1.25rem; font-weight: 700; margin-bottom: 0.25rem; }
.meta { color: #8B8680; font-size: 0.8rem; margin-bottom: 1.5rem; }
.msg { border-radius: 8px; padding: 0.75rem 1rem; margin-bottom: 0.75rem; border: 1px solid #E8E4DE; background: #fff; }
.role { font-size: 0.7rem; font-weight: 600; text-transform: uppercase; letter-spacing: 0.05em; margin-bottom: 0.4rem; }
.msg.user { background: #F0F4F8; border-color: #4A6FA5; }
.msg.user .role { color: #4A6FA5; }
.msg.assistant { background: #F5F8F6; border-color: #5B8A72; }
.msg.assistant .role { color: #5B8A72; }
.msg.tool { background: #FDF8F0; border-color: #D4A843; }
.msg.tool .role { color: #B8860B; }
.msg.compaction { background: #FFF5F0; border-color: #D97757; }
.msg.compaction .role { color: #D97757; }
.text { white-space: pre-wrap; word-break: break-word; font-size: 0.875rem; }
.toolcall { margin-top: 0.5rem; padding: 0.5rem 0.75rem; background: #F0EEE9; border-radius: 6px; font-family: ui-monospace, SFMono-Regular, Menlo, monospace; font-size: 0.78rem; white-space: pre-wrap; word-break: break-word; }
.toolcall .tname { font-weight: 600; color: #B8860B; }
.footer { text-align: center; margin-top: 2rem; color: #B0AAA4; font-size: 0.7rem; }
</style>
</head>
<body>
<div class="container">
<h1>pigo session — %s</h1>
<p class="meta">%s</p>
`, html.EscapeString(title), html.EscapeString(title), strings.Join(meta, " · "))
}
// htmlFoot closes the container and document.
func htmlFoot() string {
return `<p class="footer">Exported by pigo /export</p>
</div>
</body>
</html>
`
}
// renderEntryHTML renders one entry as a role-colored message block. Text and
// tool arguments are escaped so no session content can inject markup.
func renderEntryHTML(e Entry) string {
switch m := e.Message.(type) {
case agentcore.UserMessage:
return msgBlock("user", "User", html.EscapeString(agentcore.ContentToText(m.Content)), "")
case agentcore.AssistantMessage:
var tools strings.Builder
for _, c := range m.ToolCalls() {
args := strings.TrimSpace(string(c.Arguments))
tools.WriteString(fmt.Sprintf(`<div class="toolcall"><span class="tname">→ %s</span> %s</div>`,
html.EscapeString(c.Name), html.EscapeString(args)))
}
return msgBlock("assistant", "Assistant", html.EscapeString(agentcore.ContentToText(m.Content)), tools.String())
case agentcore.ToolResultMessage:
label := "Tool Result"
if m.ToolName != "" {
label = "Tool Result: " + m.ToolName
}
return msgBlock("tool", html.EscapeString(label), html.EscapeString(agentcore.ContentToText(m.Content)), "")
case agentcore.CompactionMessage:
return msgBlock("compaction", "Compaction", html.EscapeString(m.Summary), "")
default:
return msgBlock("assistant", html.EscapeString(e.Message.Role()), "", "")
}
}
// msgBlock assembles one .msg block with a role class, a role label, the escaped
// body text, and optional pre-rendered tool-call HTML. Callers MUST pass already
// escaped text/label; extra is trusted HTML built here from escaped parts.
func msgBlock(class, label, escapedText, extra string) string {
var b strings.Builder
b.WriteString(`<div class="msg `)
b.WriteString(class)
b.WriteString(`"><div class="role">`)
b.WriteString(label)
b.WriteString(`</div>`)
if escapedText != "" {
b.WriteString(`<div class="text">`)
b.WriteString(escapedText)
b.WriteString(`</div>`)
}
b.WriteString(extra)
b.WriteString("</div>\n")
return b.String()
}
pigo/internal/session/export_test.go
+197
View File
@@ -0,0 +1,197 @@
package session
// Tests for session export/import (US-008, #124). They cover the lossless
// JSONL round-trip (export → import preserves the header fields, message
// sequence, and entry tree), the self-contained HTML export (inline styles, no
// external network resources, HTML-escaped content), and format selection by
// file extension.
import (
"bytes"
"os"
"path/filepath"
"strings"
"testing"
"time"
"github.com/smallnest/pigo/internal/agentcore"
)
// seedSession writes a small multi-turn session and returns its header.
func seedSession(t *testing.T, s *Store) SessionHeader {
t.Helper()
now := time.Date(2026, 7, 17, 9, 0, 0, 0, time.UTC)
header := SessionHeader{
ID: NewID(now),
CreatedAt: now,
UpdatedAt: now,
Model: "anthropic/claude-opus-4",
Provider: "anthropic",
SystemPrompt: "You are pigo.",
}
if err := s.Save(header, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
return header
}
// TestExportImportJSONLRoundTrip is the core acceptance check: exporting a
// session to JSONL and importing it back yields the same message sequence and
// header carry-over, with the source id recorded as the import's parent.
func TestExportImportJSONLRoundTrip(t *testing.T) {
s := newStore(t)
header := seedSession(t, s)
out := filepath.Join(t.TempDir(), "export.jsonl")
n, err := s.Export(header.ID, out)
if err != nil {
t.Fatalf("Export: %v", err)
}
if n != len(sampleMessages()) {
t.Errorf("exported %d entries, want %d", n, len(sampleMessages()))
}
now := time.Date(2026, 7, 17, 10, 0, 0, 0, time.UTC)
newHeader, entries, err := s.Import(out, now)
if err != nil {
t.Fatalf("Import: %v", err)
}
if newHeader.ID == header.ID {
t.Errorf("import must assign a fresh id, got the source id %q", newHeader.ID)
}
if newHeader.ParentSession != header.ID {
t.Errorf("ParentSession = %q, want source id %q", newHeader.ParentSession, header.ID)
}
if newHeader.Model != header.Model || newHeader.Provider != header.Provider || newHeader.SystemPrompt != header.SystemPrompt {
t.Errorf("header fields not carried over: %+v", newHeader)
}
if len(entries) != len(sampleMessages()) {
t.Fatalf("imported %d entries, want %d", len(entries), len(sampleMessages()))
}
// The imported session must be independently loadable with the same roles.
_, gotMsgs, err := s.Load(newHeader.ID)
if err != nil {
t.Fatalf("Load imported: %v", err)
}
wantRoles := []string{agentcore.RoleUser, agentcore.RoleAssistant, agentcore.RoleToolResult, agentcore.RoleAssistant}
if len(gotMsgs) != len(wantRoles) {
t.Fatalf("imported message count = %d, want %d", len(gotMsgs), len(wantRoles))
}
for i, m := range gotMsgs {
if m.Role() != wantRoles[i] {
t.Errorf("message[%d] role = %q, want %q", i, m.Role(), wantRoles[i])
}
}
// The tool call must survive the round-trip.
a, ok := gotMsgs[1].(agentcore.AssistantMessage)
if !ok {
t.Fatalf("message[1] is not AssistantMessage: %T", gotMsgs[1])
}
if calls := a.ToolCalls(); len(calls) != 1 || calls[0].Name != "read" {
t.Errorf("tool calls = %+v, want one 'read'", calls)
}
}
// TestWriteReadJSONLPreservesTree checks the lower-level primitives: WriteJSONL
// then ReadJSONL preserves entry ids and parentIds verbatim (so the tree — and
// thus resume/PathToLeaf — behaves identically).
func TestWriteReadJSONLPreservesTree(t *testing.T) {
s := newStore(t)
header := seedSession(t, s)
_, entries, err := s.LoadEntries(header.ID)
if err != nil {
t.Fatalf("LoadEntries: %v", err)
}
var buf bytes.Buffer
if err := WriteJSONL(&buf, header, entries); err != nil {
t.Fatalf("WriteJSONL: %v", err)
}
gotHeader, gotEntries, err := ReadJSONL(&buf)
if err != nil {
t.Fatalf("ReadJSONL: %v", err)
}
if gotHeader.Version != SchemaVersion {
t.Errorf("version = %d, want %d", gotHeader.Version, SchemaVersion)
}
if len(gotEntries) != len(entries) {
t.Fatalf("entry count = %d, want %d", len(gotEntries), len(entries))
}
for i := range entries {
if gotEntries[i].ID != entries[i].ID || gotEntries[i].ParentID != entries[i].ParentID {
t.Errorf("entry[%d] tree ids changed: got {%q,%q} want {%q,%q}",
i, gotEntries[i].ID, gotEntries[i].ParentID, entries[i].ID, entries[i].ParentID)
}
}
}
// TestExportHTMLSelfContained verifies the HTML export is a self-contained
// document: it carries inline styles, has no external network resources (no
// http(s):// URLs, no <script>, no <link>), and HTML-escapes session content.
func TestExportHTMLSelfContained(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 7, 17, 9, 0, 0, 0, time.UTC)
header := SessionHeader{
ID: NewID(now),
CreatedAt: now,
UpdatedAt: now,
Model: "anthropic/claude-opus-4",
Provider: "anthropic",
}
// Include a message with HTML-significant characters to exercise escaping.
msgs := agentcore.MessageList{
agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("<script>alert('xss')</script>")}},
agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("safe & sound")}, StopReason: agentcore.StopReasonEndTurn},
}
if err := s.Save(header, msgs); err != nil {
t.Fatalf("Save: %v", err)
}
out := filepath.Join(t.TempDir(), "export.html")
if _, err := s.Export(header.ID, out); err != nil {
t.Fatalf("Export html: %v", err)
}
data, err := os.ReadFile(out)
if err != nil {
t.Fatalf("ReadFile: %v", err)
}
doc := string(data)
if !strings.Contains(doc, "<style>") {
t.Error("HTML export must inline a <style> block")
}
if strings.Contains(doc, "<script") {
t.Error("HTML export must not contain <script> tags")
}
if strings.Contains(doc, "<link") {
t.Error("HTML export must not reference external stylesheets via <link>")
}
if strings.Contains(doc, "http://") || strings.Contains(doc, "https://") {
t.Error("HTML export must not reference external network resources (http/https)")
}
// The raw script tag from the user message must be escaped, not live.
if strings.Contains(doc, "<script>alert") {
t.Error("session content must be HTML-escaped (found live <script>alert)")
}
if !strings.Contains(doc, "&lt;script&gt;alert") {
t.Error("expected escaped user content &lt;script&gt;alert in HTML export")
}
if !strings.Contains(doc, "safe &amp; sound") {
t.Error("expected escaped ampersand in assistant content")
}
}
// TestImportRejectsHTML verifies importing a non-JSONL file (e.g. an HTML
// export) fails rather than materializing garbage.
func TestImportRejectsHTML(t *testing.T) {
s := newStore(t)
header := seedSession(t, s)
out := filepath.Join(t.TempDir(), "export.html")
if _, err := s.Export(header.ID, out); err != nil {
t.Fatalf("Export html: %v", err)
}
if _, _, err := s.Import(out, time.Now()); err == nil {
t.Error("Import of an HTML file should fail, got nil error")
}
}
pigo/internal/session/inherit.go
+56
View File
@@ -0,0 +1,56 @@
package session
// Context inheritance plumbing for the "infinite context" feature (#483, built
// on the header fields added by #480). A session created to *continue* another
// session's collapsed context records two things on its header:
//
// - ContextFrom: the source session id whose checkpoint holds the collapsed
// prefix (see runtime.LoadCheckpoint / CheckpointPath).
// - ContextWatermark: the message index up to which that prefix was collapsed;
// on resume/fork the runtime (#481) can load the source checkpoint and skip
// re-summarizing messages [0, ContextWatermark).
//
// This is deliberately side-effect-free header plumbing: it never touches the
// memory dir or the runtime. It sits alongside ParentSession (which records
// lineage) — inheritance additionally records *where the collapsed context came
// from* so a forked/continued session need not re-summarize early context.
// SetContextInheritance records on header that this session continues fromID's
// collapsed context up to watermark. A non-positive watermark or an empty fromID
// clears inheritance (both fields reset to their zero/omitempty state), so the
// header round-trips as a session with no inherited checkpoint. This mirrors how
// ParentSession is a plain header field set at creation time.
func SetContextInheritance(header *SessionHeader, fromID string, watermark int) {
if header == nil {
return
}
if fromID == "" || watermark <= 0 {
header.ContextFrom = ""
header.ContextWatermark = 0
return
}
header.ContextFrom = fromID
header.ContextWatermark = watermark
}
// ContextInheritance reports the inherited-context source recorded on header:
// the source session id, the watermark, and whether inheritance is set. It is
// the read counterpart to SetContextInheritance — a convenience accessor the
// runtime (#481) uses on resume/fork to decide whether to load the source
// checkpoint. ok is true only when a non-empty source and a positive watermark
// are both present, so a header with only one of the two (a malformed or
// partially written record) reports ok=false rather than a half-configured
// inheritance.
func (h SessionHeader) ContextInheritance() (fromID string, watermark int, ok bool) {
if h.ContextFrom == "" || h.ContextWatermark <= 0 {
return "", 0, false
}
return h.ContextFrom, h.ContextWatermark, true
}
// HasContextInheritance reports whether header declares inherited collapsed
// context. It is shorthand for the ok return of ContextInheritance.
func (h SessionHeader) HasContextInheritance() bool {
_, _, ok := h.ContextInheritance()
return ok
}
pigo/internal/session/session.go
+700
View File
@@ -0,0 +1,700 @@
// Package session implements local JSONL session persistence and resume
// (US-024, #43). A session is stored as a single append-only JSONL file: the
// first line is a SessionHeader (schema version + metadata), and every
// subsequent line is one persisted message (user / assistant / toolResult),
// using the same "role"-discriminated encoding as agentcore.MessageList.
//
// The format is internally self-consistent and deliberately NOT wire-compatible
// with pi's session files (spec #16, session-format decision #5): pigo owns the schema
// and versions it via SessionHeader.Version so future migrations have a hook.
//
// A persisted session round-trips into an agentcore.AgentContext via Load, so a
// run can be resumed by feeding the reconstructed context to a fresh run and the
// transcript replays correctly in the REPL.
package session
import (
"bufio"
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"os"
"path/filepath"
"sort"
"strings"
"time"
"github.com/smallnest/pigo/internal/agentcore"
)
// SchemaVersion is the current session file schema version. It is written into
// every SessionHeader and checked on read; an unknown (higher) version is a
// hard error so a newer file is never silently misread by an older binary.
//
// v2 adds the inline "compaction" message role (US-003): a compaction
// checkpoint persisted as one message line.
//
// v3 gives every persisted entry an id/parentId, forming a tree (US-005, #121)
// — the prerequisite for fork/clone/tree navigation. Each message line is
// wrapped as {"id","parentId","timestamp","message":{…}}. v1/v2 files (bare
// message lines) remain fully readable: readSession migrates them on load by
// synthesizing ids and chaining parentId to the previous entry, so old sessions
// still load and resume.
const SchemaVersion = 3
// sessionScanBufInit / sessionScanBufMax bound the line scanner used to read a
// session file. A single line holds one message, which can be large (a long
// tool result), so the max is raised well past bufio.Scanner's 64KiB default.
const (
sessionScanBufInit = 64 * 1024
sessionScanBufMax = 16 * 1024 * 1024
)
// SessionHeader is the first line of a session file: schema version plus the
// metadata needed to list and resume a session without reading its messages.
type SessionHeader struct {
// Version is the schema version (SchemaVersion at write time).
Version int `json:"version"`
// ID is the session identifier, also the file stem (see FileName).
ID string `json:"id"`
// CreatedAt is when the session file was created (RFC 3339, UTC).
CreatedAt time.Time `json:"createdAt"`
// UpdatedAt is when the session was last appended to.
UpdatedAt time.Time `json:"updatedAt"`
// Model / Provider record what the session ran against, for display and to
// re-establish the run configuration on resume. Optional.
Model string `json:"model,omitempty"`
Provider string `json:"provider,omitempty"`
// SystemPrompt is the system prompt the session ran under. Persisted so the
// resumed context is faithful. Optional.
SystemPrompt string `json:"systemPrompt,omitempty"`
// ParentSession is the id of the session this one was forked/cloned from
// (US-006, #122). Empty for a session created from scratch. It records
// lineage only; a fork is otherwise a fully independent session file.
ParentSession string `json:"parentSession,omitempty"`
// Cwd is the absolute working directory the session ran in, recorded so a
// session can be attributed to a project (its stable project id derives from
// this path). Used by /dream distillation to select the current project's
// recent sessions (SPEC §5.3 / §11.3). Optional and additive: older schemas
// omit it and still load; a session with an empty Cwd is treated as
// unattributed and never matches a project-scoped distill.
Cwd string `json:"cwd,omitempty"`
// ContextFrom is the id of the session this one inherited its collapsed
// context from (#480, "infinite context"). Empty when the session started
// with no inherited checkpoint. Optional and additive: older schemas (v1/v2/v3)
// omit it and still load.
ContextFrom string `json:"contextFrom,omitempty"`
// ContextWatermark is the message index up to which context was collapsed
// into an inherited checkpoint (#480). Messages before this index live in the
// checkpoint summary rather than the replayed transcript. Optional/additive.
ContextWatermark int `json:"contextWatermark,omitempty"`
}
// Entry wraps one persisted message with the tree metadata introduced in schema
// v3 (US-005, #121): a stable ID plus the ParentID it descends from. A linear
// session is the degenerate tree where every entry's ParentID is the previous
// entry's ID; the first entry has an empty ParentID (a root). The wrapper is
// what lets a session fork/clone later — PathToLeaf walks the ParentID chain to
// reconstruct the linear conversation feeding any leaf.
type Entry struct {
// ID is this entry's stable identifier (see newEntryID). Unique within a file.
ID string
// ParentID is the ID this entry descends from; empty for a root entry.
ParentID string
// Timestamp is when the entry was persisted (RFC 3339, UTC).
Timestamp time.Time
// Message is the wrapped agent message (user / assistant / toolResult / …).
Message agentcore.Message
}
// entryWire is the on-disk JSON shape of an Entry: the message is carried as a
// raw object so it round-trips through MessageList's role-discriminated decoder
// (agentcore.Message is a sealed interface with no default unmarshaler).
type entryWire struct {
ID string `json:"id"`
ParentID string `json:"parentId,omitempty"`
Timestamp time.Time `json:"timestamp"`
Message json.RawMessage `json:"message"`
}
// MarshalJSON emits the entry as {"id","parentId","timestamp","message":{…}}.
func (e Entry) MarshalJSON() ([]byte, error) {
mb, err := json.Marshal(e.Message)
if err != nil {
return nil, fmt.Errorf("session: encode entry message: %w", err)
}
return json.Marshal(entryWire{ID: e.ID, ParentID: e.ParentID, Timestamp: e.Timestamp, Message: mb})
}
// UnmarshalJSON decodes an entry line, decoding the inner message with the same
// discriminated logic as agentcore.MessageList (by wrapping it in a one-element
// array).
func (e *Entry) UnmarshalJSON(data []byte) error {
var w entryWire
if err := json.Unmarshal(data, &w); err != nil {
return err
}
e.ID = w.ID
e.ParentID = w.ParentID
e.Timestamp = w.Timestamp
if len(w.Message) == 0 {
return fmt.Errorf("session: entry missing message")
}
var one agentcore.MessageList
if err := json.Unmarshal([]byte("["+string(w.Message)+"]"), &one); err != nil {
return fmt.Errorf("session: decode entry message: %w", err)
}
if len(one) != 1 {
return fmt.Errorf("session: entry decoded to %d messages, want 1", len(one))
}
e.Message = one[0]
return nil
}
// newEntryID returns a fresh 8-hex-character entry id (4 random bytes). This
// mirrors pi's generateEntryId (uuidv7().slice(-8)) in width; pigo does not need
// the time-ordering of uuidv7 because entry order is already given by the
// ParentID chain, so a simple random id suffices and collisions within a single
// file are astronomically unlikely.
func newEntryID() string {
var b [4]byte
if _, err := rand.Read(b[:]); err != nil {
// crypto/rand.Read never fails in practice; fall back to a timestamp-derived
// id so a session write never aborts on this path.
return fmt.Sprintf("%08x", time.Now().UnixNano()&0xffffffff)
}
return hex.EncodeToString(b[:])
}
// PathToLeaf walks the ParentID chain from the entry identified by leafID back
// to its root and returns the entries in root→leaf order — the linear
// conversation that feeds the given leaf (US-005 acceptance criterion c). An
// empty leafID (or an unknown one) yields nil. If the chain references a missing
// parent the walk stops at the last resolvable ancestor rather than failing, so
// a partially corrupt file still yields the recoverable prefix.
func PathToLeaf(entries []Entry, leafID string) []Entry {
if leafID == "" {
return nil
}
byID := make(map[string]Entry, len(entries))
for _, e := range entries {
byID[e.ID] = e
}
var rev []Entry
seen := make(map[string]bool, len(entries))
for id := leafID; id != ""; {
e, ok := byID[id]
if !ok || seen[id] {
break // missing parent or a cycle: stop at the last good ancestor
}
seen[id] = true
rev = append(rev, e)
id = e.ParentID
}
// rev is leaf→root; reverse to root→leaf.
for i, j := 0, len(rev)-1; i < j; i, j = i+1, j-1 {
rev[i], rev[j] = rev[j], rev[i]
}
return rev
}
// FileName returns the on-disk file name for a session id (id + ".jsonl").
func FileName(id string) string { return id + ".jsonl" }
// TreeLine pairs one entry with its rendered display line, so a caller can print
// the tree AND map a 1-based selection index back to the entry it refers to (the
// slice order is the render order — a stable pre-order DFS). See RenderTreeLines.
type TreeLine struct {
Entry Entry
Text string
}
// RenderTreeLines renders the entry forest as human-readable lines for a pure
// line REPL — no TUI, no cursor control (US-007, #123). Each entry becomes one
// line with ├─/└─ connectors showing structure; the entry whose id == leafID is
// tagged "← current" so the active branch is obvious. Roots (entries with an
// empty or dangling ParentID) anchor the forest; children are ordered by
// timestamp then id for stable output. The returned slice is in render order, so
// element i corresponds to the i-th printed line (and 1-based selector n → [n-1]).
func RenderTreeLines(entries []Entry, leafID string) []TreeLine {
present := make(map[string]bool, len(entries))
for _, e := range entries {
present[e.ID] = true
}
childrenOf := make(map[string][]Entry, len(entries))
var roots []Entry
for _, e := range entries {
if e.ParentID == "" || !present[e.ParentID] {
roots = append(roots, e)
} else {
childrenOf[e.ParentID] = append(childrenOf[e.ParentID], e)
}
}
sortEntries(roots)
for k := range childrenOf {
sortEntries(childrenOf[k])
}
var lines []TreeLine
var walk func(e Entry, prefix string, isRoot, isLast bool)
walk = func(e Entry, prefix string, isRoot, isLast bool) {
connector := ""
if !isRoot {
if isLast {
connector = "└─ "
} else {
connector = "├─ "
}
}
marker := ""
if e.ID == leafID {
marker = " ← current"
}
lines = append(lines, TreeLine{Entry: e, Text: prefix + connector + entrySummary(e) + marker})
childPrefix := prefix
if !isRoot {
if isLast {
childPrefix += " "
} else {
childPrefix += "│ "
}
}
kids := childrenOf[e.ID]
for i, k := range kids {
walk(k, childPrefix, false, i == len(kids)-1)
}
}
for i, r := range roots {
walk(r, "", true, i == len(roots)-1)
}
return lines
}
// sortEntries orders entries by timestamp, breaking ties by id so the render is
// deterministic even when entries share a timestamp (common within one turn).
func sortEntries(es []Entry) {
sort.SliceStable(es, func(i, j int) bool {
if !es[i].Timestamp.Equal(es[j].Timestamp) {
return es[i].Timestamp.Before(es[j].Timestamp)
}
return es[i].ID < es[j].ID
})
}
// entrySummary renders a one-line, role-tagged preview of an entry's message for
// the tree display (a full message would wrap and ruin the ASCII structure).
func entrySummary(e Entry) string {
switch m := e.Message.(type) {
case agentcore.UserMessage:
return "user: " + treeOneLine(agentcore.ContentToText(m.Content))
case agentcore.AssistantMessage:
text := treeOneLine(agentcore.ContentToText(m.Content))
calls := m.ToolCalls()
if len(calls) > 0 {
names := make([]string, len(calls))
for i, c := range calls {
names[i] = c.Name
}
tools := "[→ " + strings.Join(names, ", ") + "]"
if text != "" {
return "assistant: " + text + " " + tools
}
return "assistant " + tools
}
return "assistant: " + text
case agentcore.ToolResultMessage:
return "tool result: " + treeOneLine(agentcore.ContentToText(m.Content))
case agentcore.CompactionMessage:
return "compaction: " + treeOneLine(m.Summary)
default:
return e.Message.Role()
}
}
// treeOneLine collapses a possibly multi-line message into a single trimmed,
// truncated line so tree rows stay on one physical line.
func treeOneLine(s string) string {
s = strings.TrimSpace(s)
if i := strings.IndexByte(s, '\n'); i >= 0 {
s = s[:i] + " …"
}
const max = 72
if len(s) > max {
s = s[:max] + " …"
}
return s
}
// NewID returns a time-ordered session id: a UTC timestamp stem that sorts
// lexicographically by creation time (e.g. "20260710-142530-uniq"). The suffix
// disambiguates sessions created within the same second.
func NewID(now time.Time) string {
return fmt.Sprintf("%s-%06d", now.UTC().Format("20060102-150405"), now.UTC().Nanosecond()/1000%1_000_000)
}
// Store persists sessions as JSONL files under a directory (typically
// ~/.pigo/sessions). The zero value is unusable; construct with NewStore.
type Store struct {
dir string
}
// NewStore returns a Store rooted at dir, creating the directory if needed.
func NewStore(dir string) (*Store, error) {
if dir == "" {
return nil, fmt.Errorf("session: store dir must not be empty")
}
if err := os.MkdirAll(dir, 0o755); err != nil {
return nil, fmt.Errorf("session: create store dir: %w", err)
}
return &Store{dir: dir}, nil
}
// Dir returns the store's root directory.
func (s *Store) Dir() string { return s.dir }
// path returns the on-disk path for a session id.
func (s *Store) path(id string) string { return filepath.Join(s.dir, FileName(id)) }
// Save writes header and messages to a fresh session file, overwriting any
// existing file for the same id. The header's Version is forced to
// SchemaVersion; CreatedAt/UpdatedAt are left as the caller set them. Save is
// the whole-session write; Append adds messages to an existing file.
func (s *Store) Save(header SessionHeader, messages agentcore.MessageList) error {
header.Version = SchemaVersion
if header.ID == "" {
return fmt.Errorf("session: header ID must not be empty")
}
return s.atomicWrite(header.ID, func(w io.Writer) error {
return writeSession(w, header, messages)
})
}
// SaveEntries writes header plus the given entries verbatim — preserving each
// entry's id/parentId — to a fresh session file, overwriting any existing file
// for header.ID. Unlike Save (which generates fresh ids and a linear chain from
// a MessageList), SaveEntries persists an already-known tree, which is what Fork
// needs: it copies a path of existing entries into a new session file without
// disturbing their identifiers. The header's Version is forced to SchemaVersion.
func (s *Store) SaveEntries(header SessionHeader, entries []Entry) error {
header.Version = SchemaVersion
if header.ID == "" {
return fmt.Errorf("session: header ID must not be empty")
}
return s.atomicWrite(header.ID, func(w io.Writer) error {
return writeSessionEntries(w, header, entries)
})
}
// atomicWrite writes a session file for id by streaming through write into a
// temp file and atomically renaming it into place, so a concurrent reader never
// sees a half-written file. It is the shared write plumbing behind Save and
// SaveEntries.
func (s *Store) atomicWrite(id string, write func(w io.Writer) error) error {
tmp := s.path(id) + ".tmp"
f, err := os.Create(tmp)
if err != nil {
return fmt.Errorf("session: create %s: %w", tmp, err)
}
w := bufio.NewWriter(f)
if err := write(w); err != nil {
f.Close()
os.Remove(tmp)
return err
}
if err := w.Flush(); err != nil {
f.Close()
os.Remove(tmp)
return fmt.Errorf("session: flush %s: %w", tmp, err)
}
if err := f.Close(); err != nil {
os.Remove(tmp)
return fmt.Errorf("session: close %s: %w", tmp, err)
}
// Atomic replace so a reader never sees a half-written file.
if err := os.Rename(tmp, s.path(id)); err != nil {
os.Remove(tmp)
return fmt.Errorf("session: commit %s: %w", id, err)
}
return nil
}
// writeSession emits the header line followed by one entry line per message.
// Each message is wrapped in an Entry: ids are generated fresh and ParentID is
// chained to the previous entry so a linear session is persisted as a linear
// tree (schema v3). The chain is what PathToLeaf later walks.
func writeSession(w io.Writer, header SessionHeader, messages agentcore.MessageList) error {
enc := json.NewEncoder(w)
if err := enc.Encode(header); err != nil {
return fmt.Errorf("session: encode header: %w", err)
}
now := time.Now().UTC()
parentID := ""
for i, m := range messages {
e := Entry{ID: newEntryID(), ParentID: parentID, Timestamp: now, Message: m}
if err := enc.Encode(e); err != nil {
return fmt.Errorf("session: encode message[%d]: %w", i, err)
}
parentID = e.ID
}
return nil
}
// writeSessionEntries emits the header line followed by the given entries
// verbatim, preserving their ids and parentIds. It is the whole-tree counterpart
// to writeSession (which synthesizes a fresh linear chain from a MessageList).
func writeSessionEntries(w io.Writer, header SessionHeader, entries []Entry) error {
enc := json.NewEncoder(w)
if err := enc.Encode(header); err != nil {
return fmt.Errorf("session: encode header: %w", err)
}
for i, e := range entries {
if err := enc.Encode(e); err != nil {
return fmt.Errorf("session: encode entry[%d]: %w", i, err)
}
}
return nil
}
// Load reads a session file and returns its header and messages. It validates
// the schema version (an unknown version is rejected) and decodes each entry
// line, returning the messages in file order (root→leaf for a linear session).
// Old v1/v2 files (bare message lines) are migrated transparently. Load is the
// linear view; LoadEntries exposes the id/parentId tree metadata.
func (s *Store) Load(id string) (SessionHeader, agentcore.MessageList, error) {
header, entries, err := s.LoadEntries(id)
if err != nil {
return SessionHeader{}, nil, err
}
msgs := make(agentcore.MessageList, len(entries))
for i, e := range entries {
msgs[i] = e.Message
}
return header, msgs, nil
}
// LoadEntries reads a session file and returns its header plus the tree entries
// (id/parentId + message) in file order. v1/v2 files are migrated on load:
// every bare message line is wrapped in an Entry with a synthesized id and its
// ParentID chained to the previous entry, so old sessions load and resume
// exactly as they did before (US-005 acceptance criterion b/d).
func (s *Store) LoadEntries(id string) (SessionHeader, []Entry, error) {
f, err := os.Open(s.path(id))
if err != nil {
return SessionHeader{}, nil, fmt.Errorf("session: open %s: %w", id, err)
}
defer f.Close()
return readSession(f)
}
// readSession decodes a session stream: header line first, then entries. For
// schema v3 each line is an Entry ({id,parentId,timestamp,message}). For older
// v1/v2 files each line is a bare message; readSession migrates them by
// synthesizing ids and chaining parentId to the previous entry.
func readSession(r io.Reader) (SessionHeader, []Entry, error) {
sc := bufio.NewScanner(r)
// Session lines can be large (long tool results); grow the buffer well past
// the default 64KB token cap.
sc.Buffer(make([]byte, 0, sessionScanBufInit), sessionScanBufMax)
if !sc.Scan() {
if err := sc.Err(); err != nil {
return SessionHeader{}, nil, fmt.Errorf("session: read header: %w", err)
}
return SessionHeader{}, nil, fmt.Errorf("session: empty file (no header)")
}
var header SessionHeader
if err := json.Unmarshal(sc.Bytes(), &header); err != nil {
return SessionHeader{}, nil, fmt.Errorf("session: parse header: %w", err)
}
if header.Version == 0 {
return SessionHeader{}, nil, fmt.Errorf("session: header missing version")
}
if header.Version > SchemaVersion {
return SessionHeader{}, nil, fmt.Errorf("session: file schema version %d newer than supported %d", header.Version, SchemaVersion)
}
// v3+ lines are wrapped entries; v1/v2 lines are bare messages that we migrate.
wrapped := header.Version >= 3
var entries []Entry
parentID := ""
for line := 2; sc.Scan(); line++ {
raw := sc.Bytes()
if len(strings.TrimSpace(string(raw))) == 0 {
continue // tolerate blank lines
}
var e Entry
if wrapped {
if err := json.Unmarshal(raw, &e); err != nil {
return SessionHeader{}, nil, fmt.Errorf("session: parse entry line %d: %w", line, err)
}
} else {
// Migrate a bare v1/v2 message line: reuse MessageList's discriminated
// decoding, then synthesize the tree metadata (id + parentId chain).
var one agentcore.MessageList
if err := json.Unmarshal([]byte("["+string(raw)+"]"), &one); err != nil {
return SessionHeader{}, nil, fmt.Errorf("session: parse message line %d: %w", line, err)
}
if len(one) != 1 {
return SessionHeader{}, nil, fmt.Errorf("session: message line %d decoded to %d messages, want 1", line, len(one))
}
e = Entry{ID: newEntryID(), ParentID: parentID, Timestamp: header.UpdatedAt, Message: one[0]}
}
entries = append(entries, e)
parentID = e.ID
}
if err := sc.Err(); err != nil {
return SessionHeader{}, nil, fmt.Errorf("session: scan: %w", err)
}
return header, entries, nil
}
// List returns the headers of all sessions in the store, sorted by UpdatedAt
// descending (most recently used first). Files that fail to parse are skipped
// rather than failing the whole listing, so one corrupt session does not hide
// the rest.
func (s *Store) List() ([]SessionHeader, error) {
entries, err := os.ReadDir(s.dir)
if err != nil {
return nil, fmt.Errorf("session: read store dir: %w", err)
}
var headers []SessionHeader
for _, e := range entries {
if e.IsDir() || !strings.HasSuffix(e.Name(), ".jsonl") {
continue
}
id := strings.TrimSuffix(e.Name(), ".jsonl")
h, err := s.loadHeader(id)
if err != nil {
continue // skip unreadable/corrupt session
}
headers = append(headers, h)
}
sort.Slice(headers, func(i, j int) bool {
return headers[i].UpdatedAt.After(headers[j].UpdatedAt)
})
return headers, nil
}
// loadHeader reads only the header line of a session file (cheap for listing).
func (s *Store) loadHeader(id string) (SessionHeader, error) {
f, err := os.Open(s.path(id))
if err != nil {
return SessionHeader{}, err
}
defer f.Close()
sc := bufio.NewScanner(f)
sc.Buffer(make([]byte, 0, sessionScanBufInit), sessionScanBufMax)
if !sc.Scan() {
return SessionHeader{}, fmt.Errorf("session: empty file")
}
var header SessionHeader
if err := json.Unmarshal(sc.Bytes(), &header); err != nil {
return SessionHeader{}, err
}
if header.Version == 0 {
return SessionHeader{}, fmt.Errorf("session: missing version")
}
return header, nil
}
// Append adds messages to an existing session file and bumps its UpdatedAt to
// updatedAt. It is the incremental-persistence primitive: a driver that wants to
// grow a session turn-by-turn appends only the newly produced messages rather
// than rewriting the whole file itself. If the session does not exist it is an
// error — use Save to create one first.
//
// Because the header lives on the first line and JSONL is otherwise
// append-only, updating UpdatedAt requires rewriting the file; Append does a
// load-modify-save under the hood, which is simple and correct for the session
// sizes pigo produces.
func (s *Store) Append(id string, updatedAt time.Time, messages agentcore.MessageList) error {
header, existing, err := s.Load(id)
if err != nil {
return err
}
header.UpdatedAt = updatedAt
existing = append(existing, messages...)
return s.Save(header, existing)
}
// AppendBranch appends messages as a chain descending from parentLeafID,
// preserving every existing entry — and therefore any other branches — in the
// session file (US-007, #123). It is the tree-aware counterpart to Append (which
// rewrites the file as a single linear chain): where Append flattens, AppendBranch
// grows the on-disk tree, so switching the active leaf to a historical entry and
// continuing produces a real sibling branch rather than truncating history.
//
// Each message becomes a fresh entry (new id, ParentID chained to the previous
// one, first chained to parentLeafID). An empty parentLeafID roots the new chain.
// header (Version forced to SchemaVersion, UpdatedAt as the caller set it) is
// rewritten as line 1. It returns the id of the new leaf — the last appended
// entry — so the caller can track the active branch. If the session file does not
// yet exist it is created (the fresh-session first-turn case).
func (s *Store) AppendBranch(header SessionHeader, parentLeafID string, messages agentcore.MessageList) (string, error) {
if header.ID == "" {
return "", fmt.Errorf("session: header ID must not be empty")
}
var entries []Entry
if _, existing, err := s.LoadEntries(header.ID); err == nil {
entries = existing
} else if !errors.Is(err, os.ErrNotExist) {
return "", err
}
now := time.Now().UTC()
parent := parentLeafID
leaf := parentLeafID
for _, m := range messages {
e := Entry{ID: newEntryID(), ParentID: parent, Timestamp: now, Message: m}
entries = append(entries, e)
parent = e.ID
leaf = e.ID
}
if err := s.SaveEntries(header, entries); err != nil {
return "", err
}
return leaf, nil
}
// Fork creates a new session whose contents are the linear path from the root
// down to leafID in the source session, copied verbatim (each entry keeps its
// id/parentId). The new session gets a fresh id (derived from now) and its
// header carries ParentSession = sourceID so the lineage is recorded. It returns
// the new session's header and its entries.
//
// This is the primitive behind /fork and /clone (US-006, #122):
//
// - /clone passes the current leaf id → the entire current conversation is
// duplicated into an independent session (position "at").
// - /fork passes a historical user message's PARENT id → the new session holds
// everything up to but excluding that message, so the user re-prompts from
// that point on a fresh branch (position "before").
//
// Because the copy lands in a brand-new file, appending to either the source or
// the fork never touches the other — the two branches are fully isolated. An
// empty leafID copies nothing but the header (an empty new session rooted at the
// source), which is the correct behavior for forking before the very first
// message.
func (s *Store) Fork(sourceID, leafID string, now time.Time) (SessionHeader, []Entry, error) {
srcHeader, entries, err := s.LoadEntries(sourceID)
if err != nil {
return SessionHeader{}, nil, err
}
path := PathToLeaf(entries, leafID)
newHeader := SessionHeader{
ID: NewID(now),
CreatedAt: now,
UpdatedAt: now,
Model: srcHeader.Model,
Provider: srcHeader.Provider,
SystemPrompt: srcHeader.SystemPrompt,
ParentSession: sourceID,
}
if err := s.SaveEntries(newHeader, path); err != nil {
return SessionHeader{}, nil, err
}
return newHeader, path, nil
}
pigo/internal/session/session_test.go
+843
View File
@@ -0,0 +1,843 @@
package session
// Tests for local JSONL session persistence and resume (US-024, #43). They
// cover the write→read round-trip, listing order, resume into an AgentContext,
// schema-version guarding, and append — driving the real filesystem via
// t.TempDir(), the standard Go pattern for behavior tests.
import (
"os"
"path/filepath"
"strings"
"testing"
"time"
"github.com/smallnest/pigo/internal/agentcore"
)
// writeFile writes content to path (test helper for hand-crafted fixtures).
func writeFile(path, content string) error {
return os.WriteFile(path, []byte(content), 0o644)
}
// sampleMessages returns a small multi-turn transcript: user prompt, assistant
// with a tool call, tool result, then a final assistant reply.
func sampleMessages() agentcore.MessageList {
return agentcore.MessageList{
agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("read main.go")}},
agentcore.AssistantMessage{
RoleField: agentcore.RoleAssistant,
Content: agentcore.ContentList{agentcore.NewTextContent("Reading it."), agentcore.NewToolCallContent("call-1", "read", []byte(`{"path":"main.go"}`))},
StopReason: agentcore.StopReasonToolUse,
},
agentcore.ToolResultMessage{
RoleField: agentcore.RoleToolResult,
ToolCallID: "call-1",
ToolName: "read",
Content: agentcore.ContentList{agentcore.NewTextContent("package main")},
},
agentcore.AssistantMessage{
RoleField: agentcore.RoleAssistant,
Content: agentcore.ContentList{agentcore.NewTextContent("It is package main.")},
StopReason: agentcore.StopReasonEndTurn,
},
}
}
func newStore(t *testing.T) *Store {
t.Helper()
s, err := NewStore(t.TempDir())
if err != nil {
t.Fatalf("NewStore: %v", err)
}
return s
}
// TestSaveLoadRoundTrip is the core acceptance check: a saved session loads
// back with an identical header and message sequence.
func TestSaveLoadRoundTrip(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 7, 10, 14, 25, 30, 0, time.UTC)
header := SessionHeader{
ID: NewID(now),
CreatedAt: now,
UpdatedAt: now,
Model: "anthropic/claude-opus-4",
Provider: "anthropic",
SystemPrompt: "You are pigo.",
}
msgs := sampleMessages()
if err := s.Save(header, msgs); err != nil {
t.Fatalf("Save: %v", err)
}
gotHeader, gotMsgs, err := s.Load(header.ID)
if err != nil {
t.Fatalf("Load: %v", err)
}
if gotHeader.Version != SchemaVersion {
t.Errorf("version = %d, want %d", gotHeader.Version, SchemaVersion)
}
if gotHeader.Model != header.Model || gotHeader.SystemPrompt != header.SystemPrompt {
t.Errorf("header mismatch: %+v", gotHeader)
}
if len(gotMsgs) != len(msgs) {
t.Fatalf("message count = %d, want %d", len(gotMsgs), len(msgs))
}
// Roles must round-trip in order.
wantRoles := []string{agentcore.RoleUser, agentcore.RoleAssistant, agentcore.RoleToolResult, agentcore.RoleAssistant}
for i, m := range gotMsgs {
if m.Role() != wantRoles[i] {
t.Errorf("message[%d] role = %q, want %q", i, m.Role(), wantRoles[i])
}
}
// The assistant tool call must survive the round-trip.
a, ok := gotMsgs[1].(agentcore.AssistantMessage)
if !ok {
t.Fatalf("message[1] is not AssistantMessage: %T", gotMsgs[1])
}
calls := a.ToolCalls()
if len(calls) != 1 || calls[0].Name != "read" {
t.Errorf("tool calls = %+v, want one 'read'", calls)
}
}
// TestSaveLoadMalformedToolArguments verifies a transcript containing a tool
// call whose arguments are syntactically invalid JSON (as a model can stream)
// still saves and loads, rather than aborting the whole session write. This is
// the regression for the "session save failed: ... invalid character '{' after
// object key:value pair" crash.
func TestSaveLoadMalformedToolArguments(t *testing.T) {
s := newStore(t)
now := time.Now().UTC()
h := SessionHeader{ID: "malformed", CreatedAt: now, UpdatedAt: now}
msgs := agentcore.MessageList{
agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("go")}},
agentcore.AssistantMessage{
RoleField: agentcore.RoleAssistant,
Content: agentcore.ContentList{
agentcore.NewToolCallContent("c1", "todo", []byte(`{"todos": []{}"content": ""x"}`)),
},
StopReason: agentcore.StopReasonToolUse,
},
}
if err := s.Save(h, msgs); err != nil {
t.Fatalf("Save with malformed tool args: %v", err)
}
_, got, err := s.Load("malformed")
if err != nil {
t.Fatalf("Load: %v", err)
}
if len(got) != 2 {
t.Fatalf("message count = %d, want 2", len(got))
}
a, ok := got[1].(agentcore.AssistantMessage)
if !ok {
t.Fatalf("message[1] is not AssistantMessage: %T", got[1])
}
if calls := a.ToolCalls(); len(calls) != 1 || calls[0].Name != "todo" {
t.Errorf("tool calls = %+v, want one 'todo'", calls)
}
}
// TestSaveOverwrites verifies Save replaces an existing session file (same id)
// atomically rather than appending.
func TestSaveOverwrites(t *testing.T) {
s := newStore(t)
now := time.Now().UTC()
h := SessionHeader{ID: "fixed", CreatedAt: now, UpdatedAt: now}
if err := s.Save(h, sampleMessages()); err != nil {
t.Fatalf("first Save: %v", err)
}
shorter := agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}
if err := s.Save(h, shorter); err != nil {
t.Fatalf("second Save: %v", err)
}
_, msgs, err := s.Load("fixed")
if err != nil {
t.Fatalf("Load: %v", err)
}
if len(msgs) != 1 {
t.Errorf("after overwrite, message count = %d, want 1", len(msgs))
}
}
// TestListSortedByUpdatedDesc verifies List returns sessions most-recent-first.
func TestListSortedByUpdatedDesc(t *testing.T) {
s := newStore(t)
base := time.Date(2026, 7, 10, 0, 0, 0, 0, time.UTC)
for i, id := range []string{"old", "mid", "new"} {
h := SessionHeader{
ID: id,
CreatedAt: base,
UpdatedAt: base.Add(time.Duration(i) * time.Hour),
}
if err := s.Save(h, sampleMessages()); err != nil {
t.Fatalf("Save %s: %v", id, err)
}
}
headers, err := s.List()
if err != nil {
t.Fatalf("List: %v", err)
}
if len(headers) != 3 {
t.Fatalf("List count = %d, want 3", len(headers))
}
wantOrder := []string{"new", "mid", "old"}
for i, h := range headers {
if h.ID != wantOrder[i] {
t.Errorf("List[%d].ID = %q, want %q", i, h.ID, wantOrder[i])
}
}
}
// TestLoadRejectsNewerSchema verifies a file whose version is newer than the
// binary supports is rejected rather than silently misread.
func TestLoadRejectsNewerSchema(t *testing.T) {
s := newStore(t)
// Hand-write a session file with a future version.
future := `{"version":9999,"id":"future","createdAt":"2026-07-10T00:00:00Z","updatedAt":"2026-07-10T00:00:00Z"}` + "\n"
path := filepath.Join(s.Dir(), FileName("future"))
if err := writeFile(path, future); err != nil {
t.Fatalf("write fixture: %v", err)
}
if _, _, err := s.Load("future"); err == nil {
t.Error("Load must reject a newer schema version")
}
}
// TestLoadV1FileStillReadable verifies the v2 schema bump is backward-compatible:
// an old v1 session file (no compaction lines) loads without error.
func TestLoadV1FileStillReadable(t *testing.T) {
s := newStore(t)
v1 := `{"version":1,"id":"old","createdAt":"2026-07-10T00:00:00Z","updatedAt":"2026-07-10T00:00:00Z"}` + "\n" +
`{"role":"user","content":[{"type":"text","text":"hi"}],"timestamp":0}` + "\n"
path := filepath.Join(s.Dir(), FileName("old"))
if err := writeFile(path, v1); err != nil {
t.Fatalf("write fixture: %v", err)
}
h, msgs, err := s.Load("old")
if err != nil {
t.Fatalf("Load v1: %v", err)
}
if h.Version != 1 {
t.Fatalf("version: got %d, want 1", h.Version)
}
if len(msgs) != 1 || msgs[0].Role() != agentcore.RoleUser {
t.Fatalf("messages: got %+v", msgs)
}
}
// TestSaveWritesV3TreeEntries verifies Save persists each message as a wrapped
// v3 entry: the file header is version 3, every entry carries a non-empty id,
// the first entry is a root (empty parentId), and each subsequent entry's
// parentId chains to the previous entry's id — a linear tree.
func TestSaveWritesV3TreeEntries(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 7, 17, 0, 0, 0, 0, time.UTC)
header := SessionHeader{ID: NewID(now), CreatedAt: now, UpdatedAt: now}
if err := s.Save(header, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
h, entries, err := s.LoadEntries(header.ID)
if err != nil {
t.Fatalf("LoadEntries: %v", err)
}
if h.Version != 3 {
t.Fatalf("version = %d, want 3", h.Version)
}
if len(entries) != 4 {
t.Fatalf("entry count = %d, want 4", len(entries))
}
if entries[0].ParentID != "" {
t.Errorf("root entry parentId = %q, want empty", entries[0].ParentID)
}
seen := map[string]bool{}
for i, e := range entries {
if e.ID == "" {
t.Errorf("entry[%d] has empty id", i)
}
if seen[e.ID] {
t.Errorf("entry[%d] id %q is duplicated", i, e.ID)
}
seen[e.ID] = true
if i > 0 && e.ParentID != entries[i-1].ID {
t.Errorf("entry[%d] parentId = %q, want %q (previous entry)", i, e.ParentID, entries[i-1].ID)
}
}
}
// TestPathToLeaf verifies PathToLeaf walks the parentId chain from a leaf back
// to the root and returns the entries in root→leaf order.
func TestPathToLeaf(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 7, 17, 0, 0, 0, 0, time.UTC)
header := SessionHeader{ID: NewID(now), CreatedAt: now, UpdatedAt: now}
if err := s.Save(header, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
_, entries, err := s.LoadEntries(header.ID)
if err != nil {
t.Fatalf("LoadEntries: %v", err)
}
leaf := entries[len(entries)-1]
path := PathToLeaf(entries, leaf.ID)
if len(path) != len(entries) {
t.Fatalf("path length = %d, want %d (linear session)", len(path), len(entries))
}
for i := range entries {
if path[i].ID != entries[i].ID {
t.Errorf("path[%d].ID = %q, want %q", i, path[i].ID, entries[i].ID)
}
}
// A mid-chain leaf yields only its ancestors + itself.
mid := PathToLeaf(entries, entries[1].ID)
if len(mid) != 2 || mid[0].ID != entries[0].ID || mid[1].ID != entries[1].ID {
t.Errorf("PathToLeaf(entries[1]) = %+v, want [root, entries[1]]", mid)
}
// Unknown / empty leaf ids yield nil.
if got := PathToLeaf(entries, "nope"); got != nil {
t.Errorf("PathToLeaf(unknown) = %+v, want nil", got)
}
if got := PathToLeaf(entries, ""); got != nil {
t.Errorf("PathToLeaf(empty) = %+v, want nil", got)
}
}
// TestLoadV2FileMigratesToEntries verifies a v2 file (bare message lines, no
// id/parentId) still loads and resumes: readSession back-fills a synthesized id
// per line and chains parentId to the previous entry, so the migrated entries
// form a linear tree while the flat Load view is unchanged.
func TestLoadV2FileMigratesToEntries(t *testing.T) {
s := newStore(t)
v2 := `{"version":2,"id":"legacy","createdAt":"2026-07-10T00:00:00Z","updatedAt":"2026-07-10T00:00:00Z"}` + "\n" +
`{"role":"user","content":[{"type":"text","text":"hi"}]}` + "\n" +
`{"role":"assistant","content":[{"type":"text","text":"hello"}],"stopReason":"end_turn"}` + "\n"
path := filepath.Join(s.Dir(), FileName("legacy"))
if err := writeFile(path, v2); err != nil {
t.Fatalf("write fixture: %v", err)
}
// Flat Load view is unchanged by the migration.
h, msgs, err := s.Load("legacy")
if err != nil {
t.Fatalf("Load v2: %v", err)
}
if h.Version != 2 {
t.Fatalf("version = %d, want 2", h.Version)
}
if len(msgs) != 2 || msgs[0].Role() != agentcore.RoleUser || msgs[1].Role() != agentcore.RoleAssistant {
t.Fatalf("messages: got %+v", msgs)
}
// Entry view is back-filled into a linear tree.
_, entries, err := s.LoadEntries("legacy")
if err != nil {
t.Fatalf("LoadEntries v2: %v", err)
}
if len(entries) != 2 {
t.Fatalf("entry count = %d, want 2", len(entries))
}
if entries[0].ID == "" || entries[0].ParentID != "" {
t.Errorf("root entry = %+v, want non-empty id and empty parentId", entries[0])
}
if entries[1].ParentID != entries[0].ID {
t.Errorf("entry[1].parentId = %q, want %q", entries[1].ParentID, entries[0].ID)
}
// The migrated entries reconstruct the full conversation via PathToLeaf.
path2 := PathToLeaf(entries, entries[1].ID)
if len(path2) != 2 {
t.Errorf("PathToLeaf on migrated v2 = %d entries, want 2", len(path2))
}
}
// TestAppendPreservesChain verifies Append grows the linear tree: after
// appending, the file still loads with a valid root and an unbroken parentId
// chain across the combined message set.
func TestAppendPreservesChain(t *testing.T) {
s := newStore(t)
now := time.Now().UTC()
h := SessionHeader{ID: "grow", CreatedAt: now, UpdatedAt: now}
if err := s.Save(h, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
extra := agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("more")}}}
if err := s.Append("grow", now.Add(time.Minute), extra); err != nil {
t.Fatalf("Append: %v", err)
}
_, entries, err := s.LoadEntries("grow")
if err != nil {
t.Fatalf("LoadEntries: %v", err)
}
if len(entries) != 5 {
t.Fatalf("entry count = %d, want 5", len(entries))
}
if entries[0].ParentID != "" {
t.Errorf("root parentId = %q, want empty", entries[0].ParentID)
}
for i := 1; i < len(entries); i++ {
if entries[i].ParentID != entries[i-1].ID {
t.Errorf("entry[%d] parentId = %q, want %q", i, entries[i].ParentID, entries[i-1].ID)
}
}
if path := PathToLeaf(entries, entries[4].ID); len(path) != 5 {
t.Errorf("PathToLeaf after append = %d, want 5", len(path))
}
}
// TestForkClonesFullConversation verifies Fork(sourceID, lastLeaf) — the /clone
// case — copies the entire conversation verbatim into a new, independent session:
// the new header records ParentSession, the copied entries keep their ids, and
// appending to the fork does NOT touch the source (branch isolation).
func TestForkClonesFullConversation(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 7, 17, 0, 0, 0, 0, time.UTC)
src := SessionHeader{ID: NewID(now), CreatedAt: now, UpdatedAt: now, Model: "m", Provider: "p", SystemPrompt: "sp"}
if err := s.Save(src, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
_, srcEntries, err := s.LoadEntries(src.ID)
if err != nil {
t.Fatalf("LoadEntries: %v", err)
}
leaf := srcEntries[len(srcEntries)-1].ID
forkHeader, forkEntries, err := s.Fork(src.ID, leaf, now.Add(time.Hour))
if err != nil {
t.Fatalf("Fork: %v", err)
}
if forkHeader.ID == src.ID {
t.Fatal("fork must get a new id, got the source id")
}
if forkHeader.ParentSession != src.ID {
t.Errorf("ParentSession = %q, want %q", forkHeader.ParentSession, src.ID)
}
// Header metadata is inherited from the source.
if forkHeader.Model != "m" || forkHeader.Provider != "p" || forkHeader.SystemPrompt != "sp" {
t.Errorf("fork header did not inherit source metadata: %+v", forkHeader)
}
if len(forkEntries) != len(srcEntries) {
t.Fatalf("fork entry count = %d, want %d (full clone)", len(forkEntries), len(srcEntries))
}
// Copied entries keep their ids/parentIds verbatim.
for i := range srcEntries {
if forkEntries[i].ID != srcEntries[i].ID || forkEntries[i].ParentID != srcEntries[i].ParentID {
t.Errorf("entry[%d] id/parent = (%q,%q), want (%q,%q)", i,
forkEntries[i].ID, forkEntries[i].ParentID, srcEntries[i].ID, srcEntries[i].ParentID)
}
}
// Branch isolation: appending to the fork must not change the source.
extra := agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("on the fork only")}}}
if err := s.Append(forkHeader.ID, now.Add(2*time.Hour), extra); err != nil {
t.Fatalf("Append to fork: %v", err)
}
_, srcAfter, err := s.Load(src.ID)
if err != nil {
t.Fatalf("Load source: %v", err)
}
if len(srcAfter) != len(sampleMessages()) {
t.Errorf("source message count changed to %d after fork append, want %d (branches must be isolated)", len(srcAfter), len(sampleMessages()))
}
_, forkAfter, err := s.Load(forkHeader.ID)
if err != nil {
t.Fatalf("Load fork: %v", err)
}
if len(forkAfter) != len(sampleMessages())+1 {
t.Errorf("fork message count = %d, want %d", len(forkAfter), len(sampleMessages())+1)
}
}
// TestForkBeforeUserMessage verifies Fork(sourceID, parentOfUserMsg) — the
// /fork case — copies only the prefix up to (excluding) a chosen user message,
// so the branch can re-prompt from that point. Forking before the very first
// user message (empty leafID) yields an empty session rooted at the source.
func TestForkBeforeUserMessage(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 7, 17, 0, 0, 0, 0, time.UTC)
src := SessionHeader{ID: NewID(now), CreatedAt: now, UpdatedAt: now}
if err := s.Save(src, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
_, entries, err := s.LoadEntries(src.ID)
if err != nil {
t.Fatalf("LoadEntries: %v", err)
}
// sampleMessages()[0] is the first user message (a root). Forking before it
// uses its ParentID (empty) → an empty branch.
if entries[0].ParentID != "" {
t.Fatalf("precondition: first entry should be a root")
}
emptyHeader, emptyPath, err := s.Fork(src.ID, entries[0].ParentID, now.Add(time.Hour))
if err != nil {
t.Fatalf("Fork before first message: %v", err)
}
if len(emptyPath) != 0 {
t.Errorf("fork before first message = %d entries, want 0", len(emptyPath))
}
if emptyHeader.ParentSession != src.ID {
t.Errorf("ParentSession = %q, want %q", emptyHeader.ParentSession, src.ID)
}
// Reloading the empty fork yields a valid, empty session.
_, reload, err := s.LoadEntries(emptyHeader.ID)
if err != nil {
t.Fatalf("LoadEntries(empty fork): %v", err)
}
if len(reload) != 0 {
t.Errorf("reloaded empty fork = %d entries, want 0", len(reload))
}
// Forking before the SECOND-turn user message (there is only one user message
// in sampleMessages, so simulate a two-user transcript) — copy just the prefix.
// Here we fork at the parent of the last entry to get all but the last message.
lastParent := entries[len(entries)-1].ParentID
_, prefix, err := s.Fork(src.ID, lastParent, now.Add(2*time.Hour))
if err != nil {
t.Fatalf("Fork at last parent: %v", err)
}
if len(prefix) != len(entries)-1 {
t.Errorf("prefix fork = %d entries, want %d", len(prefix), len(entries)-1)
}
}
// TestAppendBranchGrowsTree verifies AppendBranch (US-007, #123) preserves all
// existing entries and chains new messages from a chosen parent leaf — so
// switching the active leaf to a historical entry and continuing produces a real
// sibling branch on disk rather than truncating history.
func TestAppendBranchGrowsTree(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 7, 17, 0, 0, 0, 0, time.UTC)
h := SessionHeader{ID: NewID(now), CreatedAt: now, UpdatedAt: now}
if err := s.Save(h, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
_, entries, err := s.LoadEntries(h.ID)
if err != nil {
t.Fatalf("LoadEntries: %v", err)
}
// Branch from the FIRST entry (the root user message): append a new user turn
// as its child. The result must keep every original entry plus the new one.
branchParent := entries[0].ID
extra := agentcore.MessageList{
agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("a different question")}},
}
leaf, err := s.AppendBranch(h, branchParent, extra)
if err != nil {
t.Fatalf("AppendBranch: %v", err)
}
_, after, err := s.LoadEntries(h.ID)
if err != nil {
t.Fatalf("LoadEntries after branch: %v", err)
}
if len(after) != len(entries)+1 {
t.Fatalf("entry count = %d, want %d (nothing dropped, one added)", len(after), len(entries)+1)
}
// The new leaf descends from the chosen parent.
var newLeaf *Entry
for i := range after {
if after[i].ID == leaf {
newLeaf = &after[i]
}
}
if newLeaf == nil {
t.Fatalf("new leaf %q not found in reloaded entries", leaf)
}
if newLeaf.ParentID != branchParent {
t.Errorf("new leaf parent = %q, want %q", newLeaf.ParentID, branchParent)
}
// The root now has two children: the original second entry and the new leaf —
// a genuine branch point.
kids := 0
for _, e := range after {
if e.ParentID == branchParent {
kids++
}
}
if kids != 2 {
t.Errorf("branch point should have 2 children, got %d", kids)
}
// PathToLeaf to the new leaf yields exactly [root, newLeaf].
path := PathToLeaf(after, leaf)
if len(path) != 2 || path[0].ID != branchParent || path[1].ID != leaf {
t.Errorf("PathToLeaf(newLeaf) = %v, want [root, newLeaf]", pathIDs(path))
}
}
// TestAppendBranchCreatesFileWhenMissing verifies AppendBranch creates the
// session file on first use (the fresh-session first-turn case) rather than
// erroring like Append does.
func TestAppendBranchCreatesFileWhenMissing(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 7, 17, 0, 0, 0, 0, time.UTC)
h := SessionHeader{ID: NewID(now), CreatedAt: now, UpdatedAt: now}
msgs := agentcore.MessageList{
agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}},
}
leaf, err := s.AppendBranch(h, "", msgs)
if err != nil {
t.Fatalf("AppendBranch on missing file: %v", err)
}
_, entries, err := s.LoadEntries(h.ID)
if err != nil {
t.Fatalf("LoadEntries: %v", err)
}
if len(entries) != 1 || entries[0].ID != leaf || entries[0].ParentID != "" {
t.Errorf("expected one root entry with id=%q, got %v", leaf, entries)
}
}
// TestRenderTreeLinesMarksCurrentAndBranches verifies the pure-text tree render
// (US-007, #123): every entry gets one numbered-able line in render order, the
// active leaf is tagged "← current", and a branch point produces two child rows.
func TestRenderTreeLinesMarksCurrentAndBranches(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 7, 17, 0, 0, 0, 0, time.UTC)
h := SessionHeader{ID: NewID(now), CreatedAt: now, UpdatedAt: now}
if err := s.Save(h, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
_, entries, err := s.LoadEntries(h.ID)
if err != nil {
t.Fatalf("LoadEntries: %v", err)
}
// Branch off the root to create a fork point.
if _, err := s.AppendBranch(h, entries[0].ID, agentcore.MessageList{
agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("sibling turn")}},
}); err != nil {
t.Fatalf("AppendBranch: %v", err)
}
_, after, err := s.LoadEntries(h.ID)
if err != nil {
t.Fatalf("LoadEntries after branch: %v", err)
}
leaf := after[len(after)-1].ID
lines := RenderTreeLines(after, leaf)
if len(lines) != len(after) {
t.Fatalf("render produced %d lines, want %d (one per entry)", len(lines), len(after))
}
// Exactly one line is tagged as current, and it is the leaf.
current := 0
for _, l := range lines {
if strings.Contains(l.Text, "← current") {
current++
if l.Entry.ID != leaf {
t.Errorf("current marker on %q, want leaf %q", l.Entry.ID, leaf)
}
}
}
if current != 1 {
t.Errorf("expected exactly one ← current line, got %d", current)
}
// Connector characters must appear (readable branch structure).
joined := ""
for _, l := range lines {
joined += l.Text + "\n"
}
if !strings.Contains(joined, "├─") && !strings.Contains(joined, "└─") {
t.Errorf("tree render lacks connectors:\n%s", joined)
}
// The first line is a root (no connector prefix) rendering the root user msg.
if !strings.HasPrefix(lines[0].Text, "user:") {
t.Errorf("first render line should be the root user message, got %q", lines[0].Text)
}
}
// pathIDs is a small test helper: the ids of a path, for readable failures.
func pathIDs(path []Entry) []string {
ids := make([]string, len(path))
for i, e := range path {
ids[i] = e.ID
}
return ids
}
// through the JSONL store as a first-class message line under schema v2.
func TestSaveLoadCompactionEntry(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 7, 17, 0, 0, 0, 0, time.UTC)
header := SessionHeader{ID: NewID(now), CreatedAt: now, UpdatedAt: now}
msgs := agentcore.MessageList{
agentcore.CompactionMessage{
RoleField: agentcore.RoleCompaction,
Summary: "## Goal\nship #119",
TokensBefore: 12345,
Details: []byte(`{"readFiles":["a.go"],"modifiedFiles":["b.go"]}`),
Timestamp: now.UnixMilli(),
},
agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("continue")}},
}
if err := s.Save(header, msgs); err != nil {
t.Fatalf("Save: %v", err)
}
loaded, back, err := s.Load(header.ID)
if err != nil {
t.Fatalf("Load: %v", err)
}
if loaded.Version != SchemaVersion {
t.Fatalf("version: got %d, want %d", loaded.Version, SchemaVersion)
}
if len(back) != 2 || back[0].Role() != agentcore.RoleCompaction {
t.Fatalf("messages: got %+v", back)
}
cm, ok := back[0].(agentcore.CompactionMessage)
if !ok {
t.Fatalf("first message is not a CompactionMessage: %T", back[0])
}
if cm.Summary != "## Goal\nship #119" || cm.TokensBefore != 12345 {
t.Fatalf("compaction fields: %+v", cm)
}
}
// TestContextHeaderFieldsRoundTrip verifies the additive "infinite context"
// header fields (#480) — ContextFrom/ContextWatermark — survive a Save→Load
// round-trip through the real session file path.
func TestContextHeaderFieldsRoundTrip(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 8, 1, 9, 0, 0, 0, time.UTC)
header := SessionHeader{
ID: NewID(now),
CreatedAt: now,
UpdatedAt: now,
ContextFrom: "parent-sess-123",
ContextWatermark: 37,
}
if err := s.Save(header, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
got, _, err := s.Load(header.ID)
if err != nil {
t.Fatalf("Load: %v", err)
}
if got.ContextFrom != "parent-sess-123" {
t.Errorf("ContextFrom = %q, want %q", got.ContextFrom, "parent-sess-123")
}
if got.ContextWatermark != 37 {
t.Errorf("ContextWatermark = %d, want 37", got.ContextWatermark)
}
}
// TestContextHeaderFieldsOmittedWhenZero verifies the fields are omitempty: a
// header without them serializes without the keys, keeping v1/v2/v3 files
// byte-compatible for the common no-inherited-context case.
func TestContextHeaderFieldsOmittedWhenZero(t *testing.T) {
s := newStore(t)
now := time.Now().UTC()
header := SessionHeader{ID: "no-ctx", CreatedAt: now, UpdatedAt: now}
if err := s.Save(header, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
raw, err := os.ReadFile(filepath.Join(s.Dir(), FileName("no-ctx")))
if err != nil {
t.Fatalf("read session file: %v", err)
}
headerLine := strings.SplitN(string(raw), "\n", 2)[0]
if strings.Contains(headerLine, "contextFrom") || strings.Contains(headerLine, "contextWatermark") {
t.Errorf("zero context fields must be omitted; header line = %s", headerLine)
}
}
// TestSetContextInheritanceRoundTrips verifies that a session created with
// inheritance set via SetContextInheritance writes both fields and exposes them
// (via the ContextInheritance accessor) when read back.
func TestSetContextInheritanceRoundTrips(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 8, 1, 10, 0, 0, 0, time.UTC)
header := SessionHeader{ID: NewID(now), CreatedAt: now, UpdatedAt: now}
SetContextInheritance(&header, "src-sess-42", 12)
if err := s.Save(header, sampleMessages()); err != nil {
t.Fatalf("Save: %v", err)
}
got, _, err := s.Load(header.ID)
if err != nil {
t.Fatalf("Load: %v", err)
}
fromID, watermark, ok := got.ContextInheritance()
if !ok {
t.Fatalf("ContextInheritance ok = false, want true")
}
if fromID != "src-sess-42" || watermark != 12 {
t.Errorf("ContextInheritance = (%q, %d), want (%q, 12)", fromID, watermark, "src-sess-42")
}
if !got.HasContextInheritance() {
t.Errorf("HasContextInheritance = false, want true")
}
}
// TestSetContextInheritanceClears verifies an empty source or non-positive
// watermark clears both fields, so the header round-trips as a no-inheritance
// session (omitempty keeps it byte-compatible with older files).
func TestSetContextInheritanceClears(t *testing.T) {
cases := []struct {
name string
fromID string
watermark int
}{
{"empty source", "", 5},
{"zero watermark", "src", 0},
{"negative watermark", "src", -3},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
header := SessionHeader{ID: "h", ContextFrom: "stale", ContextWatermark: 99}
SetContextInheritance(&header, tc.fromID, tc.watermark)
if header.ContextFrom != "" || header.ContextWatermark != 0 {
t.Errorf("expected cleared inheritance, got (%q, %d)", header.ContextFrom, header.ContextWatermark)
}
if header.HasContextInheritance() {
t.Errorf("HasContextInheritance = true after clear")
}
})
}
}
// TestContextInheritanceRejectsPartial verifies the accessor treats a header
// with only one of the two fields as "no inheritance" (ok=false) rather than
// reporting a half-configured record.
func TestContextInheritanceRejectsPartial(t *testing.T) {
onlyFrom := SessionHeader{ContextFrom: "src"}
if _, _, ok := onlyFrom.ContextInheritance(); ok {
t.Errorf("ok = true for header with source but no watermark")
}
onlyWatermark := SessionHeader{ContextWatermark: 8}
if _, _, ok := onlyWatermark.ContextInheritance(); ok {
t.Errorf("ok = true for header with watermark but no source")
}
}
// TestLoadedSessionsWithoutInheritanceReadZero verifies backward compatibility:
// v1/v2/v3 sessions written without the inheritance fields load with zero-value
// ContextFrom/ContextWatermark and report HasContextInheritance()=false.
func TestLoadedSessionsWithoutInheritanceReadZero(t *testing.T) {
s := newStore(t)
now := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
// v3 written through the normal path (no inheritance set).
v3 := SessionHeader{ID: "v3-noctx", CreatedAt: now, UpdatedAt: now}
if err := s.Save(v3, sampleMessages()); err != nil {
t.Fatalf("Save v3: %v", err)
}
// Hand-crafted v1 and v2 fixtures (bare message lines, no context fields).
msgLine := `{"role":"user","content":[{"type":"text","text":"hi"}]}`
if err := writeFile(filepath.Join(s.Dir(), FileName("v1-noctx")),
`{"version":1,"id":"v1-noctx","createdAt":"2026-01-02T03:04:05Z","updatedAt":"2026-01-02T03:04:05Z"}`+"\n"+msgLine+"\n"); err != nil {
t.Fatalf("write v1 fixture: %v", err)
}
if err := writeFile(filepath.Join(s.Dir(), FileName("v2-noctx")),
`{"version":2,"id":"v2-noctx","createdAt":"2026-01-02T03:04:05Z","updatedAt":"2026-01-02T03:04:05Z"}`+"\n"+msgLine+"\n"); err != nil {
t.Fatalf("write v2 fixture: %v", err)
}
for _, id := range []string{"v3-noctx", "v1-noctx", "v2-noctx"} {
got, _, err := s.Load(id)
if err != nil {
t.Fatalf("Load %s: %v", id, err)
}
if got.ContextFrom != "" || got.ContextWatermark != 0 {
t.Errorf("%s: expected zero inheritance, got (%q, %d)", id, got.ContextFrom, got.ContextWatermark)
}
if got.HasContextInheritance() {
t.Errorf("%s: HasContextInheritance = true, want false", id)
}
}
}