feat(pkg/kit): add SessionManager interface for custom session backends

Add SessionManager interface to allow pluggable session storage backends.
This enables users to implement custom session managers for databases,
cloud storage, or other persistence mechanisms instead of the default
JSONL file-based TreeManager.

Changes:
- Add SessionManager interface with methods for message storage,
  tree navigation, compaction, and extension data
- Add treeManagerAdapter to wrap existing TreeManager for backward compatibility
- Update Kit struct to use SessionManager interface instead of concrete type
- Add SessionManager option to Options struct
- Update all session-related methods to use interface
- Add documentation for custom SessionManager usage

The default behavior is preserved - when no SessionManager is provided,
Kit automatically uses the TreeManager via the adapter.

pkg/kit/adapter.go

@@ -0,0 +1,231 @@
+package kit
+
+import (
+	"strings"
+	"time"
+
+	"charm.land/fantasy"
+	"github.com/mark3labs/kit/internal/session"
+)
+
+// treeManagerAdapter adapts TreeManager to SessionManager interface.
+// This is unexported - users don't interact with it directly.
+type treeManagerAdapter struct {
+	inner *session.TreeManager
+}
+
+// NewTreeManagerAdapter creates an adapter (exported for use in New function).
+// This is used by the SDK when no custom SessionManager is provided.
+func NewTreeManagerAdapter(tm *session.TreeManager) SessionManager {
+	return &treeManagerAdapter{inner: tm}
+}
+
+// AppendMessage implements SessionManager.
+func (a *treeManagerAdapter) AppendMessage(msg LLMMessage) (string, error) {
+	// LLMMessage is just an alias for fantasy.Message, so no conversion needed
+	return a.inner.AppendLLMMessage(msg)
+}
+
+// GetMessages implements SessionManager.
+func (a *treeManagerAdapter) GetMessages() []LLMMessage {
+	// LLMMessage is just an alias for fantasy.Message
+	return a.inner.GetLLMMessages()
+}
+
+// BuildContext implements SessionManager.
+func (a *treeManagerAdapter) BuildContext() ([]LLMMessage, string, string) {
+	msgs, provider, modelID := a.inner.BuildContext()
+	return msgs, provider, modelID
+}
+
+// Branch implements SessionManager.
+func (a *treeManagerAdapter) Branch(entryID string) error {
+	return a.inner.Branch(entryID)
+}
+
+// GetCurrentBranch implements SessionManager.
+func (a *treeManagerAdapter) GetCurrentBranch() []BranchEntry {
+	branch := a.inner.GetBranch("")
+	var result []BranchEntry
+	for _, entry := range branch {
+		be := a.convertEntry(entry)
+		if be != nil {
+			result = append(result, *be)
+		}
+	}
+	return result
+}
+
+// GetChildren implements SessionManager.
+func (a *treeManagerAdapter) GetChildren(parentID string) []string {
+	return a.inner.GetChildren(parentID)
+}
+
+// GetEntry implements SessionManager.
+func (a *treeManagerAdapter) GetEntry(entryID string) *BranchEntry {
+	entry := a.inner.GetEntry(entryID)
+	if entry == nil {
+		return nil
+	}
+	return a.convertEntry(entry)
+}
+
+// GetSessionID implements SessionManager.
+func (a *treeManagerAdapter) GetSessionID() string {
+	return a.inner.GetSessionID()
+}
+
+// GetSessionName implements SessionManager.
+func (a *treeManagerAdapter) GetSessionName() string {
+	return a.inner.GetSessionName()
+}
+
+// SetSessionName implements SessionManager.
+func (a *treeManagerAdapter) SetSessionName(name string) error {
+	_, err := a.inner.AppendSessionInfo(name)
+	return err
+}
+
+// GetCreatedAt implements SessionManager.
+func (a *treeManagerAdapter) GetCreatedAt() time.Time {
+	return a.inner.GetHeader().Timestamp
+}
+
+// IsPersisted implements SessionManager.
+func (a *treeManagerAdapter) IsPersisted() bool {
+	return a.inner.IsPersisted()
+}
+
+// AppendCompaction implements SessionManager.
+func (a *treeManagerAdapter) AppendCompaction(summary string, firstKeptEntryID string,
+	tokensBefore, tokensAfter int, messagesRemoved int, readFiles, modifiedFiles []string) (string, error) {
+
+	return a.inner.AppendCompaction(summary, firstKeptEntryID,
+		tokensBefore, tokensAfter, messagesRemoved, readFiles, modifiedFiles)
+}
+
+// GetLastCompaction implements SessionManager.
+func (a *treeManagerAdapter) GetLastCompaction() *CompactionEntry {
+	c := a.inner.GetLastCompaction()
+	if c == nil {
+		return nil
+	}
+	return &CompactionEntry{
+		ID:               c.ID,
+		Summary:          c.Summary,
+		FirstKeptEntryID: c.FirstKeptEntryID,
+		TokensBefore:     c.TokensBefore,
+		TokensAfter:      c.TokensAfter,
+		MessagesRemoved:  c.MessagesRemoved,
+		ReadFiles:        c.ReadFiles,
+		ModifiedFiles:    c.ModifiedFiles,
+		Timestamp:        c.Timestamp,
+	}
+}
+
+// AppendExtensionData implements SessionManager.
+func (a *treeManagerAdapter) AppendExtensionData(extType, data string) (string, error) {
+	return a.inner.AppendExtensionData(extType, data)
+}
+
+// GetExtensionData implements SessionManager.
+func (a *treeManagerAdapter) GetExtensionData(extType string) []ExtensionDataEntry {
+	entries := a.inner.GetExtensionData(extType)
+	var result []ExtensionDataEntry
+	for _, e := range entries {
+		result = append(result, ExtensionDataEntry{
+			ID:        e.ID,
+			ExtType:   e.ExtType,
+			Data:      e.Data,
+			Timestamp: e.Timestamp,
+		})
+	}
+	return result
+}
+
+// AppendModelChange implements SessionManager.
+func (a *treeManagerAdapter) AppendModelChange(provider, modelID string) (string, error) {
+	return a.inner.AppendModelChange(provider, modelID)
+}
+
+// GetContextEntryIDs implements SessionManager.
+func (a *treeManagerAdapter) GetContextEntryIDs() []string {
+	return a.inner.GetContextEntryIDs()
+}
+
+// Close implements SessionManager.
+func (a *treeManagerAdapter) Close() error {
+	return a.inner.Close()
+}
+
+// Helper: Convert internal entry types to BranchEntry
+func (a *treeManagerAdapter) convertEntry(entry any) *BranchEntry {
+	switch e := entry.(type) {
+	case *session.MessageEntry:
+		msg, err := e.ToMessage()
+		if err != nil {
+			return nil
+		}
+		// Build content text from parts
+		var content strings.Builder
+		for _, part := range msg.Parts {
+			if textPart, ok := part.(TextContent); ok {
+				content.WriteString(textPart.Text)
+			}
+		}
+		return &BranchEntry{
+			ID:        e.ID,
+			ParentID:  e.ParentID,
+			Type:      EntryTypeMessage,
+			Role:      string(msg.Role),
+			Content:   content.String(),
+			Model:     e.Model,
+			Provider:  e.Provider,
+			Timestamp: e.Timestamp,
+			RawParts:  msg.Parts,
+		}
+	case *session.BranchSummaryEntry:
+		return &BranchEntry{
+			ID:        e.ID,
+			ParentID:  e.ParentID,
+			Type:      EntryTypeBranchSummary,
+			Content:   e.Summary,
+			Timestamp: e.Timestamp,
+		}
+	case *session.ModelChangeEntry:
+		return &BranchEntry{
+			ID:        e.ID,
+			ParentID:  e.ParentID,
+			Type:      EntryTypeModelChange,
+			Content:   "Model changed to " + e.Provider + "/" + e.ModelID,
+			Model:     e.ModelID,
+			Provider:  e.Provider,
+			Timestamp: e.Timestamp,
+		}
+	case *session.CompactionEntry:
+		return &BranchEntry{
+			ID:        e.ID,
+			ParentID:  e.ParentID,
+			Type:      EntryTypeCompaction,
+			Content:   e.Summary,
+			Timestamp: e.Timestamp,
+		}
+	case *session.ExtensionDataEntry:
+		return &BranchEntry{
+			ID:        e.ID,
+			ParentID:  e.ParentID,
+			Type:      EntryTypeExtensionData,
+			Content:   "Extension data: " + e.ExtType,
+			Timestamp: e.Timestamp,
+		}
+	default:
+		return nil
+	}
+}
+
+// convertKitMessagesToFantasy converts kit LLM messages to fantasy messages.
+// Since LLMMessage is an alias for fantasy.Message, this is a no-op.
+func convertKitMessagesToFantasy(msgs []LLMMessage) []fantasy.Message {
+	// LLMMessage is just an alias for fantasy.Message, so we can type convert
+	return msgs
+}

