kit/internal/session/entry.go

package session

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"time"

	"github.com/mark3labs/kit/internal/message"
)

// EntryType identifies the kind of entry stored in a JSONL session file.
// Sessions are append-only JSONL files where each line is a typed entry
// linked by id/parent_id to form a tree structure.
type EntryType string

const (
	EntryTypeSession       EntryType = "session"
	EntryTypeMessage       EntryType = "message"
	EntryTypeModelChange   EntryType = "model_change"
	EntryTypeBranchSummary EntryType = "branch_summary"
	EntryTypeLabel         EntryType = "label"
	EntryTypeSessionInfo   EntryType = "session_info"
	EntryTypeExtensionData EntryType = "extension_data"
	EntryTypeCompaction    EntryType = "compaction"
	EntryTypeSystemPrompt  EntryType = "system_prompt"
)

// CurrentVersion is the session format version for JSONL tree sessions.
const CurrentVersion = 3

// SessionHeader is the first line in a JSONL session file. It carries
// metadata about the session and does NOT participate in the tree structure
// (it has no ID or ParentID).
type SessionHeader struct {
	Type          EntryType `json:"type"`                     // always "session"
	Version       int       `json:"version"`                  // format version (3)
	ID            string    `json:"id"`                       // session UUID
	Timestamp     time.Time `json:"timestamp"`                // creation time
	Cwd           string    `json:"cwd"`                      // working directory
	ParentSession string    `json:"parent_session,omitempty"` // path to parent if forked

	// Subagent fields (set when session is created by a subagent)
	ParentSessionID string `json:"parent_session_id,omitempty"` // UUID of parent session
	SubagentTask    string `json:"subagent_task,omitempty"`     // original task prompt
}

// Entry is the common structure shared by all tree entries (everything except
// the session header). Every entry has an ID, an optional ParentID (empty for
// root entries), a type tag, and a timestamp.
type Entry struct {
	Type      EntryType `json:"type"`
	ID        string    `json:"id"`
	ParentID  string    `json:"parent_id,omitempty"`
	Timestamp time.Time `json:"timestamp"`
}

// MessageEntry stores a conversation message as a tree entry. The message
// content uses the same type-tagged parts format as the existing session
// persistence layer, enabling reuse of MarshalParts/UnmarshalParts.
type MessageEntry struct {
	Entry
	Role     string          `json:"role"`
	Parts    json.RawMessage `json:"parts"` // type-tagged parts array
	Model    string          `json:"model,omitempty"`
	Provider string          `json:"provider,omitempty"`
}

// ModelChangeEntry records a provider/model switch in the session tree.
type ModelChangeEntry struct {
	Entry
	Provider string `json:"provider"`
	ModelID  string `json:"model_id"`
}

// BranchSummaryEntry provides LLM-generated context from an abandoned branch.
// When the user navigates away from a branch, a summary of that branch's
// conversation is stored so the LLM retains context about what was explored.
type BranchSummaryEntry struct {
	Entry
	FromID  string `json:"from_id"` // leaf of the summarized branch
	Summary string `json:"summary"`
}

// LabelEntry bookmarks a specific entry with a user-defined label.
type LabelEntry struct {
	Entry
	TargetID string `json:"target_id"`
	Label    string `json:"label"`
}

// SessionInfoEntry stores a user-defined display name for the session.
type SessionInfoEntry struct {
	Entry
	Name string `json:"name"`
}

// ExtensionDataEntry stores custom extension data in the session tree.
// Extensions use this to persist state that survives across session restarts.
type ExtensionDataEntry struct {
	Entry
	ExtType string `json:"ext_type"` // Extension-defined type string (e.g. "plan-mode:state")
	Data    string `json:"data"`     // Extension-defined data (JSON or plain text)
}

// CompactionEntry records an LLM-generated summary of older messages.
// Instead of deleting old messages, the tree manager skips entries before
// FirstKeptEntryID when building the LLM context, preserving full history.
type CompactionEntry struct {
	Entry
	Summary          string   `json:"summary"`
	FirstKeptEntryID string   `json:"first_kept_entry_id"`
	TokensBefore     int      `json:"tokens_before"`
	TokensAfter      int      `json:"tokens_after"`
	MessagesRemoved  int      `json:"messages_removed"`
	ReadFiles        []string `json:"read_files,omitempty"`
	ModifiedFiles    []string `json:"modified_files,omitempty"`
}

