diff --git a/pkg/kit/adapter.go b/pkg/kit/adapter.go new file mode 100644 index 00000000..f0ab6fcb --- /dev/null +++ b/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 +} diff --git a/pkg/kit/compaction.go b/pkg/kit/compaction.go index 602549ba..2420d19f 100644 --- a/pkg/kit/compaction.go +++ b/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, diff --git a/pkg/kit/extension_api.go b/pkg/kit/extension_api.go index ea073155..d6d9c0f4 100644 --- a/pkg/kit/extension_api.go +++ b/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{ diff --git a/pkg/kit/kit.go b/pkg/kit/kit.go index 9397aca5..dbd588cf 100644 --- a/pkg/kit/kit.go +++ b/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. diff --git a/pkg/kit/session.go b/pkg/kit/session.go new file mode 100644 index 00000000..147dce9c --- /dev/null +++ b/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 +} diff --git a/pkg/kit/sessions.go b/pkg/kit/sessions.go index 33618fbc..9abdca25 100644 --- a/pkg/kit/sessions.go +++ b/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 diff --git a/skills/kit-sdk/SKILL.md b/skills/kit-sdk/SKILL.md index cdf6ac62..6b7dac77 100644 --- a/skills/kit-sdk/SKILL.md +++ b/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