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

Fix session persistence: flush all buffered entries on first assistant message

Browse source 
- SessionManager: add flushed flag, _persist() flushes all entries when first assistant seen
- SessionManager: reset() now creates session header
- MomSessionManager: same pattern, remove pendingEntries/sessionInitialized/startSession/shouldInitializeSession
This commit is contained in:
Mario Zechner 2025-12-22 02:51:56 +01:00
parent 974c8f57e5
commit 184c648334
3 changed files with 62 additions and 86 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -212,6 +212,7 @@ export class SessionManager {
private sessionDir: string; private sessionDir: string;
private cwd: string; private cwd: string;
private persist: boolean; private persist: boolean;
private flushed: boolean = false;
private inMemoryEntries: SessionEntry[] = []; private inMemoryEntries: SessionEntry[] = [];
private constructor(cwd: string, agentDir: string, sessionFile: string | null, persist: boolean) { private constructor(cwd: string, agentDir: string, sessionFile: string | null, persist: boolean) {
@ -236,9 +237,11 @@ export class SessionManager {
this.inMemoryEntries = loadEntriesFromFile(this.sessionFile); this.inMemoryEntries = loadEntriesFromFile(this.sessionFile);
const header = this.inMemoryEntries.find((e) => e.type === "session"); const header = this.inMemoryEntries.find((e) => e.type === "session");
this.sessionId = header ? (header as SessionHeader).id : uuidv4(); this.sessionId = header ? (header as SessionHeader).id : uuidv4();
this.flushed = true;
} else { } else {
this.sessionId = uuidv4(); this.sessionId = uuidv4();
this.inMemoryEntries = []; this.inMemoryEntries = [];
this.flushed = false;
const entry: SessionHeader = { const entry: SessionHeader = {
type: "session", type: "session",
id: this.sessionId, id: this.sessionId,
@ -266,14 +269,32 @@ export class SessionManager {
} }
reset(): void { reset(): void {
this.inMemoryEntries = [];
this.sessionId = uuidv4(); this.sessionId = uuidv4();
this.flushed = false;
const timestamp = new Date().toISOString().replace(/[:.]/g, "-"); const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
this.sessionFile = join(this.sessionDir, `${timestamp}_${this.sessionId}.jsonl`); this.sessionFile = join(this.sessionDir, `${timestamp}_${this.sessionId}.jsonl`);
this.inMemoryEntries = [
{
type: "session",
id: this.sessionId,
timestamp: new Date().toISOString(),
cwd: this.cwd,
},
];
} }
_persist(entry: SessionEntry): void { _persist(entry: SessionEntry): void {
if (this.persist && this.inMemoryEntries.some((e) => e.type === "message" && e.message.role === "assistant")) { if (!this.persist) return;
const hasAssistant = this.inMemoryEntries.some((e) => e.type === "message" && e.message.role === "assistant");
if (!hasAssistant) return;
if (!this.flushed) {
for (const e of this.inMemoryEntries) {
appendFileSync(this.sessionFile, `${JSON.stringify(e)}\n`);
}
this.flushed = true;
} else {
appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`); appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`);
} }
} }

7
packages/mom/src/agent.ts
View file

@ -411,12 +411,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
const systemPrompt = buildSystemPrompt(workspacePath, channelId, memory, sandboxConfig, [], [], skills); const systemPrompt = buildSystemPrompt(workspacePath, channelId, memory, sandboxConfig, [], [], skills);
// Create session manager and settings manager // Create session manager and settings manager
// Pass model info so new sessions get a header written immediately const sessionManager = new MomSessionManager(channelDir);
const sessionManager = new MomSessionManager(channelDir, {
provider: model.provider,
id: model.id,
thinkingLevel: "off",
});
const settingsManager = new MomSettingsManager(join(channelDir, "..")); const settingsManager = new MomSettingsManager(join(channelDir, ".."));
// Create agent // Create agent

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

@ -48,11 +48,10 @@ export class MomSessionManager {
private contextFile: string; private contextFile: string;
private logFile: string; private logFile: string;
private channelDir: string; private channelDir: string;
private sessionInitialized: boolean = false; private flushed: boolean = false;
private inMemoryEntries: SessionEntry[] = []; private inMemoryEntries: SessionEntry[] = [];
private pendingEntries: SessionEntry[] = [];
constructor(channelDir: string, initialModel?: { provider: string; id: string; thinkingLevel?: string }) { constructor(channelDir: string) {
this.channelDir = channelDir; this.channelDir = channelDir;
this.contextFile = join(channelDir, "context.jsonl"); this.contextFile = join(channelDir, "context.jsonl");
this.logFile = join(channelDir, "log.jsonl"); this.logFile = join(channelDir, "log.jsonl");
@ -66,30 +65,33 @@ export class MomSessionManager {
if (existsSync(this.contextFile)) { if (existsSync(this.contextFile)) {
this.inMemoryEntries = this.loadEntriesFromFile(); this.inMemoryEntries = this.loadEntriesFromFile();
this.sessionId = this.extractSessionId() || uuidv4(); this.sessionId = this.extractSessionId() || uuidv4();
this.sessionInitialized = this.inMemoryEntries.length > 0; this.flushed = true;
} else { } else {
// New session - write header immediately
this.sessionId = uuidv4(); this.sessionId = uuidv4();
if (initialModel) { this.inMemoryEntries = [
this.writeSessionHeader(); {
} type: "session",
id: this.sessionId,
timestamp: new Date().toISOString(),
cwd: this.channelDir,
},
];
} }
// Note: syncFromLog() is called explicitly from agent.ts with excludeTimestamp // Note: syncFromLog() is called explicitly from agent.ts with excludeTimestamp
} }
/** Write session header to file (called on new session creation) */ private _persist(entry: SessionEntry): void {
private writeSessionHeader(): void { const hasAssistant = this.inMemoryEntries.some((e) => e.type === "message" && e.message.role === "assistant");
this.sessionInitialized = true; if (!hasAssistant) return;
const entry: SessionHeader = { if (!this.flushed) {
type: "session", for (const e of this.inMemoryEntries) {
id: this.sessionId, appendFileSync(this.contextFile, `${JSON.stringify(e)}\n`);
timestamp: new Date().toISOString(), }
cwd: this.channelDir, this.flushed = true;
}; } else {
appendFileSync(this.contextFile, `${JSON.stringify(entry)}\n`);
this.inMemoryEntries.push(entry); }
appendFileSync(this.contextFile, `${JSON.stringify(entry)}\n`);
} }
/** /**
@ -245,44 +247,14 @@ export class MomSessionManager {
return entries; return entries;
} }
/** Initialize session with header if not already done */
startSession(): void {
if (this.sessionInitialized) return;
this.sessionInitialized = true;
const entry: SessionHeader = {
type: "session",
id: this.sessionId,
timestamp: new Date().toISOString(),
cwd: this.channelDir,
};
this.inMemoryEntries.push(entry);
for (const pending of this.pendingEntries) {
this.inMemoryEntries.push(pending);
}
this.pendingEntries = [];
// Write to file
appendFileSync(this.contextFile, `${JSON.stringify(entry)}\n`);
for (const memEntry of this.inMemoryEntries.slice(1)) {
appendFileSync(this.contextFile, `${JSON.stringify(memEntry)}\n`);
}
}
saveMessage(message: AppMessage): void { saveMessage(message: AppMessage): void {
const entry: SessionMessageEntry = { const entry: SessionMessageEntry = {
type: "message", type: "message",
timestamp: new Date().toISOString(), timestamp: new Date().toISOString(),
message, message,
}; };
this.inMemoryEntries.push(entry);
if (!this.sessionInitialized) { this._persist(entry);
this.pendingEntries.push(entry);
} else {
this.inMemoryEntries.push(entry);
appendFileSync(this.contextFile, `${JSON.stringify(entry)}\n`);
}
} }
saveThinkingLevelChange(thinkingLevel: string): void { saveThinkingLevelChange(thinkingLevel: string): void {
@ -291,13 +263,8 @@ export class MomSessionManager {
timestamp: new Date().toISOString(), timestamp: new Date().toISOString(),
thinkingLevel, thinkingLevel,
}; };
this.inMemoryEntries.push(entry);
if (!this.sessionInitialized) { this._persist(entry);
this.pendingEntries.push(entry);
} else {
this.inMemoryEntries.push(entry);
appendFileSync(this.contextFile, `${JSON.stringify(entry)}\n`);
}
} }
saveModelChange(provider: string, modelId: string): void { saveModelChange(provider: string, modelId: string): void {
@ -307,18 +274,13 @@ export class MomSessionManager {
provider, provider,
modelId, modelId,
}; };
this.inMemoryEntries.push(entry);
if (!this.sessionInitialized) { this._persist(entry);
this.pendingEntries.push(entry);
} else {
this.inMemoryEntries.push(entry);
appendFileSync(this.contextFile, `${JSON.stringify(entry)}\n`);
}
} }
saveCompaction(entry: CompactionEntry): void { saveCompaction(entry: CompactionEntry): void {
this.inMemoryEntries.push(entry); this.inMemoryEntries.push(entry);
appendFileSync(this.contextFile, `${JSON.stringify(entry)}\n`); this._persist(entry);
} }
/** Load session with compaction support */ /** Load session with compaction support */
@ -343,20 +305,18 @@ export class MomSessionManager {
return this.contextFile; return this.contextFile;
} }
/** Check if session should be initialized */
shouldInitializeSession(messages: AppMessage[]): 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;
}
/** Reset session (clears context.jsonl) */ /** Reset session (clears context.jsonl) */
reset(): void { reset(): void {
this.pendingEntries = [];
this.inMemoryEntries = [];
this.sessionInitialized = false;
this.sessionId = uuidv4(); this.sessionId = uuidv4();
this.flushed = false;
this.inMemoryEntries = [
{
type: "session",
id: this.sessionId,
timestamp: new Date().toISOString(),
cwd: this.channelDir,
},
];
// Truncate the context file // Truncate the context file
if (existsSync(this.contextFile)) { if (existsSync(this.contextFile)) {
writeFileSync(this.contextFile, ""); writeFileSync(this.contextFile, "");