From dd018b65ec5735e70b4bb66700ad35cc2f8695c1 Mon Sep 17 00:00:00 2001 From: Ed Zynda Date: Thu, 26 Feb 2026 18:47:10 +0300 Subject: [PATCH] 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 --- cmd/root.go | 179 +++++--- internal/app/app.go | 42 +- internal/app/options.go | 10 +- internal/session/entry.go | 265 ++++++++++++ internal/session/store.go | 247 +++++++++++ internal/session/tree_manager.go | 689 +++++++++++++++++++++++++++++++ internal/ui/commands.go | 28 ++ internal/ui/events.go | 17 + internal/ui/model.go | 195 ++++++++- internal/ui/model_test.go | 5 + internal/ui/tree_selector.go | 528 +++++++++++++++++++++++ 11 files changed, 2131 insertions(+), 74 deletions(-) create mode 100644 internal/session/entry.go create mode 100644 internal/session/store.go create mode 100644 internal/session/tree_manager.go create mode 100644 internal/ui/tree_selector.go diff --git a/cmd/root.go b/cmd/root.go index fb9c3aee..5f2e80ee 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -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 diff --git a/internal/app/app.go b/internal/app/app.go index 06f4de37..09f493c0 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -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 } diff --git a/internal/app/options.go b/internal/app/options.go index 19b60c18..d326dfc7 100644 --- a/internal/app/options.go +++ b/internal/app/options.go @@ -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 diff --git a/internal/session/entry.go b/internal/session/entry.go new file mode 100644 index 00000000..bbaa3f7f --- /dev/null +++ b/internal/session/entry.go @@ -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 +} diff --git a/internal/session/store.go b/internal/session/store.go new file mode 100644 index 00000000..310a8f40 --- /dev/null +++ b/internal/session/store.go @@ -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) +} diff --git a/internal/session/tree_manager.go b/internal/session/tree_manager.go new file mode 100644 index 00000000..d07c751b --- /dev/null +++ b/internal/session/tree_manager.go @@ -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/----/ +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) +} diff --git a/internal/ui/commands.go b/internal/ui/commands.go index bffaf4cd..986dbacf 100644 --- a/internal/ui/commands.go +++ b/internal/ui/commands.go @@ -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 diff --git a/internal/ui/events.go b/internal/ui/events.go index 42ff6ad1..db3d6f25 100644 --- a/internal/ui/events.go +++ b/internal/ui/events.go @@ -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{} diff --git a/internal/ui/model.go b/internal/ui/model.go index 3da4c6c3..3e9a6455 100644 --- a/internal/ui/model.go +++ b/internal/ui/model.go @@ -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 ` (not yet implemented — use the session file directly).", currentName)) + } + return m.printSystemMessage("To name this session, use: `/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 // -------------------------------------------------------------------------- diff --git a/internal/ui/model_test.go b/internal/ui/model_test.go index bc8a645a..e4d3d14b 100644 --- a/internal/ui/model_test.go +++ b/internal/ui/model_test.go @@ -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 // -------------------------------------------------------------------------- diff --git a/internal/ui/tree_selector.go b/internal/ui/tree_selector.go new file mode 100644 index 00000000..b905d399 --- /dev/null +++ b/internal/ui/tree_selector.go @@ -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") +}