mirror of
https://github.com/mark3labs/kit.git
synced 2026-06-14 03:30:26 +00:00
add tree-structured session handling and /tree command
Implement pi-style JSONL append-only session management with tree branching: - TreeManager with id/parent_id tree structure, leaf pointer, and context building that walks leaf-to-root for LLM messages - Auto-discovery by cwd in ~/.kit/sessions/ with session listing - /tree TUI overlay with ASCII art rendering, filter modes, and navigation - /fork, /new, /name, /session slash commands for tree operations - --continue, --resume, --no-session CLI flags - Default auto-creates a tree session per working directory
This commit is contained in:
Ed Zynda
+111
-68
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
tea "charm.land/bubbletea/v2"
|
||||
@@ -40,6 +41,11 @@ var (
|
||||
loadSessionPath string
|
||||
sessionPath string
|
||||
|
||||
// Tree session management (pi-style)
|
||||
continueFlag bool // --continue / -c: resume most recent session for cwd
|
||||
resumeFlag bool // --resume / -r: interactive session picker
|
||||
noSessionFlag bool // --no-session: ephemeral mode, no persistence
|
||||
|
||||
// Model generation parameters
|
||||
maxTokens int
|
||||
temperature float32
|
||||
@@ -289,6 +295,12 @@ func init() {
|
||||
StringVar(&loadSessionPath, "load-session", "", "load session from file at startup")
|
||||
rootCmd.PersistentFlags().
|
||||
StringVarP(&sessionPath, "session", "s", "", "session file to load and update")
|
||||
rootCmd.PersistentFlags().
|
||||
BoolVarP(&continueFlag, "continue", "c", false, "continue the most recent session for the current directory")
|
||||
rootCmd.PersistentFlags().
|
||||
BoolVarP(&resumeFlag, "resume", "r", false, "interactive session picker")
|
||||
rootCmd.PersistentFlags().
|
||||
BoolVar(&noSessionFlag, "no-session", false, "ephemeral mode — no session persistence")
|
||||
|
||||
flags := rootCmd.PersistentFlags()
|
||||
flags.StringVar(&providerURL, "provider-url", "", "base URL for the provider API (applies to OpenAI, Anthropic, Ollama, and Google)")
|
||||
@@ -422,86 +434,117 @@ func runNormalMode(ctx context.Context) error {
|
||||
// Main interaction logic
|
||||
var messages []fantasy.Message
|
||||
var sessionManager *session.Manager
|
||||
if sessionPath != "" {
|
||||
_, err := os.Stat(sessionPath)
|
||||
if os.IsNotExist(err) {
|
||||
content := []byte("{}")
|
||||
if err := os.WriteFile(sessionPath, content, 0664); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
var treeSession *session.TreeManager
|
||||
|
||||
cwd, _ := os.Getwd()
|
||||
|
||||
// --- Tree session handling (--continue, --resume, default) ---
|
||||
if noSessionFlag {
|
||||
// Ephemeral mode: in-memory tree session, no persistence.
|
||||
treeSession = session.InMemoryTreeSession(cwd)
|
||||
} else if continueFlag {
|
||||
// Continue the most recent session for this cwd.
|
||||
ts, err := session.ContinueRecent(cwd)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to continue session: %v", err)
|
||||
}
|
||||
treeSession = ts
|
||||
// Load existing messages into the fantasy message slice.
|
||||
messages = ts.GetFantasyMessages()
|
||||
} else if resumeFlag {
|
||||
// Interactive session picker: list sessions and let user choose.
|
||||
sessions, err := session.ListSessions(cwd)
|
||||
if err != nil || len(sessions) == 0 {
|
||||
// No sessions found — create a new one.
|
||||
ts, err := session.CreateTreeSession(cwd)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create session: %v", err)
|
||||
}
|
||||
treeSession = ts
|
||||
} else {
|
||||
// For now, pick the most recent. TODO: TUI picker.
|
||||
ts, err := session.OpenTreeSession(sessions[0].Path)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to open session: %v", err)
|
||||
}
|
||||
treeSession = ts
|
||||
messages = ts.GetFantasyMessages()
|
||||
}
|
||||
} else if sessionPath != "" {
|
||||
// Legacy --session flag: open or create a specific JSONL session.
|
||||
if strings.HasSuffix(sessionPath, ".jsonl") {
|
||||
_, statErr := os.Stat(sessionPath)
|
||||
if os.IsNotExist(statErr) {
|
||||
// Create a new tree session at the specified path.
|
||||
dir := filepath.Dir(sessionPath)
|
||||
if err := os.MkdirAll(dir, 0755); err != nil {
|
||||
return fmt.Errorf("failed to create session directory: %v", err)
|
||||
}
|
||||
ts, err := session.CreateTreeSession(cwd)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create session: %v", err)
|
||||
}
|
||||
treeSession = ts
|
||||
} else {
|
||||
ts, err := session.OpenTreeSession(sessionPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to open session: %v", err)
|
||||
}
|
||||
treeSession = ts
|
||||
messages = ts.GetFantasyMessages()
|
||||
}
|
||||
} else {
|
||||
// Legacy JSON session path handling.
|
||||
_, statErr := os.Stat(sessionPath)
|
||||
if os.IsNotExist(statErr) {
|
||||
content := []byte("{}")
|
||||
if err := os.WriteFile(sessionPath, content, 0664); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
loadSessionPath = sessionPath
|
||||
saveSessionPath = sessionPath
|
||||
}
|
||||
} else {
|
||||
// Default: auto-create a tree session for the current directory.
|
||||
ts, err := session.CreateTreeSession(cwd)
|
||||
if err != nil {
|
||||
// Non-fatal: fall back to no session.
|
||||
if debugMode {
|
||||
fmt.Fprintf(os.Stderr, "Warning: could not create tree session: %v\n", err)
|
||||
}
|
||||
} else {
|
||||
treeSession = ts
|
||||
}
|
||||
loadSessionPath = sessionPath
|
||||
saveSessionPath = sessionPath
|
||||
}
|
||||
|
||||
// Load existing session if specified
|
||||
if loadSessionPath != "" {
|
||||
loadedSession, err := session.LoadFromFile(loadSessionPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load session: %v", err)
|
||||
}
|
||||
|
||||
// Convert session messages to fantasy messages
|
||||
for _, msg := range loadedSession.Messages {
|
||||
messages = append(messages, msg.ToFantasyMessages()...)
|
||||
}
|
||||
|
||||
// If we're also saving, use the loaded session with the session manager
|
||||
if saveSessionPath != "" {
|
||||
sessionManager = session.NewManagerWithSession(loadedSession, saveSessionPath)
|
||||
}
|
||||
|
||||
if !quietFlag && cli != nil {
|
||||
// Build a map of tool call IDs to tool calls for quick lookup
|
||||
type toolCallInfo struct {
|
||||
Name string
|
||||
Input string
|
||||
// --- Legacy JSON session handling (--load-session, --save-session) ---
|
||||
if treeSession == nil {
|
||||
if loadSessionPath != "" {
|
||||
loadedSession, err := session.LoadFromFile(loadSessionPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load session: %v", err)
|
||||
}
|
||||
toolCallMap := make(map[string]toolCallInfo)
|
||||
for _, sessionMsg := range loadedSession.Messages {
|
||||
for _, tc := range sessionMsg.ToolCalls() {
|
||||
toolCallMap[tc.ID] = toolCallInfo{Name: tc.Name, Input: tc.Input}
|
||||
}
|
||||
for _, msg := range loadedSession.Messages {
|
||||
messages = append(messages, msg.ToFantasyMessages()...)
|
||||
}
|
||||
|
||||
// Display all previous messages as they would have appeared
|
||||
for _, sessionMsg := range loadedSession.Messages {
|
||||
switch sessionMsg.Role {
|
||||
case "user":
|
||||
cli.DisplayUserMessage(sessionMsg.Content())
|
||||
case "assistant":
|
||||
// Tool calls are rendered as part of the unified tool result
|
||||
// block, so we skip displaying them here separately.
|
||||
|
||||
// Display assistant response (only if there's content)
|
||||
if text := sessionMsg.Content(); text != "" {
|
||||
_ = cli.DisplayAssistantMessage(text)
|
||||
}
|
||||
case "tool":
|
||||
// Display tool results
|
||||
for _, result := range sessionMsg.ToolResults() {
|
||||
if tc, exists := toolCallMap[result.ToolCallID]; exists {
|
||||
cli.DisplayToolMessage(tc.Name, tc.Input, result.Content, result.IsError)
|
||||
}
|
||||
}
|
||||
}
|
||||
if saveSessionPath != "" {
|
||||
sessionManager = session.NewManagerWithSession(loadedSession, saveSessionPath)
|
||||
}
|
||||
} else if saveSessionPath != "" {
|
||||
sessionManager = session.NewManager(saveSessionPath)
|
||||
_ = sessionManager.SetMetadata(session.Metadata{
|
||||
KitVersion: "dev",
|
||||
Provider: parsedProvider,
|
||||
Model: modelName,
|
||||
})
|
||||
}
|
||||
} else if saveSessionPath != "" {
|
||||
// Only saving, create new session manager
|
||||
sessionManager = session.NewManager(saveSessionPath)
|
||||
|
||||
// Set metadata
|
||||
_ = sessionManager.SetMetadata(session.Metadata{
|
||||
KitVersion: "dev", // TODO: Get actual version
|
||||
Provider: parsedProvider,
|
||||
Model: modelName,
|
||||
})
|
||||
}
|
||||
|
||||
// Create the app.App instance now that session messages are loaded.
|
||||
appOpts := BuildAppOptions(mcpAgent, mcpConfig, modelName, serverNames, toolNames)
|
||||
appOpts.SessionManager = sessionManager
|
||||
appOpts.TreeSession = treeSession
|
||||
|
||||
// Create a usage tracker that is shared between the app layer (for recording
|
||||
// usage after each step) and the TUI (for /usage display). For non-interactive
|
||||
|
||||
+40
-2
@@ -9,6 +9,7 @@ import (
|
||||
"charm.land/fantasy"
|
||||
|
||||
"github.com/mark3labs/kit/internal/agent"
|
||||
"github.com/mark3labs/kit/internal/session"
|
||||
)
|
||||
|
||||
// App is the application-layer orchestrator. It owns the agentic loop,
|
||||
@@ -151,11 +152,20 @@ func (a *App) ClearQueue() {
|
||||
a.mu.Unlock()
|
||||
}
|
||||
|
||||
// ClearMessages empties the conversation history.
|
||||
// ClearMessages empties the conversation history. When a tree session is
|
||||
// active the leaf pointer is reset to the root, creating an implicit branch.
|
||||
//
|
||||
// Satisfies ui.AppController.
|
||||
func (a *App) ClearMessages() {
|
||||
a.store.Clear()
|
||||
if a.opts.TreeSession != nil {
|
||||
a.opts.TreeSession.ResetLeaf()
|
||||
}
|
||||
}
|
||||
|
||||
// GetTreeSession returns the tree session manager, or nil if not configured.
|
||||
func (a *App) GetTreeSession() *session.TreeManager {
|
||||
return a.opts.TreeSession
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
@@ -250,6 +260,11 @@ func (a *App) Close() {
|
||||
|
||||
// Wait for background goroutines.
|
||||
a.wg.Wait()
|
||||
|
||||
// Close tree session file handle.
|
||||
if a.opts.TreeSession != nil {
|
||||
_ = a.opts.TreeSession.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
@@ -352,8 +367,23 @@ func (a *App) executeStep(ctx context.Context, prompt string, eventFn func(tea.M
|
||||
userMsg := fantasy.NewUserMessage(prompt)
|
||||
a.store.Add(userMsg)
|
||||
|
||||
// Persist user message to tree session if configured.
|
||||
if a.opts.TreeSession != nil {
|
||||
_, _ = a.opts.TreeSession.AppendFantasyMessage(userMsg)
|
||||
}
|
||||
|
||||
// Build the full message slice for the agent call.
|
||||
msgs := a.store.GetAll()
|
||||
// When a tree session is active, build context from the tree (walks
|
||||
// leaf-to-root) so that only the current branch is sent to the LLM.
|
||||
var msgs []fantasy.Message
|
||||
if a.opts.TreeSession != nil {
|
||||
msgs = a.opts.TreeSession.GetFantasyMessages()
|
||||
} else {
|
||||
msgs = a.store.GetAll()
|
||||
}
|
||||
|
||||
// Track message count before agent runs so we can diff new messages.
|
||||
sentCount := len(msgs)
|
||||
|
||||
// Signal spinner start.
|
||||
sendFn(SpinnerEvent{Show: true})
|
||||
@@ -398,6 +428,14 @@ func (a *App) executeStep(ctx context.Context, prompt string, eventFn func(tea.M
|
||||
// (includes tool call/result messages added during the step).
|
||||
a.store.Replace(result.ConversationMessages)
|
||||
|
||||
// Persist new messages (tool calls, tool results, assistant response)
|
||||
// to the tree session. Only append messages beyond what we sent.
|
||||
if a.opts.TreeSession != nil && len(result.ConversationMessages) > sentCount {
|
||||
for _, msg := range result.ConversationMessages[sentCount:] {
|
||||
_, _ = a.opts.TreeSession.AppendFantasyMessage(msg)
|
||||
}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -51,10 +51,16 @@ type Options struct {
|
||||
// *agent.Agent satisfies this interface; tests may supply stubs.
|
||||
Agent AgentRunner
|
||||
|
||||
// SessionManager is the optional session manager for persisting conversation
|
||||
// history to disk. When non-nil, the MessageStore calls it on every mutation.
|
||||
// SessionManager is the optional legacy session manager for persisting
|
||||
// conversation history to disk using the old JSON format.
|
||||
// Deprecated: Use TreeSession instead.
|
||||
SessionManager *session.Manager
|
||||
|
||||
// TreeSession is the tree-structured JSONL session manager. When non-nil,
|
||||
// conversation history is persisted as an append-only JSONL tree and tree
|
||||
// navigation (/tree, /fork) is enabled. Takes precedence over SessionManager.
|
||||
TreeSession *session.TreeManager
|
||||
|
||||
// MCPConfig is the full MCP configuration used for session continuation and
|
||||
// slash command resolution.
|
||||
MCPConfig *config.Config
|
||||
|
||||
@@ -0,0 +1,265 @@
|
||||
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.
|
||||
// Following pi's design, 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"
|
||||
)
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// 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"`
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// NewMessageEntryFromRaw creates a MessageEntry with pre-marshaled parts.
|
||||
func NewMessageEntryFromRaw(parentID, role string, parts json.RawMessage, model, provider string) *MessageEntry {
|
||||
return &MessageEntry{
|
||||
Entry: NewEntry(EntryTypeMessage, parentID),
|
||||
Role: role,
|
||||
Parts: parts,
|
||||
Model: model,
|
||||
Provider: provider,
|
||||
}
|
||||
}
|
||||
|
||||
// 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,
|
||||
}
|
||||
}
|
||||
|
||||
// --- 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
|
||||
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SessionInfo contains metadata about a discovered session, used for listing
|
||||
// and session picker display. Follows pi's SessionInfo design.
|
||||
type SessionInfo struct {
|
||||
// Path is the absolute path to the JSONL session file.
|
||||
Path string
|
||||
|
||||
// ID is the session UUID from the header.
|
||||
ID string
|
||||
|
||||
// Cwd is the working directory the session was created in.
|
||||
Cwd string
|
||||
|
||||
// Name is the user-defined display name (from session_info entries).
|
||||
Name string
|
||||
|
||||
// ParentSessionPath is the parent session path if this session was forked.
|
||||
ParentSessionPath string
|
||||
|
||||
// Created is when the session was first created.
|
||||
Created time.Time
|
||||
|
||||
// Modified is the timestamp of the last activity (latest message).
|
||||
Modified time.Time
|
||||
|
||||
// MessageCount is the number of message entries in the session.
|
||||
MessageCount int
|
||||
|
||||
// FirstMessage is a preview of the first user message.
|
||||
FirstMessage string
|
||||
}
|
||||
|
||||
// ListSessions finds all sessions for a given working directory, sorted by
|
||||
// modification time (newest first).
|
||||
func ListSessions(cwd string) ([]SessionInfo, error) {
|
||||
sessionDir := DefaultSessionDir(cwd)
|
||||
return listSessionsInDir(sessionDir)
|
||||
}
|
||||
|
||||
// ListAllSessions finds all sessions across all working directories, sorted
|
||||
// by modification time (newest first).
|
||||
func ListAllSessions() ([]SessionInfo, error) {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to find home directory: %w", err)
|
||||
}
|
||||
|
||||
sessionsRoot := filepath.Join(home, ".kit", "sessions")
|
||||
if _, err := os.Stat(sessionsRoot); os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var allSessions []SessionInfo
|
||||
|
||||
dirs, err := os.ReadDir(sessionsRoot)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read sessions directory: %w", err)
|
||||
}
|
||||
|
||||
for _, dir := range dirs {
|
||||
if !dir.IsDir() {
|
||||
continue
|
||||
}
|
||||
dirPath := filepath.Join(sessionsRoot, dir.Name())
|
||||
sessions, err := listSessionsInDir(dirPath)
|
||||
if err != nil {
|
||||
continue // skip unreadable directories
|
||||
}
|
||||
allSessions = append(allSessions, sessions...)
|
||||
}
|
||||
|
||||
// Sort by modification time, newest first.
|
||||
sort.Slice(allSessions, func(i, j int) bool {
|
||||
return allSessions[i].Modified.After(allSessions[j].Modified)
|
||||
})
|
||||
|
||||
return allSessions, nil
|
||||
}
|
||||
|
||||
// listSessionsInDir reads all .jsonl files in a directory and extracts session info.
|
||||
func listSessionsInDir(dir string) ([]SessionInfo, error) {
|
||||
if _, err := os.Stat(dir); os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read directory %s: %w", dir, err)
|
||||
}
|
||||
|
||||
var sessions []SessionInfo
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".jsonl") {
|
||||
continue
|
||||
}
|
||||
|
||||
path := filepath.Join(dir, entry.Name())
|
||||
info, err := extractSessionInfo(path)
|
||||
if err != nil {
|
||||
continue // skip malformed session files
|
||||
}
|
||||
sessions = append(sessions, *info)
|
||||
}
|
||||
|
||||
// Sort by modification time, newest first.
|
||||
sort.Slice(sessions, func(i, j int) bool {
|
||||
return sessions[i].Modified.After(sessions[j].Modified)
|
||||
})
|
||||
|
||||
return sessions, nil
|
||||
}
|
||||
|
||||
// extractSessionInfo reads a JSONL session file and extracts metadata.
|
||||
// It only reads enough of the file to get the header and scan for messages.
|
||||
func extractSessionInfo(path string) (*SessionInfo, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
info := &SessionInfo{
|
||||
Path: path,
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(f)
|
||||
// Increase scanner buffer for large lines.
|
||||
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
|
||||
lineNum := 0
|
||||
var lastTimestamp time.Time
|
||||
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
lineNum++
|
||||
|
||||
if lineNum == 1 {
|
||||
// Parse header.
|
||||
var h SessionHeader
|
||||
if err := json.Unmarshal([]byte(line), &h); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse header: %w", err)
|
||||
}
|
||||
if h.Type != EntryTypeSession {
|
||||
return nil, fmt.Errorf("first line is not a session header")
|
||||
}
|
||||
info.ID = h.ID
|
||||
info.Cwd = h.Cwd
|
||||
info.Created = h.Timestamp
|
||||
info.Modified = h.Timestamp
|
||||
info.ParentSessionPath = h.ParentSession
|
||||
continue
|
||||
}
|
||||
|
||||
// For subsequent lines, only parse enough to get type and timestamp.
|
||||
var env struct {
|
||||
Type EntryType `json:"type"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Role string `json:"role,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(line), &env); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
if !env.Timestamp.IsZero() && env.Timestamp.After(lastTimestamp) {
|
||||
lastTimestamp = env.Timestamp
|
||||
}
|
||||
|
||||
switch env.Type {
|
||||
case EntryTypeMessage:
|
||||
info.MessageCount++
|
||||
// Capture first user message as preview.
|
||||
if env.Role == "user" && info.FirstMessage == "" {
|
||||
var msgEntry struct {
|
||||
Parts json.RawMessage `json:"parts"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(line), &msgEntry); err == nil {
|
||||
info.FirstMessage = extractTextPreview(msgEntry.Parts)
|
||||
}
|
||||
}
|
||||
case EntryTypeSessionInfo:
|
||||
if env.Name != "" {
|
||||
info.Name = env.Name
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !lastTimestamp.IsZero() {
|
||||
info.Modified = lastTimestamp
|
||||
}
|
||||
|
||||
// Fall back to file modification time if no timestamps found.
|
||||
if info.Modified.IsZero() {
|
||||
fi, err := os.Stat(path)
|
||||
if err == nil {
|
||||
info.Modified = fi.ModTime()
|
||||
}
|
||||
}
|
||||
|
||||
return info, nil
|
||||
}
|
||||
|
||||
// extractTextPreview extracts a short text preview from type-tagged parts JSON.
|
||||
func extractTextPreview(partsJSON json.RawMessage) string {
|
||||
var parts []struct {
|
||||
Type string `json:"type"`
|
||||
Data json.RawMessage `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(partsJSON, &parts); err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
for _, p := range parts {
|
||||
if p.Type == "text" {
|
||||
var text struct {
|
||||
Text string `json:"text"`
|
||||
}
|
||||
if err := json.Unmarshal(p.Data, &text); err == nil && text.Text != "" {
|
||||
preview := text.Text
|
||||
if len(preview) > 100 {
|
||||
preview = preview[:100] + "..."
|
||||
}
|
||||
return preview
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// DeleteSession removes a session file from disk.
|
||||
func DeleteSession(path string) error {
|
||||
return os.Remove(path)
|
||||
}
|
||||
@@ -0,0 +1,689 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"charm.land/fantasy"
|
||||
|
||||
"github.com/mark3labs/kit/internal/message"
|
||||
)
|
||||
|
||||
// TreeNode represents a node in the session tree for display purposes.
|
||||
// It mirrors pi's SessionTreeNode design.
|
||||
type TreeNode struct {
|
||||
Entry any // the underlying entry (*MessageEntry, *ModelChangeEntry, etc.)
|
||||
ID string // entry ID
|
||||
ParentID string // parent entry ID
|
||||
Children []*TreeNode // child nodes
|
||||
}
|
||||
|
||||
// TreeManager manages a tree-structured JSONL session. It is the replacement
|
||||
// for the linear session.Manager, following pi's design decisions:
|
||||
//
|
||||
// - JSONL append-only format (one JSON object per line)
|
||||
// - Tree structure via id/parent_id on every entry
|
||||
// - Leaf pointer tracking current position
|
||||
// - Context building walks from leaf to root
|
||||
// - Auto-discovery by working directory
|
||||
type TreeManager struct {
|
||||
mu sync.RWMutex
|
||||
|
||||
// header is the session header (first line of the JSONL file).
|
||||
header SessionHeader
|
||||
|
||||
// entries is the ordered list of all entries (excluding header).
|
||||
entries []any
|
||||
|
||||
// index maps entry ID to the entry for O(1) lookup.
|
||||
index map[string]any
|
||||
|
||||
// childIndex maps parent ID to child entry IDs for tree traversal.
|
||||
childIndex map[string][]string
|
||||
|
||||
// labels maps entry ID to user-defined label string.
|
||||
labels map[string]string
|
||||
|
||||
// leafID is the current position in the tree. Empty string means
|
||||
// the session is at the root (before any entries).
|
||||
leafID string
|
||||
|
||||
// sessionName is the latest user-defined display name.
|
||||
sessionName string
|
||||
|
||||
// filePath is the JSONL file path. Empty for in-memory sessions.
|
||||
filePath string
|
||||
|
||||
// file is the open file handle for appending entries. Nil for in-memory.
|
||||
file *os.File
|
||||
}
|
||||
|
||||
// --- Constructors ---
|
||||
|
||||
// CreateTreeSession creates a new tree session persisted at the default
|
||||
// location for the given working directory.
|
||||
func CreateTreeSession(cwd string) (*TreeManager, error) {
|
||||
sessionDir := DefaultSessionDir(cwd)
|
||||
if err := os.MkdirAll(sessionDir, 0755); err != nil {
|
||||
return nil, fmt.Errorf("failed to create session directory: %w", err)
|
||||
}
|
||||
|
||||
now := time.Now().UTC()
|
||||
fileName := fmt.Sprintf("%s_%s.jsonl",
|
||||
now.Format("2006-01-02T15-04-05-000Z"),
|
||||
GenerateSessionID()[:12],
|
||||
)
|
||||
filePath := filepath.Join(sessionDir, fileName)
|
||||
|
||||
header := SessionHeader{
|
||||
Type: EntryTypeSession,
|
||||
Version: CurrentVersion,
|
||||
ID: GenerateSessionID(),
|
||||
Timestamp: now,
|
||||
Cwd: cwd,
|
||||
}
|
||||
|
||||
tm := &TreeManager{
|
||||
header: header,
|
||||
entries: make([]any, 0),
|
||||
index: make(map[string]any),
|
||||
childIndex: make(map[string][]string),
|
||||
labels: make(map[string]string),
|
||||
filePath: filePath,
|
||||
}
|
||||
|
||||
// Create the file and write the header.
|
||||
f, err := os.Create(filePath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create session file: %w", err)
|
||||
}
|
||||
tm.file = f
|
||||
|
||||
if err := tm.writeEntry(&header); err != nil {
|
||||
f.Close()
|
||||
return nil, fmt.Errorf("failed to write session header: %w", err)
|
||||
}
|
||||
|
||||
return tm, nil
|
||||
}
|
||||
|
||||
// OpenTreeSession opens an existing JSONL session file.
|
||||
func OpenTreeSession(path string) (*TreeManager, error) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read session file: %w", err)
|
||||
}
|
||||
|
||||
tm := &TreeManager{
|
||||
entries: make([]any, 0),
|
||||
index: make(map[string]any),
|
||||
childIndex: make(map[string][]string),
|
||||
labels: make(map[string]string),
|
||||
filePath: path,
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(strings.NewReader(string(data)))
|
||||
lineNum := 0
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
lineNum++
|
||||
|
||||
entry, err := UnmarshalEntry([]byte(line))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("line %d: %w", lineNum, err)
|
||||
}
|
||||
|
||||
if lineNum == 1 {
|
||||
h, ok := entry.(*SessionHeader)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("first line must be a session header, got %T", entry)
|
||||
}
|
||||
tm.header = *h
|
||||
continue
|
||||
}
|
||||
|
||||
tm.addEntryToIndex(entry)
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, fmt.Errorf("failed to scan session file: %w", err)
|
||||
}
|
||||
|
||||
// Set leaf to the last entry.
|
||||
if len(tm.entries) > 0 {
|
||||
tm.leafID = tm.entryID(tm.entries[len(tm.entries)-1])
|
||||
}
|
||||
|
||||
// Open file for appending.
|
||||
f, err := os.OpenFile(path, os.O_WRONLY|os.O_APPEND, 0644)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to open session file for append: %w", err)
|
||||
}
|
||||
tm.file = f
|
||||
|
||||
return tm, nil
|
||||
}
|
||||
|
||||
// ContinueRecent finds the most recently modified session for the given cwd,
|
||||
// or creates a new one if none exists.
|
||||
func ContinueRecent(cwd string) (*TreeManager, error) {
|
||||
sessions, err := ListSessions(cwd)
|
||||
if err != nil || len(sessions) == 0 {
|
||||
return CreateTreeSession(cwd)
|
||||
}
|
||||
// sessions are sorted by modified time (newest first).
|
||||
return OpenTreeSession(sessions[0].Path)
|
||||
}
|
||||
|
||||
// InMemoryTreeSession creates a tree session that is not persisted to disk.
|
||||
func InMemoryTreeSession(cwd string) *TreeManager {
|
||||
return &TreeManager{
|
||||
header: SessionHeader{
|
||||
Type: EntryTypeSession,
|
||||
Version: CurrentVersion,
|
||||
ID: GenerateSessionID(),
|
||||
Timestamp: time.Now().UTC(),
|
||||
Cwd: cwd,
|
||||
},
|
||||
entries: make([]any, 0),
|
||||
index: make(map[string]any),
|
||||
childIndex: make(map[string][]string),
|
||||
labels: make(map[string]string),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Append operations (all return entry ID) ---
|
||||
|
||||
// AppendMessage adds a message entry to the tree and persists it.
|
||||
func (tm *TreeManager) AppendMessage(msg message.Message) (string, error) {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
|
||||
entry, err := NewMessageEntry(tm.leafID, msg)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
tm.leafID = entry.ID
|
||||
return entry.ID, nil
|
||||
}
|
||||
|
||||
// AppendFantasyMessage converts a fantasy.Message and appends it.
|
||||
func (tm *TreeManager) AppendFantasyMessage(msg fantasy.Message) (string, error) {
|
||||
return tm.AppendMessage(message.FromFantasyMessage(msg))
|
||||
}
|
||||
|
||||
// AppendModelChange records a model/provider change.
|
||||
func (tm *TreeManager) AppendModelChange(provider, modelID string) (string, error) {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
|
||||
entry := NewModelChangeEntry(tm.leafID, provider, modelID)
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
tm.leafID = entry.ID
|
||||
return entry.ID, nil
|
||||
}
|
||||
|
||||
// AppendBranchSummary adds a summary of an abandoned branch.
|
||||
func (tm *TreeManager) AppendBranchSummary(fromID, summary string) (string, error) {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
|
||||
entry := NewBranchSummaryEntry(tm.leafID, fromID, summary)
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
tm.leafID = entry.ID
|
||||
return entry.ID, nil
|
||||
}
|
||||
|
||||
// AppendLabel sets a label on a target entry.
|
||||
func (tm *TreeManager) AppendLabel(targetID, label string) (string, error) {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
|
||||
entry := NewLabelEntry(tm.leafID, targetID, label)
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
tm.labels[targetID] = label
|
||||
tm.leafID = entry.ID
|
||||
return entry.ID, nil
|
||||
}
|
||||
|
||||
// AppendSessionInfo sets a display name for the session.
|
||||
func (tm *TreeManager) AppendSessionInfo(name string) (string, error) {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
|
||||
entry := NewSessionInfoEntry(tm.leafID, name)
|
||||
if err := tm.appendAndPersist(entry); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
tm.sessionName = name
|
||||
tm.leafID = entry.ID
|
||||
return entry.ID, nil
|
||||
}
|
||||
|
||||
// --- Tree navigation ---
|
||||
|
||||
// Branch moves the leaf pointer to the given entry ID, creating a branch
|
||||
// point. Subsequent appends will extend from this new position.
|
||||
func (tm *TreeManager) Branch(entryID string) error {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
|
||||
if entryID == "" {
|
||||
tm.leafID = ""
|
||||
return nil
|
||||
}
|
||||
|
||||
if _, ok := tm.index[entryID]; !ok {
|
||||
return fmt.Errorf("entry %q not found", entryID)
|
||||
}
|
||||
tm.leafID = entryID
|
||||
return nil
|
||||
}
|
||||
|
||||
// ResetLeaf moves the leaf pointer to before the first entry (empty conversation).
|
||||
func (tm *TreeManager) ResetLeaf() {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
tm.leafID = ""
|
||||
}
|
||||
|
||||
// GetLeafID returns the current leaf position.
|
||||
func (tm *TreeManager) GetLeafID() string {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
return tm.leafID
|
||||
}
|
||||
|
||||
// GetEntry returns the entry with the given ID, or nil if not found.
|
||||
func (tm *TreeManager) GetEntry(id string) any {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
return tm.index[id]
|
||||
}
|
||||
|
||||
// GetEntries returns all entries (excluding the session header).
|
||||
func (tm *TreeManager) GetEntries() []any {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
|
||||
cp := make([]any, len(tm.entries))
|
||||
copy(cp, tm.entries)
|
||||
return cp
|
||||
}
|
||||
|
||||
// GetChildren returns direct child entry IDs for a given parent.
|
||||
func (tm *TreeManager) GetChildren(parentID string) []string {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
|
||||
cp := make([]string, len(tm.childIndex[parentID]))
|
||||
copy(cp, tm.childIndex[parentID])
|
||||
return cp
|
||||
}
|
||||
|
||||
// GetBranch returns the path of entries from the given entry to the root,
|
||||
// ordered from root to the entry. If fromID is empty, uses the current leaf.
|
||||
func (tm *TreeManager) GetBranch(fromID string) []any {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
|
||||
if fromID == "" {
|
||||
fromID = tm.leafID
|
||||
}
|
||||
if fromID == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
var path []any
|
||||
current := fromID
|
||||
for current != "" {
|
||||
entry, ok := tm.index[current]
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
path = append(path, entry)
|
||||
current = tm.entryParentID(entry)
|
||||
}
|
||||
|
||||
// Reverse to get root-to-leaf order.
|
||||
for i, j := 0, len(path)-1; i < j; i, j = i+1, j-1 {
|
||||
path[i], path[j] = path[j], path[i]
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
// GetLabel returns the label for an entry, or empty string if none.
|
||||
func (tm *TreeManager) GetLabel(id string) string {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
return tm.labels[id]
|
||||
}
|
||||
|
||||
// GetTree builds the full tree structure from root entries.
|
||||
func (tm *TreeManager) GetTree() []*TreeNode {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
|
||||
// Find root entries (entries with empty ParentID).
|
||||
var roots []*TreeNode
|
||||
rootIDs := tm.childIndex[""]
|
||||
|
||||
for _, id := range rootIDs {
|
||||
node := tm.buildTreeNode(id)
|
||||
if node != nil {
|
||||
roots = append(roots, node)
|
||||
}
|
||||
}
|
||||
return roots
|
||||
}
|
||||
|
||||
// --- Context building ---
|
||||
|
||||
// BuildContext walks from the current leaf to the root and returns the
|
||||
// conversation messages suitable for sending to the LLM. Branch summaries
|
||||
// are converted to user messages to provide context from abandoned branches.
|
||||
// Also returns the latest model/provider settings encountered on the path.
|
||||
func (tm *TreeManager) BuildContext() (messages []fantasy.Message, provider string, modelID string) {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
|
||||
if tm.leafID == "" {
|
||||
return nil, "", ""
|
||||
}
|
||||
|
||||
// Walk from leaf to root collecting entries.
|
||||
branch := tm.getBranchLocked(tm.leafID)
|
||||
|
||||
for _, entry := range branch {
|
||||
switch e := entry.(type) {
|
||||
case *MessageEntry:
|
||||
msg, err := e.ToMessage()
|
||||
if err != nil {
|
||||
continue // skip malformed entries
|
||||
}
|
||||
msgs := msg.ToFantasyMessages()
|
||||
messages = append(messages, msgs...)
|
||||
|
||||
case *BranchSummaryEntry:
|
||||
// Convert branch summary to a user message for context.
|
||||
if e.Summary != "" {
|
||||
messages = append(messages, fantasy.Message{
|
||||
Role: fantasy.MessageRoleUser,
|
||||
Content: []fantasy.MessagePart{
|
||||
fantasy.TextPart{
|
||||
Text: fmt.Sprintf("[Branch context: %s]", e.Summary),
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
case *ModelChangeEntry:
|
||||
provider = e.Provider
|
||||
modelID = e.ModelID
|
||||
}
|
||||
}
|
||||
|
||||
return messages, provider, modelID
|
||||
}
|
||||
|
||||
// --- Session info ---
|
||||
|
||||
// GetSessionID returns the session UUID.
|
||||
func (tm *TreeManager) GetSessionID() string {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
return tm.header.ID
|
||||
}
|
||||
|
||||
// GetSessionName returns the user-defined display name, or empty string.
|
||||
func (tm *TreeManager) GetSessionName() string {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
return tm.sessionName
|
||||
}
|
||||
|
||||
// GetCwd returns the working directory this session was created in.
|
||||
func (tm *TreeManager) GetCwd() string {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
return tm.header.Cwd
|
||||
}
|
||||
|
||||
// GetFilePath returns the JSONL file path, or empty for in-memory sessions.
|
||||
func (tm *TreeManager) GetFilePath() string {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
return tm.filePath
|
||||
}
|
||||
|
||||
// GetHeader returns a copy of the session header.
|
||||
func (tm *TreeManager) GetHeader() SessionHeader {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
return tm.header
|
||||
}
|
||||
|
||||
// IsPersisted returns true if this session writes to disk.
|
||||
func (tm *TreeManager) IsPersisted() bool {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
return tm.filePath != ""
|
||||
}
|
||||
|
||||
// EntryCount returns the number of entries (excluding header).
|
||||
func (tm *TreeManager) EntryCount() int {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
return len(tm.entries)
|
||||
}
|
||||
|
||||
// MessageCount returns the number of message entries.
|
||||
func (tm *TreeManager) MessageCount() int {
|
||||
tm.mu.RLock()
|
||||
defer tm.mu.RUnlock()
|
||||
count := 0
|
||||
for _, e := range tm.entries {
|
||||
if _, ok := e.(*MessageEntry); ok {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
// Close closes the underlying file handle.
|
||||
func (tm *TreeManager) Close() error {
|
||||
tm.mu.Lock()
|
||||
defer tm.mu.Unlock()
|
||||
if tm.file != nil {
|
||||
err := tm.file.Close()
|
||||
tm.file = nil
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- Legacy bridge ---
|
||||
|
||||
// AddFantasyMessages appends multiple fantasy messages as entries. This is
|
||||
// used when syncing from the agent's ConversationMessages after a step.
|
||||
func (tm *TreeManager) AddFantasyMessages(msgs []fantasy.Message) error {
|
||||
for _, msg := range msgs {
|
||||
if _, err := tm.AppendFantasyMessage(msg); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetFantasyMessages builds the context and returns just the messages.
|
||||
// This satisfies the same conceptual role as the old Manager.GetMessages().
|
||||
func (tm *TreeManager) GetFantasyMessages() []fantasy.Message {
|
||||
msgs, _, _ := tm.BuildContext()
|
||||
return msgs
|
||||
}
|
||||
|
||||
// --- Internal helpers ---
|
||||
|
||||
// addEntryToIndex adds an entry to the in-memory indices.
|
||||
func (tm *TreeManager) addEntryToIndex(entry any) {
|
||||
tm.entries = append(tm.entries, entry)
|
||||
|
||||
id := tm.entryID(entry)
|
||||
parentID := tm.entryParentID(entry)
|
||||
|
||||
if id != "" {
|
||||
tm.index[id] = entry
|
||||
tm.childIndex[parentID] = append(tm.childIndex[parentID], id)
|
||||
}
|
||||
|
||||
// Track labels and session names.
|
||||
switch e := entry.(type) {
|
||||
case *LabelEntry:
|
||||
tm.labels[e.TargetID] = e.Label
|
||||
case *SessionInfoEntry:
|
||||
tm.sessionName = e.Name
|
||||
}
|
||||
}
|
||||
|
||||
// appendAndPersist adds an entry to indices and writes it to the JSONL file.
|
||||
func (tm *TreeManager) appendAndPersist(entry any) error {
|
||||
tm.addEntryToIndex(entry)
|
||||
if tm.file != nil {
|
||||
return tm.writeEntry(entry)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// writeEntry serializes an entry and appends it as a line to the file.
|
||||
func (tm *TreeManager) writeEntry(entry any) error {
|
||||
data, err := json.Marshal(entry)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal entry: %w", err)
|
||||
}
|
||||
data = append(data, '\n')
|
||||
_, err = tm.file.Write(data)
|
||||
return err
|
||||
}
|
||||
|
||||
// entryID extracts the ID from any entry type.
|
||||
func (tm *TreeManager) entryID(entry any) string {
|
||||
switch e := entry.(type) {
|
||||
case *MessageEntry:
|
||||
return e.ID
|
||||
case *ModelChangeEntry:
|
||||
return e.ID
|
||||
case *BranchSummaryEntry:
|
||||
return e.ID
|
||||
case *LabelEntry:
|
||||
return e.ID
|
||||
case *SessionInfoEntry:
|
||||
return e.ID
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// entryParentID extracts the ParentID from any entry type.
|
||||
func (tm *TreeManager) entryParentID(entry any) string {
|
||||
switch e := entry.(type) {
|
||||
case *MessageEntry:
|
||||
return e.ParentID
|
||||
case *ModelChangeEntry:
|
||||
return e.ParentID
|
||||
case *BranchSummaryEntry:
|
||||
return e.ParentID
|
||||
case *LabelEntry:
|
||||
return e.ParentID
|
||||
case *SessionInfoEntry:
|
||||
return e.ParentID
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// getBranchLocked walks from an entry to the root (must hold at least RLock).
|
||||
func (tm *TreeManager) getBranchLocked(fromID string) []any {
|
||||
var path []any
|
||||
visited := make(map[string]bool) // prevent cycles
|
||||
current := fromID
|
||||
for current != "" {
|
||||
if visited[current] {
|
||||
break
|
||||
}
|
||||
visited[current] = true
|
||||
entry, ok := tm.index[current]
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
path = append(path, entry)
|
||||
current = tm.entryParentID(entry)
|
||||
}
|
||||
|
||||
// Reverse to root-first order.
|
||||
for i, j := 0, len(path)-1; i < j; i, j = i+1, j-1 {
|
||||
path[i], path[j] = path[j], path[i]
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
// buildTreeNode recursively builds a TreeNode from an entry ID.
|
||||
func (tm *TreeManager) buildTreeNode(id string) *TreeNode {
|
||||
entry, ok := tm.index[id]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
node := &TreeNode{
|
||||
Entry: entry,
|
||||
ID: id,
|
||||
ParentID: tm.entryParentID(entry),
|
||||
}
|
||||
|
||||
for _, childID := range tm.childIndex[id] {
|
||||
child := tm.buildTreeNode(childID)
|
||||
if child != nil {
|
||||
node.Children = append(node.Children, child)
|
||||
}
|
||||
}
|
||||
|
||||
return node
|
||||
}
|
||||
|
||||
// --- Path conventions ---
|
||||
|
||||
// DefaultSessionDir returns the default session storage directory for a cwd.
|
||||
// Following pi's convention: ~/.kit/sessions/--<cwd-path>--/
|
||||
func DefaultSessionDir(cwd string) string {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
home = "."
|
||||
}
|
||||
// Convert path separators to double dashes.
|
||||
safeCwd := strings.ReplaceAll(cwd, string(filepath.Separator), "--")
|
||||
// Remove leading separator replacement.
|
||||
safeCwd = strings.TrimPrefix(safeCwd, "--")
|
||||
return filepath.Join(home, ".kit", "sessions", safeCwd)
|
||||
}
|
||||
@@ -65,6 +65,34 @@ var SlashCommands = []SlashCommand{
|
||||
Category: "System",
|
||||
Aliases: []string{"/q", "/exit"},
|
||||
},
|
||||
|
||||
// Navigation commands (tree sessions)
|
||||
{
|
||||
Name: "/tree",
|
||||
Description: "Navigate session tree (switch branches)",
|
||||
Category: "Navigation",
|
||||
},
|
||||
{
|
||||
Name: "/fork",
|
||||
Description: "Branch from an earlier message",
|
||||
Category: "Navigation",
|
||||
},
|
||||
{
|
||||
Name: "/new",
|
||||
Description: "Start a new session",
|
||||
Category: "Navigation",
|
||||
Aliases: []string{"/n"},
|
||||
},
|
||||
{
|
||||
Name: "/name",
|
||||
Description: "Set a display name for this session",
|
||||
Category: "Navigation",
|
||||
},
|
||||
{
|
||||
Name: "/session",
|
||||
Description: "Show session info and statistics",
|
||||
Category: "Info",
|
||||
},
|
||||
}
|
||||
|
||||
// GetCommandByName looks up a slash command by its primary name or any of its
|
||||
|
||||
@@ -11,3 +11,20 @@ type submitMsg struct {
|
||||
// presses ESC once during stateWorking. If this message arrives before the user
|
||||
// presses ESC a second time, the canceling state is reset to false.
|
||||
type cancelTimerExpiredMsg struct{}
|
||||
|
||||
// --- Tree session events ---
|
||||
|
||||
// TreeNodeSelectedMsg is sent when the user selects a node in the tree selector.
|
||||
type TreeNodeSelectedMsg struct {
|
||||
// ID is the entry ID of the selected node.
|
||||
ID string
|
||||
// Entry is the underlying entry object.
|
||||
Entry any
|
||||
// IsUser is true if the selected entry is a user message.
|
||||
IsUser bool
|
||||
// UserText is the user message text (only set when IsUser is true).
|
||||
UserText string
|
||||
}
|
||||
|
||||
// TreeCancelledMsg is sent when the user cancels the tree selector (ESC).
|
||||
type TreeCancelledMsg struct{}
|
||||
|
||||
+193
-2
@@ -8,6 +8,7 @@ import (
|
||||
tea "charm.land/bubbletea/v2"
|
||||
"charm.land/lipgloss/v2"
|
||||
"github.com/mark3labs/kit/internal/app"
|
||||
"github.com/mark3labs/kit/internal/session"
|
||||
)
|
||||
|
||||
// appState represents the current state of the parent TUI model.
|
||||
@@ -21,6 +22,9 @@ const (
|
||||
// stateWorking means the agent is running. The stream component is active.
|
||||
// The input component remains visible and editable for queueing messages.
|
||||
stateWorking
|
||||
|
||||
// stateTreeSelector means the /tree viewer is active.
|
||||
stateTreeSelector
|
||||
)
|
||||
|
||||
// AppController is the interface the parent TUI model uses to interact with the
|
||||
@@ -45,6 +49,9 @@ type AppController interface {
|
||||
ClearQueue()
|
||||
// ClearMessages clears the conversation history.
|
||||
ClearMessages()
|
||||
// GetTreeSession returns the tree session manager, or nil if tree sessions
|
||||
// are not enabled. Used by slash commands like /tree, /fork, /session.
|
||||
GetTreeSession() *session.TreeManager
|
||||
}
|
||||
|
||||
// AppModelOptions holds configuration passed to NewAppModel.
|
||||
@@ -147,6 +154,9 @@ type AppModel struct {
|
||||
// May be nil when usage tracking is unavailable.
|
||||
usageTracker *UsageTracker
|
||||
|
||||
// treeSelector is the tree navigation overlay, active in stateTreeSelector.
|
||||
treeSelector *TreeSelectorComponent
|
||||
|
||||
// width and height track the terminal dimensions.
|
||||
width int
|
||||
height int
|
||||
@@ -273,6 +283,50 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
|
||||
switch msg := msg.(type) {
|
||||
|
||||
// ── Tree selector events ─────────────────────────────────────────────────
|
||||
case TreeNodeSelectedMsg:
|
||||
// User selected a node in the tree. Branch to it and return to input.
|
||||
if ts := m.appCtrl.GetTreeSession(); ts != nil {
|
||||
// For user messages: branch to parent (so user can resubmit).
|
||||
// For other entries: branch directly to the selected entry.
|
||||
targetID := msg.ID
|
||||
if msg.IsUser {
|
||||
// Branch to parent of user message, place text in editor.
|
||||
if node := ts.GetEntry(msg.ID); node != nil {
|
||||
if me, ok := node.(*session.MessageEntry); ok {
|
||||
targetID = me.ParentID
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = ts.Branch(targetID)
|
||||
m.appCtrl.ClearMessages()
|
||||
|
||||
// If it was a user message, populate the input with the text.
|
||||
if msg.IsUser && msg.UserText != "" {
|
||||
if ic, ok := m.input.(*InputComponent); ok {
|
||||
ic.textarea.SetValue(msg.UserText)
|
||||
ic.textarea.CursorEnd()
|
||||
}
|
||||
}
|
||||
|
||||
cmds = append(cmds, m.printSystemMessage(
|
||||
fmt.Sprintf("Navigated to branch point. %s",
|
||||
func() string {
|
||||
if msg.IsUser {
|
||||
return "Edit and resubmit to create a new branch."
|
||||
}
|
||||
return "Continue from this point."
|
||||
}())))
|
||||
}
|
||||
m.treeSelector = nil
|
||||
m.state = stateInput
|
||||
return m, tea.Batch(cmds...)
|
||||
|
||||
case TreeCancelledMsg:
|
||||
m.treeSelector = nil
|
||||
m.state = stateInput
|
||||
return m, nil
|
||||
|
||||
// ── Window resize ────────────────────────────────────────────────────────
|
||||
case tea.WindowSizeMsg:
|
||||
m.width = msg.Width
|
||||
@@ -294,7 +348,17 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
case "ctrl+c":
|
||||
// Graceful quit: app.Close() is deferred in cmd/root.go.
|
||||
return m, tea.Quit
|
||||
}
|
||||
|
||||
// Route to tree selector when active.
|
||||
if m.state == stateTreeSelector && m.treeSelector != nil {
|
||||
updated, cmd := m.treeSelector.Update(msg)
|
||||
m.treeSelector = updated.(*TreeSelectorComponent)
|
||||
cmds = append(cmds, cmd)
|
||||
return m, tea.Batch(cmds...)
|
||||
}
|
||||
|
||||
switch msg.String() {
|
||||
case "esc":
|
||||
if m.state == stateWorking {
|
||||
if m.canceling {
|
||||
@@ -483,7 +547,13 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
|
||||
// View implements tea.Model. It renders the stacked layout:
|
||||
// stream region + [usage info] + separator + [queued messages] + input region.
|
||||
// When the tree selector is active, it replaces the stream region.
|
||||
func (m *AppModel) View() tea.View {
|
||||
// Tree selector overlay replaces the normal layout.
|
||||
if m.state == stateTreeSelector && m.treeSelector != nil {
|
||||
return m.treeSelector.View()
|
||||
}
|
||||
|
||||
streamView := m.renderStream()
|
||||
separator := m.renderSeparator()
|
||||
inputView := m.renderInput()
|
||||
@@ -691,6 +761,18 @@ func (m *AppModel) handleSlashCommand(sc *SlashCommand) tea.Cmd {
|
||||
m.queuedMessages = m.queuedMessages[:0]
|
||||
m.distributeHeight()
|
||||
return nil
|
||||
|
||||
case "/tree":
|
||||
return m.handleTreeCommand()
|
||||
case "/fork":
|
||||
return m.handleForkCommand()
|
||||
case "/new":
|
||||
return m.handleNewCommand()
|
||||
case "/name":
|
||||
return m.handleNameCommand()
|
||||
case "/session":
|
||||
return m.handleSessionInfoCommand()
|
||||
|
||||
default:
|
||||
return m.printSystemMessage(fmt.Sprintf("Unknown command: %s", sc.Name))
|
||||
}
|
||||
@@ -712,13 +794,21 @@ func (m *AppModel) printSystemMessage(text string) tea.Cmd {
|
||||
// printHelpMessage renders the help text listing all available slash commands.
|
||||
func (m *AppModel) printHelpMessage() tea.Cmd {
|
||||
help := "## Available Commands\n\n" +
|
||||
"**Info:**\n" +
|
||||
"- `/help`: Show this help message\n" +
|
||||
"- `/tools`: List all available tools\n" +
|
||||
"- `/servers`: List configured MCP servers\n" +
|
||||
"- `/usage`: Show token usage and cost statistics\n" +
|
||||
"- `/reset-usage`: Reset usage statistics\n" +
|
||||
"- `/session`: Show session info and statistics\n\n" +
|
||||
"**Navigation:**\n" +
|
||||
"- `/tree`: Navigate session tree (switch branches)\n" +
|
||||
"- `/fork`: Branch from an earlier message\n" +
|
||||
"- `/new`: Start a new branch (preserves history)\n\n" +
|
||||
"**System:**\n" +
|
||||
"- `/clear`: Clear message history\n" +
|
||||
"- `/quit`: Exit the application\n" +
|
||||
"- `/reset-usage`: Reset usage statistics\n" +
|
||||
"- `/quit`: Exit the application\n\n" +
|
||||
"**Keys:**\n" +
|
||||
"- `Ctrl+C`: Exit at any time\n" +
|
||||
"- `ESC` (x2): Cancel ongoing LLM generation\n\n" +
|
||||
"You can also just type your message to chat with the AI assistant."
|
||||
@@ -840,6 +930,107 @@ func repeatRune(r rune, n int) string {
|
||||
return string(runes)
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Tree session command handlers
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
// handleTreeCommand opens the tree selector overlay.
|
||||
func (m *AppModel) handleTreeCommand() tea.Cmd {
|
||||
ts := m.appCtrl.GetTreeSession()
|
||||
if ts == nil {
|
||||
return m.printSystemMessage("No tree session active. Start with `--continue` or `--resume` to enable tree sessions.")
|
||||
}
|
||||
if ts.EntryCount() == 0 {
|
||||
return m.printSystemMessage("No entries in session yet.")
|
||||
}
|
||||
|
||||
m.treeSelector = NewTreeSelector(ts, m.width, m.height)
|
||||
m.state = stateTreeSelector
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleForkCommand creates a branch from the current position. Like /tree
|
||||
// but opens the selector directly for fork semantics.
|
||||
func (m *AppModel) handleForkCommand() tea.Cmd {
|
||||
ts := m.appCtrl.GetTreeSession()
|
||||
if ts == nil {
|
||||
return m.printSystemMessage("No tree session active. Start with `--continue` or `--resume` to enable tree sessions.")
|
||||
}
|
||||
if ts.EntryCount() == 0 {
|
||||
return m.printSystemMessage("No entries to fork from.")
|
||||
}
|
||||
|
||||
m.treeSelector = NewTreeSelector(ts, m.width, m.height)
|
||||
m.state = stateTreeSelector
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleNewCommand starts a fresh session by resetting the tree leaf.
|
||||
func (m *AppModel) handleNewCommand() tea.Cmd {
|
||||
ts := m.appCtrl.GetTreeSession()
|
||||
if ts == nil {
|
||||
// No tree session — just clear messages.
|
||||
if m.appCtrl != nil {
|
||||
m.appCtrl.ClearMessages()
|
||||
}
|
||||
return m.printSystemMessage("Conversation cleared. Starting fresh.")
|
||||
}
|
||||
|
||||
ts.ResetLeaf()
|
||||
if m.appCtrl != nil {
|
||||
m.appCtrl.ClearMessages()
|
||||
}
|
||||
return m.printSystemMessage("New branch started. Previous conversation is preserved in the tree.")
|
||||
}
|
||||
|
||||
// handleNameCommand sets a display name for the current session.
|
||||
func (m *AppModel) handleNameCommand() tea.Cmd {
|
||||
ts := m.appCtrl.GetTreeSession()
|
||||
if ts == nil {
|
||||
return m.printSystemMessage("No tree session active.")
|
||||
}
|
||||
// For now, prompt user to provide name via input. We print instructions
|
||||
// and the next non-command input starting with "name:" will be captured.
|
||||
// TODO: inline input dialog like pi's implementation.
|
||||
currentName := ts.GetSessionName()
|
||||
if currentName != "" {
|
||||
return m.printSystemMessage(fmt.Sprintf("Current session name: %q\nTo rename, type: `/name <new name>` (not yet implemented — use the session file directly).", currentName))
|
||||
}
|
||||
return m.printSystemMessage("To name this session, use: `/name <new name>` (not yet implemented — use the session file directly).")
|
||||
}
|
||||
|
||||
// handleSessionInfoCommand shows session statistics.
|
||||
func (m *AppModel) handleSessionInfoCommand() tea.Cmd {
|
||||
ts := m.appCtrl.GetTreeSession()
|
||||
if ts == nil {
|
||||
return m.printSystemMessage("No tree session active.")
|
||||
}
|
||||
|
||||
header := ts.GetHeader()
|
||||
info := fmt.Sprintf("## Session Info\n\n"+
|
||||
"- **ID:** `%s`\n"+
|
||||
"- **File:** `%s`\n"+
|
||||
"- **Working Dir:** `%s`\n"+
|
||||
"- **Created:** %s\n"+
|
||||
"- **Entries:** %d\n"+
|
||||
"- **Messages:** %d\n"+
|
||||
"- **Current Leaf:** `%s`\n",
|
||||
header.ID,
|
||||
ts.GetFilePath(),
|
||||
header.Cwd,
|
||||
header.Timestamp.Format(time.RFC3339),
|
||||
ts.EntryCount(),
|
||||
ts.MessageCount(),
|
||||
ts.GetLeafID(),
|
||||
)
|
||||
|
||||
if name := ts.GetSessionName(); name != "" {
|
||||
info += fmt.Sprintf("- **Name:** %s\n", name)
|
||||
}
|
||||
|
||||
return m.printSystemMessage(info)
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Cancel timer command
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
tea "charm.land/bubbletea/v2"
|
||||
"charm.land/fantasy"
|
||||
"github.com/mark3labs/kit/internal/app"
|
||||
"github.com/mark3labs/kit/internal/session"
|
||||
)
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
@@ -45,6 +46,10 @@ func (s *stubAppController) ClearMessages() {
|
||||
s.clearMsgCalled++
|
||||
}
|
||||
|
||||
func (s *stubAppController) GetTreeSession() *session.TreeManager {
|
||||
return nil
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Stub child components
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,528 @@
|
||||
package ui
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"charm.land/bubbles/v2/key"
|
||||
tea "charm.land/bubbletea/v2"
|
||||
"charm.land/lipgloss/v2"
|
||||
|
||||
"github.com/mark3labs/kit/internal/session"
|
||||
)
|
||||
|
||||
// TreeFilterMode controls which entries are visible in the tree selector.
|
||||
type TreeFilterMode int
|
||||
|
||||
const (
|
||||
TreeFilterDefault TreeFilterMode = iota // hide settings entries
|
||||
TreeFilterNoTools // hide tool results
|
||||
TreeFilterUserOnly // show only user messages
|
||||
TreeFilterLabelOnly // show only labeled entries
|
||||
TreeFilterAll // show everything
|
||||
)
|
||||
|
||||
func (m TreeFilterMode) String() string {
|
||||
switch m {
|
||||
case TreeFilterDefault:
|
||||
return "default"
|
||||
case TreeFilterNoTools:
|
||||
return "no-tools"
|
||||
case TreeFilterUserOnly:
|
||||
return "user-only"
|
||||
case TreeFilterLabelOnly:
|
||||
return "labeled"
|
||||
case TreeFilterAll:
|
||||
return "all"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// FlatNode is a tree entry flattened for list rendering with indentation info.
|
||||
type FlatNode struct {
|
||||
Entry any // the underlying entry
|
||||
ID string // entry ID
|
||||
ParentID string
|
||||
Depth int // indentation level
|
||||
IsLast bool // last child at this depth
|
||||
Prefix string // computed prefix string (├─, └─, etc.)
|
||||
Label string // user-defined label, if any
|
||||
}
|
||||
|
||||
// TreeSelectorComponent is a Bubble Tea component that renders the session
|
||||
// tree as an ASCII art list with navigation and selection. It follows pi's
|
||||
// tree selector design.
|
||||
type TreeSelectorComponent struct {
|
||||
tm *session.TreeManager
|
||||
flatNodes []FlatNode
|
||||
cursor int
|
||||
filter TreeFilterMode
|
||||
leafID string // real leaf for "active" marker
|
||||
width int
|
||||
height int
|
||||
search string
|
||||
active bool
|
||||
selectedID string // set when user selects a node
|
||||
cancelled bool
|
||||
}
|
||||
|
||||
// NewTreeSelector creates a tree selector from a TreeManager.
|
||||
func NewTreeSelector(tm *session.TreeManager, width, height int) *TreeSelectorComponent {
|
||||
ts := &TreeSelectorComponent{
|
||||
tm: tm,
|
||||
filter: TreeFilterDefault,
|
||||
leafID: tm.GetLeafID(),
|
||||
width: width,
|
||||
height: height,
|
||||
active: true,
|
||||
}
|
||||
ts.rebuildFlatList()
|
||||
// Position cursor at the active leaf.
|
||||
for i, node := range ts.flatNodes {
|
||||
if node.ID == ts.leafID {
|
||||
ts.cursor = i
|
||||
break
|
||||
}
|
||||
}
|
||||
return ts
|
||||
}
|
||||
|
||||
// Init implements tea.Model.
|
||||
func (ts *TreeSelectorComponent) Init() tea.Cmd {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Update implements tea.Model.
|
||||
func (ts *TreeSelectorComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
switch msg := msg.(type) {
|
||||
case tea.WindowSizeMsg:
|
||||
ts.width = msg.Width
|
||||
ts.height = msg.Height
|
||||
return ts, nil
|
||||
|
||||
case tea.KeyPressMsg:
|
||||
switch {
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("up", "k"))):
|
||||
if ts.cursor > 0 {
|
||||
ts.cursor--
|
||||
}
|
||||
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("down", "j"))):
|
||||
if ts.cursor < len(ts.flatNodes)-1 {
|
||||
ts.cursor++
|
||||
}
|
||||
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("left", "pgup"))):
|
||||
// Page up.
|
||||
ts.cursor -= ts.visibleHeight()
|
||||
if ts.cursor < 0 {
|
||||
ts.cursor = 0
|
||||
}
|
||||
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("right", "pgdown"))):
|
||||
// Page down.
|
||||
ts.cursor += ts.visibleHeight()
|
||||
if ts.cursor >= len(ts.flatNodes) {
|
||||
ts.cursor = len(ts.flatNodes) - 1
|
||||
}
|
||||
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("home"))):
|
||||
ts.cursor = 0
|
||||
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("end"))):
|
||||
ts.cursor = len(ts.flatNodes) - 1
|
||||
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("enter"))):
|
||||
if ts.cursor < len(ts.flatNodes) {
|
||||
ts.selectedID = ts.flatNodes[ts.cursor].ID
|
||||
ts.active = false
|
||||
return ts, func() tea.Msg {
|
||||
return TreeNodeSelectedMsg{
|
||||
ID: ts.selectedID,
|
||||
Entry: ts.flatNodes[ts.cursor].Entry,
|
||||
IsUser: ts.isUserMessage(ts.flatNodes[ts.cursor].Entry),
|
||||
UserText: ts.extractUserText(ts.flatNodes[ts.cursor].Entry),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("esc"))):
|
||||
if ts.search != "" {
|
||||
ts.search = ""
|
||||
ts.rebuildFlatList()
|
||||
} else {
|
||||
ts.cancelled = true
|
||||
ts.active = false
|
||||
return ts, func() tea.Msg {
|
||||
return TreeCancelledMsg{}
|
||||
}
|
||||
}
|
||||
|
||||
// Filter cycle with ctrl+o.
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("ctrl+o"))):
|
||||
ts.filter = (ts.filter + 1) % 5
|
||||
ts.rebuildFlatList()
|
||||
|
||||
// Direct filter shortcuts.
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("ctrl+d"))):
|
||||
ts.filter = TreeFilterDefault
|
||||
ts.rebuildFlatList()
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("ctrl+t"))):
|
||||
ts.filter = TreeFilterNoTools
|
||||
ts.rebuildFlatList()
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("ctrl+u"))):
|
||||
ts.filter = TreeFilterUserOnly
|
||||
ts.rebuildFlatList()
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("ctrl+l"))):
|
||||
ts.filter = TreeFilterLabelOnly
|
||||
ts.rebuildFlatList()
|
||||
case key.Matches(msg, key.NewBinding(key.WithKeys("ctrl+a"))):
|
||||
ts.filter = TreeFilterAll
|
||||
ts.rebuildFlatList()
|
||||
|
||||
default:
|
||||
// Typing search.
|
||||
if msg.Text != "" && len(msg.Text) == 1 {
|
||||
ch := msg.Text[0]
|
||||
if ch >= 32 && ch < 127 {
|
||||
ts.search += string(ch)
|
||||
ts.rebuildFlatList()
|
||||
}
|
||||
}
|
||||
if key.Matches(msg, key.NewBinding(key.WithKeys("backspace"))) && len(ts.search) > 0 {
|
||||
ts.search = ts.search[:len(ts.search)-1]
|
||||
ts.rebuildFlatList()
|
||||
}
|
||||
}
|
||||
}
|
||||
return ts, nil
|
||||
}
|
||||
|
||||
// View implements tea.Model.
|
||||
func (ts *TreeSelectorComponent) View() tea.View {
|
||||
theme := GetTheme()
|
||||
|
||||
headerStyle := lipgloss.NewStyle().
|
||||
Bold(true).
|
||||
Foreground(theme.Accent).
|
||||
PaddingLeft(2)
|
||||
|
||||
helpStyle := lipgloss.NewStyle().
|
||||
Foreground(theme.Muted).
|
||||
PaddingLeft(2)
|
||||
|
||||
var b strings.Builder
|
||||
|
||||
// Header.
|
||||
b.WriteString(headerStyle.Render("Session Tree"))
|
||||
b.WriteString("\n")
|
||||
b.WriteString(helpStyle.Render("↑/↓: move ←/→: page enter: select esc: cancel ^O: cycle filter"))
|
||||
b.WriteString("\n")
|
||||
|
||||
if ts.search != "" {
|
||||
searchStyle := lipgloss.NewStyle().Foreground(theme.Info).PaddingLeft(2)
|
||||
b.WriteString(searchStyle.Render(fmt.Sprintf("Search: %s", ts.search)))
|
||||
b.WriteString("\n")
|
||||
}
|
||||
|
||||
b.WriteString(lipgloss.NewStyle().Foreground(theme.Muted).Render(strings.Repeat("─", ts.width)))
|
||||
b.WriteString("\n")
|
||||
|
||||
if len(ts.flatNodes) == 0 {
|
||||
emptyStyle := lipgloss.NewStyle().Foreground(theme.Muted).PaddingLeft(2)
|
||||
b.WriteString(emptyStyle.Render("No entries in session"))
|
||||
b.WriteString("\n")
|
||||
} else {
|
||||
// Compute visible window.
|
||||
visH := ts.visibleHeight()
|
||||
startIdx := 0
|
||||
if ts.cursor >= visH {
|
||||
startIdx = ts.cursor - visH + 1
|
||||
}
|
||||
endIdx := startIdx + visH
|
||||
if endIdx > len(ts.flatNodes) {
|
||||
endIdx = len(ts.flatNodes)
|
||||
}
|
||||
|
||||
for i := startIdx; i < endIdx; i++ {
|
||||
node := ts.flatNodes[i]
|
||||
line := ts.renderNode(node, i == ts.cursor, node.ID == ts.leafID)
|
||||
b.WriteString(line)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
}
|
||||
|
||||
// Footer.
|
||||
b.WriteString(lipgloss.NewStyle().Foreground(theme.Muted).Render(strings.Repeat("─", ts.width)))
|
||||
b.WriteString("\n")
|
||||
|
||||
footerStyle := lipgloss.NewStyle().Foreground(theme.Muted).PaddingLeft(2)
|
||||
footer := fmt.Sprintf("(%d/%d) [%s]", ts.cursor+1, len(ts.flatNodes), ts.filter)
|
||||
b.WriteString(footerStyle.Render(footer))
|
||||
|
||||
return tea.NewView(b.String())
|
||||
}
|
||||
|
||||
// IsActive returns whether the tree selector is still accepting input.
|
||||
func (ts *TreeSelectorComponent) IsActive() bool {
|
||||
return ts.active
|
||||
}
|
||||
|
||||
// --- Internal helpers ---
|
||||
|
||||
func (ts *TreeSelectorComponent) visibleHeight() int {
|
||||
// Reserve lines for header(3) + search(1) + separator(1) + footer(2).
|
||||
h := ts.height/2 - 7
|
||||
if h < 5 {
|
||||
h = 5
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
func (ts *TreeSelectorComponent) rebuildFlatList() {
|
||||
tree := ts.tm.GetTree()
|
||||
ts.flatNodes = ts.flatNodes[:0]
|
||||
for i, root := range tree {
|
||||
isLast := i == len(tree)-1
|
||||
ts.flattenNode(root, 0, isLast, "")
|
||||
}
|
||||
|
||||
// Apply search filter.
|
||||
if ts.search != "" {
|
||||
query := strings.ToLower(ts.search)
|
||||
filtered := make([]FlatNode, 0)
|
||||
for _, node := range ts.flatNodes {
|
||||
text := ts.entryDisplayText(node.Entry)
|
||||
if strings.Contains(strings.ToLower(text), query) {
|
||||
filtered = append(filtered, node)
|
||||
}
|
||||
}
|
||||
ts.flatNodes = filtered
|
||||
}
|
||||
|
||||
// Clamp cursor.
|
||||
if ts.cursor >= len(ts.flatNodes) {
|
||||
ts.cursor = max(len(ts.flatNodes)-1, 0)
|
||||
}
|
||||
}
|
||||
|
||||
func (ts *TreeSelectorComponent) flattenNode(node *session.TreeNode, depth int, isLast bool, gutterPrefix string) {
|
||||
if !ts.passesFilter(node) {
|
||||
// Still recurse into children in case they pass.
|
||||
for i, child := range node.Children {
|
||||
childIsLast := i == len(node.Children)-1
|
||||
ts.flattenNode(child, depth, childIsLast, gutterPrefix)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
var prefix string
|
||||
if depth == 0 {
|
||||
prefix = ""
|
||||
} else if isLast {
|
||||
prefix = gutterPrefix + "└─ "
|
||||
} else {
|
||||
prefix = gutterPrefix + "├─ "
|
||||
}
|
||||
|
||||
label := ts.tm.GetLabel(node.ID)
|
||||
|
||||
ts.flatNodes = append(ts.flatNodes, FlatNode{
|
||||
Entry: node.Entry,
|
||||
ID: node.ID,
|
||||
ParentID: node.ParentID,
|
||||
Depth: depth,
|
||||
IsLast: isLast,
|
||||
Prefix: prefix,
|
||||
Label: label,
|
||||
})
|
||||
|
||||
// Build gutter prefix for children.
|
||||
var childGutter string
|
||||
if depth == 0 {
|
||||
childGutter = ""
|
||||
} else if isLast {
|
||||
childGutter = gutterPrefix + " "
|
||||
} else {
|
||||
childGutter = gutterPrefix + "│ "
|
||||
}
|
||||
|
||||
for i, child := range node.Children {
|
||||
childIsLast := i == len(node.Children)-1
|
||||
ts.flattenNode(child, depth+1, childIsLast, childGutter)
|
||||
}
|
||||
}
|
||||
|
||||
func (ts *TreeSelectorComponent) passesFilter(node *session.TreeNode) bool {
|
||||
switch ts.filter {
|
||||
case TreeFilterAll:
|
||||
return true
|
||||
|
||||
case TreeFilterDefault:
|
||||
// Hide settings entries.
|
||||
switch node.Entry.(type) {
|
||||
case *session.ModelChangeEntry, *session.LabelEntry, *session.SessionInfoEntry:
|
||||
return false
|
||||
}
|
||||
// Hide tool messages unless they're the leaf.
|
||||
if me, ok := node.Entry.(*session.MessageEntry); ok {
|
||||
if me.Role == "tool" && node.ID != ts.leafID {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
|
||||
case TreeFilterNoTools:
|
||||
if me, ok := node.Entry.(*session.MessageEntry); ok {
|
||||
return me.Role != "tool"
|
||||
}
|
||||
return true
|
||||
|
||||
case TreeFilterUserOnly:
|
||||
if me, ok := node.Entry.(*session.MessageEntry); ok {
|
||||
return me.Role == "user"
|
||||
}
|
||||
return false
|
||||
|
||||
case TreeFilterLabelOnly:
|
||||
return ts.tm.GetLabel(node.ID) != ""
|
||||
|
||||
default:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
func (ts *TreeSelectorComponent) renderNode(node FlatNode, isCursor, isLeaf bool) string {
|
||||
theme := GetTheme()
|
||||
maxWidth := ts.width - 4
|
||||
|
||||
// Cursor indicator.
|
||||
var cursor string
|
||||
if isCursor {
|
||||
cursor = lipgloss.NewStyle().Foreground(theme.Accent).Render("› ")
|
||||
} else {
|
||||
cursor = " "
|
||||
}
|
||||
|
||||
// Role-colored content.
|
||||
text := ts.entryDisplayText(node.Entry)
|
||||
if len(text) > maxWidth-len(node.Prefix)-10 {
|
||||
trimLen := maxWidth - len(node.Prefix) - 13
|
||||
if trimLen > 0 && trimLen < len(text) {
|
||||
text = text[:trimLen] + "..."
|
||||
}
|
||||
}
|
||||
|
||||
var style lipgloss.Style
|
||||
switch e := node.Entry.(type) {
|
||||
case *session.MessageEntry:
|
||||
switch e.Role {
|
||||
case "user":
|
||||
style = lipgloss.NewStyle().Foreground(theme.Accent)
|
||||
case "assistant":
|
||||
style = lipgloss.NewStyle().Foreground(theme.Success)
|
||||
default:
|
||||
style = lipgloss.NewStyle().Foreground(theme.Muted)
|
||||
}
|
||||
case *session.BranchSummaryEntry:
|
||||
style = lipgloss.NewStyle().Foreground(theme.Warning).Italic(true)
|
||||
default:
|
||||
style = lipgloss.NewStyle().Foreground(theme.Muted)
|
||||
}
|
||||
|
||||
if isCursor {
|
||||
style = style.Bold(true)
|
||||
}
|
||||
|
||||
content := style.Render(text)
|
||||
|
||||
// Label badge.
|
||||
var labelBadge string
|
||||
if node.Label != "" {
|
||||
labelBadge = " " + lipgloss.NewStyle().Foreground(theme.Warning).Render("["+node.Label+"]")
|
||||
}
|
||||
|
||||
// Active marker.
|
||||
var activeMarker string
|
||||
if isLeaf {
|
||||
activeMarker = lipgloss.NewStyle().Foreground(theme.Accent).Bold(true).Render(" ← active")
|
||||
}
|
||||
|
||||
// Prefix (tree lines).
|
||||
prefixStyle := lipgloss.NewStyle().Foreground(theme.Muted)
|
||||
renderedPrefix := prefixStyle.Render(node.Prefix)
|
||||
|
||||
return cursor + renderedPrefix + content + labelBadge + activeMarker
|
||||
}
|
||||
|
||||
func (ts *TreeSelectorComponent) entryDisplayText(entry any) string {
|
||||
switch e := entry.(type) {
|
||||
case *session.MessageEntry:
|
||||
role := e.Role
|
||||
text := extractTextFromParts(e.Parts)
|
||||
if len(text) > 80 {
|
||||
text = text[:80] + "..."
|
||||
}
|
||||
if text == "" {
|
||||
// Tool call messages may not have text.
|
||||
text = "(tool interaction)"
|
||||
}
|
||||
return fmt.Sprintf("%s: %s", role, text)
|
||||
|
||||
case *session.ModelChangeEntry:
|
||||
return fmt.Sprintf("model: %s/%s", e.Provider, e.ModelID)
|
||||
|
||||
case *session.BranchSummaryEntry:
|
||||
summary := e.Summary
|
||||
if len(summary) > 60 {
|
||||
summary = summary[:60] + "..."
|
||||
}
|
||||
return fmt.Sprintf("branch summary: %s", summary)
|
||||
|
||||
case *session.LabelEntry:
|
||||
return fmt.Sprintf("label: %s", e.Label)
|
||||
|
||||
case *session.SessionInfoEntry:
|
||||
return fmt.Sprintf("name: %s", e.Name)
|
||||
|
||||
default:
|
||||
return "(unknown entry)"
|
||||
}
|
||||
}
|
||||
|
||||
func (ts *TreeSelectorComponent) isUserMessage(entry any) bool {
|
||||
if me, ok := entry.(*session.MessageEntry); ok {
|
||||
return me.Role == "user"
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (ts *TreeSelectorComponent) extractUserText(entry any) string {
|
||||
if me, ok := entry.(*session.MessageEntry); ok && me.Role == "user" {
|
||||
return extractTextFromParts(me.Parts)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// extractTextFromParts extracts text content from type-tagged parts JSON.
|
||||
func extractTextFromParts(partsJSON []byte) string {
|
||||
// Quick extraction without full unmarshal.
|
||||
var parts []struct {
|
||||
Type string `json:"type"`
|
||||
Data struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(partsJSON, &parts); err != nil {
|
||||
return ""
|
||||
}
|
||||
var texts []string
|
||||
for _, p := range parts {
|
||||
if p.Type == "text" && p.Data.Text != "" {
|
||||
texts = append(texts, p.Data.Text)
|
||||
}
|
||||
}
|
||||
return strings.Join(texts, "\n")
|
||||
}
|
||||
Reference in New Issue
Block a user