// SystemPromptEntry records the system prompt and model used for the session.
// This is primarily for sharing/debugging to see what instructions were
// active during the conversation. It does NOT participate in the tree
// structure (no ParentID) and is not used when building LLM context.
type SystemPromptEntry struct {
	Type      EntryType `json:"type"`      // always "system_prompt"
	ID        string    `json:"id"`        // unique entry ID
	Timestamp time.Time `json:"timestamp"` // when captured
	Content   string    `json:"content"`   // the system prompt text
	Model     string    `json:"model"`     // the model used (e.g., "claude-sonnet-4-5")
	Provider  string    `json:"provider"`  // the provider used (e.g., "anthropic")
}

// GenerateEntryID creates a unique entry identifier (16 hex chars).
func GenerateEntryID() string {
	bytes := make([]byte, 8)
	_, _ = rand.Read(bytes)
	return hex.EncodeToString(bytes)
}

// GenerateSessionID creates a unique session identifier (32 hex chars).
func GenerateSessionID() string {
	bytes := make([]byte, 16)
	_, _ = rand.Read(bytes)
	return hex.EncodeToString(bytes)
}

// NewEntry creates a base Entry with a generated ID and current timestamp.
func NewEntry(entryType EntryType, parentID string) Entry {
	return Entry{
		Type:      entryType,
		ID:        GenerateEntryID(),
		ParentID:  parentID,
		Timestamp: time.Now(),
	}
}

// --- Entry constructors ---

// NewMessageEntry creates a MessageEntry from a message.Message, linking it
// to the given parent entry in the tree.
func NewMessageEntry(parentID string, msg message.Message) (*MessageEntry, error) {
	parts, err := message.MarshalParts(msg.Parts)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal message parts: %w", err)
	}
	return &MessageEntry{
		Entry:    NewEntry(EntryTypeMessage, parentID),
		Role:     string(msg.Role),
		Parts:    parts,
		Model:    msg.Model,
		Provider: msg.Provider,
	}, nil
}

// NewModelChangeEntry creates a ModelChangeEntry.
func NewModelChangeEntry(parentID, provider, modelID string) *ModelChangeEntry {
	return &ModelChangeEntry{
		Entry:    NewEntry(EntryTypeModelChange, parentID),
		Provider: provider,
		ModelID:  modelID,
	}
}

// NewBranchSummaryEntry creates a BranchSummaryEntry.
func NewBranchSummaryEntry(parentID, fromID, summary string) *BranchSummaryEntry {
	return &BranchSummaryEntry{
		Entry:   NewEntry(EntryTypeBranchSummary, parentID),
		FromID:  fromID,
		Summary: summary,
	}
}

// NewLabelEntry creates a LabelEntry.
func NewLabelEntry(parentID, targetID, label string) *LabelEntry {
	return &LabelEntry{
		Entry:    NewEntry(EntryTypeLabel, parentID),
		TargetID: targetID,
		Label:    label,
	}
}

// NewSessionInfoEntry creates a SessionInfoEntry.
func NewSessionInfoEntry(parentID, name string) *SessionInfoEntry {
	return &SessionInfoEntry{
		Entry: NewEntry(EntryTypeSessionInfo, parentID),
		Name:  name,
	}
}

// NewExtensionDataEntry creates an ExtensionDataEntry.
func NewExtensionDataEntry(parentID, extType, data string) *ExtensionDataEntry {
	return &ExtensionDataEntry{
		Entry:   NewEntry(EntryTypeExtensionData, parentID),
		ExtType: extType,
		Data:    data,
	}
}

// NewCompactionEntry creates a CompactionEntry.
func NewCompactionEntry(parentID, summary, firstKeptEntryID string, tokensBefore, tokensAfter, messagesRemoved int, readFiles, modifiedFiles []string) *CompactionEntry {
	return &CompactionEntry{
		Entry:            NewEntry(EntryTypeCompaction, parentID),
		Summary:          summary,
		FirstKeptEntryID: firstKeptEntryID,
		TokensBefore:     tokensBefore,
		TokensAfter:      tokensAfter,
		MessagesRemoved:  messagesRemoved,
		ReadFiles:        readFiles,
		ModifiedFiles:    modifiedFiles,
	}
}

