diff --git a/src/agents/pi-embedded-runner.ts b/src/agents/pi-embedded-runner.ts index d153a5802..86b790a44 100644 --- a/src/agents/pi-embedded-runner.ts +++ b/src/agents/pi-embedded-runner.ts @@ -1,5 +1,7 @@ import fs from "node:fs/promises"; import os from "node:os"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; import type { AgentMessage, @@ -40,7 +42,11 @@ import { markAuthProfileUsed, } from "./auth-profiles.js"; import type { BashElevatedDefaults } from "./bash-tools.js"; -import { DEFAULT_MODEL, DEFAULT_PROVIDER } from "./defaults.js"; +import { + DEFAULT_CONTEXT_TOKENS, + DEFAULT_MODEL, + DEFAULT_PROVIDER, +} from "./defaults.js"; import { ensureAuthProfileStore, getApiKeyForModel, @@ -67,6 +73,9 @@ import { extractAssistantThinking, formatReasoningMarkdown, } from "./pi-embedded-utils.js"; +import { setContextPruningRuntime } from "./pi-extensions/context-pruning/runtime.js"; +import { computeEffectiveSettings } from "./pi-extensions/context-pruning/settings.js"; +import { makeToolPrunablePredicate } from "./pi-extensions/context-pruning/tools.js"; import { toToolDefinitions } from "./pi-tool-definition-adapter.js"; import { createClawdbotCodingTools } from "./pi-tools.js"; import { resolveSandboxContext } from "./sandbox.js"; @@ -82,6 +91,84 @@ import { buildAgentSystemPromptAppend } from "./system-prompt.js"; import { normalizeUsage, type UsageLike } from "./usage.js"; import { loadWorkspaceBootstrapFiles } from "./workspace.js"; +// Optional features can be implemented as Pi extensions that run in the same Node process. +// We configure context pruning per-session via a WeakMap registry keyed by the SessionManager instance. + +function resolvePiExtensionPath(id: string): string { + const self = fileURLToPath(import.meta.url); + const dir = path.dirname(self); + // In dev this file is `.ts` (tsx), in production it's `.js`. + const ext = path.extname(self) === ".ts" ? "ts" : "js"; + return path.join(dir, "pi-extensions", `${id}.${ext}`); +} + +function resolveContextWindowTokens(params: { + cfg: ClawdbotConfig | undefined; + provider: string; + modelId: string; + model: Model | undefined; +}): number { + const fromModel = + typeof params.model?.contextWindow === "number" && + Number.isFinite(params.model.contextWindow) && + params.model.contextWindow > 0 + ? params.model.contextWindow + : undefined; + if (fromModel) return fromModel; + + const fromModelsConfig = (() => { + const providers = params.cfg?.models?.providers as + | Record< + string, + { models?: Array<{ id?: string; contextWindow?: number }> } + > + | undefined; + const providerEntry = providers?.[params.provider]; + const models = Array.isArray(providerEntry?.models) + ? providerEntry.models + : []; + const match = models.find((m) => m?.id === params.modelId); + return typeof match?.contextWindow === "number" && match.contextWindow > 0 + ? match.contextWindow + : undefined; + })(); + if (fromModelsConfig) return fromModelsConfig; + + const fromAgentConfig = + typeof params.cfg?.agent?.contextTokens === "number" && + Number.isFinite(params.cfg.agent.contextTokens) && + params.cfg.agent.contextTokens > 0 + ? Math.floor(params.cfg.agent.contextTokens) + : undefined; + if (fromAgentConfig) return fromAgentConfig; + + return DEFAULT_CONTEXT_TOKENS; +} + +function buildContextPruningExtension(params: { + cfg: ClawdbotConfig | undefined; + sessionManager: SessionManager; + provider: string; + modelId: string; + model: Model | undefined; +}): { additionalExtensionPaths?: string[] } { + const raw = params.cfg?.agent?.contextPruning; + if (raw?.mode !== "adaptive" && raw?.mode !== "aggressive") return {}; + + const settings = computeEffectiveSettings(raw); + if (!settings) return {}; + + setContextPruningRuntime(params.sessionManager, { + settings, + contextWindowTokens: resolveContextWindowTokens(params), + isToolPrunable: makeToolPrunablePredicate(settings.tools), + }); + + return { + additionalExtensionPaths: [resolvePiExtensionPath("context-pruning")], + }; +} + export type EmbeddedPiAgentMeta = { sessionId: string; provider: string; @@ -578,13 +665,22 @@ export async function compactEmbeddedPiSession(params: { effectiveWorkspace, agentDir, ); + const pruning = buildContextPruningExtension({ + cfg: params.config, + sessionManager, + provider, + modelId, + model, + }); + const additionalExtensionPaths = pruning.additionalExtensionPaths; const { builtInTools, customTools } = splitSdkTools({ tools, sandboxEnabled: !!sandbox?.enabled, }); - const { session } = await createAgentSession({ + let session: Awaited>["session"]; + ({ session } = await createAgentSession({ cwd: resolvedWorkspace, agentDir, authStorage, @@ -598,7 +694,8 @@ export async function compactEmbeddedPiSession(params: { settingsManager, skills: promptSkills, contextFiles, - }); + additionalExtensionPaths, + })); try { const prior = await sanitizeSessionMessagesImages( @@ -887,13 +984,24 @@ export async function runEmbeddedPiAgent(params: { effectiveWorkspace, agentDir, ); + const pruning = buildContextPruningExtension({ + cfg: params.config, + sessionManager, + provider, + modelId, + model, + }); + const additionalExtensionPaths = pruning.additionalExtensionPaths; const { builtInTools, customTools } = splitSdkTools({ tools, sandboxEnabled: !!sandbox?.enabled, }); - const { session } = await createAgentSession({ + let session: Awaited< + ReturnType + >["session"]; + ({ session } = await createAgentSession({ cwd: resolvedWorkspace, agentDir, authStorage, @@ -909,14 +1017,20 @@ export async function runEmbeddedPiAgent(params: { settingsManager, skills: promptSkills, contextFiles, - }); + additionalExtensionPaths, + })); - const prior = await sanitizeSessionMessagesImages( - session.messages, - "session:history", - ); - if (prior.length > 0) { - session.agent.replaceMessages(prior); + try { + const prior = await sanitizeSessionMessagesImages( + session.messages, + "session:history", + ); + if (prior.length > 0) { + session.agent.replaceMessages(prior); + } + } catch (err) { + session.dispose(); + throw err; } let aborted = Boolean(params.abortSignal?.aborted); let timedOut = false; @@ -925,21 +1039,27 @@ export async function runEmbeddedPiAgent(params: { if (isTimeout) timedOut = true; void session.abort(); }; - const subscription = subscribeEmbeddedPiSession({ - session, - runId: params.runId, - verboseLevel: params.verboseLevel, - reasoningMode: params.reasoningLevel ?? "off", - shouldEmitToolResult: params.shouldEmitToolResult, - onToolResult: params.onToolResult, - onReasoningStream: params.onReasoningStream, - onBlockReply: params.onBlockReply, - blockReplyBreak: params.blockReplyBreak, - blockReplyChunking: params.blockReplyChunking, - onPartialReply: params.onPartialReply, - onAgentEvent: params.onAgentEvent, - enforceFinalTag: params.enforceFinalTag, - }); + let subscription: ReturnType; + try { + subscription = subscribeEmbeddedPiSession({ + session, + runId: params.runId, + verboseLevel: params.verboseLevel, + reasoningMode: params.reasoningLevel ?? "off", + shouldEmitToolResult: params.shouldEmitToolResult, + onToolResult: params.onToolResult, + onReasoningStream: params.onReasoningStream, + onBlockReply: params.onBlockReply, + blockReplyBreak: params.blockReplyBreak, + blockReplyChunking: params.blockReplyChunking, + onPartialReply: params.onPartialReply, + onAgentEvent: params.onAgentEvent, + enforceFinalTag: params.enforceFinalTag, + }); + } catch (err) { + session.dispose(); + throw err; + } const { assistantTexts, toolMetas, diff --git a/src/agents/pi-extensions/context-pruning.ts b/src/agents/pi-extensions/context-pruning.ts new file mode 100644 index 000000000..b80addb9d --- /dev/null +++ b/src/agents/pi-extensions/context-pruning.ts @@ -0,0 +1,19 @@ +/** + * Opt-in context pruning (“microcompact”-style) for Pi sessions. + * + * This only affects the in-memory context for the current request; it does not rewrite session + * history persisted on disk. + */ + +export { default } from "./context-pruning/extension.js"; + +export { pruneContextMessages } from "./context-pruning/pruner.js"; +export type { + ContextPruningConfig, + ContextPruningToolMatch, + EffectiveContextPruningSettings, +} from "./context-pruning/settings.js"; +export { + computeEffectiveSettings, + DEFAULT_CONTEXT_PRUNING_SETTINGS, +} from "./context-pruning/settings.js"; diff --git a/src/agents/pi-extensions/context-pruning/extension.ts b/src/agents/pi-extensions/context-pruning/extension.ts new file mode 100644 index 000000000..13b9a8d4b --- /dev/null +++ b/src/agents/pi-extensions/context-pruning/extension.ts @@ -0,0 +1,27 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import type { + ContextEvent, + ExtensionAPI, + ExtensionContext, +} from "@mariozechner/pi-coding-agent"; + +import { pruneContextMessages } from "./pruner.js"; +import { getContextPruningRuntime } from "./runtime.js"; + +export default function contextPruningExtension(api: ExtensionAPI): void { + api.on("context", (event: ContextEvent, ctx: ExtensionContext) => { + const runtime = getContextPruningRuntime(ctx.sessionManager); + if (!runtime) return undefined; + + const next = pruneContextMessages({ + messages: event.messages as AgentMessage[], + settings: runtime.settings, + ctx, + isToolPrunable: runtime.isToolPrunable, + contextWindowTokensOverride: runtime.contextWindowTokens ?? undefined, + }); + + if (next === event.messages) return undefined; + return { messages: next }; + }); +} diff --git a/src/agents/pi-extensions/context-pruning/pruner.ts b/src/agents/pi-extensions/context-pruning/pruner.ts new file mode 100644 index 000000000..0341b2bbf --- /dev/null +++ b/src/agents/pi-extensions/context-pruning/pruner.ts @@ -0,0 +1,310 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import type { + ImageContent, + TextContent, + ToolResultMessage, +} from "@mariozechner/pi-ai"; +import type { ExtensionContext } from "@mariozechner/pi-coding-agent"; + +import type { EffectiveContextPruningSettings } from "./settings.js"; +import { makeToolPrunablePredicate } from "./tools.js"; + +const CHARS_PER_TOKEN_ESTIMATE = 4; +// We currently skip pruning tool results that contain images. Still, we count them (approx.) so +// we start trimming prunable tool results earlier when image-heavy context is consuming the window. +const IMAGE_CHAR_ESTIMATE = 8_000; + +function asText(text: string): TextContent { + return { type: "text", text }; +} + +function collectTextSegments( + content: ReadonlyArray, +): string[] { + const parts: string[] = []; + for (const block of content) { + if (block.type === "text") parts.push(block.text); + } + return parts; +} + +function estimateJoinedTextLength(parts: string[]): number { + if (parts.length === 0) return 0; + let len = 0; + for (const p of parts) len += p.length; + // Joined with "\n" separators between blocks. + len += Math.max(0, parts.length - 1); + return len; +} + +function takeHeadFromJoinedText(parts: string[], maxChars: number): string { + if (maxChars <= 0 || parts.length === 0) return ""; + let remaining = maxChars; + let out = ""; + for (let i = 0; i < parts.length && remaining > 0; i++) { + if (i > 0) { + out += "\n"; + remaining -= 1; + if (remaining <= 0) break; + } + const p = parts[i]; + if (p.length <= remaining) { + out += p; + remaining -= p.length; + } else { + out += p.slice(0, remaining); + remaining = 0; + } + } + return out; +} + +function takeTailFromJoinedText(parts: string[], maxChars: number): string { + if (maxChars <= 0 || parts.length === 0) return ""; + let remaining = maxChars; + const out: string[] = []; + for (let i = parts.length - 1; i >= 0 && remaining > 0; i--) { + const p = parts[i]; + if (p.length <= remaining) { + out.push(p); + remaining -= p.length; + } else { + out.push(p.slice(p.length - remaining)); + remaining = 0; + break; + } + if (remaining > 0 && i > 0) { + out.push("\n"); + remaining -= 1; + } + } + out.reverse(); + return out.join(""); +} + +function hasImageBlocks( + content: ReadonlyArray, +): boolean { + for (const block of content) { + if (block.type === "image") return true; + } + return false; +} + +function estimateMessageChars(message: AgentMessage): number { + if (message.role === "user") { + const content = message.content; + if (typeof content === "string") return content.length; + let chars = 0; + for (const b of content) { + if (b.type === "text") chars += b.text.length; + if (b.type === "image") chars += IMAGE_CHAR_ESTIMATE; + } + return chars; + } + + if (message.role === "assistant") { + let chars = 0; + for (const b of message.content) { + if (b.type === "text") chars += b.text.length; + if (b.type === "thinking") chars += b.thinking.length; + if (b.type === "toolCall") { + try { + chars += JSON.stringify(b.arguments ?? {}).length; + } catch { + chars += 128; + } + } + } + return chars; + } + + if (message.role === "toolResult") { + let chars = 0; + for (const b of message.content) { + if (b.type === "text") chars += b.text.length; + if (b.type === "image") chars += IMAGE_CHAR_ESTIMATE; + } + return chars; + } + + return 256; +} + +function estimateContextChars(messages: AgentMessage[]): number { + return messages.reduce((sum, m) => sum + estimateMessageChars(m), 0); +} + +function findAssistantCutoffIndex( + messages: AgentMessage[], + keepLastAssistants: number, +): number | null { + // keepLastAssistants <= 0 => everything is potentially prunable. + if (keepLastAssistants <= 0) return messages.length; + + let remaining = keepLastAssistants; + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i]?.role !== "assistant") continue; + remaining--; + if (remaining === 0) return i; + } + + // Not enough assistant messages to establish a protected tail. + return null; +} + +function softTrimToolResultMessage(params: { + msg: ToolResultMessage; + settings: EffectiveContextPruningSettings; +}): ToolResultMessage | null { + const { msg, settings } = params; + // Ignore image tool results for now: these are often directly relevant and hard to partially prune safely. + if (hasImageBlocks(msg.content)) return null; + + const parts = collectTextSegments(msg.content); + const rawLen = estimateJoinedTextLength(parts); + if (rawLen <= settings.softTrim.maxChars) return null; + + const headChars = Math.max(0, settings.softTrim.headChars); + const tailChars = Math.max(0, settings.softTrim.tailChars); + if (headChars + tailChars >= rawLen) return null; + + const head = takeHeadFromJoinedText(parts, headChars); + const tail = takeTailFromJoinedText(parts, tailChars); + const trimmed = `${head} +... +${tail}`; + + const note = ` + +[Tool result trimmed: kept first ${headChars} chars and last ${tailChars} chars of ${rawLen} chars.]`; + + return { ...msg, content: [asText(trimmed + note)] }; +} + +export function pruneContextMessages(params: { + messages: AgentMessage[]; + settings: EffectiveContextPruningSettings; + ctx: Pick; + isToolPrunable?: (toolName: string) => boolean; + contextWindowTokensOverride?: number; +}): AgentMessage[] { + const { messages, settings, ctx } = params; + const contextWindowTokens = + typeof params.contextWindowTokensOverride === "number" && + Number.isFinite(params.contextWindowTokensOverride) && + params.contextWindowTokensOverride > 0 + ? params.contextWindowTokensOverride + : ctx.model?.contextWindow; + if (!contextWindowTokens || contextWindowTokens <= 0) return messages; + + const charWindow = contextWindowTokens * CHARS_PER_TOKEN_ESTIMATE; + if (charWindow <= 0) return messages; + + const cutoffIndex = findAssistantCutoffIndex( + messages, + settings.keepLastAssistants, + ); + if (cutoffIndex === null) return messages; + + const isToolPrunable = + params.isToolPrunable ?? makeToolPrunablePredicate(settings.tools); + + if (settings.mode === "aggressive") { + let next: AgentMessage[] | null = null; + + for (let i = 0; i < cutoffIndex; i++) { + const msg = messages[i]; + if (!msg || msg.role !== "toolResult") continue; + if (!isToolPrunable(msg.toolName)) continue; + if (hasImageBlocks(msg.content)) { + continue; + } + + const alreadyCleared = + msg.content.length === 1 && + msg.content[0]?.type === "text" && + msg.content[0].text === settings.hardClear.placeholder; + if (alreadyCleared) continue; + + const cleared: ToolResultMessage = { + ...msg, + content: [asText(settings.hardClear.placeholder)], + }; + if (!next) next = messages.slice(); + next[i] = cleared as unknown as AgentMessage; + } + + return next ?? messages; + } + + const totalCharsBefore = estimateContextChars(messages); + let totalChars = totalCharsBefore; + let ratio = totalChars / charWindow; + if (ratio < settings.softTrimRatio) { + return messages; + } + + const prunableToolIndexes: number[] = []; + let next: AgentMessage[] | null = null; + + for (let i = 0; i < cutoffIndex; i++) { + const msg = messages[i]; + if (!msg || msg.role !== "toolResult") continue; + if (!isToolPrunable(msg.toolName)) continue; + if (hasImageBlocks(msg.content)) { + continue; + } + prunableToolIndexes.push(i); + + const updated = softTrimToolResultMessage({ + msg: msg as unknown as ToolResultMessage, + settings, + }); + if (!updated) continue; + + const beforeChars = estimateMessageChars(msg); + const afterChars = estimateMessageChars(updated as unknown as AgentMessage); + totalChars += afterChars - beforeChars; + if (!next) next = messages.slice(); + next[i] = updated as unknown as AgentMessage; + } + + const outputAfterSoftTrim = next ?? messages; + ratio = totalChars / charWindow; + if (ratio < settings.hardClearRatio) { + return outputAfterSoftTrim; + } + if (!settings.hardClear.enabled) { + return outputAfterSoftTrim; + } + + let prunableToolChars = 0; + for (const i of prunableToolIndexes) { + const msg = outputAfterSoftTrim[i]; + if (!msg || msg.role !== "toolResult") continue; + prunableToolChars += estimateMessageChars(msg); + } + if (prunableToolChars < settings.minPrunableToolChars) { + return outputAfterSoftTrim; + } + + for (const i of prunableToolIndexes) { + if (ratio < settings.hardClearRatio) break; + const msg = (next ?? messages)[i]; + if (!msg || msg.role !== "toolResult") continue; + + const beforeChars = estimateMessageChars(msg); + const cleared: ToolResultMessage = { + ...msg, + content: [asText(settings.hardClear.placeholder)], + }; + if (!next) next = messages.slice(); + next[i] = cleared as unknown as AgentMessage; + const afterChars = estimateMessageChars(cleared as unknown as AgentMessage); + totalChars += afterChars - beforeChars; + ratio = totalChars / charWindow; + } + + return next ?? messages; +} diff --git a/src/agents/pi-extensions/context-pruning/runtime.ts b/src/agents/pi-extensions/context-pruning/runtime.ts new file mode 100644 index 000000000..b497e6383 --- /dev/null +++ b/src/agents/pi-extensions/context-pruning/runtime.ts @@ -0,0 +1,39 @@ +import type { EffectiveContextPruningSettings } from "./settings.js"; + +export type ContextPruningRuntimeValue = { + settings: EffectiveContextPruningSettings; + contextWindowTokens?: number | null; + isToolPrunable: (toolName: string) => boolean; +}; + +// Session-scoped runtime registry keyed by object identity. +// Important: this relies on Pi passing the same SessionManager object instance into +// ExtensionContext (ctx.sessionManager) that we used when calling setContextPruningRuntime. +const REGISTRY = new WeakMap(); + +export function setContextPruningRuntime( + sessionManager: unknown, + value: ContextPruningRuntimeValue | null, +): void { + if (!sessionManager || typeof sessionManager !== "object") { + return; + } + + const key = sessionManager as object; + if (value === null) { + REGISTRY.delete(key); + return; + } + + REGISTRY.set(key, value); +} + +export function getContextPruningRuntime( + sessionManager: unknown, +): ContextPruningRuntimeValue | null { + if (!sessionManager || typeof sessionManager !== "object") { + return null; + } + + return REGISTRY.get(sessionManager as object) ?? null; +} diff --git a/src/agents/pi-extensions/context-pruning/settings.ts b/src/agents/pi-extensions/context-pruning/settings.ts new file mode 100644 index 000000000..f3bb6de83 --- /dev/null +++ b/src/agents/pi-extensions/context-pruning/settings.ts @@ -0,0 +1,135 @@ +export type ContextPruningToolMatch = { + allow?: string[]; + deny?: string[]; +}; + +export type ContextPruningMode = "off" | "adaptive" | "aggressive"; + +export type ContextPruningConfig = { + mode?: ContextPruningMode; + keepLastAssistants?: number; + softTrimRatio?: number; + hardClearRatio?: number; + minPrunableToolChars?: number; + tools?: ContextPruningToolMatch; + softTrim?: { + maxChars?: number; + headChars?: number; + tailChars?: number; + }; + hardClear?: { + enabled?: boolean; + placeholder?: string; + }; +}; + +export type EffectiveContextPruningSettings = { + mode: Exclude; + keepLastAssistants: number; + softTrimRatio: number; + hardClearRatio: number; + minPrunableToolChars: number; + tools: ContextPruningToolMatch; + softTrim: { + maxChars: number; + headChars: number; + tailChars: number; + }; + hardClear: { + enabled: boolean; + placeholder: string; + }; +}; + +export const DEFAULT_CONTEXT_PRUNING_SETTINGS: EffectiveContextPruningSettings = + { + mode: "adaptive", + keepLastAssistants: 3, + softTrimRatio: 0.3, + hardClearRatio: 0.5, + minPrunableToolChars: 50_000, + tools: {}, + softTrim: { + maxChars: 4_000, + headChars: 1_500, + tailChars: 1_500, + }, + hardClear: { + enabled: true, + placeholder: "[Old tool result content cleared]", + }, + }; + +export function computeEffectiveSettings( + raw: unknown, +): EffectiveContextPruningSettings | null { + if (!raw || typeof raw !== "object") return null; + const cfg = raw as ContextPruningConfig; + if (cfg.mode !== "adaptive" && cfg.mode !== "aggressive") return null; + + const s: EffectiveContextPruningSettings = structuredClone( + DEFAULT_CONTEXT_PRUNING_SETTINGS, + ); + s.mode = cfg.mode; + + if ( + typeof cfg.keepLastAssistants === "number" && + Number.isFinite(cfg.keepLastAssistants) + ) { + s.keepLastAssistants = Math.max(0, Math.floor(cfg.keepLastAssistants)); + } + if ( + typeof cfg.softTrimRatio === "number" && + Number.isFinite(cfg.softTrimRatio) + ) { + s.softTrimRatio = Math.min(1, Math.max(0, cfg.softTrimRatio)); + } + if ( + typeof cfg.hardClearRatio === "number" && + Number.isFinite(cfg.hardClearRatio) + ) { + s.hardClearRatio = Math.min(1, Math.max(0, cfg.hardClearRatio)); + } + if ( + typeof cfg.minPrunableToolChars === "number" && + Number.isFinite(cfg.minPrunableToolChars) + ) { + s.minPrunableToolChars = Math.max(0, Math.floor(cfg.minPrunableToolChars)); + } + if (cfg.tools) { + s.tools = cfg.tools; + } + if (cfg.softTrim) { + if ( + typeof cfg.softTrim.maxChars === "number" && + Number.isFinite(cfg.softTrim.maxChars) + ) { + s.softTrim.maxChars = Math.max(0, Math.floor(cfg.softTrim.maxChars)); + } + if ( + typeof cfg.softTrim.headChars === "number" && + Number.isFinite(cfg.softTrim.headChars) + ) { + s.softTrim.headChars = Math.max(0, Math.floor(cfg.softTrim.headChars)); + } + if ( + typeof cfg.softTrim.tailChars === "number" && + Number.isFinite(cfg.softTrim.tailChars) + ) { + s.softTrim.tailChars = Math.max(0, Math.floor(cfg.softTrim.tailChars)); + } + } + if (cfg.hardClear) { + if (s.mode === "adaptive" && typeof cfg.hardClear.enabled === "boolean") { + s.hardClear.enabled = cfg.hardClear.enabled; + } + if ( + typeof cfg.hardClear.placeholder === "string" && + cfg.hardClear.placeholder.trim() + ) { + s.hardClear.placeholder = cfg.hardClear.placeholder.trim(); + } + } + + return s; +} diff --git a/src/agents/pi-extensions/context-pruning/tools.ts b/src/agents/pi-extensions/context-pruning/tools.ts new file mode 100644 index 000000000..81b064767 --- /dev/null +++ b/src/agents/pi-extensions/context-pruning/tools.ts @@ -0,0 +1,46 @@ +import type { ContextPruningToolMatch } from "./settings.js"; + +function normalizePatterns(patterns?: string[]): string[] { + if (!Array.isArray(patterns)) return []; + return patterns.map((p) => String(p ?? "").trim()).filter(Boolean); +} + +type CompiledPattern = + | { kind: "all" } + | { kind: "exact"; value: string } + | { kind: "regex"; value: RegExp }; + +function compilePattern(pattern: string): CompiledPattern { + if (pattern === "*") return { kind: "all" }; + if (!pattern.includes("*")) return { kind: "exact", value: pattern }; + + const escaped = pattern.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); + const re = new RegExp(`^${escaped.replaceAll("\\*", ".*")}$`); + return { kind: "regex", value: re }; +} + +function compilePatterns(patterns?: string[]): CompiledPattern[] { + return normalizePatterns(patterns).map(compilePattern); +} + +function matchesAny(toolName: string, patterns: CompiledPattern[]): boolean { + for (const p of patterns) { + if (p.kind === "all") return true; + if (p.kind === "exact" && toolName === p.value) return true; + if (p.kind === "regex" && p.value.test(toolName)) return true; + } + return false; +} + +export function makeToolPrunablePredicate( + match: ContextPruningToolMatch, +): (toolName: string) => boolean { + const deny = compilePatterns(match.deny); + const allow = compilePatterns(match.allow); + + return (toolName: string) => { + if (matchesAny(toolName, deny)) return false; + if (allow.length === 0) return true; + return matchesAny(toolName, allow); + }; +} diff --git a/src/config/types.ts b/src/config/types.ts index e8a16f23d..c00a635cd 100644 --- a/src/config/types.ts +++ b/src/config/types.ts @@ -850,6 +850,27 @@ export type AgentModelListConfig = { fallbacks?: string[]; }; +export type AgentContextPruningConfig = { + mode?: "off" | "adaptive" | "aggressive"; + keepLastAssistants?: number; + softTrimRatio?: number; + hardClearRatio?: number; + minPrunableToolChars?: number; + tools?: { + allow?: string[]; + deny?: string[]; + }; + softTrim?: { + maxChars?: number; + headChars?: number; + tailChars?: number; + }; + hardClear?: { + enabled?: boolean; + placeholder?: string; + }; +}; + export type ClawdbotConfig = { auth?: AuthConfig; env?: { @@ -895,6 +916,8 @@ export type ClawdbotConfig = { userTimezone?: string; /** Optional display-only context window override (used for % in status UIs). */ contextTokens?: number; + /** Opt-in: prune old tool results from the LLM context to reduce token usage. */ + contextPruning?: AgentContextPruningConfig; /** Default thinking level when no /think directive is present. */ thinkingDefault?: "off" | "minimal" | "low" | "medium" | "high"; /** Default verbose level when no /verbose directive is present. */ diff --git a/src/config/zod-schema.ts b/src/config/zod-schema.ts index b3dfef5ab..bcf3249ff 100644 --- a/src/config/zod-schema.ts +++ b/src/config/zod-schema.ts @@ -513,6 +513,40 @@ export const ClawdbotSchema = z.object({ skipBootstrap: z.boolean().optional(), userTimezone: z.string().optional(), contextTokens: z.number().int().positive().optional(), + contextPruning: z + .object({ + mode: z + .union([ + z.literal("off"), + z.literal("adaptive"), + z.literal("aggressive"), + ]) + .optional(), + keepLastAssistants: z.number().int().nonnegative().optional(), + softTrimRatio: z.number().min(0).max(1).optional(), + hardClearRatio: z.number().min(0).max(1).optional(), + minPrunableToolChars: z.number().int().nonnegative().optional(), + tools: z + .object({ + allow: z.array(z.string()).optional(), + deny: z.array(z.string()).optional(), + }) + .optional(), + softTrim: z + .object({ + maxChars: z.number().int().nonnegative().optional(), + headChars: z.number().int().nonnegative().optional(), + tailChars: z.number().int().nonnegative().optional(), + }) + .optional(), + hardClear: z + .object({ + enabled: z.boolean().optional(), + placeholder: z.string().optional(), + }) + .optional(), + }) + .optional(), tools: z .object({ allow: z.array(z.string()).optional(),