pkg/kit/compaction.go

@@ -21,9 +21,9 @@ type ContextStats struct {
 const defaultReserveTokens = 16384
 
 // EstimateContextTokens returns the estimated token count of the current
-// conversation based on tree session messages.
+// conversation based on session messages.
 func (m *Kit) EstimateContextTokens() int {
-	messages := m.treeSession.GetLLMMessages()
+	messages := m.session.GetMessages()
 	return compaction.EstimateMessageTokens(messages)
 }
 
@@ -42,8 +42,8 @@ func (m *Kit) ShouldCompact() bool {
 		reserveTokens = m.compactionOpts.ReserveTokens
 	}
 
-	messages := m.treeSession.GetLLMMessages()
-	return compaction.ShouldCompact(messages, info.Limit.Context, reserveTokens)
+	messages := m.session.GetMessages()
+	return compaction.ShouldCompact(convertKitMessagesToFantasy(messages), info.Limit.Context, reserveTokens)
 }
 
 // GetContextStats returns current context usage statistics including
@@ -55,7 +55,7 @@ func (m *Kit) ShouldCompact() bool {
 // because it includes system prompts, tool definitions, and other overhead
 // that the heuristic cannot account for.
 func (m *Kit) GetContextStats() ContextStats {
-	messages := m.treeSession.GetLLMMessages()
+	messages := m.session.GetMessages()
 
 	// Prefer the real API-reported input token count when available.
 	m.lastInputTokensMu.RLock()
@@ -114,7 +114,7 @@ func (m *Kit) compactInternal(ctx context.Context, opts *CompactionOptions, cust
 		}
 	}
 
-	messages := m.treeSession.GetLLMMessages()
+	messages := m.session.GetMessages()
 	if len(messages) < 2 {
 		return nil, fmt.Errorf("cannot compact: need at least 2 messages")
 	}
@@ -145,7 +145,7 @@ func (m *Kit) compactInternal(ctx context.Context, opts *CompactionOptions, cust
 
 	// Carry forward file tracking from previous compaction.
 	var prev *compaction.PreviousCompaction
-	if lastCompaction := m.treeSession.GetLastCompaction(); lastCompaction != nil {
+	if lastCompaction := m.session.GetLastCompaction(); lastCompaction != nil {
 		prev = &compaction.PreviousCompaction{
 			ReadFiles:     lastCompaction.ReadFiles,
 			ModifiedFiles: lastCompaction.ModifiedFiles,
@@ -171,7 +171,7 @@ func (m *Kit) compactInternal(ctx context.Context, opts *CompactionOptions, cust
 
 	// Non-destructive: append a CompactionEntry to the session tree instead
 	// of clearing and rewriting messages.
-	entryIDs := m.treeSession.GetContextEntryIDs()
+	entryIDs := m.session.GetContextEntryIDs()
 	firstKeptEntryID := ""
 	if result.CutPoint >= 0 && result.CutPoint < len(entryIDs) {
 		firstKeptEntryID = entryIDs[result.CutPoint]
@@ -188,9 +188,9 @@ func (m *Kit) compactInternal(ctx context.Context, opts *CompactionOptions, cust
 // custom summary. It still determines the cut point and persists a
 // CompactionEntry.
 func (m *Kit) applyCustomCompaction(summary string, messages []LLMMessage, opts *CompactionOptions) (*CompactionResult, error) {
-	originalTokens := compaction.EstimateMessageTokens(messages)
+	originalTokens := compaction.EstimateMessageTokens(convertKitMessagesToFantasy(messages))
 
-	cutPoint := compaction.FindCutPoint(messages, opts.KeepRecentTokens)
+	cutPoint := compaction.FindCutPoint(convertKitMessagesToFantasy(messages), opts.KeepRecentTokens)
 	if cutPoint == 0 {
 		cutPoint = len(messages) - 1
 		if cutPoint < 1 {
@@ -198,7 +198,7 @@ func (m *Kit) applyCustomCompaction(summary string, messages []LLMMessage, opts
 		}
 	}
 
-	entryIDs := m.treeSession.GetContextEntryIDs()
+	entryIDs := m.session.GetContextEntryIDs()
 	firstKeptEntryID := ""
 	if cutPoint >= 0 && cutPoint < len(entryIDs) {
 		firstKeptEntryID = entryIDs[cutPoint]
@@ -234,7 +234,7 @@ func (m *Kit) persistAndEmitCompaction(
 	originalTokens, compactedTokens, messagesRemoved int,
 	readFiles, modifiedFiles []string,
 ) error {
-	if _, err := m.treeSession.AppendCompaction(
+	if _, err := m.session.AppendCompaction(
 		summary,
 		firstKeptEntryID,
 		originalTokens,

pkg/kit/extension_api.go

@@ -227,28 +227,51 @@ func (e *extensionAPI) GetMessageRenderer(name string) *extensions.MessageRender
 // Session data
 
 func (e *extensionAPI) GetSessionMessages() []extensions.SessionMessage {
-	return iterBranchMessages(e.kit.treeSession, func(me *session.MessageEntry, msg message.Message) extensions.SessionMessage {
-		return extensions.SessionMessage{
-			ID:        me.ID,
-			Role:      string(msg.Role),
-			Content:   msg.Content(),
-			Timestamp: me.Timestamp.Format("2006-01-02T15:04:05Z07:00"),
+	if e.kit.session == nil {
+		return nil
+	}
+
+	// Try to use the legacy iterBranchMessages for backward compatibility
+	// with the default TreeManager adapter
+	if adapter, ok := e.kit.session.(*treeManagerAdapter); ok {
+		return iterBranchMessages(adapter.inner, func(me *session.MessageEntry, msg message.Message) extensions.SessionMessage {
+			return extensions.SessionMessage{
+				ID:        me.ID,
+				Role:      string(msg.Role),
+				Content:   msg.Content(),
+				Timestamp: me.Timestamp.Format("2006-01-02T15:04:05Z07:00"),
+			}
+		})
+	}
+
+	// For custom SessionManagers, use the public interface
+	branch := e.kit.session.GetCurrentBranch()
+	var result []extensions.SessionMessage
+	for _, entry := range branch {
+		if entry.Type == EntryTypeMessage {
+			result = append(result, extensions.SessionMessage{
+				ID:        entry.ID,
+				Role:      entry.Role,
+				Content:   entry.Content,
+				Timestamp: entry.Timestamp.Format("2006-01-02T15:04:05Z07:00"),
+			})
 		}
-	})
+	}
+	return result
 }
 
 func (e *extensionAPI) AppendEntry(extType, data string) (string, error) {
-	if e.kit.treeSession == nil {
+	if e.kit.session == nil {
 		return "", fmt.Errorf("no session available")
 	}
-	return e.kit.treeSession.AppendExtensionData(extType, data)
+	return e.kit.session.AppendExtensionData(extType, data)
 }
 
 func (e *extensionAPI) GetEntries(extType string) []extensions.ExtensionEntry {
-	if e.kit.treeSession == nil {
+	if e.kit.session == nil {
 		return nil
 	}
-	entries := e.kit.treeSession.GetExtensionData(extType)
+	entries := e.kit.session.GetExtensionData(extType)
 	result := make([]extensions.ExtensionEntry, 0, len(entries))
 	for _, e := range entries {
 		result = append(result, extensions.ExtensionEntry{

pkg/kit/kit.go

@@ -39,7 +39,7 @@ type ContextFile struct {
 // agents, sessions, and model configurations.
 type Kit struct {
 	agent          *agent.Agent
-	treeSession    *session.TreeManager
+	session        SessionManager
 	modelString    string
 	events         *eventBus
 	autoCompact    bool
@@ -172,27 +172,39 @@ type StructuredMessage struct {
 // flattens all content to a single text string, this preserves tool calls,
 // tool results, reasoning blocks, and finish markers as distinct typed parts.
 func (m *Kit) GetStructuredMessages() []StructuredMessage {
-	return iterBranchMessages(m.treeSession, func(me *session.MessageEntry, msg message.Message) StructuredMessage {
-		return StructuredMessage{
-			ID:        me.ID,
-			ParentID:  me.ParentID,
-			Role:      msg.Role,
-			Parts:     msg.Parts,
-			Model:     msg.Model,
-			Provider:  msg.Provider,
-			Timestamp: me.Timestamp.Format("2006-01-02T15:04:05Z07:00"),
+	if m.session == nil {
+		return nil
+	}
+
+	branch := m.session.GetCurrentBranch()
+	var results []StructuredMessage
+	for _, entry := range branch {
+		if entry.Type != EntryTypeMessage {
+			continue
 		}
-	})
+		results = append(results, StructuredMessage{
+			ID:        entry.ID,
+			ParentID:  entry.ParentID,
+			Role:      MessageRole(entry.Role),
+			Parts:     entry.RawParts,
+			Model:     entry.Model,
+			Provider:  entry.Provider,
+			Timestamp: entry.Timestamp.Format("2006-01-02T15:04:05Z07:00"),
+		})
+	}
+	return results
 }
 
 // iterBranchMessages iterates over the current branch's MessageEntry items,
 // converting each to a message.Message and calling fn to build the result.
-// Returns nil if there is no tree session. Skips entries that are not
+// Returns nil if there is no session. Skips entries that are not
 // MessageEntry or that fail conversion.
+// Deprecated: Use SessionManager.GetCurrentBranch() directly.
 func iterBranchMessages[T any](tm *session.TreeManager, fn func(*session.MessageEntry, message.Message) T) []T {
 	if tm == nil {
 		return nil
 	}
+
 	branch := tm.GetBranch("")
 	var results []T
 	for _, entry := range branch {
@@ -494,6 +506,11 @@ type Options struct {
 
 	// CLI is optional CLI-specific configuration. SDK users leave this nil.
 	CLI *CLIOptions
+
+	// SessionManager allows custom session storage backends.
+	// If nil (default), Kit uses the built-in file-based TreeManager.
+	// When provided, SessionPath, Continue, and NoSession options are ignored.
+	SessionManager SessionManager
 }
 
 // CLIOptions holds fields only relevant to the CLI binary. SDK users should
@@ -740,16 +757,25 @@ func New(ctx context.Context, opts *Options) (*Kit, error) {
 		return nil, err
 	}
 
-	// Initialize tree session.
-	treeSession, err := InitTreeSession(opts)
-	if err != nil {
-		_ = agentResult.Agent.Close()
-		return nil, fmt.Errorf("failed to initialize session: %w", err)
+	// Initialize session manager.
+	var sessionManager SessionManager
+	if opts.SessionManager != nil {
+		// Use custom session manager provided by user.
+		sessionManager = opts.SessionManager
+	} else {
+		// DEFAULT: Use built-in TreeManager (existing behavior).
+		treeSession, err := InitTreeSession(opts)
+		if err != nil {
+			_ = agentResult.Agent.Close()
+			return nil, fmt.Errorf("failed to initialize session: %w", err)
+		}
+		// Wrap TreeManager in adapter to satisfy SessionManager interface.
+		sessionManager = NewTreeManagerAdapter(treeSession)
 	}
 
 	k := &Kit{
 		agent:           agentResult.Agent,
-		treeSession:     treeSession,
+		session:         sessionManager,
 		modelString:     modelString,
 		events:          newEventBus(),
 		autoCompact:     opts.AutoCompact,
@@ -1357,9 +1383,9 @@ func (m *Kit) runTurn(ctx context.Context, promptLabel string, prompt string, pr
 		}
 	}
 
-	// Persist pre-generation messages to tree session.
+	// Persist pre-generation messages to session.
 	for _, msg := range preMessages {
-		_, _ = m.treeSession.AppendLLMMessage(msg)
+		_, _ = m.session.AppendMessage(msg)
 	}
 
 	// Auto-compact if enabled and conversation is near the context limit.
@@ -1367,8 +1393,8 @@ func (m *Kit) runTurn(ctx context.Context, promptLabel string, prompt string, pr
 		_, _ = m.compactInternal(ctx, m.compactionOpts, "", true) // best-effort, automatic
 	}
 
-	// Build context from the tree so only the current branch is sent.
-	messages := m.treeSession.GetLLMMessages()
+	// Build context from the session so only the current branch is sent.
+	messages, _, _ := m.session.BuildContext()
 
 	// Run ContextPrepare hooks — extensions can filter, reorder, or inject messages.
 	if hookResult := m.contextPrepare.run(ContextPrepareHook{Messages: messages}); hookResult != nil && hookResult.Messages != nil {
@@ -1391,7 +1417,7 @@ func (m *Kit) runTurn(ctx context.Context, promptLabel string, prompt string, pr
 		// (pending) message or tool call is discarded.
 		if result != nil && len(result.ConversationMessages) > sentCount {
 			for _, msg := range result.ConversationMessages[sentCount:] {
-				_, _ = m.treeSession.AppendLLMMessage(msg)
+				_, _ = m.session.AppendMessage(msg)
 			}
 		}
 		m.events.emit(TurnEndEvent{Error: err})
@@ -1407,7 +1433,7 @@ func (m *Kit) runTurn(ctx context.Context, promptLabel string, prompt string, pr
 	// GetContextStats() see up-to-date token counts.
 	if len(result.ConversationMessages) > sentCount {
 		for _, msg := range result.ConversationMessages[sentCount:] {
-			_, _ = m.treeSession.AppendLLMMessage(msg)
+			_, _ = m.session.AppendMessage(msg)
 		}
 	}
 
@@ -1489,7 +1515,7 @@ func (m *Kit) Steer(ctx context.Context, instruction string) (string, error) {
 // Returns an error if there are no previous messages in the session.
 func (m *Kit) FollowUp(ctx context.Context, text string) (string, error) {
 	// Verify there is conversation history to follow up on.
-	if len(m.treeSession.GetLLMMessages()) == 0 {
+	if len(m.session.GetMessages()) == 0 {
 		return "", fmt.Errorf("cannot follow up: no previous messages")
 	}
 
@@ -1645,10 +1671,12 @@ func (m *Kit) PromptResultWithMessages(ctx context.Context, messages []string) (
 	return m.runTurn(ctx, promptLabel, messages[len(messages)-1], preMessages)
 }
 
-// ClearSession resets the tree session's leaf pointer to the root, starting
+// ClearSession resets the session's leaf pointer to the root, starting
 // a fresh conversation branch.
 func (m *Kit) ClearSession() {
-	m.treeSession.ResetLeaf()
+	if m.session != nil {
+		_ = m.session.Branch("")
+	}
 }
 
 // GetModelString returns the current model string identifier (e.g.,
@@ -1717,8 +1745,8 @@ func (m *Kit) Close() error {
 	if m.extRunner != nil && m.extRunner.HasHandlers(extensions.SessionShutdown) {
 		_, _ = m.extRunner.Emit(extensions.SessionShutdownEvent{})
 	}
-	if m.treeSession != nil {
-		_ = m.treeSession.Close()
+	if m.session != nil {
+		_ = m.session.Close()
 	}
 	// Release the OAuth callback port if we own the handler.
 	if closer, ok := m.authHandler.(interface{ Close() error }); ok {
@@ -1726,3 +1754,5 @@ func (m *Kit) Close() error {
 	}
 	return m.agent.Close()
 }
+
+// Conversion helpers are defined in adapter.go.

pkg/kit/session.go

@@ -0,0 +1,132 @@
+package kit
+
+import (
+	"time"
+)
+
+// SessionManager defines the contract for conversation storage backends.
+// Implementations can use files (default), databases, cloud storage, etc.
+type SessionManager interface {
+	// AppendMessage adds a message to the current branch and returns its entry ID.
+	// The entry ID is used for tree navigation and must be unique within the session.
+	AppendMessage(msg LLMMessage) (entryID string, err error)
+
+	// GetMessages returns all messages on the current branch (from root to leaf),
+	// including any compaction summaries at the appropriate positions.
+	GetMessages() []LLMMessage
+
+	// BuildContext returns the message history to send to the LLM, applying
+	// compaction rules and branch summaries as needed.
+	// Returns: messages, currentProvider, currentModelID
+	BuildContext() (messages []LLMMessage, provider string, modelID string)
+
+	// Branch moves the leaf pointer to the given entry ID, creating a branch point.
+	// Subsequent AppendMessage calls extend from this new position.
+	// entryID can be empty to reset to root (new conversation branch).
+	Branch(entryID string) error
+
+	// GetCurrentBranch returns the path from root to current leaf as entry metadata.
+	// Used for UI display and navigation.
+	GetCurrentBranch() []BranchEntry
+
+	// GetChildren returns direct child entry IDs for a given parent entry.
+	// Used to display branch points in the conversation tree.
+	GetChildren(parentID string) []string
+
+	// GetEntry returns a specific entry by ID, or nil if not found.
+	GetEntry(entryID string) *BranchEntry
+
+	// GetSessionID returns the unique session identifier (UUID).
+	GetSessionID() string
+
+	// GetSessionName returns the user-defined display name, or empty.
+	GetSessionName() string
+
+	// SetSessionName sets a display name for the session.
+	SetSessionName(name string) error
+
+	// GetCreatedAt returns when the session was created.
+	GetCreatedAt() time.Time
+
+	// IsPersisted returns true if this session writes to durable storage.
+	IsPersisted() bool
+
+	// AppendCompaction adds a compaction entry that summarizes older messages.
+	// firstKeptEntryID is the ID of the first message to preserve in context.
+	// readFiles and modifiedFiles track file changes for the compaction summary.
+	AppendCompaction(summary string, firstKeptEntryID string,
+		tokensBefore, tokensAfter int, messagesRemoved int, readFiles, modifiedFiles []string) (string, error)
+
+	// GetLastCompaction returns the most recent compaction entry on the current
+	// branch, or nil if none exists.
+	GetLastCompaction() *CompactionEntry
+
+	// AppendExtensionData stores custom extension data in the session tree.
+	// Extensions use this to persist state across restarts.
+	AppendExtensionData(extType, data string) (string, error)
+
+	// GetExtensionData returns all extension data entries of the given type
+	// on the current branch. If extType is empty, returns all extension data.
+	GetExtensionData(extType string) []ExtensionDataEntry
+
+	// AppendModelChange records a provider/model switch in the session.
+	AppendModelChange(provider, modelID string) (string, error)
+
+	// GetContextEntryIDs returns the entry IDs corresponding to the messages
+	// returned by BuildContext, in the same order. Used by compaction to
+	// determine which entries to summarize.
+	GetContextEntryIDs() []string
+
+	// Close releases resources (database connections, file handles, etc.).
+	Close() error
+}
+
+// BranchEntry represents a single node in the conversation tree.
+// This is a SDK-friendly struct (not the internal entry types).
+type BranchEntry struct {
+	ID        string
+	ParentID  string
+	Type      EntryType // "message", "branch_summary", "model_change", "compaction", "extension_data"
+	Role      string    // for messages: "user", "assistant", "system", "tool"
+	Content   string    // text content or summary
+	Model     string    // model used (for messages and model_change)
+	Provider  string    // provider used
+	Timestamp time.Time
+	Children  []string // child entry IDs (for tree display)
+
+	// RawParts contains the full typed content parts for structured access.
+	// Only populated for message entries.
+	RawParts []ContentPart
+}
+
+// EntryType identifies the kind of entry in the session tree.
+type EntryType string
+
+const (
+	EntryTypeMessage       EntryType = "message"
+	EntryTypeBranchSummary EntryType = "branch_summary"
+	EntryTypeModelChange   EntryType = "model_change"
+	EntryTypeCompaction    EntryType = "compaction"
+	EntryTypeExtensionData EntryType = "extension_data"
+)
+
+// CompactionEntry represents a context compaction/summarization event.
+type CompactionEntry struct {
+	ID               string
+	Summary          string
+	FirstKeptEntryID string
+	TokensBefore     int
+	TokensAfter      int
+	MessagesRemoved  int
+	ReadFiles        []string
+	ModifiedFiles    []string
+	Timestamp        time.Time
+}
+
+// ExtensionDataEntry represents custom extension data stored in the session.
+type ExtensionDataEntry struct {
+	ID        string
+	ExtType   string
+	Data      string
+	Timestamp time.Time
+}

pkg/kit/sessions.go

@@ -8,7 +8,6 @@ import (
 	"time"
 
 	"github.com/mark3labs/kit/internal/extensions"
-	"github.com/mark3labs/kit/internal/message"
 	"github.com/mark3labs/kit/internal/session"
 )
 
@@ -47,49 +46,73 @@ func OpenTreeSession(path string) (*TreeManager, error) {
 
 // --- Instance methods on Kit ---
 
+// GetSessionManager returns the session manager, or nil if not configured.
+func (m *Kit) GetSessionManager() SessionManager {
+	return m.session
+}
+
 // GetTreeSession returns the tree session manager, or nil if not configured.
+// Deprecated: Use GetSessionManager instead.
 func (m *Kit) GetTreeSession() *TreeManager {
-	return m.treeSession
+	// Try to unwrap the adapter if using default implementation
+	if adapter, ok := m.session.(*treeManagerAdapter); ok {
+		return adapter.inner
+	}
+	return nil
+}
+
+// SetSessionManager replaces the session manager on a Kit instance.
+func (m *Kit) SetSessionManager(sm SessionManager) {
+	m.session = sm
 }
 
 // SetTreeSession replaces the tree session on a Kit instance. This is used by
 // the CLI when it handles session creation externally (e.g. --resume with a
 // TUI picker) and needs to inject the result into a Kit-like workflow.
+// Deprecated: Use SetSessionManager instead.
 func (m *Kit) SetTreeSession(ts *TreeManager) {
-	m.treeSession = ts
+	m.session = NewTreeManagerAdapter(ts)
 }
 
-// GetSessionPath returns the file path of the active tree session, or empty
-// for in-memory sessions or when no tree session is configured.
+// GetSessionPath returns the file path of the active session, or empty
+// for in-memory sessions or when no file-based session is configured.
 func (m *Kit) GetSessionPath() string {
-	if m.treeSession != nil {
-		return m.treeSession.GetFilePath()
+	// Only file-based sessions have a path
+	// Try to get it from the underlying TreeManager if using default adapter
+	if m.session == nil {
+		return ""
+	}
+	// Check if it's the default adapter
+	if adapter, ok := m.session.(*treeManagerAdapter); ok {
+		return adapter.inner.GetFilePath()
 	}
 	return ""
 }
 
-// GetSessionID returns the UUID of the active tree session, or empty when no
-// tree session is configured.
+// GetSessionID returns the UUID of the active session, or empty when no
+// session is configured.
 func (m *Kit) GetSessionID() string {
-	if m.treeSession != nil {
-		return m.treeSession.GetSessionID()
+	if m.session == nil {
+		return ""
 	}
-	return ""
+	return m.session.GetSessionID()
 }
 
-// Branch moves the tree session's leaf pointer to the given entry ID, creating
+// Branch moves the session's leaf pointer to the given entry ID, creating
 // a branch point. Subsequent Prompt() calls will extend from the new position.
 func (m *Kit) Branch(entryID string) error {
-	return m.treeSession.Branch(entryID)
+	if m.session == nil {
+		return fmt.Errorf("no session available")
+	}
+	return m.session.Branch(entryID)
 }
 
-// SetSessionName sets a user-defined display name for the active tree session.
+// SetSessionName sets a user-defined display name for the active session.
 func (m *Kit) SetSessionName(name string) error {
-	if m.treeSession == nil {
-		return fmt.Errorf("session naming requires a tree session")
+	if m.session == nil {
+		return fmt.Errorf("session naming requires a session")
 	}
-	_, err := m.treeSession.AppendSessionInfo(name)
-	return err
+	return m.session.SetSessionName(name)
 }
 
 // ---------------------------------------------------------------------------
@@ -97,27 +120,27 @@ func (m *Kit) SetSessionName(name string) error {
 // ---------------------------------------------------------------------------
 
 // GetTreeNode returns a node by ID with full metadata and children.
-// Returns nil if entry not found or no tree session.
+// Returns nil if entry not found or no session.
 func (m *Kit) GetTreeNode(entryID string) *TreeNode {
-	if m.treeSession == nil {
+	if m.session == nil {
 		return nil
 	}
-	entry := m.treeSession.GetEntry(entryID)
+	entry := m.session.GetEntry(entryID)
 	if entry == nil {
 		return nil
 	}
-	return m.entryToTreeNode(entry)
+	return m.branchEntryToTreeNode(entry)
 }
 
 // GetCurrentBranch returns the path from root to current leaf as TreeNodes.
 func (m *Kit) GetCurrentBranch() []TreeNode {
-	if m.treeSession == nil {
+	if m.session == nil {
 		return nil
 	}
-	branch := m.treeSession.GetBranch("")
+	branch := m.session.GetCurrentBranch()
 	var nodes []TreeNode
 	for _, entry := range branch {
-		node := m.entryToTreeNode(entry)
+		node := m.branchEntryToTreeNode(&entry)
 		if node != nil {
 			nodes = append(nodes, *node)
 		}
@@ -127,34 +150,34 @@ func (m *Kit) GetCurrentBranch() []TreeNode {
 
 // GetChildren returns direct child IDs of an entry.
 func (m *Kit) GetChildren(parentID string) []string {
-	if m.treeSession == nil {
+	if m.session == nil {
 		return nil
 	}
-	return m.treeSession.GetChildren(parentID)
+	return m.session.GetChildren(parentID)
 }
 
 // NavigateTo branches/forks the session to the specified entry ID.
 // Returns an error if the session is unavailable or the entry ID is not found.
 func (m *Kit) NavigateTo(entryID string) error {
-	if m.treeSession == nil {
-		return fmt.Errorf("no tree session available")
+	if m.session == nil {
+		return fmt.Errorf("no session available")
 	}
-	return m.treeSession.Branch(entryID)
+	return m.session.Branch(entryID)
 }
 
 // SummarizeBranch uses the LLM to summarize the conversation between two
 // entry IDs. Returns the summary text, or an error if the range is invalid,
 // the session is unavailable, or the LLM call fails.
 func (m *Kit) SummarizeBranch(fromID, toID string) (string, error) {
-	if m.treeSession == nil {
-		return "", fmt.Errorf("no tree session available")
+	if m.session == nil {
+		return "", fmt.Errorf("no session available")
 	}
 
 	// Get the branch and find the range
-	branch := m.treeSession.GetBranch("")
+	branch := m.session.GetCurrentBranch()
 	var startIdx, endIdx = -1, -1
 	for i, entry := range branch {
-		id := m.treeSession.EntryID(entry)
+		id := entry.ID
 		if id == fromID {
 			startIdx = i
 		}
@@ -170,7 +193,7 @@ func (m *Kit) SummarizeBranch(fromID, toID string) (string, error) {
 	// Build text to summarize
 	var content strings.Builder
 	for i := startIdx; i <= endIdx; i++ {
-		node := m.entryToTreeNode(branch[i])
+		node := m.branchEntryToTreeNode(&branch[i])
 		if node != nil && node.Content != "" {
 			fmt.Fprintf(&content, "[%s] %s\n\n", node.Role, node.Content)
 		}
@@ -195,73 +218,81 @@ func (m *Kit) SummarizeBranch(fromID, toID string) (string, error) {
 // CollapseBranch replaces a branch range with a summary entry.
 // Returns an error if the session is unavailable or the operation fails.
 func (m *Kit) CollapseBranch(fromID, toID, summary string) error {
-	if m.treeSession == nil {
-		return fmt.Errorf("no tree session available")
+	if m.session == nil {
+		return fmt.Errorf("no session available")
 	}
-	_, err := m.treeSession.AppendBranchSummary(fromID, summary)
-	return err
+	// Note: This operation is not directly supported by SessionManager interface
+	// as it requires AppendBranchSummary which is TreeManager-specific.
+	// For custom SessionManagers, this would need to be implemented differently.
+	// For now, we try to use the underlying TreeManager if available.
+	if adapter, ok := m.session.(*treeManagerAdapter); ok {
+		_, err := adapter.inner.AppendBranchSummary(fromID, summary)
+		return err
+	}
+	return fmt.Errorf("CollapseBranch not supported by custom session manager")
 }
 
-// entryToTreeNode converts a session entry to a TreeNode.
-func (m *Kit) entryToTreeNode(entry any) *TreeNode {
-	switch e := entry.(type) {
-	case *session.MessageEntry:
-		msg, err := e.ToMessage()
-		if err != nil {
-			return nil
-		}
+// branchEntryToTreeNode converts a BranchEntry to a TreeNode.
+func (m *Kit) branchEntryToTreeNode(entry *BranchEntry) *TreeNode {
+	if entry == nil {
+		return nil
+	}
+
+	switch entry.Type {
+	case EntryTypeMessage:
+		// Build content from RawParts
 		var content strings.Builder
-		for _, p := range msg.Parts {
+		for _, p := range entry.RawParts {
 			switch pt := p.(type) {
-			case message.TextContent:
+			case TextContent:
 				content.WriteString(pt.Text)
-			case message.ReasoningContent:
+			case ReasoningContent:
 				content.WriteString(pt.Thinking)
-			case message.ToolCall:
+			case ToolCall:
 				fmt.Fprintf(&content, "[tool_call: %s]", pt.Name)
-			case message.ToolResult:
+			case ToolResult:
 				fmt.Fprintf(&content, "[tool_result: %s]", pt.Content)
 			}
 		}
 		return &TreeNode{
-			ID:        e.ID,
-			ParentID:  e.ParentID,
+			ID:        entry.ID,
+			ParentID:  entry.ParentID,
 			Type:      "message",
-			Role:      string(msg.Role),
+			Role:      entry.Role,
 			Content:   content.String(),
-			Model:     msg.Model,
-			Provider:  msg.Provider,
-			Timestamp: e.Timestamp.Format(time.RFC3339),
-			Children:  m.treeSession.GetChildren(e.ID),
+			Model:     entry.Model,
+			Provider:  entry.Provider,
+			Timestamp: entry.Timestamp.Format(time.RFC3339),
+			Children:  m.session.GetChildren(entry.ID),
 		}
-	case *session.BranchSummaryEntry:
+	case EntryTypeBranchSummary:
 		return &TreeNode{
-			ID:        e.ID,
-			ParentID:  e.ParentID,
+			ID:        entry.ID,
+			ParentID:  entry.ParentID,
 			Type:      "branch_summary",
-			Content:   e.Summary,
-			Timestamp: e.Timestamp.Format(time.RFC3339),
-			Children:  m.treeSession.GetChildren(e.ID),
+			Content:   entry.Content,
+			Timestamp: entry.Timestamp.Format(time.RFC3339),
+			Children:  m.session.GetChildren(entry.ID),
 		}
-	case *session.ModelChangeEntry:
+	case EntryTypeModelChange:
 		return &TreeNode{
-			ID:        e.ID,
-			ParentID:  e.ParentID,
+			ID:        entry.ID,
+			ParentID:  entry.ParentID,
 			Type:      "model_change",
-			Content:   fmt.Sprintf("Model changed to %s/%s", e.Provider, e.ModelID),
-			Model:     e.Provider + "/" + e.ModelID,
-			Provider:  e.Provider,
-			Timestamp: e.Timestamp.Format(time.RFC3339),
-			Children:  m.treeSession.GetChildren(e.ID),
+			Content:   entry.Content,
+			Model:     entry.Model,
+			Provider:  entry.Provider,
+			Timestamp: entry.Timestamp.Format(time.RFC3339),
+			Children:  m.session.GetChildren(entry.ID),
 		}
-	case *session.ExtensionDataEntry:
+	case EntryTypeExtensionData:
 		return &TreeNode{
-			ID:        e.ID,
-			ParentID:  e.ParentID,
+			ID:        entry.ID,
+			ParentID:  entry.ParentID,
 			Type:      "extension_data",
-			Content:   fmt.Sprintf("Extension data: %s", e.ExtType),
-			Timestamp: e.Timestamp.Format(time.RFC3339),
-			Children:  m.treeSession.GetChildren(e.ID),
+			Content:   entry.Content,
+			Timestamp: entry.Timestamp.Format(time.RFC3339),
+			Children:  m.session.GetChildren(entry.ID),
 		}
 	default:
 		return nil

skills/kit-sdk/SKILL.md

@@ -85,6 +85,7 @@ host, err := kit.New(ctx, &kit.Options{
     SessionPath: "/path/to/session.jsonl", // open specific session file
     Continue:    true,                // resume most recent session for SessionDir
     NoSession:   true,                // ephemeral in-memory session, no disk persistence
+    SessionManager: myCustomSession,  // custom SessionManager implementation (advanced)
 
     // Tools
     Tools:            []kit.Tool{kit.NewBashTool()}, // REPLACES entire default tool set
@@ -435,6 +436,72 @@ kit.DeleteSession("/path/to/session.jsonl")
 tm, _ := kit.OpenTreeSession("/path/to/session.jsonl") // open for direct access
 ```
 
+### Custom Session Manager (Advanced)
+
+You can provide a custom session manager to store conversation history in your own backend (database, cloud storage, etc.) instead of the default JSONL files.
+
+```go
+// Implement the SessionManager interface
+type MyDatabaseSessionManager struct {
+    db *sql.DB
+    // ... other fields
+}
+
+func (s *MyDatabaseSessionManager) AppendMessage(msg kit.LLMMessage) (string, error) {
+    // Store message in your database
+}
+
+func (s *MyDatabaseSessionManager) GetMessages() []kit.LLMMessage {
+    // Retrieve messages from your database
+}
+
+// ... implement all other SessionManager methods
+
+// Use with Kit
+host, _ := kit.New(ctx, &kit.Options{
+    SessionManager: myCustomSession,  // Your custom implementation
+    Model: "anthropic/claude-sonnet-latest",
+})
+```
+
+**SessionManager Interface:**
+
+```go
+type SessionManager interface {
+    AppendMessage(msg kit.LLMMessage) (entryID string, err error)
+    GetMessages() []kit.LLMMessage
+    BuildContext() (messages []kit.LLMMessage, provider string, modelID string)
+    Branch(entryID string) error
+    GetCurrentBranch() []kit.BranchEntry
+    GetChildren(parentID string) []string
+    GetEntry(entryID string) *kit.BranchEntry
+    GetSessionID() string
+    GetSessionName() string
+    SetSessionName(name string) error
+    GetCreatedAt() time.Time
+    IsPersisted() bool
+    AppendCompaction(summary string, firstKeptEntryID string,
+        tokensBefore, tokensAfter int, messagesRemoved int, readFiles, modifiedFiles []string) (string, error)
+    GetLastCompaction() *kit.CompactionEntry
+    AppendExtensionData(extType, data string) (string, error)
+    GetExtensionData(extType string) []kit.ExtensionDataEntry
+    AppendModelChange(provider, modelID string) (string, error)
+    GetContextEntryIDs() []string
+    Close() error
+}
+```
+
+**Use Cases:**
+- **PocketBase integration**: Store sessions as PocketBase records
+- **Cloud storage**: Persist sessions to S3, GCS, or Azure Blob
+- **Multi-user apps**: Store sessions per user in a database
+- **Custom retention**: Implement your own session cleanup policies
+
+**Note:** When using a custom SessionManager, the following Options are ignored:
+- `SessionPath` - your manager handles its own storage
+- `Continue` - your manager handles session selection
+- `NoSession` - use an in-memory implementation instead
+
 ---
 
 ## Model Management