diff --git a/src/auto-reply/reply/agent-runner.ts b/src/auto-reply/reply/agent-runner.ts index 227e6f17e..c78de2769 100644 --- a/src/auto-reply/reply/agent-runner.ts +++ b/src/auto-reply/reply/agent-runner.ts @@ -30,6 +30,7 @@ import { } from "./agent-runner-helpers.js"; import { runMemoryFlushIfNeeded } from "./agent-runner-memory.js"; import { buildReplyPayloads } from "./agent-runner-payloads.js"; +import { checkContextRecoveryNeeded, recoverContext } from "./context-recovery.js"; import { appendUsageLine, formatResponseUsageLine } from "./agent-runner-utils.js"; import { createAudioAsVoiceBuffer, createBlockReplyPipeline } from "./block-reply-pipeline.js"; import { resolveBlockStreamingCoalescing } from "./block-streaming.js"; @@ -213,6 +214,42 @@ export async function runReplyAgent(params: { 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({ opts, typing, diff --git a/src/auto-reply/reply/context-recovery.test.ts b/src/auto-reply/reply/context-recovery.test.ts new file mode 100644 index 000000000..b46bf1cb9 --- /dev/null +++ b/src/auto-reply/reply/context-recovery.test.ts @@ -0,0 +1,179 @@ +import { describe, expect, it, vi } from "vitest"; +import { + checkContextRecoveryNeeded, + resolveContextRecoverySettings, + type RecoveredMessage, +} 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); + }); + }); +}); diff --git a/src/auto-reply/reply/context-recovery.ts b/src/auto-reply/reply/context-recovery.ts new file mode 100644 index 000000000..cb2013504 --- /dev/null +++ b/src/auto-reply/reply/context-recovery.ts @@ -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 { + 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).role; + if (role !== "user" && role !== "assistant") continue; + + // Extract text content + const content = (msg as Record).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).type === "text" + ) { + const blockText = (block as Record).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).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 { + 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, + }; + } +} diff --git a/src/config/sessions/types.ts b/src/config/sessions/types.ts index 48ce428c1..02817873f 100644 --- a/src/config/sessions/types.ts +++ b/src/config/sessions/types.ts @@ -77,6 +77,8 @@ export type SessionEntry = { compactionCount?: number; memoryFlushAt?: number; memoryFlushCompactionCount?: number; + /** Compaction count at which context recovery was last performed. */ + lastContextRecoveryCompactionCount?: number; cliSessionIds?: Record; claudeCliSessionId?: string; label?: string; diff --git a/src/config/types.agent-defaults.ts b/src/config/types.agent-defaults.ts index 9c6ce0211..0de7c69a6 100644 --- a/src/config/types.agent-defaults.ts +++ b/src/config/types.agent-defaults.ts @@ -248,6 +248,8 @@ export type AgentCompactionConfig = { maxHistoryShare?: number; /** Pre-compaction memory flush (agentic turn). Default: enabled. */ memoryFlush?: AgentCompactionMemoryFlushConfig; + /** Auto-recover recent messages from session history after compaction. */ + contextRecovery?: AgentCompactionContextRecoveryConfig; }; export type AgentCompactionMemoryFlushConfig = { @@ -260,3 +262,8 @@ export type AgentCompactionMemoryFlushConfig = { /** System prompt appended for the memory flush turn. */ systemPrompt?: string; }; + +export type AgentCompactionContextRecoveryConfig = { + /** Number of recent messages to recover from session history after compaction (0 = disabled). Default: 0. */ + messages?: number; +}; diff --git a/src/config/zod-schema.agent-defaults.ts b/src/config/zod-schema.agent-defaults.ts index a849078ed..5fb8116f0 100644 --- a/src/config/zod-schema.agent-defaults.ts +++ b/src/config/zod-schema.agent-defaults.ts @@ -100,6 +100,12 @@ export const AgentDefaultsSchema = z }) .strict() .optional(), + contextRecovery: z + .object({ + messages: z.number().int().nonnegative().max(50).optional(), + }) + .strict() + .optional(), }) .strict() .optional(),