// NewSystemPromptEntry creates a SystemPromptEntry.
func NewSystemPromptEntry(content, model, provider string) *SystemPromptEntry {
	return &SystemPromptEntry{
		Type:      EntryTypeSystemPrompt,
		ID:        GenerateEntryID(),
		Timestamp: time.Now(),
		Content:   content,
		Model:     model,
		Provider:  provider,
	}
}

// --- JSONL marshaling helpers ---

// MarshalEntry serializes any entry to a JSON line (no trailing newline).
func MarshalEntry(entry any) ([]byte, error) {
	return json.Marshal(entry)
}

// entryEnvelope is used for initial unmarshaling to determine the entry type.
type entryEnvelope struct {
	Type EntryType `json:"type"`
}

// UnmarshalEntry deserializes a JSON line into the appropriate entry type.
// Returns one of: *SessionHeader, *MessageEntry, *ModelChangeEntry,
// *BranchSummaryEntry, *LabelEntry, *SessionInfoEntry.
func UnmarshalEntry(data []byte) (any, error) {
	var env entryEnvelope
	if err := json.Unmarshal(data, &env); err != nil {
		return nil, fmt.Errorf("failed to detect entry type: %w", err)
	}

	switch env.Type {
	case EntryTypeSession:
		var h SessionHeader
		if err := json.Unmarshal(data, &h); err != nil {
			return nil, fmt.Errorf("failed to unmarshal session header: %w", err)
		}
		return &h, nil

	case EntryTypeMessage:
		var e MessageEntry
		if err := json.Unmarshal(data, &e); err != nil {
			return nil, fmt.Errorf("failed to unmarshal message entry: %w", err)
		}
		return &e, nil

	case EntryTypeModelChange:
		var e ModelChangeEntry
		if err := json.Unmarshal(data, &e); err != nil {
			return nil, fmt.Errorf("failed to unmarshal model_change entry: %w", err)
		}
		return &e, nil

	case EntryTypeBranchSummary:
		var e BranchSummaryEntry
		if err := json.Unmarshal(data, &e); err != nil {
			return nil, fmt.Errorf("failed to unmarshal branch_summary entry: %w", err)
		}
		return &e, nil

	case EntryTypeLabel:
		var e LabelEntry
		if err := json.Unmarshal(data, &e); err != nil {
			return nil, fmt.Errorf("failed to unmarshal label entry: %w", err)
		}
		return &e, nil

	case EntryTypeSessionInfo:
		var e SessionInfoEntry
		if err := json.Unmarshal(data, &e); err != nil {
			return nil, fmt.Errorf("failed to unmarshal session_info entry: %w", err)
		}
		return &e, nil

	case EntryTypeExtensionData:
		var e ExtensionDataEntry
		if err := json.Unmarshal(data, &e); err != nil {
			return nil, fmt.Errorf("failed to unmarshal extension_data entry: %w", err)
		}
		return &e, nil

	case EntryTypeCompaction:
		var e CompactionEntry
		if err := json.Unmarshal(data, &e); err != nil {
			return nil, fmt.Errorf("failed to unmarshal compaction entry: %w", err)
		}
		return &e, nil

	case EntryTypeSystemPrompt:
		var e SystemPromptEntry
		if err := json.Unmarshal(data, &e); err != nil {
			return nil, fmt.Errorf("failed to unmarshal system_prompt entry: %w", err)
		}
		return &e, nil

	default:
		return nil, fmt.Errorf("unknown entry type: %q", env.Type)
	}
}

// ToMessage converts a MessageEntry back to a message.Message by
// unmarshaling the type-tagged parts.
func (e *MessageEntry) ToMessage() (message.Message, error) {
	parts, err := message.UnmarshalParts(e.Parts)
	if err != nil {
		return message.Message{}, fmt.Errorf("failed to unmarshal parts: %w", err)
	}
	return message.Message{
		ID:        e.ID,
		Role:      message.MessageRole(e.Role),
		Parts:     parts,
		Model:     e.Model,
		Provider:  e.Provider,
		CreatedAt: e.Timestamp,
		UpdatedAt: e.Timestamp,
	}, nil
}