From 3ae28c6ac32f8b2d4c198e1dc3e7d5e9d3201831 Mon Sep 17 00:00:00 2001 From: s4na Date: Thu, 29 Jan 2026 23:50:03 +0900 Subject: [PATCH] feat(agents): add session_compact tool for agent-initiated context compaction Add a tool that allows agents to programmatically trigger context compaction without requiring user intervention via `/compact`. Currently agents have no way to proactively manage their context window before memory-intensive tasks, and automatic compaction doesn't allow custom instructions for what to preserve in the summary. - Add `session_compact` tool with optional `instructions` and `sessionKey` params - Support cross-session compaction with A2A policy enforcement - Disable `bashElevated` during compaction for safety - Add 7 unit tests covering success, failure, and policy scenarios --- .../moltbot-tools.session-compact.test.ts | 262 ++++++++++++++++++ src/agents/moltbot-tools.ts | 5 + src/agents/system-prompt.ts | 3 + src/agents/tools/session-compact-tool.ts | 247 +++++++++++++++++ 4 files changed, 517 insertions(+) create mode 100644 src/agents/moltbot-tools.session-compact.test.ts create mode 100644 src/agents/tools/session-compact-tool.ts diff --git a/src/agents/moltbot-tools.session-compact.test.ts b/src/agents/moltbot-tools.session-compact.test.ts new file mode 100644 index 000000000..bdd3e34ae --- /dev/null +++ b/src/agents/moltbot-tools.session-compact.test.ts @@ -0,0 +1,262 @@ +import { describe, expect, it, vi } from "vitest"; + +const loadSessionStoreMock = vi.fn(); + +vi.mock("../config/sessions.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + loadSessionStore: (storePath: string) => loadSessionStoreMock(storePath), + resolveStorePath: (_store: string | undefined, opts?: { agentId?: string }) => + opts?.agentId === "support" ? "/tmp/support/sessions.json" : "/tmp/main/sessions.json", + resolveSessionFilePath: (sessionId: string) => `/tmp/sessions/${sessionId}.json`, + }; +}); + +vi.mock("../config/config.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + loadConfig: () => ({ + session: { mainKey: "main", scope: "per-sender" }, + agents: { + defaults: { + model: { primary: "anthropic/claude-sonnet-4-5" }, + models: {}, + workspaceDir: "/tmp/workspace", + }, + }, + }), + }; +}); + +const compactEmbeddedPiSessionMock = vi.fn(); + +vi.mock("../agents/pi-embedded.js", () => ({ + compactEmbeddedPiSession: (params: unknown) => compactEmbeddedPiSessionMock(params), +})); + +vi.mock("../auto-reply/thinking.js", () => ({ + resolveDefaultThinkingLevel: async () => "off", +})); + +import "./test-helpers/fast-core-tools.js"; +import { createMoltbotTools } from "./moltbot-tools.js"; + +describe("session_compact tool", () => { + it("compacts the current session successfully", async () => { + loadSessionStoreMock.mockReset(); + compactEmbeddedPiSessionMock.mockReset(); + + loadSessionStoreMock.mockReturnValue({ + main: { + sessionId: "s1", + updatedAt: 10, + }, + }); + + compactEmbeddedPiSessionMock.mockResolvedValue({ + ok: true, + compacted: true, + result: { + tokensBefore: 50000, + tokensAfter: 15000, + }, + }); + + const tool = createMoltbotTools({ agentSessionKey: "main" }).find( + (candidate) => candidate.name === "session_compact", + ); + expect(tool).toBeDefined(); + if (!tool) throw new Error("missing session_compact tool"); + + const result = await tool.execute("call1", {}); + const details = result.details as { + ok?: boolean; + compacted?: boolean; + tokensBefore?: number; + tokensAfter?: number; + }; + + expect(details.ok).toBe(true); + expect(details.compacted).toBe(true); + expect(details.tokensBefore).toBe(50000); + expect(details.tokensAfter).toBe(15000); + expect(compactEmbeddedPiSessionMock).toHaveBeenCalledWith( + expect.objectContaining({ + sessionId: "s1", + }), + ); + }); + + it("passes custom instructions to compaction", async () => { + loadSessionStoreMock.mockReset(); + compactEmbeddedPiSessionMock.mockReset(); + + loadSessionStoreMock.mockReturnValue({ + main: { + sessionId: "s1", + updatedAt: 10, + }, + }); + + compactEmbeddedPiSessionMock.mockResolvedValue({ + ok: true, + compacted: true, + result: { + tokensBefore: 30000, + tokensAfter: 10000, + }, + }); + + const tool = createMoltbotTools({ agentSessionKey: "main" }).find( + (candidate) => candidate.name === "session_compact", + ); + expect(tool).toBeDefined(); + if (!tool) throw new Error("missing session_compact tool"); + + await tool.execute("call2", { instructions: "Keep all TODOs" }); + + expect(compactEmbeddedPiSessionMock).toHaveBeenCalledWith( + expect.objectContaining({ + customInstructions: "Keep all TODOs", + }), + ); + }); + + it("errors for unknown session keys", async () => { + loadSessionStoreMock.mockReset(); + compactEmbeddedPiSessionMock.mockReset(); + + loadSessionStoreMock.mockReturnValue({ + main: { sessionId: "s1", updatedAt: 10 }, + }); + + const tool = createMoltbotTools({ agentSessionKey: "main" }).find( + (candidate) => candidate.name === "session_compact", + ); + expect(tool).toBeDefined(); + if (!tool) throw new Error("missing session_compact tool"); + + await expect(tool.execute("call3", { sessionKey: "nope" })).rejects.toThrow( + "Unknown sessionId", + ); + expect(compactEmbeddedPiSessionMock).not.toHaveBeenCalled(); + }); + + it("errors when session has no sessionId", async () => { + loadSessionStoreMock.mockReset(); + compactEmbeddedPiSessionMock.mockReset(); + + loadSessionStoreMock.mockReturnValue({ + main: { + updatedAt: 10, + // no sessionId + }, + }); + + const tool = createMoltbotTools({ agentSessionKey: "main" }).find( + (candidate) => candidate.name === "session_compact", + ); + expect(tool).toBeDefined(); + if (!tool) throw new Error("missing session_compact tool"); + + await expect(tool.execute("call4", {})).rejects.toThrow( + "Compaction unavailable (missing session id)", + ); + expect(compactEmbeddedPiSessionMock).not.toHaveBeenCalled(); + }); + + it("blocks cross-agent session_compact without agent-to-agent access", async () => { + loadSessionStoreMock.mockReset(); + compactEmbeddedPiSessionMock.mockReset(); + + loadSessionStoreMock.mockReturnValue({ + "agent:other:main": { + sessionId: "s2", + updatedAt: 10, + }, + }); + + const tool = createMoltbotTools({ agentSessionKey: "agent:main:main" }).find( + (candidate) => candidate.name === "session_compact", + ); + expect(tool).toBeDefined(); + if (!tool) throw new Error("missing session_compact tool"); + + await expect(tool.execute("call5", { sessionKey: "agent:other:main" })).rejects.toThrow( + "Agent-to-agent compact is disabled", + ); + expect(compactEmbeddedPiSessionMock).not.toHaveBeenCalled(); + }); + + it("handles compaction failure gracefully", async () => { + loadSessionStoreMock.mockReset(); + compactEmbeddedPiSessionMock.mockReset(); + + loadSessionStoreMock.mockReturnValue({ + main: { + sessionId: "s1", + updatedAt: 10, + }, + }); + + compactEmbeddedPiSessionMock.mockResolvedValue({ + ok: false, + compacted: false, + reason: "API error", + }); + + const tool = createMoltbotTools({ agentSessionKey: "main" }).find( + (candidate) => candidate.name === "session_compact", + ); + expect(tool).toBeDefined(); + if (!tool) throw new Error("missing session_compact tool"); + + const result = await tool.execute("call6", {}); + const details = result.details as { + ok?: boolean; + compacted?: boolean; + reason?: string; + }; + + expect(details.ok).toBe(false); + expect(details.compacted).toBe(false); + expect(details.reason).toBe("API error"); + }); + + it("handles compaction skipped", async () => { + loadSessionStoreMock.mockReset(); + compactEmbeddedPiSessionMock.mockReset(); + + loadSessionStoreMock.mockReturnValue({ + main: { + sessionId: "s1", + updatedAt: 10, + }, + }); + + compactEmbeddedPiSessionMock.mockResolvedValue({ + ok: true, + compacted: false, + reason: "Context too small", + }); + + const tool = createMoltbotTools({ agentSessionKey: "main" }).find( + (candidate) => candidate.name === "session_compact", + ); + expect(tool).toBeDefined(); + if (!tool) throw new Error("missing session_compact tool"); + + const result = await tool.execute("call7", {}); + const details = result.details as { + ok?: boolean; + compacted?: boolean; + reason?: string; + }; + + expect(details.ok).toBe(true); + expect(details.compacted).toBe(false); + expect(details.reason).toBe("Context too small"); + }); +}); diff --git a/src/agents/moltbot-tools.ts b/src/agents/moltbot-tools.ts index c10a55190..94dc1a515 100644 --- a/src/agents/moltbot-tools.ts +++ b/src/agents/moltbot-tools.ts @@ -11,6 +11,7 @@ import { createGatewayTool } from "./tools/gateway-tool.js"; import { createImageTool } from "./tools/image-tool.js"; import { createMessageTool } from "./tools/message-tool.js"; import { createNodesTool } from "./tools/nodes-tool.js"; +import { createSessionCompactTool } from "./tools/session-compact-tool.js"; import { createSessionStatusTool } from "./tools/session-status-tool.js"; import { createSessionsHistoryTool } from "./tools/sessions-history-tool.js"; import { createSessionsListTool } from "./tools/sessions-list-tool.js"; @@ -134,6 +135,10 @@ export function createMoltbotTools(options?: { agentSessionKey: options?.agentSessionKey, config: options?.config, }), + createSessionCompactTool({ + agentSessionKey: options?.agentSessionKey, + config: options?.config, + }), ...(webSearchTool ? [webSearchTool] : []), ...(webFetchTool ? [webFetchTool] : []), ...(imageTool ? [imageTool] : []), diff --git a/src/agents/system-prompt.ts b/src/agents/system-prompt.ts index ed97fd539..70d62db7b 100644 --- a/src/agents/system-prompt.ts +++ b/src/agents/system-prompt.ts @@ -205,6 +205,8 @@ export function buildAgentSystemPrompt(params: { sessions_spawn: "Spawn a sub-agent session", session_status: "Show a /status-equivalent status card (usage + time + Reasoning/Verbose/Elevated); use for model-use questions (📊 session_status); optional per-session model override", + session_compact: + "Compact the session context by summarizing conversation history; use when context is getting large or before complex tasks", image: "Analyze an image with the configured image model", }; @@ -231,6 +233,7 @@ export function buildAgentSystemPrompt(params: { "sessions_history", "sessions_send", "session_status", + "session_compact", "image", ]; diff --git a/src/agents/tools/session-compact-tool.ts b/src/agents/tools/session-compact-tool.ts new file mode 100644 index 000000000..3d2f46691 --- /dev/null +++ b/src/agents/tools/session-compact-tool.ts @@ -0,0 +1,247 @@ +import { Type } from "@sinclair/typebox"; + +import { compactEmbeddedPiSession } from "../pi-embedded.js"; +import { + resolveSessionFilePath, + loadSessionStore, + resolveStorePath, +} from "../../config/sessions.js"; +import type { MoltbotConfig } from "../../config/config.js"; +import { loadConfig } from "../../config/config.js"; +import { formatTokenCount, formatContextUsageShort } from "../../auto-reply/status.js"; +import { + buildAgentMainSessionKey, + resolveAgentIdFromSessionKey, + DEFAULT_AGENT_ID, +} from "../../routing/session-key.js"; +import { resolveDefaultModelForAgent } from "../model-selection.js"; +import type { ThinkLevel } from "../../auto-reply/thinking.js"; +import type { AnyAgentTool } from "./common.js"; +import { readStringParam } from "./common.js"; +import { + shouldResolveSessionIdInput, + resolveInternalSessionKey, + resolveMainSessionAlias, + createAgentToAgentPolicy, +} from "./sessions-helpers.js"; +import { loadCombinedSessionStoreForGateway } from "../../gateway/session-utils.js"; +import type { SessionEntry } from "../../config/sessions.js"; + +const SessionCompactToolSchema = Type.Object({ + sessionKey: Type.Optional(Type.String()), + instructions: Type.Optional( + Type.String({ + description: "Extra compaction instructions to guide the summarization", + }), + ), +}); + +function resolveSessionEntry(params: { + store: Record; + keyRaw: string; + alias: string; + mainKey: string; +}): { key: string; entry: SessionEntry } | null { + const keyRaw = params.keyRaw.trim(); + if (!keyRaw) return null; + const internal = resolveInternalSessionKey({ + key: keyRaw, + alias: params.alias, + mainKey: params.mainKey, + }); + + const candidates = new Set([keyRaw, internal]); + if (!keyRaw.startsWith("agent:")) { + candidates.add(`agent:${DEFAULT_AGENT_ID}:${keyRaw}`); + candidates.add(`agent:${DEFAULT_AGENT_ID}:${internal}`); + } + if (keyRaw === "main") { + candidates.add( + buildAgentMainSessionKey({ + agentId: DEFAULT_AGENT_ID, + mainKey: params.mainKey, + }), + ); + } + + for (const key of candidates) { + const entry = params.store[key]; + if (entry) return { key, entry }; + } + + return null; +} + +function resolveSessionKeyFromSessionId(params: { + cfg: MoltbotConfig; + sessionId: string; + agentId?: string; +}): string | null { + const trimmed = params.sessionId.trim(); + if (!trimmed) return null; + const { store } = loadCombinedSessionStoreForGateway(params.cfg); + const match = Object.entries(store).find(([key, entry]) => { + if (entry?.sessionId !== trimmed) return false; + if (!params.agentId) return true; + return resolveAgentIdFromSessionKey(key) === params.agentId; + }); + return match?.[0] ?? null; +} + +export function createSessionCompactTool(opts?: { + agentSessionKey?: string; + config?: MoltbotConfig; +}): AnyAgentTool { + return { + label: "Session Compact", + name: "session_compact", + description: + "Compact the session context by summarizing conversation history. " + + "Use when context is getting large or before complex tasks. " + + "Optional instructions can guide what to preserve in the summary.", + parameters: SessionCompactToolSchema, + execute: async (_toolCallId, args) => { + const params = args as Record; + const cfg = opts?.config ?? loadConfig(); + const { mainKey, alias } = resolveMainSessionAlias(cfg); + const a2aPolicy = createAgentToAgentPolicy(cfg); + + const requestedKeyParam = readStringParam(params, "sessionKey"); + let requestedKeyRaw = requestedKeyParam ?? opts?.agentSessionKey; + if (!requestedKeyRaw?.trim()) { + throw new Error("sessionKey required"); + } + + const requesterAgentId = resolveAgentIdFromSessionKey( + opts?.agentSessionKey ?? requestedKeyRaw, + ); + const ensureAgentAccess = (targetAgentId: string) => { + if (targetAgentId === requesterAgentId) return; + if (!a2aPolicy.enabled) { + throw new Error( + "Agent-to-agent compact is disabled. Set tools.agentToAgent.enabled=true to allow cross-agent access.", + ); + } + if (!a2aPolicy.isAllowed(requesterAgentId, targetAgentId)) { + throw new Error("Agent-to-agent session compact denied by tools.agentToAgent.allow."); + } + }; + + if (requestedKeyRaw.startsWith("agent:")) { + ensureAgentAccess(resolveAgentIdFromSessionKey(requestedKeyRaw)); + } + + const isExplicitAgentKey = requestedKeyRaw.startsWith("agent:"); + let agentId = isExplicitAgentKey + ? resolveAgentIdFromSessionKey(requestedKeyRaw) + : requesterAgentId; + let storePath = resolveStorePath(cfg.session?.store, { agentId }); + let store = loadSessionStore(storePath); + + let resolved = resolveSessionEntry({ + store, + keyRaw: requestedKeyRaw, + alias, + mainKey, + }); + + if (!resolved && shouldResolveSessionIdInput(requestedKeyRaw)) { + const resolvedKey = resolveSessionKeyFromSessionId({ + cfg, + sessionId: requestedKeyRaw, + agentId: a2aPolicy.enabled ? undefined : requesterAgentId, + }); + if (resolvedKey) { + ensureAgentAccess(resolveAgentIdFromSessionKey(resolvedKey)); + requestedKeyRaw = resolvedKey; + agentId = resolveAgentIdFromSessionKey(resolvedKey); + storePath = resolveStorePath(cfg.session?.store, { agentId }); + store = loadSessionStore(storePath); + resolved = resolveSessionEntry({ + store, + keyRaw: requestedKeyRaw, + alias, + mainKey, + }); + } + } + + if (!resolved) { + const kind = shouldResolveSessionIdInput(requestedKeyRaw) ? "sessionId" : "sessionKey"; + throw new Error(`Unknown ${kind}: ${requestedKeyRaw}`); + } + + if (!resolved.entry.sessionId) { + throw new Error("Compaction unavailable (missing session id)"); + } + + const customInstructions = readStringParam(params, "instructions"); + const configured = resolveDefaultModelForAgent({ cfg, agentId }); + const provider = resolved.entry.providerOverride?.trim() || configured.provider; + const model = resolved.entry.modelOverride?.trim() || configured.model; + // Use session's thinking level or default to "off" for compaction + const thinkLevel: ThinkLevel = (resolved.entry.thinkingLevel as ThinkLevel) ?? "off"; + + const result = await compactEmbeddedPiSession({ + sessionId: resolved.entry.sessionId, + sessionKey: resolved.key, + messageChannel: resolved.entry.channel ?? resolved.entry.lastChannel, + groupId: resolved.entry.groupId, + groupChannel: resolved.entry.groupChannel, + groupSpace: resolved.entry.space, + spawnedBy: resolved.entry.spawnedBy, + sessionFile: resolveSessionFilePath(resolved.entry.sessionId, resolved.entry), + workspaceDir: cfg.agents?.defaults?.workspace ?? process.cwd(), + config: cfg, + skillsSnapshot: resolved.entry.skillsSnapshot, + provider, + model, + thinkLevel, + bashElevated: { + enabled: false, + allowed: false, + defaultLevel: "off", + }, + customInstructions, + }); + + const tokensBefore = result.result?.tokensBefore; + const tokensAfter = result.result?.tokensAfter; + + let statusText: string; + if (result.ok) { + if (result.compacted) { + const tokenInfo = + tokensBefore != null && tokensAfter != null + ? `${formatTokenCount(tokensBefore)} → ${formatTokenCount(tokensAfter)}` + : tokensBefore != null + ? `${formatTokenCount(tokensBefore)} before` + : ""; + const contextSummary = formatContextUsageShort( + tokensAfter ?? null, + resolved.entry.contextTokens ?? null, + ); + statusText = tokenInfo + ? `⚙️ Compacted (${tokenInfo}) • ${contextSummary}` + : `⚙️ Compacted • ${contextSummary}`; + } else { + statusText = `⚙️ Compaction skipped${result.reason ? `: ${result.reason}` : ""}`; + } + } else { + statusText = `❌ Compaction failed${result.reason ? `: ${result.reason}` : ""}`; + } + + return { + content: [{ type: "text", text: statusText }], + details: { + ok: result.ok, + compacted: result.compacted ?? false, + sessionKey: resolved.key, + tokensBefore, + tokensAfter, + reason: result.reason, + }, + }; + }, + }; +}