mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-16 14:01:06 +00:00
Session tree: simplify types, add branching API, comprehensive tests
Types: - SessionEntryBase with type field, extended by all entry types - CustomEntry for hooks (type: 'custom', customType, data) - Remove XXXContent types and TreeNode (redundant) API: - Rename saveXXX to appendXXX with JSDoc explaining tree semantics - Rename branchInPlace to branch() with better docs - Add createBranchedSession(leafId) replacing index-based version - Add getTree() returning SessionTreeNode[] for tree traversal - Add appendCustomEntry(customType, data) for hooks Tests: - tree-traversal.test.ts: 28 tests covering append, getPath, getTree, branch, branchWithSummary, createBranchedSession - save-entry.test.ts: custom entry integration Docs: - Class-level JSDoc explaining append-only tree model - Method docs explaining leaf advancement and branching - CHANGELOG.md entry for all changes
This commit is contained in:
Mario Zechner
parent
beb70f126d
commit
6f94e24629
8 changed files with 779 additions and 135 deletions
|
|
@ -2,13 +2,41 @@
|
|||
|
||||
## [Unreleased]
|
||||
|
||||
### Breaking Changes
|
||||
|
||||
- **Session tree structure (v2)**: Sessions now store entries as a tree with `id`/`parentId` fields, enabling in-place branching without creating new files. Existing v1 sessions are auto-migrated on load.
|
||||
- **SessionManager API**:
|
||||
- `saveXXX()` renamed to `appendXXX()` (e.g., `appendMessage`, `appendCompaction`)
|
||||
- `branchInPlace()` renamed to `branch()`
|
||||
- `reset()` renamed to `newSession()`
|
||||
- `createBranchedSessionFromEntries(entries, index)` replaced with `createBranchedSession(leafId)`
|
||||
- `saveCompaction(entry)` replaced with `appendCompaction(summary, firstKeptEntryId, tokensBefore)`
|
||||
- `getEntries()` now excludes the session header (use `getHeader()` separately)
|
||||
- New methods: `getTree()`, `getPath()`, `getLeafUuid()`, `getLeafEntry()`, `getEntry()`, `branchWithSummary()`
|
||||
- New `appendCustomEntry(customType, data)` for hooks to store custom data
|
||||
- **Compaction API**:
|
||||
- `compact()` now returns `CompactionResult` (`{ summary, firstKeptEntryId, tokensBefore }`) instead of `CompactionEntry`
|
||||
- `CompactionEntry.firstKeptEntryIndex` replaced with `firstKeptEntryId`
|
||||
- `prepareCompaction()` now returns `firstKeptEntryId` in its result
|
||||
- **Hook types**:
|
||||
- `SessionEventResult.compactionEntry` replaced with `SessionEventResult.compaction` (content only, SessionManager adds id/parentId)
|
||||
- `before_compact` event now includes `firstKeptEntryId` field for hooks that return custom compaction
|
||||
|
||||
### Added
|
||||
|
||||
- **`enabledModels` setting**: Configure whitelisted models in `settings.json` (same format as `--models` CLI flag). CLI `--models` takes precedence over the setting.
|
||||
|
||||
### Changed
|
||||
|
||||
- **Entry IDs**: Session entries now use short 8-character hex IDs instead of full UUIDs
|
||||
- **API key priority**: `ANTHROPIC_OAUTH_TOKEN` now takes precedence over `ANTHROPIC_API_KEY`
|
||||
- **New entry types**: `BranchSummaryEntry` for branch context, `CustomEntry` for hook data
|
||||
|
||||
### Fixed
|
||||
|
||||
- **Edit tool fails on Windows due to CRLF line endings**: Files with CRLF line endings now match correctly when LLMs send LF-only text. Line endings are normalized before matching and restored to original style on write. ([#355](https://github.com/badlogic/pi-mono/issues/355))
|
||||
- **Session file validation**: `findMostRecentSession()` now validates session headers before returning, preventing non-session JSONL files from being loaded
|
||||
- **Compaction error handling**: `generateSummary()` and `generateTurnPrefixSummary()` now throw on LLM errors instead of returning empty strings
|
||||
|
||||
## [0.30.2] - 2025-12-26
|
||||
|
||||
|
|
|
|||
|
|
@ -211,7 +211,7 @@ export class AgentSession {
|
|||
|
||||
// Handle session persistence
|
||||
if (event.type === "message_end") {
|
||||
this.sessionManager.saveMessage(event.message);
|
||||
this.sessionManager.appendMessage(event.message);
|
||||
|
||||
// Track assistant message for auto-compaction (checked on agent_end)
|
||||
if (event.message.role === "assistant") {
|
||||
|
|
@ -535,7 +535,7 @@ export class AgentSession {
|
|||
this._disconnectFromAgent();
|
||||
await this.abort();
|
||||
this.agent.reset();
|
||||
this.sessionManager.reset();
|
||||
this.sessionManager.newSession();
|
||||
this._queuedMessages = [];
|
||||
this._reconnectToAgent();
|
||||
|
||||
|
|
@ -572,7 +572,7 @@ export class AgentSession {
|
|||
}
|
||||
|
||||
this.agent.setModel(model);
|
||||
this.sessionManager.saveModelChange(model.provider, model.id);
|
||||
this.sessionManager.appendModelChange(model.provider, model.id);
|
||||
this.settingsManager.setDefaultModelAndProvider(model.provider, model.id);
|
||||
|
||||
// Re-clamp thinking level for new model's capabilities
|
||||
|
|
@ -611,7 +611,7 @@ export class AgentSession {
|
|||
|
||||
// Apply model
|
||||
this.agent.setModel(next.model);
|
||||
this.sessionManager.saveModelChange(next.model.provider, next.model.id);
|
||||
this.sessionManager.appendModelChange(next.model.provider, next.model.id);
|
||||
this.settingsManager.setDefaultModelAndProvider(next.model.provider, next.model.id);
|
||||
|
||||
// Apply thinking level (setThinkingLevel clamps to model capabilities)
|
||||
|
|
@ -638,7 +638,7 @@ export class AgentSession {
|
|||
}
|
||||
|
||||
this.agent.setModel(nextModel);
|
||||
this.sessionManager.saveModelChange(nextModel.provider, nextModel.id);
|
||||
this.sessionManager.appendModelChange(nextModel.provider, nextModel.id);
|
||||
this.settingsManager.setDefaultModelAndProvider(nextModel.provider, nextModel.id);
|
||||
|
||||
// Re-clamp thinking level for new model's capabilities
|
||||
|
|
@ -671,7 +671,7 @@ export class AgentSession {
|
|||
effectiveLevel = "high";
|
||||
}
|
||||
this.agent.setThinkingLevel(effectiveLevel);
|
||||
this.sessionManager.saveThinkingLevelChange(effectiveLevel);
|
||||
this.sessionManager.appendThinkingLevelChange(effectiveLevel);
|
||||
this.settingsManager.setDefaultThinkingLevel(effectiveLevel);
|
||||
}
|
||||
|
||||
|
|
@ -831,7 +831,7 @@ export class AgentSession {
|
|||
throw new Error("Compaction cancelled");
|
||||
}
|
||||
|
||||
this.sessionManager.saveCompaction(summary, firstKeptEntryId, tokensBefore);
|
||||
this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore);
|
||||
const newEntries = this.sessionManager.getEntries();
|
||||
const sessionContext = this.sessionManager.buildSessionContext();
|
||||
this.agent.replaceMessages(sessionContext.messages);
|
||||
|
|
@ -1013,7 +1013,7 @@ export class AgentSession {
|
|||
return;
|
||||
}
|
||||
|
||||
this.sessionManager.saveCompaction(summary, firstKeptEntryId, tokensBefore);
|
||||
this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore);
|
||||
const newEntries = this.sessionManager.getEntries();
|
||||
const sessionContext = this.sessionManager.buildSessionContext();
|
||||
this.agent.replaceMessages(sessionContext.messages);
|
||||
|
|
@ -1271,7 +1271,7 @@ export class AgentSession {
|
|||
this.agent.appendMessage(bashMessage);
|
||||
|
||||
// Save to session
|
||||
this.sessionManager.saveMessage(bashMessage);
|
||||
this.sessionManager.appendMessage(bashMessage);
|
||||
}
|
||||
|
||||
return result;
|
||||
|
|
@ -1309,7 +1309,7 @@ export class AgentSession {
|
|||
this.agent.appendMessage(bashMessage);
|
||||
|
||||
// Save to session
|
||||
this.sessionManager.saveMessage(bashMessage);
|
||||
this.sessionManager.appendMessage(bashMessage);
|
||||
}
|
||||
|
||||
this._pendingBashMessages = [];
|
||||
|
|
@ -1431,8 +1431,12 @@ export class AgentSession {
|
|||
skipConversationRestore = result?.skipConversationRestore ?? false;
|
||||
}
|
||||
|
||||
// Create branched session (returns null in --no-session mode)
|
||||
const newSessionFile = this.sessionManager.createBranchedSessionFromEntries(entries, entryIndex);
|
||||
// Create branched session ending before the selected message (returns null in --no-session mode)
|
||||
// User will re-enter/edit the selected message
|
||||
if (!selectedEntry.parentId) {
|
||||
throw new Error("Cannot branch from first message");
|
||||
}
|
||||
const newSessionFile = this.sessionManager.createBranchedSession(selectedEntry.parentId);
|
||||
|
||||
// Update session file if we have one (file-based mode)
|
||||
if (newSessionFile !== null) {
|
||||
|
|
|
|||
|
|
@ -26,54 +26,47 @@ export interface SessionHeader {
|
|||
branchedFrom?: string;
|
||||
}
|
||||
|
||||
export interface MessageContent {
|
||||
type: "message";
|
||||
message: AppMessage;
|
||||
}
|
||||
|
||||
export interface ThinkingLevelContent {
|
||||
type: "thinking_level_change";
|
||||
thinkingLevel: string;
|
||||
}
|
||||
|
||||
export interface ModelChangeContent {
|
||||
type: "model_change";
|
||||
provider: string;
|
||||
modelId: string;
|
||||
}
|
||||
|
||||
export interface CompactionContent {
|
||||
type: "compaction";
|
||||
summary: string;
|
||||
firstKeptEntryId: string;
|
||||
tokensBefore: number;
|
||||
}
|
||||
|
||||
export interface BranchSummaryContent {
|
||||
type: "branch_summary";
|
||||
summary: string;
|
||||
}
|
||||
|
||||
/** Union of all content types (for "write" methods in SessionManager) */
|
||||
export type SessionContent =
|
||||
| MessageContent
|
||||
| ThinkingLevelContent
|
||||
| ModelChangeContent
|
||||
| CompactionContent
|
||||
| BranchSummaryContent;
|
||||
|
||||
export interface TreeNode {
|
||||
export interface SessionEntryBase {
|
||||
type: string;
|
||||
id: string;
|
||||
parentId: string | null;
|
||||
timestamp: string;
|
||||
}
|
||||
|
||||
export type SessionMessageEntry = TreeNode & MessageContent;
|
||||
export type ThinkingLevelChangeEntry = TreeNode & ThinkingLevelContent;
|
||||
export type ModelChangeEntry = TreeNode & ModelChangeContent;
|
||||
export type CompactionEntry = TreeNode & CompactionContent;
|
||||
export type BranchSummaryEntry = TreeNode & BranchSummaryContent;
|
||||
export interface SessionMessageEntry extends SessionEntryBase {
|
||||
type: "message";
|
||||
message: AppMessage;
|
||||
}
|
||||
|
||||
export interface ThinkingLevelChangeEntry extends SessionEntryBase {
|
||||
type: "thinking_level_change";
|
||||
thinkingLevel: string;
|
||||
}
|
||||
|
||||
export interface ModelChangeEntry extends SessionEntryBase {
|
||||
type: "model_change";
|
||||
provider: string;
|
||||
modelId: string;
|
||||
}
|
||||
|
||||
export interface CompactionEntry extends SessionEntryBase {
|
||||
type: "compaction";
|
||||
summary: string;
|
||||
firstKeptEntryId: string;
|
||||
tokensBefore: number;
|
||||
}
|
||||
|
||||
export interface BranchSummaryEntry extends SessionEntryBase {
|
||||
type: "branch_summary";
|
||||
summary: string;
|
||||
}
|
||||
|
||||
/** Custom entry for hooks. Use customType to identify your hook's entries. */
|
||||
export interface CustomEntry extends SessionEntryBase {
|
||||
type: "custom";
|
||||
customType: string;
|
||||
data?: unknown;
|
||||
}
|
||||
|
||||
/** Session entry - has id/parentId for tree structure (returned by "read" methods in SessionManager) */
|
||||
export type SessionEntry =
|
||||
|
|
@ -81,11 +74,18 @@ export type SessionEntry =
|
|||
| ThinkingLevelChangeEntry
|
||||
| ModelChangeEntry
|
||||
| CompactionEntry
|
||||
| BranchSummaryEntry;
|
||||
| BranchSummaryEntry
|
||||
| CustomEntry;
|
||||
|
||||
/** Raw file entry (includes header) */
|
||||
export type FileEntry = SessionHeader | SessionEntry;
|
||||
|
||||
/** Tree node for getTree() - defensive copy of session structure */
|
||||
export interface SessionTreeNode {
|
||||
entry: SessionEntry;
|
||||
children: SessionTreeNode[];
|
||||
}
|
||||
|
||||
export interface SessionContext {
|
||||
messages: AppMessage[];
|
||||
thinkingLevel: string;
|
||||
|
|
@ -387,6 +387,17 @@ export function findMostRecentSession(sessionDir: string): string | null {
|
|||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Manages conversation sessions as append-only trees stored in JSONL files.
|
||||
*
|
||||
* Each session entry has an id and parentId forming a tree structure. The "leaf"
|
||||
* pointer tracks the current position. Appending creates a child of the current leaf.
|
||||
* Branching moves the leaf to an earlier entry, allowing new branches without
|
||||
* modifying history.
|
||||
*
|
||||
* Use buildSessionContext() to get the resolved message list for the LLM, which
|
||||
* handles compaction summaries and follows the path from root to current leaf.
|
||||
*/
|
||||
export class SessionManager {
|
||||
private sessionId: string = "";
|
||||
private sessionFile: string = "";
|
||||
|
|
@ -394,7 +405,7 @@ export class SessionManager {
|
|||
private cwd: string;
|
||||
private persist: boolean;
|
||||
private flushed: boolean = false;
|
||||
private inMemoryEntries: FileEntry[] = [];
|
||||
private fileEntries: FileEntry[] = [];
|
||||
private byId: Map<string, SessionEntry> = new Map();
|
||||
private leafId: string = "";
|
||||
|
||||
|
|
@ -409,7 +420,7 @@ export class SessionManager {
|
|||
if (sessionFile) {
|
||||
this.setSessionFile(sessionFile);
|
||||
} else {
|
||||
this._initNewSession();
|
||||
this.newSession();
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -417,22 +428,22 @@ export class SessionManager {
|
|||
setSessionFile(sessionFile: string): void {
|
||||
this.sessionFile = resolve(sessionFile);
|
||||
if (existsSync(this.sessionFile)) {
|
||||
this.inMemoryEntries = loadEntriesFromFile(this.sessionFile);
|
||||
const header = this.inMemoryEntries.find((e) => e.type === "session") as SessionHeader | undefined;
|
||||
this.fileEntries = loadEntriesFromFile(this.sessionFile);
|
||||
const header = this.fileEntries.find((e) => e.type === "session") as SessionHeader | undefined;
|
||||
this.sessionId = header?.id ?? randomUUID();
|
||||
|
||||
if (migrateToCurrentVersion(this.inMemoryEntries)) {
|
||||
if (migrateToCurrentVersion(this.fileEntries)) {
|
||||
this._rewriteFile();
|
||||
}
|
||||
|
||||
this._buildIndex();
|
||||
this.flushed = true;
|
||||
} else {
|
||||
this._initNewSession();
|
||||
this.newSession();
|
||||
}
|
||||
}
|
||||
|
||||
private _initNewSession(): void {
|
||||
newSession(): void {
|
||||
this.sessionId = randomUUID();
|
||||
const timestamp = new Date().toISOString();
|
||||
const header: SessionHeader = {
|
||||
|
|
@ -442,7 +453,7 @@ export class SessionManager {
|
|||
timestamp,
|
||||
cwd: this.cwd,
|
||||
};
|
||||
this.inMemoryEntries = [header];
|
||||
this.fileEntries = [header];
|
||||
this.byId.clear();
|
||||
this.leafId = "";
|
||||
this.flushed = false;
|
||||
|
|
@ -456,7 +467,7 @@ export class SessionManager {
|
|||
private _buildIndex(): void {
|
||||
this.byId.clear();
|
||||
this.leafId = "";
|
||||
for (const entry of this.inMemoryEntries) {
|
||||
for (const entry of this.fileEntries) {
|
||||
if (entry.type === "session") continue;
|
||||
this.byId.set(entry.id, entry);
|
||||
this.leafId = entry.id;
|
||||
|
|
@ -465,7 +476,7 @@ export class SessionManager {
|
|||
|
||||
private _rewriteFile(): void {
|
||||
if (!this.persist) return;
|
||||
const content = `${this.inMemoryEntries.map((e) => JSON.stringify(e)).join("\n")}\n`;
|
||||
const content = `${this.fileEntries.map((e) => JSON.stringify(e)).join("\n")}\n`;
|
||||
writeFileSync(this.sessionFile, content);
|
||||
}
|
||||
|
||||
|
|
@ -489,18 +500,14 @@ export class SessionManager {
|
|||
return this.sessionFile;
|
||||
}
|
||||
|
||||
reset(): void {
|
||||
this._initNewSession();
|
||||
}
|
||||
|
||||
_persist(entry: SessionEntry): void {
|
||||
if (!this.persist) return;
|
||||
|
||||
const hasAssistant = this.inMemoryEntries.some((e) => e.type === "message" && e.message.role === "assistant");
|
||||
const hasAssistant = this.fileEntries.some((e) => e.type === "message" && e.message.role === "assistant");
|
||||
if (!hasAssistant) return;
|
||||
|
||||
if (!this.flushed) {
|
||||
for (const e of this.inMemoryEntries) {
|
||||
for (const e of this.fileEntries) {
|
||||
appendFileSync(this.sessionFile, `${JSON.stringify(e)}\n`);
|
||||
}
|
||||
this.flushed = true;
|
||||
|
|
@ -510,13 +517,14 @@ export class SessionManager {
|
|||
}
|
||||
|
||||
private _appendEntry(entry: SessionEntry): void {
|
||||
this.inMemoryEntries.push(entry);
|
||||
this.fileEntries.push(entry);
|
||||
this.byId.set(entry.id, entry);
|
||||
this.leafId = entry.id;
|
||||
this._persist(entry);
|
||||
}
|
||||
|
||||
saveMessage(message: AppMessage): string {
|
||||
/** Append a message as child of current leaf, then advance leaf. Returns entry id. */
|
||||
appendMessage(message: AppMessage): string {
|
||||
const entry: SessionMessageEntry = {
|
||||
type: "message",
|
||||
id: generateId(this.byId),
|
||||
|
|
@ -528,7 +536,8 @@ export class SessionManager {
|
|||
return entry.id;
|
||||
}
|
||||
|
||||
saveThinkingLevelChange(thinkingLevel: string): string {
|
||||
/** Append a thinking level change as child of current leaf, then advance leaf. Returns entry id. */
|
||||
appendThinkingLevelChange(thinkingLevel: string): string {
|
||||
const entry: ThinkingLevelChangeEntry = {
|
||||
type: "thinking_level_change",
|
||||
id: generateId(this.byId),
|
||||
|
|
@ -540,7 +549,8 @@ export class SessionManager {
|
|||
return entry.id;
|
||||
}
|
||||
|
||||
saveModelChange(provider: string, modelId: string): string {
|
||||
/** Append a model change as child of current leaf, then advance leaf. Returns entry id. */
|
||||
appendModelChange(provider: string, modelId: string): string {
|
||||
const entry: ModelChangeEntry = {
|
||||
type: "model_change",
|
||||
id: generateId(this.byId),
|
||||
|
|
@ -553,7 +563,8 @@ export class SessionManager {
|
|||
return entry.id;
|
||||
}
|
||||
|
||||
saveCompaction(summary: string, firstKeptEntryId: string, tokensBefore: number): string {
|
||||
/** Append a compaction summary as child of current leaf, then advance leaf. Returns entry id. */
|
||||
appendCompaction(summary: string, firstKeptEntryId: string, tokensBefore: number): string {
|
||||
const entry: CompactionEntry = {
|
||||
type: "compaction",
|
||||
id: generateId(this.byId),
|
||||
|
|
@ -567,6 +578,20 @@ export class SessionManager {
|
|||
return entry.id;
|
||||
}
|
||||
|
||||
/** Append a custom entry (for hooks) as child of current leaf, then advance leaf. Returns entry id. */
|
||||
appendCustomEntry(customType: string, data?: unknown): string {
|
||||
const entry: CustomEntry = {
|
||||
type: "custom",
|
||||
customType,
|
||||
data,
|
||||
id: generateId(this.byId),
|
||||
parentId: this.leafId || null,
|
||||
timestamp: new Date().toISOString(),
|
||||
};
|
||||
this._appendEntry(entry);
|
||||
return entry.id;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Tree Traversal
|
||||
// =========================================================================
|
||||
|
|
@ -575,11 +600,19 @@ export class SessionManager {
|
|||
return this.leafId;
|
||||
}
|
||||
|
||||
getLeafEntry(): SessionEntry | undefined {
|
||||
return this.byId.get(this.leafId);
|
||||
}
|
||||
|
||||
getEntry(id: string): SessionEntry | undefined {
|
||||
return this.byId.get(id);
|
||||
}
|
||||
|
||||
/** Walk from entry to root, returning path (conversation entries only) */
|
||||
/**
|
||||
* Walk from entry to root, returning all entries in path order.
|
||||
* Includes all entry types (messages, compaction, model changes, etc.).
|
||||
* Use buildSessionContext() to get the resolved messages for the LLM.
|
||||
*/
|
||||
getPath(fromId?: string): SessionEntry[] {
|
||||
const path: SessionEntry[] = [];
|
||||
let current = this.byId.get(fromId ?? this.leafId);
|
||||
|
|
@ -602,31 +635,75 @@ export class SessionManager {
|
|||
* Get session header.
|
||||
*/
|
||||
getHeader(): SessionHeader | null {
|
||||
const h = this.inMemoryEntries.find((e) => e.type === "session");
|
||||
const h = this.fileEntries.find((e) => e.type === "session");
|
||||
return h ? (h as SessionHeader) : null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all session entries (excludes header). Returns a defensive copy.
|
||||
* Use buildSessionContext() if you need the messages for the LLM.
|
||||
* Get all session entries (excludes header). Returns a shallow copy.
|
||||
* The session is append-only: use appendXXX() to add entries, branch() to
|
||||
* change the leaf pointer. Entries cannot be modified or deleted.
|
||||
*/
|
||||
getEntries(): SessionEntry[] {
|
||||
return this.inMemoryEntries.filter((e): e is SessionEntry => e.type !== "session");
|
||||
return this.fileEntries.filter((e): e is SessionEntry => e.type !== "session");
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the session as a tree structure. Returns a shallow defensive copy of all entries.
|
||||
* A well-formed session has exactly one root (first entry with parentId === null).
|
||||
* Orphaned entries (broken parent chain) are also returned as roots.
|
||||
*/
|
||||
getTree(): SessionTreeNode[] {
|
||||
const entries = this.getEntries();
|
||||
const nodeMap = new Map<string, SessionTreeNode>();
|
||||
const roots: SessionTreeNode[] = [];
|
||||
|
||||
// Create nodes
|
||||
for (const entry of entries) {
|
||||
nodeMap.set(entry.id, { entry, children: [] });
|
||||
}
|
||||
|
||||
// Build tree
|
||||
for (const entry of entries) {
|
||||
const node = nodeMap.get(entry.id)!;
|
||||
if (entry.parentId === null) {
|
||||
roots.push(node);
|
||||
} else {
|
||||
const parent = nodeMap.get(entry.parentId);
|
||||
if (parent) {
|
||||
parent.children.push(node);
|
||||
} else {
|
||||
// Orphan - treat as root
|
||||
roots.push(node);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return roots;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Branching
|
||||
// =========================================================================
|
||||
|
||||
/** Branch in-place by changing the leaf pointer */
|
||||
branchInPlace(branchFromId: string): void {
|
||||
/**
|
||||
* Start a new branch from an earlier entry.
|
||||
* Moves the leaf pointer to the specified entry. The next appendXXX() call
|
||||
* will create a child of that entry, forming a new branch. Existing entries
|
||||
* are not modified or deleted.
|
||||
*/
|
||||
branch(branchFromId: string): void {
|
||||
if (!this.byId.has(branchFromId)) {
|
||||
throw new Error(`Entry ${branchFromId} not found`);
|
||||
}
|
||||
this.leafId = branchFromId;
|
||||
}
|
||||
|
||||
/** Branch with a summary of the abandoned path */
|
||||
/**
|
||||
* Start a new branch with a summary of the abandoned path.
|
||||
* Same as branch(), but also appends a branch_summary entry that captures
|
||||
* context from the abandoned conversation path.
|
||||
*/
|
||||
branchWithSummary(branchFromId: string, summary: string): string {
|
||||
if (!this.byId.has(branchFromId)) {
|
||||
throw new Error(`Entry ${branchFromId} not found`);
|
||||
|
|
@ -643,35 +720,41 @@ export class SessionManager {
|
|||
return entry.id;
|
||||
}
|
||||
|
||||
createBranchedSessionFromEntries(entries: FileEntry[], branchBeforeIndex: number): string | null {
|
||||
const newSessionId = randomUUID();
|
||||
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
|
||||
const newSessionFile = join(this.getSessionDir(), `${timestamp}_${newSessionId}.jsonl`);
|
||||
|
||||
const newEntries: FileEntry[] = [];
|
||||
for (let i = 0; i < branchBeforeIndex; i++) {
|
||||
const entry = entries[i];
|
||||
|
||||
if (entry.type === "session") {
|
||||
newEntries.push({
|
||||
...entry,
|
||||
version: CURRENT_SESSION_VERSION,
|
||||
id: newSessionId,
|
||||
timestamp: new Date().toISOString(),
|
||||
branchedFrom: this.persist ? this.sessionFile : undefined,
|
||||
});
|
||||
} else {
|
||||
newEntries.push(entry);
|
||||
}
|
||||
/**
|
||||
* Create a new session file containing only the path from root to the specified leaf.
|
||||
* Useful for extracting a single conversation path from a branched session.
|
||||
* Returns the new session file path, or null if not persisting.
|
||||
*/
|
||||
createBranchedSession(leafId: string): string | null {
|
||||
const path = this.getPath(leafId);
|
||||
if (path.length === 0) {
|
||||
throw new Error(`Entry ${leafId} not found`);
|
||||
}
|
||||
|
||||
const newSessionId = randomUUID();
|
||||
const timestamp = new Date().toISOString();
|
||||
const fileTimestamp = timestamp.replace(/[:.]/g, "-");
|
||||
const newSessionFile = join(this.getSessionDir(), `${fileTimestamp}_${newSessionId}.jsonl`);
|
||||
|
||||
const header: SessionHeader = {
|
||||
type: "session",
|
||||
version: CURRENT_SESSION_VERSION,
|
||||
id: newSessionId,
|
||||
timestamp,
|
||||
cwd: this.cwd,
|
||||
branchedFrom: this.persist ? this.sessionFile : undefined,
|
||||
};
|
||||
|
||||
if (this.persist) {
|
||||
for (const entry of newEntries) {
|
||||
appendFileSync(newSessionFile, `${JSON.stringify(header)}\n`);
|
||||
for (const entry of path) {
|
||||
appendFileSync(newSessionFile, `${JSON.stringify(entry)}\n`);
|
||||
}
|
||||
return newSessionFile;
|
||||
}
|
||||
this.inMemoryEntries = newEntries;
|
||||
|
||||
// In-memory mode: replace current session with the path
|
||||
this.fileEntries = [header, ...path];
|
||||
this.sessionId = newSessionId;
|
||||
this._buildIndex();
|
||||
return null;
|
||||
|
|
|
|||
|
|
@ -107,23 +107,19 @@ export {
|
|||
readOnlyTools,
|
||||
} from "./core/sdk.js";
|
||||
export {
|
||||
type BranchSummaryContent,
|
||||
type BranchSummaryEntry,
|
||||
buildSessionContext,
|
||||
type CompactionContent,
|
||||
type CompactionEntry,
|
||||
CURRENT_SESSION_VERSION,
|
||||
createSummaryMessage,
|
||||
type FileEntry,
|
||||
getLatestCompactionEntry,
|
||||
type MessageContent,
|
||||
type ModelChangeContent,
|
||||
type ModelChangeEntry,
|
||||
migrateSessionEntries,
|
||||
parseSessionEntries,
|
||||
type SessionContent as ConversationContent,
|
||||
type SessionContext as LoadedSession,
|
||||
type SessionContext,
|
||||
type SessionEntry,
|
||||
type SessionEntryBase,
|
||||
type SessionHeader,
|
||||
type SessionInfo,
|
||||
SessionManager,
|
||||
|
|
@ -131,9 +127,6 @@ export {
|
|||
SUMMARY_PREFIX,
|
||||
SUMMARY_SUFFIX,
|
||||
type ThinkingLevelChangeEntry,
|
||||
type ThinkingLevelContent,
|
||||
// Tree types (v2)
|
||||
type TreeNode,
|
||||
} from "./core/session-manager.js";
|
||||
export {
|
||||
type CompactionSettings,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,55 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import { type CustomEntry, SessionManager } from "../../src/core/session-manager.js";
|
||||
|
||||
describe("SessionManager.saveCustomEntry", () => {
|
||||
it("saves custom entries and includes them in tree traversal", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
// Save a message
|
||||
const msgId = session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
|
||||
|
||||
// Save a custom entry
|
||||
const customId = session.appendCustomEntry("my_hook", { foo: "bar" });
|
||||
|
||||
// Save another message
|
||||
const msg2Id = session.appendMessage({
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "hi" }],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "test",
|
||||
usage: {
|
||||
input: 1,
|
||||
output: 1,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 2,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: 2,
|
||||
});
|
||||
|
||||
// Custom entry should be in entries
|
||||
const entries = session.getEntries();
|
||||
expect(entries).toHaveLength(3);
|
||||
|
||||
const customEntry = entries.find((e) => e.type === "custom") as CustomEntry;
|
||||
expect(customEntry).toBeDefined();
|
||||
expect(customEntry.customType).toBe("my_hook");
|
||||
expect(customEntry.data).toEqual({ foo: "bar" });
|
||||
expect(customEntry.id).toBe(customId);
|
||||
expect(customEntry.parentId).toBe(msgId);
|
||||
|
||||
// Tree structure should be correct
|
||||
const path = session.getPath();
|
||||
expect(path).toHaveLength(3);
|
||||
expect(path[0].id).toBe(msgId);
|
||||
expect(path[1].id).toBe(customId);
|
||||
expect(path[2].id).toBe(msg2Id);
|
||||
|
||||
// buildSessionContext should work (custom entries skipped in messages)
|
||||
const ctx = session.buildSessionContext();
|
||||
expect(ctx.messages).toHaveLength(2); // only message entries
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,483 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import { type CustomEntry, SessionManager } from "../../src/core/session-manager.js";
|
||||
|
||||
function userMsg(text: string) {
|
||||
return { role: "user" as const, content: text, timestamp: Date.now() };
|
||||
}
|
||||
|
||||
function assistantMsg(text: string) {
|
||||
return {
|
||||
role: "assistant" as const,
|
||||
content: [{ type: "text" as const, text }],
|
||||
api: "anthropic-messages" as const,
|
||||
provider: "anthropic",
|
||||
model: "test",
|
||||
usage: {
|
||||
input: 1,
|
||||
output: 1,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 2,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop" as const,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
describe("SessionManager append and tree traversal", () => {
|
||||
describe("append operations", () => {
|
||||
it("appendMessage creates entry with correct parentId chain", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const id1 = session.appendMessage(userMsg("first"));
|
||||
const id2 = session.appendMessage(assistantMsg("second"));
|
||||
const id3 = session.appendMessage(userMsg("third"));
|
||||
|
||||
const entries = session.getEntries();
|
||||
expect(entries).toHaveLength(3);
|
||||
|
||||
expect(entries[0].id).toBe(id1);
|
||||
expect(entries[0].parentId).toBeNull();
|
||||
expect(entries[0].type).toBe("message");
|
||||
|
||||
expect(entries[1].id).toBe(id2);
|
||||
expect(entries[1].parentId).toBe(id1);
|
||||
|
||||
expect(entries[2].id).toBe(id3);
|
||||
expect(entries[2].parentId).toBe(id2);
|
||||
});
|
||||
|
||||
it("appendThinkingLevelChange integrates into tree", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const msgId = session.appendMessage(userMsg("hello"));
|
||||
const thinkingId = session.appendThinkingLevelChange("high");
|
||||
const _msg2Id = session.appendMessage(assistantMsg("response"));
|
||||
|
||||
const entries = session.getEntries();
|
||||
expect(entries).toHaveLength(3);
|
||||
|
||||
const thinkingEntry = entries.find((e) => e.type === "thinking_level_change");
|
||||
expect(thinkingEntry).toBeDefined();
|
||||
expect(thinkingEntry!.id).toBe(thinkingId);
|
||||
expect(thinkingEntry!.parentId).toBe(msgId);
|
||||
|
||||
expect(entries[2].parentId).toBe(thinkingId);
|
||||
});
|
||||
|
||||
it("appendModelChange integrates into tree", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const msgId = session.appendMessage(userMsg("hello"));
|
||||
const modelId = session.appendModelChange("openai", "gpt-4");
|
||||
const _msg2Id = session.appendMessage(assistantMsg("response"));
|
||||
|
||||
const entries = session.getEntries();
|
||||
const modelEntry = entries.find((e) => e.type === "model_change");
|
||||
expect(modelEntry).toBeDefined();
|
||||
expect(modelEntry?.id).toBe(modelId);
|
||||
expect(modelEntry?.parentId).toBe(msgId);
|
||||
if (modelEntry?.type === "model_change") {
|
||||
expect(modelEntry.provider).toBe("openai");
|
||||
expect(modelEntry.modelId).toBe("gpt-4");
|
||||
}
|
||||
|
||||
expect(entries[2].parentId).toBe(modelId);
|
||||
});
|
||||
|
||||
it("appendCompaction integrates into tree", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
const id2 = session.appendMessage(assistantMsg("2"));
|
||||
const compactionId = session.appendCompaction("summary", id1, 1000);
|
||||
const _id3 = session.appendMessage(userMsg("3"));
|
||||
|
||||
const entries = session.getEntries();
|
||||
const compactionEntry = entries.find((e) => e.type === "compaction");
|
||||
expect(compactionEntry).toBeDefined();
|
||||
expect(compactionEntry?.id).toBe(compactionId);
|
||||
expect(compactionEntry?.parentId).toBe(id2);
|
||||
if (compactionEntry?.type === "compaction") {
|
||||
expect(compactionEntry.summary).toBe("summary");
|
||||
expect(compactionEntry.firstKeptEntryId).toBe(id1);
|
||||
expect(compactionEntry.tokensBefore).toBe(1000);
|
||||
}
|
||||
|
||||
expect(entries[3].parentId).toBe(compactionId);
|
||||
});
|
||||
|
||||
it("appendCustomEntry integrates into tree", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const msgId = session.appendMessage(userMsg("hello"));
|
||||
const customId = session.appendCustomEntry("my_hook", { key: "value" });
|
||||
const _msg2Id = session.appendMessage(assistantMsg("response"));
|
||||
|
||||
const entries = session.getEntries();
|
||||
const customEntry = entries.find((e) => e.type === "custom") as CustomEntry;
|
||||
expect(customEntry).toBeDefined();
|
||||
expect(customEntry.id).toBe(customId);
|
||||
expect(customEntry.parentId).toBe(msgId);
|
||||
expect(customEntry.customType).toBe("my_hook");
|
||||
expect(customEntry.data).toEqual({ key: "value" });
|
||||
|
||||
expect(entries[2].parentId).toBe(customId);
|
||||
});
|
||||
|
||||
it("leaf pointer advances after each append", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
expect(session.getLeafUuid()).toBe("");
|
||||
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
expect(session.getLeafUuid()).toBe(id1);
|
||||
|
||||
const id2 = session.appendMessage(assistantMsg("2"));
|
||||
expect(session.getLeafUuid()).toBe(id2);
|
||||
|
||||
const id3 = session.appendThinkingLevelChange("high");
|
||||
expect(session.getLeafUuid()).toBe(id3);
|
||||
});
|
||||
});
|
||||
|
||||
describe("getPath", () => {
|
||||
it("returns empty array for empty session", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
expect(session.getPath()).toEqual([]);
|
||||
});
|
||||
|
||||
it("returns single entry path", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
const id = session.appendMessage(userMsg("hello"));
|
||||
|
||||
const path = session.getPath();
|
||||
expect(path).toHaveLength(1);
|
||||
expect(path[0].id).toBe(id);
|
||||
});
|
||||
|
||||
it("returns full path from root to leaf", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
const id2 = session.appendMessage(assistantMsg("2"));
|
||||
const id3 = session.appendThinkingLevelChange("high");
|
||||
const id4 = session.appendMessage(userMsg("3"));
|
||||
|
||||
const path = session.getPath();
|
||||
expect(path).toHaveLength(4);
|
||||
expect(path.map((e) => e.id)).toEqual([id1, id2, id3, id4]);
|
||||
});
|
||||
|
||||
it("returns path from specified entry to root", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
const id2 = session.appendMessage(assistantMsg("2"));
|
||||
const _id3 = session.appendMessage(userMsg("3"));
|
||||
const _id4 = session.appendMessage(assistantMsg("4"));
|
||||
|
||||
const path = session.getPath(id2);
|
||||
expect(path).toHaveLength(2);
|
||||
expect(path.map((e) => e.id)).toEqual([id1, id2]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("getTree", () => {
|
||||
it("returns empty array for empty session", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
expect(session.getTree()).toEqual([]);
|
||||
});
|
||||
|
||||
it("returns single root for linear session", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
const id2 = session.appendMessage(assistantMsg("2"));
|
||||
const id3 = session.appendMessage(userMsg("3"));
|
||||
|
||||
const tree = session.getTree();
|
||||
expect(tree).toHaveLength(1);
|
||||
|
||||
const root = tree[0];
|
||||
expect(root.entry.id).toBe(id1);
|
||||
expect(root.children).toHaveLength(1);
|
||||
expect(root.children[0].entry.id).toBe(id2);
|
||||
expect(root.children[0].children).toHaveLength(1);
|
||||
expect(root.children[0].children[0].entry.id).toBe(id3);
|
||||
expect(root.children[0].children[0].children).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("returns tree with branches after branch", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
// Build: 1 -> 2 -> 3
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
const id2 = session.appendMessage(assistantMsg("2"));
|
||||
const id3 = session.appendMessage(userMsg("3"));
|
||||
|
||||
// Branch from id2, add new path: 2 -> 4
|
||||
session.branch(id2);
|
||||
const id4 = session.appendMessage(userMsg("4-branch"));
|
||||
|
||||
const tree = session.getTree();
|
||||
expect(tree).toHaveLength(1);
|
||||
|
||||
const root = tree[0];
|
||||
expect(root.entry.id).toBe(id1);
|
||||
expect(root.children).toHaveLength(1);
|
||||
|
||||
const node2 = root.children[0];
|
||||
expect(node2.entry.id).toBe(id2);
|
||||
expect(node2.children).toHaveLength(2); // id3 and id4 are siblings
|
||||
|
||||
const childIds = node2.children.map((c) => c.entry.id).sort();
|
||||
expect(childIds).toEqual([id3, id4].sort());
|
||||
});
|
||||
|
||||
it("handles multiple branches at same point", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const _id1 = session.appendMessage(userMsg("root"));
|
||||
const id2 = session.appendMessage(assistantMsg("response"));
|
||||
|
||||
// Branch A
|
||||
session.branch(id2);
|
||||
const idA = session.appendMessage(userMsg("branch-A"));
|
||||
|
||||
// Branch B
|
||||
session.branch(id2);
|
||||
const idB = session.appendMessage(userMsg("branch-B"));
|
||||
|
||||
// Branch C
|
||||
session.branch(id2);
|
||||
const idC = session.appendMessage(userMsg("branch-C"));
|
||||
|
||||
const tree = session.getTree();
|
||||
const node2 = tree[0].children[0];
|
||||
expect(node2.entry.id).toBe(id2);
|
||||
expect(node2.children).toHaveLength(3);
|
||||
|
||||
const branchIds = node2.children.map((c) => c.entry.id).sort();
|
||||
expect(branchIds).toEqual([idA, idB, idC].sort());
|
||||
});
|
||||
|
||||
it("handles deep branching", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
// Main path: 1 -> 2 -> 3 -> 4
|
||||
const _id1 = session.appendMessage(userMsg("1"));
|
||||
const id2 = session.appendMessage(assistantMsg("2"));
|
||||
const id3 = session.appendMessage(userMsg("3"));
|
||||
const _id4 = session.appendMessage(assistantMsg("4"));
|
||||
|
||||
// Branch from 2: 2 -> 5 -> 6
|
||||
session.branch(id2);
|
||||
const id5 = session.appendMessage(userMsg("5"));
|
||||
const _id6 = session.appendMessage(assistantMsg("6"));
|
||||
|
||||
// Branch from 5: 5 -> 7
|
||||
session.branch(id5);
|
||||
const _id7 = session.appendMessage(userMsg("7"));
|
||||
|
||||
const tree = session.getTree();
|
||||
|
||||
// Verify structure
|
||||
const node2 = tree[0].children[0];
|
||||
expect(node2.children).toHaveLength(2); // id3 and id5
|
||||
|
||||
const node5 = node2.children.find((c) => c.entry.id === id5)!;
|
||||
expect(node5.children).toHaveLength(2); // id6 and id7
|
||||
|
||||
const node3 = node2.children.find((c) => c.entry.id === id3)!;
|
||||
expect(node3.children).toHaveLength(1); // id4
|
||||
});
|
||||
});
|
||||
|
||||
describe("branch", () => {
|
||||
it("moves leaf pointer to specified entry", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
const _id2 = session.appendMessage(assistantMsg("2"));
|
||||
const id3 = session.appendMessage(userMsg("3"));
|
||||
|
||||
expect(session.getLeafUuid()).toBe(id3);
|
||||
|
||||
session.branch(id1);
|
||||
expect(session.getLeafUuid()).toBe(id1);
|
||||
});
|
||||
|
||||
it("throws for non-existent entry", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
session.appendMessage(userMsg("hello"));
|
||||
|
||||
expect(() => session.branch("nonexistent")).toThrow("Entry nonexistent not found");
|
||||
});
|
||||
|
||||
it("new appends become children of branch point", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
const _id2 = session.appendMessage(assistantMsg("2"));
|
||||
|
||||
session.branch(id1);
|
||||
const id3 = session.appendMessage(userMsg("branched"));
|
||||
|
||||
const entries = session.getEntries();
|
||||
const branchedEntry = entries.find((e) => e.id === id3)!;
|
||||
expect(branchedEntry.parentId).toBe(id1); // sibling of id2
|
||||
});
|
||||
});
|
||||
|
||||
describe("branchWithSummary", () => {
|
||||
it("inserts branch summary and advances leaf", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
const _id2 = session.appendMessage(assistantMsg("2"));
|
||||
const _id3 = session.appendMessage(userMsg("3"));
|
||||
|
||||
const summaryId = session.branchWithSummary(id1, "Summary of abandoned work");
|
||||
|
||||
expect(session.getLeafUuid()).toBe(summaryId);
|
||||
|
||||
const entries = session.getEntries();
|
||||
const summaryEntry = entries.find((e) => e.type === "branch_summary");
|
||||
expect(summaryEntry).toBeDefined();
|
||||
expect(summaryEntry?.parentId).toBe(id1);
|
||||
if (summaryEntry?.type === "branch_summary") {
|
||||
expect(summaryEntry.summary).toBe("Summary of abandoned work");
|
||||
}
|
||||
});
|
||||
|
||||
it("throws for non-existent entry", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
session.appendMessage(userMsg("hello"));
|
||||
|
||||
expect(() => session.branchWithSummary("nonexistent", "summary")).toThrow("Entry nonexistent not found");
|
||||
});
|
||||
});
|
||||
|
||||
describe("getLeafEntry", () => {
|
||||
it("returns undefined for empty session", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
expect(session.getLeafEntry()).toBeUndefined();
|
||||
});
|
||||
|
||||
it("returns current leaf entry", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
session.appendMessage(userMsg("1"));
|
||||
const id2 = session.appendMessage(assistantMsg("2"));
|
||||
|
||||
const leaf = session.getLeafEntry();
|
||||
expect(leaf).toBeDefined();
|
||||
expect(leaf!.id).toBe(id2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("getEntry", () => {
|
||||
it("returns undefined for non-existent id", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
expect(session.getEntry("nonexistent")).toBeUndefined();
|
||||
});
|
||||
|
||||
it("returns entry by id", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
const id1 = session.appendMessage(userMsg("first"));
|
||||
const id2 = session.appendMessage(assistantMsg("second"));
|
||||
|
||||
const entry1 = session.getEntry(id1);
|
||||
expect(entry1).toBeDefined();
|
||||
expect(entry1?.type).toBe("message");
|
||||
if (entry1?.type === "message" && entry1.message.role === "user") {
|
||||
expect(entry1.message.content).toBe("first");
|
||||
}
|
||||
|
||||
const entry2 = session.getEntry(id2);
|
||||
expect(entry2).toBeDefined();
|
||||
if (entry2?.type === "message" && entry2.message.role === "assistant") {
|
||||
expect((entry2.message.content as any)[0].text).toBe("second");
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildSessionContext with branches", () => {
|
||||
it("returns messages from current branch only", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
// Main: 1 -> 2 -> 3
|
||||
session.appendMessage(userMsg("msg1"));
|
||||
const id2 = session.appendMessage(assistantMsg("msg2"));
|
||||
session.appendMessage(userMsg("msg3"));
|
||||
|
||||
// Branch from 2: 2 -> 4
|
||||
session.branch(id2);
|
||||
session.appendMessage(assistantMsg("msg4-branch"));
|
||||
|
||||
const ctx = session.buildSessionContext();
|
||||
expect(ctx.messages).toHaveLength(3); // msg1, msg2, msg4-branch (not msg3)
|
||||
|
||||
expect((ctx.messages[0] as any).content).toBe("msg1");
|
||||
expect((ctx.messages[1] as any).content[0].text).toBe("msg2");
|
||||
expect((ctx.messages[2] as any).content[0].text).toBe("msg4-branch");
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("createBranchedSession", () => {
|
||||
it("throws for non-existent entry", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
session.appendMessage(userMsg("hello"));
|
||||
|
||||
expect(() => session.createBranchedSession("nonexistent")).toThrow("Entry nonexistent not found");
|
||||
});
|
||||
|
||||
it("creates new session with path to specified leaf (in-memory)", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
// Build: 1 -> 2 -> 3 -> 4
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
const id2 = session.appendMessage(assistantMsg("2"));
|
||||
const id3 = session.appendMessage(userMsg("3"));
|
||||
session.appendMessage(assistantMsg("4"));
|
||||
|
||||
// Branch from 3: 3 -> 5
|
||||
session.branch(id3);
|
||||
const _id5 = session.appendMessage(userMsg("5"));
|
||||
|
||||
// Create branched session from id2 (should only have 1 -> 2)
|
||||
const result = session.createBranchedSession(id2);
|
||||
expect(result).toBeNull(); // in-memory returns null
|
||||
|
||||
// Session should now only have entries 1 and 2
|
||||
const entries = session.getEntries();
|
||||
expect(entries).toHaveLength(2);
|
||||
expect(entries[0].id).toBe(id1);
|
||||
expect(entries[1].id).toBe(id2);
|
||||
});
|
||||
|
||||
it("extracts correct path from branched tree", () => {
|
||||
const session = SessionManager.inMemory();
|
||||
|
||||
// Build: 1 -> 2 -> 3
|
||||
const id1 = session.appendMessage(userMsg("1"));
|
||||
const id2 = session.appendMessage(assistantMsg("2"));
|
||||
session.appendMessage(userMsg("3"));
|
||||
|
||||
// Branch from 2: 2 -> 4 -> 5
|
||||
session.branch(id2);
|
||||
const id4 = session.appendMessage(userMsg("4"));
|
||||
const id5 = session.appendMessage(assistantMsg("5"));
|
||||
|
||||
// Create branched session from id5 (should have 1 -> 2 -> 4 -> 5)
|
||||
session.createBranchedSession(id5);
|
||||
|
||||
const entries = session.getEntries();
|
||||
expect(entries).toHaveLength(4);
|
||||
expect(entries.map((e) => e.id)).toEqual([id1, id2, id4, id5]);
|
||||
});
|
||||
});
|
||||
|
|
@ -441,7 +441,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
|
|||
});
|
||||
|
||||
// Load existing messages
|
||||
const loadedSession = sessionManager.loadSession();
|
||||
const loadedSession = sessionManager.buildSessionContex();
|
||||
if (loadedSession.messages.length > 0) {
|
||||
agent.replaceMessages(loadedSession.messages);
|
||||
log.logInfo(`[${channelId}] Loaded ${loadedSession.messages.length} messages from context.jsonl`);
|
||||
|
|
@ -628,7 +628,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
|
|||
|
||||
// Reload messages from context.jsonl
|
||||
// This picks up any messages synced from log.jsonl before this run
|
||||
const reloadedSession = sessionManager.loadSession();
|
||||
const reloadedSession = sessionManager.buildSessionContex();
|
||||
if (reloadedSession.messages.length > 0) {
|
||||
agent.replaceMessages(reloadedSession.messages);
|
||||
log.logInfo(`[${channelId}] Reloaded ${reloadedSession.messages.length} messages from context`);
|
||||
|
|
|
|||
|
|
@ -15,15 +15,12 @@ import {
|
|||
buildSessionContext,
|
||||
type CompactionEntry,
|
||||
type FileEntry,
|
||||
type LoadedSession,
|
||||
type MessageContent,
|
||||
type ModelChangeContent,
|
||||
type ModelChangeEntry,
|
||||
type SessionContext,
|
||||
type SessionEntry,
|
||||
type SessionEntryBase,
|
||||
type SessionMessageEntry,
|
||||
type ThinkingLevelChangeEntry,
|
||||
type ThinkingLevelContent,
|
||||
type TreeNode,
|
||||
} from "@mariozechner/pi-coding-agent";
|
||||
import { randomBytes } from "crypto";
|
||||
import { appendFileSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs";
|
||||
|
|
@ -98,15 +95,15 @@ export class MomSessionManager {
|
|||
this.leafId = null;
|
||||
}
|
||||
|
||||
private _createTreeNode(): Omit<TreeNode, "type"> {
|
||||
private _createEntryBase(): Omit<SessionEntryBase, "type"> {
|
||||
const id = uuidv4();
|
||||
const node = {
|
||||
const base = {
|
||||
id,
|
||||
parentId: this.leafId,
|
||||
timestamp: new Date().toISOString(),
|
||||
};
|
||||
this.leafId = id;
|
||||
return node;
|
||||
return base;
|
||||
}
|
||||
|
||||
private _persist(entry: SessionEntry): void {
|
||||
|
|
@ -281,22 +278,23 @@ export class MomSessionManager {
|
|||
}
|
||||
|
||||
saveMessage(message: AppMessage): void {
|
||||
const content: MessageContent = { type: "message", message };
|
||||
const entry: SessionMessageEntry = { ...this._createTreeNode(), ...content };
|
||||
const entry: SessionMessageEntry = { ...this._createEntryBase(), type: "message", message };
|
||||
this.inMemoryEntries.push(entry);
|
||||
this._persist(entry);
|
||||
}
|
||||
|
||||
saveThinkingLevelChange(thinkingLevel: string): void {
|
||||
const content: ThinkingLevelContent = { type: "thinking_level_change", thinkingLevel };
|
||||
const entry: ThinkingLevelChangeEntry = { ...this._createTreeNode(), ...content };
|
||||
const entry: ThinkingLevelChangeEntry = {
|
||||
...this._createEntryBase(),
|
||||
type: "thinking_level_change",
|
||||
thinkingLevel,
|
||||
};
|
||||
this.inMemoryEntries.push(entry);
|
||||
this._persist(entry);
|
||||
}
|
||||
|
||||
saveModelChange(provider: string, modelId: string): void {
|
||||
const content: ModelChangeContent = { type: "model_change", provider, modelId };
|
||||
const entry: ModelChangeEntry = { ...this._createTreeNode(), ...content };
|
||||
const entry: ModelChangeEntry = { ...this._createEntryBase(), type: "model_change", provider, modelId };
|
||||
this.inMemoryEntries.push(entry);
|
||||
this._persist(entry);
|
||||
}
|
||||
|
|
@ -307,7 +305,7 @@ export class MomSessionManager {
|
|||
}
|
||||
|
||||
/** Load session with compaction support */
|
||||
loadSession(): LoadedSession {
|
||||
buildSessionContex(): SessionContext {
|
||||
const entries = this.loadEntries();
|
||||
return buildSessionContext(entries);
|
||||
}
|
||||
|
|
@ -354,15 +352,15 @@ export class MomSessionManager {
|
|||
}
|
||||
|
||||
loadModel(): { provider: string; modelId: string } | null {
|
||||
return this.loadSession().model;
|
||||
return this.buildSessionContex().model;
|
||||
}
|
||||
|
||||
loadThinkingLevel(): string {
|
||||
return this.loadSession().thinkingLevel;
|
||||
return this.buildSessionContex().thinkingLevel;
|
||||
}
|
||||
|
||||
/** Not used by mom but required by AgentSession interface */
|
||||
createBranchedSessionFromEntries(_entries: SessionEntry[], _branchBeforeIndex: number): string | null {
|
||||
createBranchedSession(_leafId: string): string | null {
|
||||
return null; // Mom doesn't support branching
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue