This commit is contained in:
vbozhkov 2026-01-30 14:09:01 +03:00 committed by GitHub
commit bf881d9e1b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 455 additions and 0 deletions

View File

@ -30,6 +30,7 @@ import {
} from "./agent-runner-helpers.js"; } from "./agent-runner-helpers.js";
import { runMemoryFlushIfNeeded } from "./agent-runner-memory.js"; import { runMemoryFlushIfNeeded } from "./agent-runner-memory.js";
import { buildReplyPayloads } from "./agent-runner-payloads.js"; import { buildReplyPayloads } from "./agent-runner-payloads.js";
import { checkContextRecoveryNeeded, recoverContext } from "./context-recovery.js";
import { appendUsageLine, formatResponseUsageLine } from "./agent-runner-utils.js"; import { appendUsageLine, formatResponseUsageLine } from "./agent-runner-utils.js";
import { createAudioAsVoiceBuffer, createBlockReplyPipeline } from "./block-reply-pipeline.js"; import { createAudioAsVoiceBuffer, createBlockReplyPipeline } from "./block-reply-pipeline.js";
import { resolveBlockStreamingCoalescing } from "./block-streaming.js"; import { resolveBlockStreamingCoalescing } from "./block-streaming.js";
@ -213,6 +214,42 @@ export async function runReplyAgent(params: {
isHeartbeat, isHeartbeat,
}); });
// Context recovery: inject recent messages after compaction
if (sessionKey && !isHeartbeat) {
const recoveryCheck = checkContextRecoveryNeeded({
cfg,
sessionEntry: activeSessionEntry,
});
if (recoveryCheck.needed) {
const recoveryResult = await recoverContext({
sessionKey,
messageCount: recoveryCheck.messageCount,
});
if (recoveryResult.ok && recoveryResult.contextBlock) {
// Prepend recovered context to extraSystemPrompt
const existingPrompt = followupRun.run.extraSystemPrompt;
followupRun.run.extraSystemPrompt = existingPrompt
? `${recoveryResult.contextBlock}\n\n${existingPrompt}`
: recoveryResult.contextBlock;
// Update session to track that we've recovered for this compaction
if (activeSessionEntry && activeSessionStore && storePath) {
activeSessionEntry.lastContextRecoveryCompactionCount = recoveryCheck.compactionCount;
activeSessionStore[sessionKey] = activeSessionEntry;
await updateSessionStoreEntry({
storePath,
sessionKey,
update: async () => ({
lastContextRecoveryCompactionCount: recoveryCheck.compactionCount,
}),
});
}
}
}
}
const runFollowupTurn = createFollowupRunner({ const runFollowupTurn = createFollowupRunner({
opts, opts,
typing, typing,

View File

@ -0,0 +1,175 @@
import { describe, expect, it } from "vitest";
import { checkContextRecoveryNeeded, resolveContextRecoverySettings } from "./context-recovery.js";
import type { MoltbotConfig } from "../../config/config.js";
import type { SessionEntry } from "../../config/sessions.js";
describe("context-recovery", () => {
describe("resolveContextRecoverySettings", () => {
it("returns disabled when not configured", () => {
const cfg = {} as MoltbotConfig;
const settings = resolveContextRecoverySettings(cfg);
expect(settings.enabled).toBe(false);
expect(settings.messages).toBe(0);
});
it("returns disabled when messages is 0", () => {
const cfg = {
agents: {
defaults: {
compaction: {
contextRecovery: { messages: 0 },
},
},
},
} as MoltbotConfig;
const settings = resolveContextRecoverySettings(cfg);
expect(settings.enabled).toBe(false);
expect(settings.messages).toBe(0);
});
it("returns enabled with correct message count", () => {
const cfg = {
agents: {
defaults: {
compaction: {
contextRecovery: { messages: 10 },
},
},
},
} as MoltbotConfig;
const settings = resolveContextRecoverySettings(cfg);
expect(settings.enabled).toBe(true);
expect(settings.messages).toBe(10);
});
it("caps message count at 50", () => {
const cfg = {
agents: {
defaults: {
compaction: {
contextRecovery: { messages: 100 },
},
},
},
} as MoltbotConfig;
const settings = resolveContextRecoverySettings(cfg);
expect(settings.enabled).toBe(true);
expect(settings.messages).toBe(50);
});
});
describe("checkContextRecoveryNeeded", () => {
it("returns not needed when disabled", () => {
const cfg = {} as MoltbotConfig;
const sessionEntry = {
sessionId: "test",
updatedAt: Date.now(),
compactionCount: 5,
} as SessionEntry;
const result = checkContextRecoveryNeeded({ cfg, sessionEntry });
expect(result.needed).toBe(false);
});
it("returns not needed when no session entry", () => {
const cfg = {
agents: {
defaults: {
compaction: {
contextRecovery: { messages: 10 },
},
},
},
} as MoltbotConfig;
const result = checkContextRecoveryNeeded({ cfg, sessionEntry: undefined });
expect(result.needed).toBe(false);
});
it("returns not needed when compaction count is 0", () => {
const cfg = {
agents: {
defaults: {
compaction: {
contextRecovery: { messages: 10 },
},
},
},
} as MoltbotConfig;
const sessionEntry = {
sessionId: "test",
updatedAt: Date.now(),
compactionCount: 0,
} as SessionEntry;
const result = checkContextRecoveryNeeded({ cfg, sessionEntry });
expect(result.needed).toBe(false);
});
it("returns not needed when already recovered for current compaction", () => {
const cfg = {
agents: {
defaults: {
compaction: {
contextRecovery: { messages: 10 },
},
},
},
} as MoltbotConfig;
const sessionEntry = {
sessionId: "test",
updatedAt: Date.now(),
compactionCount: 3,
lastContextRecoveryCompactionCount: 3,
} as SessionEntry;
const result = checkContextRecoveryNeeded({ cfg, sessionEntry });
expect(result.needed).toBe(false);
});
it("returns needed when compaction count exceeds last recovery", () => {
const cfg = {
agents: {
defaults: {
compaction: {
contextRecovery: { messages: 10 },
},
},
},
} as MoltbotConfig;
const sessionEntry = {
sessionId: "test",
updatedAt: Date.now(),
compactionCount: 3,
lastContextRecoveryCompactionCount: 2,
} as SessionEntry;
const result = checkContextRecoveryNeeded({ cfg, sessionEntry });
expect(result.needed).toBe(true);
expect(result.messageCount).toBe(10);
expect(result.compactionCount).toBe(3);
});
it("returns needed on first compaction when never recovered", () => {
const cfg = {
agents: {
defaults: {
compaction: {
contextRecovery: { messages: 5 },
},
},
},
} as MoltbotConfig;
const sessionEntry = {
sessionId: "test",
updatedAt: Date.now(),
compactionCount: 1,
} as SessionEntry;
const result = checkContextRecoveryNeeded({ cfg, sessionEntry });
expect(result.needed).toBe(true);
expect(result.messageCount).toBe(5);
expect(result.compactionCount).toBe(1);
});
});
});

View File

@ -0,0 +1,228 @@
/**
* Context recovery after compaction.
*
* When enabled via `compaction.contextRecovery.messages`, this module fetches recent
* messages from the session transcript after a compaction event and injects them
* as context for the agent's next turn.
*/
import type { MoltbotConfig } from "../../config/config.js";
import type { SessionEntry } from "../../config/sessions.js";
import { callGateway } from "../../gateway/call.js";
import { logVerbose } from "../../globals.js";
/** Result of checking whether context recovery is needed. */
export type ContextRecoveryCheck = {
/** Whether context recovery should be performed. */
needed: boolean;
/** Number of messages to recover. */
messageCount: number;
/** Current compaction count (for updating session after recovery). */
compactionCount: number;
};
/** Recovered message for context injection. */
export type RecoveredMessage = {
role: "user" | "assistant";
text: string;
timestamp?: string;
};
/** Result of context recovery. */
export type ContextRecoveryResult = {
/** Whether recovery was successful. */
ok: boolean;
/** Recovered messages (empty if none found or recovery failed). */
messages: RecoveredMessage[];
/** Formatted context block for system prompt injection. */
contextBlock: string | null;
/** Error message if recovery failed. */
error?: string;
};
/**
* Resolve the context recovery settings from config.
*/
export function resolveContextRecoverySettings(cfg: MoltbotConfig): {
enabled: boolean;
messages: number;
} {
const contextRecovery = cfg.agents?.defaults?.compaction?.contextRecovery;
const messages = contextRecovery?.messages ?? 0;
return {
enabled: messages > 0,
messages: Math.min(messages, 50), // Cap at 50 messages max
};
}
/**
* Check if context recovery is needed for a session.
*
* Recovery is needed when:
* 1. Context recovery is enabled in config
* 2. The session has been compacted (compactionCount > 0)
* 3. We haven't already recovered for this compaction (compactionCount > lastContextRecoveryCompactionCount)
*/
export function checkContextRecoveryNeeded(params: {
cfg: MoltbotConfig;
sessionEntry?: SessionEntry;
}): ContextRecoveryCheck {
const settings = resolveContextRecoverySettings(params.cfg);
if (!settings.enabled || !params.sessionEntry) {
return { needed: false, messageCount: 0, compactionCount: 0 };
}
const compactionCount = params.sessionEntry.compactionCount ?? 0;
const lastRecoveryAt = params.sessionEntry.lastContextRecoveryCompactionCount ?? 0;
// Recovery needed if we've had a compaction since last recovery
const needed = compactionCount > 0 && compactionCount > lastRecoveryAt;
return {
needed,
messageCount: settings.messages,
compactionCount,
};
}
/**
* Fetch recent messages from session history via Gateway.
*/
async function fetchSessionMessages(params: {
sessionKey: string;
limit: number;
}): Promise<RecoveredMessage[]> {
try {
const result = (await callGateway({
method: "chat.history",
params: {
sessionKey: params.sessionKey,
limit: params.limit + 10, // Fetch extra to account for system/tool messages
},
})) as { messages?: unknown[] };
const rawMessages = Array.isArray(result?.messages) ? result.messages : [];
// Filter and transform to RecoveredMessage format
const messages: RecoveredMessage[] = [];
for (const msg of rawMessages) {
if (!msg || typeof msg !== "object") continue;
const role = (msg as Record<string, unknown>).role;
if (role !== "user" && role !== "assistant") continue;
// Extract text content
const content = (msg as Record<string, unknown>).content;
let text = "";
if (typeof content === "string") {
text = content;
} else if (Array.isArray(content)) {
// Handle content blocks (e.g., [{type: "text", text: "..."}])
for (const block of content) {
if (
block &&
typeof block === "object" &&
(block as Record<string, unknown>).type === "text"
) {
const blockText = (block as Record<string, unknown>).text;
if (typeof blockText === "string") {
text += blockText;
}
}
}
}
if (!text.trim()) continue;
// Skip tool-related messages and system injections
if (text.startsWith("[tool:") || text.startsWith("[System Event]")) continue;
const timestamp = (msg as Record<string, unknown>).timestamp;
messages.push({
role: role as "user" | "assistant",
text: text.trim(),
timestamp: typeof timestamp === "string" ? timestamp : undefined,
});
}
// Return the most recent N messages
return messages.slice(-params.limit);
} catch (err) {
logVerbose(`Context recovery: failed to fetch session history: ${String(err)}`);
return [];
}
}
/**
* Format recovered messages as a context block for system prompt injection.
*/
function formatContextBlock(messages: RecoveredMessage[]): string | null {
if (messages.length === 0) return null;
const lines: string[] = [
"## Recent Conversation (Recovered After Context Compaction)",
"",
"The following is a summary of the recent conversation before compaction. Use this to maintain continuity:",
"",
];
for (const msg of messages) {
const roleLabel = msg.role === "user" ? "User" : "Assistant";
// Truncate very long messages to avoid bloating context
const text = msg.text.length > 500 ? `${msg.text.slice(0, 500)}...` : msg.text;
lines.push(`**${roleLabel}:** ${text}`);
lines.push("");
}
lines.push("---");
lines.push("");
return lines.join("\n");
}
/**
* Perform context recovery: fetch recent messages and format for injection.
*/
export async function recoverContext(params: {
sessionKey: string;
messageCount: number;
}): Promise<ContextRecoveryResult> {
try {
const messages = await fetchSessionMessages({
sessionKey: params.sessionKey,
limit: params.messageCount,
});
if (messages.length === 0) {
return {
ok: true,
messages: [],
contextBlock: null,
};
}
const contextBlock = formatContextBlock(messages);
logVerbose(
`Context recovery: recovered ${messages.length} messages for session ${params.sessionKey}`,
);
return {
ok: true,
messages,
contextBlock,
};
} catch (err) {
const error = err instanceof Error ? err.message : String(err);
logVerbose(`Context recovery failed: ${error}`);
return {
ok: false,
messages: [],
contextBlock: null,
error,
};
}
}

View File

@ -77,6 +77,8 @@ export type SessionEntry = {
compactionCount?: number; compactionCount?: number;
memoryFlushAt?: number; memoryFlushAt?: number;
memoryFlushCompactionCount?: number; memoryFlushCompactionCount?: number;
/** Compaction count at which context recovery was last performed. */
lastContextRecoveryCompactionCount?: number;
cliSessionIds?: Record<string, string>; cliSessionIds?: Record<string, string>;
claudeCliSessionId?: string; claudeCliSessionId?: string;
label?: string; label?: string;

View File

@ -248,6 +248,8 @@ export type AgentCompactionConfig = {
maxHistoryShare?: number; maxHistoryShare?: number;
/** Pre-compaction memory flush (agentic turn). Default: enabled. */ /** Pre-compaction memory flush (agentic turn). Default: enabled. */
memoryFlush?: AgentCompactionMemoryFlushConfig; memoryFlush?: AgentCompactionMemoryFlushConfig;
/** Auto-recover recent messages from session history after compaction. */
contextRecovery?: AgentCompactionContextRecoveryConfig;
}; };
export type AgentCompactionMemoryFlushConfig = { export type AgentCompactionMemoryFlushConfig = {
@ -260,3 +262,8 @@ export type AgentCompactionMemoryFlushConfig = {
/** System prompt appended for the memory flush turn. */ /** System prompt appended for the memory flush turn. */
systemPrompt?: string; systemPrompt?: string;
}; };
export type AgentCompactionContextRecoveryConfig = {
/** Number of recent messages to recover from session history after compaction (0 = disabled). Default: 0. */
messages?: number;
};

View File

@ -100,6 +100,12 @@ export const AgentDefaultsSchema = z
}) })
.strict() .strict()
.optional(), .optional(),
contextRecovery: z
.object({
messages: z.number().int().nonnegative().max(50).optional(),
})
.strict()
.optional(),
}) })
.strict() .strict()
.optional(), .optional(),