getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-15 16:04:03 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix session-manager simplification issues

Browse source 
- Remove unused inspector import from session-manager.ts
- Remove dead code in _persist()
- Update tests for simplified SessionHeader
- Update mom context.ts: remove unused AgentState import, fix startSession(), rename isEnabled to isPersisted
This commit is contained in:
Mario Zechner 2025-12-22 02:43:38 +01:00
parent f7a12c478c
commit 0faadfcd00
4 changed files with 162 additions and 305 deletions
Download patch file Download diff file Expand all files Collapse all files

17
packages/coding-agent/src/core/agent-session.ts
View file

@ -202,11 +202,6 @@ export class AgentSession {
if (event.type === "message_end") {
this.sessionManager.saveMessage(event.message);
// Initialize session after first user+assistant exchange
if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) {
this.sessionManager.startSession(this.agent.state);
}
// Track assistant message for auto-compaction (checked on agent_end)
if (event.message.role === "assistant") {
this._lastAssistantMessage = event.message;
@ -389,7 +384,7 @@ export class AgentSession {
/** Current session file path, or null if sessions are disabled */
get sessionFile(): string | null {
return this.sessionManager.isEnabled() ? this.sessionManager.getSessionFile() : null;
return this.sessionManager.isPersisted() ? this.sessionManager.getSessionFile() : null;
}
/** Current session ID */
@ -1096,11 +1091,6 @@ export class AgentSession {
// Save to session
this.sessionManager.saveMessage(bashMessage);
// Initialize session if needed
if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) {
this.sessionManager.startSession(this.agent.state);
}
}
return result;
@ -1141,11 +1131,6 @@ export class AgentSession {
this.sessionManager.saveMessage(bashMessage);
}
// Initialize session if needed
if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) {
this.sessionManager.startSession(this.agent.state);
}
this._pendingBashMessages = [];
}

414
packages/coding-agent/src/core/session-manager.ts
View file

