let5see/clawdbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(agent): opt-in tool-result context pruning

Browse Source
This commit is contained in:
Max Sumrall
2026-01-07 12:02:46 +01:00
committed by Peter Steinberger
parent 937e0265a3
commit eeaa6ea46f
9 changed files with 779 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

172
src/agents/pi-embedded-runner.ts
View File

@@ -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<Api> | 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<Api> | 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<ReturnType<typeof createAgentSession>>["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<typeof createAgentSession>
>["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<typeof subscribeEmbeddedPiSession>;
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,