getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-16 12:03:23 +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 cwd: string;
private persist: boolean;
private flushed: boolean = false;
private inMemoryEntries: SessionEntry[] = [];
private constructor(cwd: string, agentDir: string, sessionFile: string | null, persist: boolean) {
@ -236,9 +237,11 @@ export class SessionManager {
this.inMemoryEntries = loadEntriesFromFile(this.sessionFile);
const header = this.inMemoryEntries.find((e) => e.type === "session");
this.sessionId = header ? (header as SessionHeader).id : uuidv4();
this.flushed = true;
} else {
this.sessionId = uuidv4();
this.inMemoryEntries = [];
this.flushed = false;
const entry: SessionHeader = {
type: "session",
id: this.sessionId,
@ -266,14 +269,32 @@ export class SessionManager {
}
reset(): void {
this.inMemoryEntries = [];
this.sessionId = uuidv4();
this.flushed = false;
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
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 {
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`);
}
}

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);
// Create session manager and settings manager
// Pass model info so new sessions get a header written immediately
const sessionManager = new MomSessionManager(channelDir, {
provider: model.provider,
id: model.id,
thinkingLevel: "off",
});
const sessionManager = new MomSessionManager(channelDir);
const settingsManager = new MomSettingsManager(join(channelDir, ".."));
// Create agent

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

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