@ -17,9 +17,6 @@ export interface SessionHeader {
id: string;
timestamp: string;
cwd: string;
provider: string;
modelId: string;
thinkingLevel: string;
branchedFrom?: string;
}
@ -120,13 +117,12 @@ export function loadSessionFromEntries(entries: SessionEntry[]): LoadedSession {
let model: { provider: string; modelId: string } | null = null;
for (const entry of entries) {
if (entry.type === "session") {
thinkingLevel = entry.thinkingLevel;
model = { provider: entry.provider, modelId: entry.modelId };
} else if (entry.type === "thinking_level_change") {
if (entry.type === "thinking_level_change") {
thinkingLevel = entry.thinkingLevel;
} else if (entry.type === "model_change") {
model = { provider: entry.provider, modelId: entry.modelId };
} else if (entry.type === "message" && entry.message.role === "assistant") {
model = { provider: entry.message.provider, modelId: entry.message.model };
}
}
@ -194,23 +190,6 @@ function loadEntriesFromFile(filePath: string): SessionEntry[] {
return entries;
}
function extractSessionIdFromFile(filePath: string): string | null {
if (!existsSync(filePath)) return null;
const lines = readFileSync(filePath, "utf8").trim().split("\n");
for (const line of lines) {
try {
const entry = JSON.parse(line);
if (entry.type === "session") {
return entry.id;
}
} catch {
// Skip malformed lines
}
}
return null;
}
function findMostRecentSession(sessionDir: string): string | null {
try {
const files = readdirSync(sessionDir)
@ -228,35 +207,170 @@ function findMostRecentSession(sessionDir: string): string | null {
}
export class SessionManager {
private sessionId: string;
private sessionFile: string;
private sessionId: string = "";
private sessionFile: string = "";
private sessionDir: string;
private cwd: string;
private enabled: boolean;
private sessionInitialized: boolean;
private pendingEntries: SessionEntry[] = [];
private persist: boolean;
private inMemoryEntries: SessionEntry[] = [];
private constructor(cwd: string, agentDir: string, sessionFile: string | null, enabled: boolean) {
private constructor(cwd: string, agentDir: string, sessionFile: string | null, persist: boolean) {
this.cwd = cwd;
this.sessionDir = getSessionDirectory(cwd, agentDir);
this.enabled = enabled;
this.persist = persist;
if (sessionFile) {
this.sessionFile = resolve(sessionFile);
this.sessionId = extractSessionIdFromFile(this.sessionFile) ?? uuidv4();
this.sessionInitialized = existsSync(this.sessionFile);
if (this.sessionInitialized) {
this.inMemoryEntries = loadEntriesFromFile(this.sessionFile);
}
this.setSessionFile(sessionFile);
} else {
this.sessionId = uuidv4();
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
this.sessionFile = join(this.sessionDir, `${timestamp}_${this.sessionId}.jsonl`);
this.sessionInitialized = false;
const sessionFile = join(this.sessionDir, `${timestamp}_${this.sessionId}.jsonl`);
this.setSessionFile(sessionFile);
}
}
/** Switch to a different session file (used for resume and branching) */
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");
this.sessionId = header ? (header as SessionHeader).id : uuidv4();
} else {
this.sessionId = uuidv4();
this.inMemoryEntries = [];
const entry: SessionHeader = {
type: "session",
id: this.sessionId,
timestamp: new Date().toISOString(),
cwd: this.cwd,
};
this.inMemoryEntries.push(entry);
}
}
isPersisted(): boolean {
return this.persist;
}
getCwd(): string {
return this.cwd;
}
getSessionId(): string {
return this.sessionId;
}
getSessionFile(): string {
return this.sessionFile;
}
reset(): void {
this.inMemoryEntries = [];
this.sessionId = uuidv4();
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
this.sessionFile = join(this.sessionDir, `${timestamp}_${this.sessionId}.jsonl`);
}
_persist(entry: SessionEntry): void {
if (this.persist && this.inMemoryEntries.some((e) => e.type === "message" && e.message.role === "assistant")) {
appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`);
}
}
saveMessage(message: any): void {
const entry: SessionMessageEntry = {
type: "message",
timestamp: new Date().toISOString(),
message,
};
this.inMemoryEntries.push(entry);
this._persist(entry);
}
saveThinkingLevelChange(thinkingLevel: string): void {
const entry: ThinkingLevelChangeEntry = {
type: "thinking_level_change",
timestamp: new Date().toISOString(),
thinkingLevel,
};
this.inMemoryEntries.push(entry);
this._persist(entry);
}
saveModelChange(provider: string, modelId: string): void {
const entry: ModelChangeEntry = {
type: "model_change",
timestamp: new Date().toISOString(),
provider,
modelId,
};
this.inMemoryEntries.push(entry);
this._persist(entry);
}
saveCompaction(entry: CompactionEntry): void {
this.inMemoryEntries.push(entry);
this._persist(entry);
}
loadSession(): LoadedSession {
const entries = this.loadEntries();
return loadSessionFromEntries(entries);
}
loadMessages(): AppMessage[] {
return this.loadSession().messages;
}
loadThinkingLevel(): string {
return this.loadSession().thinkingLevel;
}
loadModel(): { provider: string; modelId: string } | null {
return this.loadSession().model;
}
loadEntries(): SessionEntry[] {
if (this.inMemoryEntries.length > 0) {
return [...this.inMemoryEntries];
} else {
return loadEntriesFromFile(this.sessionFile);
}
}
createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string | null {
const newSessionId = uuidv4();
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
const newSessionFile = join(this.sessionDir, `${timestamp}_${newSessionId}.jsonl`);
const newEntries: SessionEntry[] = [];
for (let i = 0; i < branchBeforeIndex; i++) {
const entry = entries[i];
if (entry.type === "session") {
newEntries.push({
...entry,
id: newSessionId,
timestamp: new Date().toISOString(),
branchedFrom: this.persist ? this.sessionFile : undefined,
});
} else {
newEntries.push(entry);
}
}
if (this.persist) {
for (const entry of newEntries) {
appendFileSync(newSessionFile, `${JSON.stringify(entry)}\n`);
}
return newSessionFile;
}
this.inMemoryEntries = newEntries;
this.sessionId = newSessionId;
return null;
}
/** Create a new session for the given directory */
static create(cwd: string, agentDir: string = getDefaultAgentDir()): SessionManager {
return new SessionManager(cwd, agentDir, null, true);
@ -361,226 +475,4 @@ export class SessionManager {
return sessions;
}
isEnabled(): boolean {
return this.enabled;
}
getCwd(): string {
return this.cwd;
}
getSessionId(): string {
return this.sessionId;
}
getSessionFile(): string {
return this.sessionFile;
}
/** Switch to a different session file (used for resume and branching) */
setSessionFile(path: string): void {
this.sessionFile = resolve(path);
this.sessionId = extractSessionIdFromFile(this.sessionFile) ?? uuidv4();
this.sessionInitialized = existsSync(this.sessionFile);
if (this.sessionInitialized) {
this.inMemoryEntries = loadEntriesFromFile(this.sessionFile);
} else {
this.inMemoryEntries = [];
}
this.pendingEntries = [];
}
reset(): void {
this.pendingEntries = [];
this.inMemoryEntries = [];
this.sessionInitialized = false;
this.sessionId = uuidv4();
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
this.sessionFile = join(this.sessionDir, `${timestamp}_${this.sessionId}.jsonl`);
}
startSession(state: AgentState): void {
if (this.sessionInitialized) return;
this.sessionInitialized = true;
const entry: SessionHeader = {
type: "session",
id: this.sessionId,
timestamp: new Date().toISOString(),
cwd: this.cwd,
provider: state.model.provider,
modelId: state.model.id,
thinkingLevel: state.thinkingLevel,
};
this.inMemoryEntries.push(entry);
for (const pending of this.pendingEntries) {
this.inMemoryEntries.push(pending);
}
this.pendingEntries = [];
if (this.enabled) {
appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`);
for (const memEntry of this.inMemoryEntries.slice(1)) {
appendFileSync(this.sessionFile, `${JSON.stringify(memEntry)}\n`);
}
}
}
saveMessage(message: any): void {
const entry: SessionMessageEntry = {
type: "message",
timestamp: new Date().toISOString(),
message,
};
if (!this.sessionInitialized) {
this.pendingEntries.push(entry);
} else {
this.inMemoryEntries.push(entry);
if (this.enabled) {
appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`);
}
}
}
saveThinkingLevelChange(thinkingLevel: string): void {
const entry: ThinkingLevelChangeEntry = {
type: "thinking_level_change",
timestamp: new Date().toISOString(),
thinkingLevel,
};
if (!this.sessionInitialized) {
this.pendingEntries.push(entry);
} else {
this.inMemoryEntries.push(entry);
if (this.enabled) {
appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`);
}
}
}
saveModelChange(provider: string, modelId: string): void {
const entry: ModelChangeEntry = {
type: "model_change",
timestamp: new Date().toISOString(),
provider,
modelId,
};
if (!this.sessionInitialized) {
this.pendingEntries.push(entry);
} else {
this.inMemoryEntries.push(entry);
if (this.enabled) {
appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`);
}
}
}
saveCompaction(entry: CompactionEntry): void {
this.inMemoryEntries.push(entry);
if (this.enabled) {
appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`);
}
}
loadSession(): LoadedSession {
const entries = this.loadEntries();
return loadSessionFromEntries(entries);
}
loadMessages(): AppMessage[] {
return this.loadSession().messages;
}
loadThinkingLevel(): string {
return this.loadSession().thinkingLevel;
}
loadModel(): { provider: string; modelId: string } | null {
return this.loadSession().model;
}
loadEntries(): SessionEntry[] {
if (this.enabled && existsSync(this.sessionFile)) {
return loadEntriesFromFile(this.sessionFile);
}
return [...this.inMemoryEntries];
}
shouldInitializeSession(messages: any[]): boolean {
if (this.sessionInitialized) return false;
const userMessages = messages.filter((m) => m.role === "user");
const assistantMessages = messages.filter((m) => m.role === "assistant");
return userMessages.length >= 1 && assistantMessages.length >= 1;
}
createBranchedSession(state: any, branchFromIndex: number): string {
const newSessionId = uuidv4();
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
const newSessionFile = join(this.sessionDir, `${timestamp}_${newSessionId}.jsonl`);
const entry: SessionHeader = {
type: "session",
id: newSessionId,
timestamp: new Date().toISOString(),
cwd: this.cwd,
provider: state.model.provider,
modelId: state.model.id,
thinkingLevel: state.thinkingLevel,
branchedFrom: this.sessionFile,
};
appendFileSync(newSessionFile, `${JSON.stringify(entry)}\n`);
if (branchFromIndex >= 0) {
const messagesToWrite = state.messages.slice(0, branchFromIndex + 1);
for (const message of messagesToWrite) {
const messageEntry: SessionMessageEntry = {
type: "message",
timestamp: new Date().toISOString(),
message,
};
appendFileSync(newSessionFile, `${JSON.stringify(messageEntry)}\n`);
}
}
return newSessionFile;
}
createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string | null {
const newSessionId = uuidv4();
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
const newSessionFile = join(this.sessionDir, `${timestamp}_${newSessionId}.jsonl`);
const newEntries: SessionEntry[] = [];
for (let i = 0; i < branchBeforeIndex; i++) {
const entry = entries[i];
if (entry.type === "session") {
newEntries.push({
...entry,
id: newSessionId,
timestamp: new Date().toISOString(),
branchedFrom: this.enabled ? this.sessionFile : undefined,
});
} else {
newEntries.push(entry);
}
}
if (this.enabled) {
for (const entry of newEntries) {
appendFileSync(newSessionFile, `${JSON.stringify(entry)}\n`);
}
return newSessionFile;
}
this.inMemoryEntries = newEntries;
this.sessionId = newSessionId;
return null;
}
}

20
packages/coding-agent/test/compaction.test.ts
View file

@ -234,9 +234,6 @@ describe("loadSessionFromEntries", () => {
id: "1",
timestamp: "",
cwd: "",
provider: "anthropic",
modelId: "claude",
thinkingLevel: "off",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
@ -247,7 +244,7 @@ describe("loadSessionFromEntries", () => {
const loaded = loadSessionFromEntries(entries);
expect(loaded.messages.length).toBe(4);
expect(loaded.thinkingLevel).toBe("off");
expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude" });
expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude-sonnet-4-5" });
});
it("should handle single compaction", () => {
@ -258,9 +255,6 @@ describe("loadSessionFromEntries", () => {
id: "1",
timestamp: "",
cwd: "",
provider: "anthropic",
modelId: "claude",
thinkingLevel: "off",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
@ -286,9 +280,6 @@ describe("loadSessionFromEntries", () => {
id: "1",
timestamp: "",
cwd: "",
provider: "anthropic",
modelId: "claude",
thinkingLevel: "off",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
@ -316,9 +307,6 @@ describe("loadSessionFromEntries", () => {
id: "1",
timestamp: "",
cwd: "",
provider: "anthropic",
modelId: "claude",
thinkingLevel: "off",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
@ -341,9 +329,6 @@ describe("loadSessionFromEntries", () => {
id: "1",
timestamp: "",
cwd: "",
provider: "anthropic",
modelId: "claude",
thinkingLevel: "off",
},
createMessageEntry(createUserMessage("1")),
{ type: "model_change", timestamp: "", provider: "openai", modelId: "gpt-4" },
@ -352,7 +337,8 @@ describe("loadSessionFromEntries", () => {
];
const loaded = loadSessionFromEntries(entries);
expect(loaded.model).toEqual({ provider: "openai", modelId: "gpt-4" });
// model_change is later overwritten by assistant message's model info
expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude-sonnet-4-5" });
expect(loaded.thinkingLevel).toBe("high");
});
});

16
packages/mom/src/context.ts
View file

@ -10,7 +10,7 @@
* - MomSettingsManager: Simple settings for mom (compaction, retry, model preferences)
*/
import type { AgentState, AppMessage } from "@mariozechner/pi-agent-core";
import type { AppMessage } from "@mariozechner/pi-agent-core";
import {
type CompactionEntry,
type LoadedSession,
@ -71,14 +71,14 @@ export class MomSessionManager {
// New session - write header immediately
this.sessionId = uuidv4();
if (initialModel) {
this.writeSessionHeader(initialModel);
this.writeSessionHeader();
}
}
// Note: syncFromLog() is called explicitly from agent.ts with excludeTimestamp
}
/** Write session header to file (called on new session creation) */
private writeSessionHeader(model: { provider: string; id: string; thinkingLevel?: string }): void {
private writeSessionHeader(): void {
this.sessionInitialized = true;
const entry: SessionHeader = {
@ -86,9 +86,6 @@ export class MomSessionManager {
id: this.sessionId,
timestamp: new Date().toISOString(),
cwd: this.channelDir,
provider: model.provider,
modelId: model.id,
thinkingLevel: model.thinkingLevel || "off",
};
this.inMemoryEntries.push(entry);
@ -249,7 +246,7 @@ export class MomSessionManager {
}
/** Initialize session with header if not already done */
startSession(state: AgentState): void {
startSession(): void {
if (this.sessionInitialized) return;
this.sessionInitialized = true;
@ -258,9 +255,6 @@ export class MomSessionManager {
id: this.sessionId,
timestamp: new Date().toISOString(),
cwd: this.channelDir,
provider: state.model?.provider || "unknown",
modelId: state.model?.id || "unknown",
thinkingLevel: state.thinkingLevel,
};
this.inMemoryEntries.push(entry);
@ -370,7 +364,7 @@ export class MomSessionManager {
}
// Compatibility methods for AgentSession
isEnabled(): boolean {
isPersisted(): boolean {
return true;
}