Pi agent extension - prune old images from the context

A developer released an extension for the Pi coding agent that prunes old images from the context to manage token usage. The extension, called image-context, allows users to configure per-model image limits and replaces pruned images with a placeholder. It is designed to work with Pi agent's message history and supports both new and legacy state formats.

| import { readFileSync } from "node:fs"; | | import { homedir } from "node:os"; | | import { join } from "node:path"; | | import type { AgentMessage } from "@earendil-works/pi-agent-core"; | | import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; | | | | type JsonObject = Record<string, unknown ; | | | | type ImageLimitRule = { | | provider: string; | | models: string ; | | maxImages: number; | | }; | | | | export type ImageContextConfig = { | | defaultMaxImages: number | null; | | rules: ImageLimitRule ; | | placeholder: string; | | }; | | | | type ImageBlock = { | | type: "image"; | | }; | | | | type ContentBlock = { | | type?: string; | | text?: string; | | }; | | | | type MessageWithContent = AgentMessage & { | | content?: string | ContentBlock ; | | timestamp?: number; | | }; | | | | type PruneResult<T = { | | messages: T ; | | totalImages: number; | | keptImages: number; | | prunedImages: number; | | }; | | | | const CONFIG PATH = join homedir , ".pi", "agent", "image-context.json" ; | | const STATE TYPE = "image-context"; | | const LEGACY STATE TYPE = "qwen-image-context"; | | const STATUS KEY = "image-context"; | | | | function isObject value: unknown : value is JsonObject { | | return typeof value === "object" && value == null && Array.isArray value ; | | } | | | | function parseImageLimit value: unknown, name: string : number | null { | | if value === null return null; | | if typeof value == "number" || Number.isInteger value || value < 0 { | | throw new Error image-context: ${name} must be null or a non-negative integer ; | | } | | return value; | | } | | | | function parseRule value: unknown, index: number : ImageLimitRule { | | if isObject value || typeof value.provider == "string" { | | throw new Error image-context: rules ${index} .provider must be a string ; | | } | | const models = Array.isArray value.models | | ? value.models.filter model : model is string = typeof model === "string" | | : ; | | if models.length === 0 { | | throw new Error image-context: rules ${index} .models must contain a model id or ; | | } | | const maxImages = parseImageLimit value.maxImages, rules ${index} .maxImages ; | | if maxImages === null { | | throw new Error image-context: rules ${index} .maxImages cannot be null ; | | } | | return { provider: value.provider, models, maxImages }; | | } | | | | function loadConfig : ImageContextConfig { | | const raw: unknown = JSON.parse readFileSync CONFIG PATH, "utf8" ; | | if isObject raw throw new Error "image-context: config must be an object" ; | | if typeof raw.placeholder == "string" || raw.placeholder.length === 0 { | | throw new Error "image-context: placeholder must be a non-empty string" ; | | } | | if Array.isArray raw.rules { | | throw new Error "image-context: rules must be an array" ; | | } | | return { | | defaultMaxImages: parseImageLimit raw.defaultMaxImages, "defaultMaxImages" , | | rules: raw.rules.map parseRule , | | placeholder: raw.placeholder, | | }; | | } | | | | export function resolveImageLimit | | config: ImageContextConfig, | | provider: string | undefined, | | model: string | undefined, | | : number | null { | | if provider || model return config.defaultMaxImages; | | const rule = config.rules.find | | candidate = | | candidate.provider === " " || candidate.provider === provider && | | candidate.models.includes " " || candidate.models.includes model , | | ; | | return rule?.maxImages ?? config.defaultMaxImages; | | } | | | | function isImageBlock value: unknown : value is ImageBlock { | | return isObject value && value.type === "image"; | | } | | | | function messageTimestamp message: MessageWithContent : number | undefined { | | return typeof message.timestamp === "number" ? message.timestamp : undefined; | | } | | | | export function pruneImageContext<T extends MessageWithContent | | messages: T , | | maxImages: number, | | clearBeforeTimestamp: number, | | placeholder: string, | | : PruneResult<T { | | let totalImages = 0; | | let keptImages = 0; | | let prunedImages = 0; | | let remaining = maxImages; | | | | for let messageIndex = messages.length - 1; messageIndex = 0; messageIndex-- { | | const message = messages messageIndex ; | | if message || Array.isArray message.content continue; | | const timestamp = messageTimestamp message ; | | const explicitlyCleared = | | clearBeforeTimestamp 0 && | | timestamp === undefined || timestamp <= clearBeforeTimestamp ; | | | | for let blockIndex = 0; blockIndex < message.content.length; blockIndex++ { | | const block = message.content blockIndex ; | | if isImageBlock block continue; | | totalImages++; | | if explicitlyCleared && remaining 0 { | | remaining--; | | keptImages++; | | continue; | | } | | message.content blockIndex = { type: "text", text: placeholder }; | | prunedImages++; | | } | | } | | | | return { messages, totalImages, keptImages, prunedImages }; | | } | | | | export default function imageContext pi: ExtensionAPI : void { | | let config = loadConfig ; | | let clearBeforeTimestamp = 0; | | | | pi.on "session start", async event, ctx = { | | clearBeforeTimestamp = 0; | | for const entry of ctx.sessionManager.getEntries { | | if | | entry.type === "custom" && | | entry.customType === STATE TYPE || entry.customType === LEGACY STATE TYPE && | | isObject entry.data && | | typeof entry.data.clearBeforeTimestamp === "number" | | { | | clearBeforeTimestamp = entry.data.clearBeforeTimestamp; | | } | | } | | ctx.ui.setStatus STATUS KEY, undefined ; | | } ; | | | | pi.on "context", async event, ctx = { | | try { | | config = loadConfig ; | | } catch error { | | console.error error ; | | } | | | | const limit = resolveImageLimit config, ctx.model?.provider, ctx.model?.id ; | | if limit === null && clearBeforeTimestamp === 0 return; | | | | const result = pruneImageContext | | event.messages as MessageWithContent , | | limit ?? Number.MAX SAFE INTEGER, | | clearBeforeTimestamp, | | config.placeholder, | | ; | | if result.totalImages === 0 { | | ctx.ui.setStatus STATUS KEY, undefined ; | | return { messages: result.messages }; | | } | | | | const status = | | result.prunedImages 0 | | ? images ${result.keptImages}/${result.totalImages} | | : images ${result.keptImages} ; | | ctx.ui.setStatus STATUS KEY, ctx.ui.theme.fg "muted", status ; | | return { messages: result.messages }; | | } ; | | | | pi.registerCommand "clear-images", { | | description: "Exclude every image already in this session from future LLM context", | | handler: async args, ctx = { | | clearBeforeTimestamp = Date.now ; | | pi.appendEntry STATE TYPE, { clearBeforeTimestamp } ; | | ctx.ui.notify | | "Existing session images will no longer be sent to the model. New images follow the active model limit.", | | "info", | | ; | | }, | | } ; | | | | pi.registerCommand "image-context", { | | description: "Show the active model image-history limit and clear cutoff", | | handler: async args, ctx = { | | config = loadConfig ; | | const provider = ctx.model?.provider; | | const model = ctx.model?.id; | | const limit = resolveImageLimit config, provider, model ; | | ctx.ui.notify | | ${provider ?? "?"}/${model ?? "?"}: + | | maxImages=${limit ?? "unlimited"}; + | | clearBefore=${clearBeforeTimestamp || "none"} , | | "info", | | ; | | }, | | } ; | | } |