From 475ea69f2c356104f4b601f4976788081e79b8b4 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 31 Jul 2026 19:05:33 -0700 Subject: [PATCH] refactor(agents): consolidate context budgets and compaction recovery (#117149) --- .../compaction-safeguard-quality.ts | 76 +-- .../agent-hooks/compaction-safeguard.ts | 511 ++++++------------ src/agents/compaction-planning-worker.test.ts | 33 +- src/agents/compaction-planning-worker.ts | 209 ++----- src/agents/compaction-planning.worker.ts | 159 ++---- .../context-engine-maintenance.ts | 389 ++++--------- .../run/compaction-runtime.ts | 168 +++++- .../run/overflow-context-recovery.ts | 177 +----- .../run/timeout-context-recovery.ts | 158 +----- .../tool-result-char-estimator.ts | 7 - .../tool-result-context-guard.ts | 129 +---- .../tool-result-truncation.ts | 386 ++++++------- src/agents/harness/tool-result-middleware.ts | 164 ++---- src/agents/session-tool-result-guard.ts | 235 ++++---- src/context-engine/registry.ts | 322 ++++------- 15 files changed, 1050 insertions(+), 2073 deletions(-) diff --git a/src/agents/agent-hooks/compaction-safeguard-quality.ts b/src/agents/agent-hooks/compaction-safeguard-quality.ts index 03d0af54da74..0d99dd74910d 100644 --- a/src/agents/agent-hooks/compaction-safeguard-quality.ts +++ b/src/agents/agent-hooks/compaction-safeguard-quality.ts @@ -39,19 +39,16 @@ function resolveExactIdentifierSectionInstruction( if (policy === "off") { return POLICY_OFF_EXACT_IDENTIFIERS_INSTRUCTION; } - if (policy === "custom") { - const custom = summarizationInstructions?.identifierInstructions?.trim(); - if (custom) { - // Operator text is runtime data, not prompt authority. Wrap it as - // untrusted data before inserting it into compaction instructions. - const customBlock = wrapUntrustedInstructionBlock( + const custom = + policy === "custom" ? summarizationInstructions?.identifierInstructions?.trim() : undefined; + if (custom) { + // Operator text is runtime data, never prompt authority. + return ( + wrapUntrustedInstructionBlock( "For ## Exact identifiers, apply this operator-defined policy text", custom, - ); - if (customBlock) { - return customBlock; - } - } + ) || STRICT_EXACT_IDENTIFIERS_INSTRUCTION + ); } return STRICT_EXACT_IDENTIFIERS_INSTRUCTION; } @@ -71,14 +68,9 @@ export function buildCompactionStructureInstructions( "When prior compaction summaries are present, re-distill them with new messages and remove stale duplicate detail.", ].join("\n"); const custom = customInstructions?.trim(); - if (!custom) { - return sectionsTemplate; - } - const customBlock = wrapUntrustedInstructionBlock("Additional context from /compact", custom); - if (!customBlock) { - return sectionsTemplate; - } - return `${sectionsTemplate}\n\n${customBlock}`; + const customBlock = + custom && wrapUntrustedInstructionBlock("Additional context from /compact", custom); + return customBlock ? `${sectionsTemplate}\n\n${customBlock}` : sectionsTemplate; } function normalizedSummaryLines(summary: string): string[] { @@ -110,26 +102,18 @@ export function buildStructuredFallbackSummary( if (trimmedPreviousSummary && hasRequiredSummarySections(trimmedPreviousSummary)) { return trimmedPreviousSummary; } - const exactIdentifiersSummary = "None captured."; - return [ - "## Decisions", + const values = [ trimmedPreviousSummary || "No prior history.", - "", - "## Open TODOs", "None.", - "", - "## Constraints/Rules", "None.", - "", - "## Pending user asks", "None.", - "", - "## Exact identifiers", - exactIdentifiersSummary, - ].join("\n"); + "None captured.", + ]; + return REQUIRED_SUMMARY_SECTIONS.map((heading, index) => `${heading}\n${values[index]}`).join( + "\n\n", + ); } -/** Append an already-formatted summary section without disturbing empty summaries. */ /** Appends a bounded post-compaction section to an existing summary. */ export function appendSummarySection(summary: string, section: string): string { if (!section) { @@ -172,8 +156,7 @@ export function extractOpaqueIdentifiers(text: string): string[] { return Array.from( new Set( matches - .map((value) => sanitizeExtractedIdentifier(value)) - .map((value) => normalizeOpaqueIdentifier(value)) + .map((value) => normalizeOpaqueIdentifier(sanitizeExtractedIdentifier(value))) .filter((value) => value.length >= 4), ), ).slice(0, MAX_EXTRACTED_IDENTIFIERS); @@ -205,31 +188,16 @@ function hasAskOverlap(summary: string, latestAsk: string | null): boolean { if (askTokens.length === 0) { return true; } - const meaningfulAskTokens = askTokens.filter((token) => { - if (token.length <= 1) { - return false; - } - if (isQueryStopWordToken(token)) { - return false; - } - return true; - }); + const meaningfulAskTokens = askTokens.filter( + (token) => token.length > 1 && !isQueryStopWordToken(token), + ); const tokensToCheck = meaningfulAskTokens.length > 0 ? meaningfulAskTokens : askTokens; - if (tokensToCheck.length === 0) { - return true; - } const summaryTokens = new Set(tokenizeAskOverlapText(summary)); - let overlapCount = 0; - for (const token of tokensToCheck) { - if (summaryTokens.has(token)) { - overlapCount += 1; - } - } + const overlapCount = tokensToCheck.filter((token) => summaryTokens.has(token)).length; const requiredMatches = tokensToCheck.length >= MIN_ASK_OVERLAP_TOKENS_FOR_DOUBLE_MATCH ? 2 : 1; return overlapCount >= requiredMatches; } -/** Audit summary structure, exact identifier preservation, and latest-ask coverage. */ /** Audits a candidate summary for required sections, pending asks, and identifier preservation. */ export function auditSummaryQuality(params: { summary: string; diff --git a/src/agents/agent-hooks/compaction-safeguard.ts b/src/agents/agent-hooks/compaction-safeguard.ts index c9f295525eb7..3eb63ff3cd5f 100644 --- a/src/agents/agent-hooks/compaction-safeguard.ts +++ b/src/agents/agent-hooks/compaction-safeguard.ts @@ -124,16 +124,8 @@ type SessionBranchEntry = { }; function coerceTimestamp(value: unknown): number { - if (typeof value === "number" && Number.isFinite(value)) { - return value; - } - if (typeof value === "string") { - const parsed = Date.parse(value); - if (Number.isFinite(parsed)) { - return parsed; - } - } - return 0; + const timestamp = typeof value === "string" ? Date.parse(value) : value; + return typeof timestamp === "number" && Number.isFinite(timestamp) ? timestamp : 0; } function sessionBranchEntryToMessage(entry: SessionBranchEntry): unknown { @@ -166,22 +158,20 @@ function collectSessionBranchMessages(sessionManager: unknown): AgentMessage[] { if (typeof getBranch !== "function") { return []; } - let entries: unknown; try { - entries = getBranch.call(sessionManager); + const entries: unknown = getBranch.call(sessionManager); + return Array.isArray(entries) + ? entries.flatMap((entry) => { + const message = + entry && typeof entry === "object" + ? sessionBranchEntryToMessage(entry as SessionBranchEntry) + : undefined; + return message ? [message as AgentMessage] : []; + }) + : []; } catch { return []; } - if (!Array.isArray(entries)) { - return []; - } - return entries - .map((entry) => - entry && typeof entry === "object" - ? sessionBranchEntryToMessage(entry as SessionBranchEntry) - : undefined, - ) - .filter((message): message is AgentMessage => Boolean(message)); } function isReplayUnsafeInterSessionInput(message: AgentMessage): boolean { @@ -206,26 +196,16 @@ function isSessionsSendToolName(value: unknown): boolean { function sanitizeSourceSessionSends(messages: AgentMessage[]): AgentMessage[] { const sendCallIds = new Set(); - const resolvedCallIds = new Set(); const resultTextByCallId = new Map(); for (const message of messages) { - if (message.role !== "assistant" || !Array.isArray(message.content)) { + if (message.role !== "assistant") { continue; } - for (const block of message.content) { - if (!block || typeof block !== "object") { - continue; - } - const record = block as { type?: unknown; id?: unknown; name?: unknown }; - if ( - typeof record.type === "string" && - TOOL_CALL_BLOCK_TYPES.has(record.type) && - isSessionsSendToolName(record.name) && - typeof record.id === "string" && - record.id.trim() - ) { - sendCallIds.add(record.id.trim()); + for (const call of extractToolCallsFromAssistant(message)) { + const callId = call.id.trim(); + if (callId && isSessionsSendToolName(call.name)) { + sendCallIds.add(callId); } } } @@ -238,11 +218,10 @@ function sanitizeSourceSessionSends(messages: AgentMessage[]): AgentMessage[] { if (!callId || !sendCallIds.has(callId)) { continue; } - resolvedCallIds.add(callId); - const resultText = extractMessageText(message) || formatNonTextPlaceholder(message.content); - if (resultText) { - resultTextByCallId.set(callId, resultText); - } + resultTextByCallId.set( + callId, + extractMessageText(message) || formatNonTextPlaceholder(message.content) || "", + ); } return messages.flatMap((message) => { @@ -268,7 +247,7 @@ function sanitizeSourceSessionSends(messages: AgentMessage[]): AgentMessage[] { replaced = true; const callId = typeof record.id === "string" ? record.id.trim() : ""; const resultText = callId ? resultTextByCallId.get(callId) : undefined; - const resolved = Boolean(callId && resolvedCallIds.has(callId)); + const resolved = Boolean(callId && resultTextByCallId.has(callId)); const requestText = JSON.stringify({ callId: callId || undefined, args: record.arguments }); return { type: "text", @@ -338,85 +317,15 @@ function containsRealConversation(messages: AgentMessage[]): boolean { ); } -/** - * Attempt provider-based summarization. Returns the summary string on success, - * or `undefined` when the caller should fall back to built-in LLM summarization. - * Rethrows abort/timeout errors so cancellation is always respected. - */ -async function tryProviderSummarize( - provider: CompactionProvider, - params: { - messages: unknown[]; - signal?: AbortSignal; - customInstructions?: string; - summarizationInstructions?: { - identifierPolicy?: "strict" | "off" | "custom"; - identifierInstructions?: string; - }; - previousSummary?: string; - }, -): Promise { - try { - const result = await provider.summarize(params); - if (typeof result === "string" && result.trim()) { - return result; - } - log.warn(`Compaction provider "${provider.id}" returned empty result, falling back to LLM.`); - return undefined; - } catch (err) { - // Propagate only when the caller explicitly cancelled. Provider-side - // AbortErrors (signal not aborted) fall through to LLM summarization. - if (params.signal?.aborted) { - throw err; - } - // Real non-abort transport timeouts (e.g. ETIMEDOUT) still propagate. - if (!isAbortError(err) && isTimeoutError(err)) { - throw err; - } - log.warn( - `Compaction provider "${provider.id}" failed, falling back to LLM: ${err instanceof Error ? err.message : String(err)}`, - ); - return undefined; - } -} - /** * Summarize via the built-in LLM pipeline (summarizeInStages). * Only called when no compaction provider is available or the provider failed. */ -async function summarizeViaLLM(params: { - messages: AgentMessage[]; - model: NonNullable[0]["model"]>; - apiKey: string; - headers?: Record; - signal: AbortSignal; - reserveTokens: number; - maxChunkTokens: number; - contextWindow: number; - customInstructions?: string; - summarizationInstructions?: Parameters[0]["summarizationInstructions"]; - previousSummary?: string; - thinkingLevel?: Parameters[0]["thinkingLevel"]; - streamFn?: Parameters[0]["streamFn"]; -}): Promise { - const messages = prependPreviousSummaryForRedistill({ - messages: params.messages, - previousSummary: params.previousSummary, - }); +async function summarizeViaLLM(params: Parameters[0]): Promise { const result = await compactionSafeguardDeps.summarizeInStages({ - messages, - model: params.model, - apiKey: params.apiKey, - headers: params.headers, - signal: params.signal, - reserveTokens: params.reserveTokens, - maxChunkTokens: params.maxChunkTokens, - contextWindow: params.contextWindow, - customInstructions: params.customInstructions, - summarizationInstructions: params.summarizationInstructions, + ...params, + messages: prependPreviousSummaryForRedistill(params), previousSummary: undefined, - thinkingLevel: params.thinkingLevel, - streamFn: params.streamFn, }); if (result.kind === "summary") { return result.text; @@ -439,12 +348,10 @@ function assembleSuffix(parts: { fileOpsSummary?: string; workspaceContext?: string; }): string { - let suffix = ""; - suffix = appendSummarySection(suffix, parts.splitTurnSection ?? ""); - suffix = appendSummarySection(suffix, parts.preservedTurnsSection ?? ""); - suffix = appendSummarySection(suffix, parts.toolFailureSection ?? ""); - suffix = appendSummarySection(suffix, parts.fileOpsSummary ?? ""); - suffix = appendSummarySection(suffix, parts.workspaceContext ?? ""); + let suffix = Object.values(parts).reduce( + (summary, section) => appendSummarySection(summary, section ?? ""), + "", + ); // Ensure leading separator so suffix does not merge with body (e.g. when body // ends without newline: "...## Exact identifiers## Tool Failures"). if (suffix && !/^\s/.test(suffix)) { @@ -580,18 +487,11 @@ function formatToolFailureMeta(details: unknown): string | undefined { typeof record.exitCode === "number" && Number.isFinite(record.exitCode) ? record.exitCode : undefined; - const parts: string[] = []; - if (status) { - parts.push(`status=${status}`); - } - if (exitCode !== undefined) { - parts.push(`exitCode=${exitCode}`); - } - return parts.length > 0 ? parts.join(" ") : undefined; -} - -function extractToolResultText(content: unknown): string { - return collectTextContentBlocks(content).join("\n"); + const parts = [ + status ? `status=${status}` : "", + exitCode !== undefined ? `exitCode=${exitCode}` : "", + ]; + return parts.filter(Boolean).join(" ") || undefined; } function collectToolFailures(messages: AgentMessage[]): ToolFailure[] { @@ -599,11 +499,7 @@ function collectToolFailures(messages: AgentMessage[]): ToolFailure[] { const seen = new Set(); for (const message of messages) { - if (!message || typeof message !== "object") { - continue; - } - const role = (message as { role?: unknown }).role; - if (role !== "toolResult") { + if (message.role !== "toolResult" || !message.isError) { continue; } const toolResult = message as { @@ -613,9 +509,6 @@ function collectToolFailures(messages: AgentMessage[]): ToolFailure[] { details?: unknown; isError?: unknown; }; - if (toolResult.isError !== true) { - continue; - } // Accepted sessions_spawn launches are successes, not failures, even when a legacy // transcript persisted them with isError:true. Mirror the observer's detection // (toolName + accepted child-run identity, see embedded-agent-subscribe.handlers.tools) @@ -637,7 +530,7 @@ function collectToolFailures(messages: AgentMessage[]): ToolFailure[] { typeof toolResult.toolName === "string" && toolResult.toolName.trim() ? toolResult.toolName : "tool"; - const rawText = extractToolResultText(toolResult.content); + const rawText = collectTextContentBlocks(toolResult.content).join("\n"); const meta = formatToolFailureMeta(toolResult.details); const normalized = normalizeFailureText(rawText); const summary = truncateFailureText( @@ -703,24 +596,13 @@ function formatFileOperations(readFiles: string[], modifiedFiles: string[]): str return lines.length > 0 ? `${openTag}${lines.join("")}${closeTag}` : ""; } - const sections: string[] = []; - const readSection = formatBoundedFileList("read-files", readFiles, MAX_FILE_OPS_LIST_CHARS); - const modifiedSection = formatBoundedFileList( - "modified-files", - modifiedFiles, - MAX_FILE_OPS_LIST_CHARS, - ); - if (readSection) { - sections.push(readSection); - } - if (modifiedSection) { - sections.push(modifiedSection); - } - if (sections.length === 0) { - return ""; - } - const combined = `\n\n${sections.join("\n\n")}`; - return capCompactionSummary(combined, MAX_FILE_OPS_SECTION_CHARS); + const sections = [ + formatBoundedFileList("read-files", readFiles, MAX_FILE_OPS_LIST_CHARS), + formatBoundedFileList("modified-files", modifiedFiles, MAX_FILE_OPS_LIST_CHARS), + ].filter(Boolean); + return sections.length > 0 + ? capCompactionSummary(`\n\n${sections.join("\n\n")}`, MAX_FILE_OPS_SECTION_CHARS) + : ""; } function capCompactionSummary(summary: string, maxChars = MAX_COMPACTION_SUMMARY_CHARS): string { @@ -777,20 +659,15 @@ function extractMessageText(message: AgentMessage): string { if (typeof content === "string") { return content.trim(); } - if (!Array.isArray(content)) { - return ""; - } - const parts: string[] = []; - for (const block of content) { - if (!block || typeof block !== "object") { - continue; - } - const text = (block as { text?: unknown }).text; - if (typeof text === "string" && text.trim().length > 0) { - parts.push(text.trim()); - } - } - return parts.join("\n").trim(); + return Array.isArray(content) + ? content + .flatMap((block) => { + const text = + block && typeof block === "object" ? (block as { text?: unknown }).text : undefined; + return typeof text === "string" && text.trim() ? [text.trim()] : []; + }) + .join("\n") + : ""; } function formatNonTextPlaceholder(content: unknown): string | null { @@ -835,94 +712,68 @@ function splitPreservedRecentTurns(params: { if (preserveTurns <= 0) { return { summarizableMessages: params.messages, preservedMessages: [] }; } - const conversationIndexes: number[] = []; - const userIndexes: number[] = []; - for (const [i, message] of params.messages.entries()) { - const role = message.role; - if (role === "user" || role === "assistant") { - conversationIndexes.push(i); - if (role === "user") { - userIndexes.push(i); - } - } - } + const conversationIndexes = params.messages.flatMap((message, index) => + message.role === "user" || message.role === "assistant" ? [index] : [], + ); if (conversationIndexes.length === 0) { return { summarizableMessages: params.messages, preservedMessages: [] }; } + const userIndexes = conversationIndexes.filter( + (index) => params.messages[index]?.role === "user", + ); const preservedIndexSet = new Set(); if (userIndexes.length >= preserveTurns) { const boundaryStartIndex = userIndexes[userIndexes.length - preserveTurns] ?? -1; - if (boundaryStartIndex >= 0) { - for (const index of conversationIndexes) { - if (index >= boundaryStartIndex) { - preservedIndexSet.add(index); - } + for (const index of conversationIndexes) { + if (index >= boundaryStartIndex) { + preservedIndexSet.add(index); } } } else { - const fallbackMessageCount = preserveTurns * 2; for (const userIndex of userIndexes) { preservedIndexSet.add(userIndex); } - for (let i = conversationIndexes.length - 1; i >= 0; i -= 1) { - const index = conversationIndexes[i]; - if (index === undefined) { - continue; - } + for (const index of conversationIndexes.toReversed()) { preservedIndexSet.add(index); - if (preservedIndexSet.size >= fallbackMessageCount) { + if (preservedIndexSet.size >= preserveTurns * 2) { break; } } } - if (preservedIndexSet.size === 0) { - return { summarizableMessages: params.messages, preservedMessages: [] }; - } const preservedToolCallIds = new Set(); - for (const [i, message] of params.messages.entries()) { - if (!preservedIndexSet.has(i)) { - continue; - } - if (message.role !== "assistant") { - continue; - } - const toolCalls = extractToolCallsFromAssistant(message); - for (const toolCall of toolCalls) { - preservedToolCallIds.add(toolCall.id); + for (const index of preservedIndexSet) { + const message = params.messages[index]; + if (message?.role === "assistant") { + for (const toolCall of extractToolCallsFromAssistant(message)) { + preservedToolCallIds.add(toolCall.id); + } } } if (preservedToolCallIds.size > 0) { - let preservedStartIndex = -1; - for (let i = 0; i < params.messages.length; i += 1) { - if (preservedIndexSet.has(i)) { - preservedStartIndex = i; - break; + const preservedStartIndex = conversationIndexes.find((index) => preservedIndexSet.has(index))!; + for (let index = preservedStartIndex; index < params.messages.length; index += 1) { + const message = params.messages[index]; + if (message?.role !== "toolResult") { + continue; } - } - if (preservedStartIndex >= 0) { - for (const [offset, message] of params.messages.slice(preservedStartIndex).entries()) { - if (message.role !== "toolResult") { - continue; - } - const toolResultId = extractToolResultId(message); - if (toolResultId && preservedToolCallIds.has(toolResultId)) { - preservedIndexSet.add(preservedStartIndex + offset); - } + const toolResultId = extractToolResultId(message); + if (toolResultId && preservedToolCallIds.has(toolResultId)) { + preservedIndexSet.add(index); } } } - const summarizableMessages = params.messages.filter((_, idx) => !preservedIndexSet.has(idx)); + const summarizableMessages: AgentMessage[] = []; + const preservedMessages: AgentMessage[] = []; + for (const [index, message] of params.messages.entries()) { + (preservedIndexSet.has(index) ? preservedMessages : summarizableMessages).push(message); + } // Preserving recent assistant turns can orphan downstream toolResult messages. // Repair pairings here so compaction summarization doesn't trip strict providers. - const repairedSummarizableMessages = repairToolUseResultPairing(summarizableMessages).messages; - const preservedMessages = params.messages - .filter((_, idx) => preservedIndexSet.has(idx)) - .filter((msg) => { - const role = (msg as { role?: unknown }).role; - return role === "user" || role === "assistant" || role === "toolResult"; - }); - return { summarizableMessages: repairedSummarizableMessages, preservedMessages }; + return { + summarizableMessages: repairToolUseResultPairing(summarizableMessages).messages, + preservedMessages, + }; } function formatContextMessages(messages: AgentMessage[]): string[] { @@ -958,26 +809,17 @@ function formatContextMessages(messages: AgentMessage[]): string[] { .filter((line): line is string => Boolean(line)); } -function formatPreservedTurnsSection(messages: AgentMessage[]): string { - if (messages.length === 0) { - return ""; - } +function formatContextSection(messages: AgentMessage[], heading: string): string { const lines = formatContextMessages(messages); - if (lines.length === 0) { - return ""; - } - return `\n\n## Recent turns preserved verbatim\n${lines.join("\n")}`; + return lines.length > 0 ? `${heading}\n${lines.join("\n")}` : ""; +} + +function formatPreservedTurnsSection(messages: AgentMessage[]): string { + return formatContextSection(messages, "\n\n## Recent turns preserved verbatim"); } function formatSplitTurnContextSection(messages: AgentMessage[]): string { - if (messages.length === 0) { - return ""; - } - const lines = formatContextMessages(messages); - if (lines.length === 0) { - return ""; - } - return `**Turn Context (split turn):**\n\n${lines.join("\n")}`; + return formatContextSection(messages, "**Turn Context (split turn):**\n"); } function extractLatestUserAsk(messages: AgentMessage[]): string | null { @@ -1140,75 +982,67 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { const providerId = runtime?.provider; const turnPrefixMessages = baseTurnPrefixMessages; const recentTurnsPreserve = resolveRecentTurnsPreserve(runtime?.recentTurnsPreserve); - const { preservedMessages: providerPreservedMessages } = splitPreservedRecentTurns({ - messages: baseMessagesToSummarize, - recentTurnsPreserve, - }); - const preservedTurnsSection = formatPreservedTurnsSection(providerPreservedMessages); - const splitTurnSection = preparation.isSplitTurn - ? formatSplitTurnContextSection(turnPrefixMessages) - : ""; const structuredInstructions = buildCompactionStructureInstructions( customInstructions, summarizationInstructions, ); + const finalizeSummary = async ( + body: string, + sections: { splitTurnSection?: string; preservedTurnsSection?: string }, + ) => ({ + compaction: { + summary: capCompactionSummaryPreservingSuffix( + body, + assembleSuffix({ + ...sections, + toolFailureSection, + fileOpsSummary, + workspaceContext: await readWorkspaceContextForSummary( + runtime?.postCompactionSections, + runtime?.workspaceDir, + ), + }), + ), + firstKeptEntryId: preparation.firstKeptEntryId, + tokensBefore: preparation.tokensBefore, + details: { readFiles, modifiedFiles }, + }, + }); - // ----------------------------------------------------------------------- - // Provider path — one call with all messages, no LLM-specific prep. - // Falls through to the LLM path below on failure. - // ----------------------------------------------------------------------- if (providerId) { - const compactionProvider = getCompactionProvider(providerId); + const compactionProvider: CompactionProvider | undefined = getCompactionProvider(providerId); if (compactionProvider) { try { // Give the provider ALL messages — no pruning, no chunking, no split-turn splitting. - // The provider handles its own context management. - const allMessages = [...baseMessagesToSummarize, ...turnPrefixMessages]; - const providerResult = await tryProviderSummarize(compactionProvider, { - messages: allMessages, + const providerResult = await compactionProvider.summarize({ + messages: [...baseMessagesToSummarize, ...turnPrefixMessages], signal, customInstructions: structuredInstructions, summarizationInstructions, previousSummary: preparation.previousSummary, }); - - if (providerResult !== undefined) { - // Provider succeeded — assemble suffix metadata and return. - // No quality guard: the provider is trusted. - const workspaceContext = await readWorkspaceContextForSummary( - runtime?.postCompactionSections, - runtime?.workspaceDir, - ); - const suffix = assembleSuffix({ - splitTurnSection, - preservedTurnsSection, - toolFailureSection, - fileOpsSummary, - workspaceContext, + if (typeof providerResult === "string" && providerResult.trim()) { + const { preservedMessages } = splitPreservedRecentTurns({ + messages: baseMessagesToSummarize, + recentTurnsPreserve, + }); + return await finalizeSummary(providerResult, { + splitTurnSection: preparation.isSplitTurn + ? formatSplitTurnContextSection(turnPrefixMessages) + : "", + preservedTurnsSection: formatPreservedTurnsSection(preservedMessages), }); - const summary = capCompactionSummaryPreservingSuffix(providerResult, suffix); - return { - compaction: { - summary, - firstKeptEntryId: preparation.firstKeptEntryId, - tokensBefore: preparation.tokensBefore, - details: { readFiles, modifiedFiles }, - }, - }; } - // Provider returned empty — fall through to LLM path. - log.info("Compaction provider did not produce a result; falling back to LLM path."); + log.warn( + `Compaction provider "${compactionProvider.id}" returned empty result, falling back to LLM.`, + ); } catch (err) { - // tryProviderSummarize rethrows on caller cancellation; reaching here - // means an unexpected error in the assembly step. Fall through to LLM. - if (signal?.aborted) { - throw err; - } - if (!isAbortError(err) && isTimeoutError(err)) { + // Caller cancellation and real transport timeouts remain terminal. + if (signal?.aborted || (!isAbortError(err) && isTimeoutError(err))) { throw err; } log.warn( - `Compaction provider path failed unexpectedly: ${err instanceof Error ? err.message : String(err)}`, + `Compaction provider "${compactionProvider.id}" failed, falling back to LLM: ${formatErrorMessage(err)}`, ); } } else { @@ -1218,9 +1052,6 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { } } - // ----------------------------------------------------------------------- - // LLM path — resolve model + auth, prune, chunk, quality guard. - // ----------------------------------------------------------------------- const model = ctx.model ?? runtime?.model; if (!model) { if (!ctx.model && !runtime?.model && !missedModelWarningSessions.has(ctx.sessionManager)) { @@ -1243,9 +1074,6 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { setCompactionSafeguardCancelReason(ctx.sessionManager, authResult.reason); return { cancel: true }; } - const apiKey = authResult.apiKey ?? ""; - const authHeaders = authResult.headers; - try { const modelContextWindow = resolveContextWindowTokens(model); const contextWindowTokens = runtime?.contextWindowTokens ?? modelContextWindow; @@ -1253,8 +1081,19 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { const headers = buildCompactionSummaryHeaders({ model, messages: messagesToSummarize, - headers: authHeaders, + headers: authResult.headers, }); + const llmSummaryParams = { + model, + apiKey: authResult.apiKey ?? "", + headers, + signal, + reserveTokens: resolveSummaryReserveTokens(preparation.settings.reserveTokens, model), + contextWindow: contextWindowTokens, + summarizationInstructions, + thinkingLevel, + streamFn, + }; const qualityGuardEnabled = runtime?.qualityGuardEnabled ?? false; const qualityGuardMaxRetries = resolveQualityGuardMaxRetries(runtime?.qualityGuardMaxRetries); @@ -1304,22 +1143,11 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { SUMMARIZATION_OVERHEAD_TOKENS, ); droppedSummary = await summarizeViaLLM({ + ...llmSummaryParams, messages: pruned.droppedMessagesList, - model, - apiKey, - headers, - signal, - reserveTokens: resolveSummaryReserveTokens( - preparation.settings.reserveTokens, - model, - ), maxChunkTokens: droppedMaxChunkTokens, - contextWindow: contextWindowTokens, customInstructions: structuredInstructions, - summarizationInstructions, previousSummary: preparation.previousSummary, - thinkingLevel, - streamFn, }); } catch (droppedError) { if (signal?.aborted) { @@ -1364,8 +1192,6 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { 1, Math.floor(contextWindowTokens * adaptiveRatio) - SUMMARIZATION_OVERHEAD_TOKENS, ); - const reserveTokens = resolveSummaryReserveTokens(preparation.settings.reserveTokens, model); - // Feed dropped-messages summary as previousSummary so the main summarization // incorporates context from pruned messages instead of losing it entirely. const effectivePreviousSummary = droppedSummary ?? preparation.previousSummary; @@ -1386,41 +1212,25 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { historySummary = messagesToSummarize.length > 0 ? await summarizeViaLLM({ + ...llmSummaryParams, messages: messagesToSummarize, - model, - apiKey, - headers, - signal, - reserveTokens, maxChunkTokens, - contextWindow: contextWindowTokens, customInstructions: currentInstructions, - summarizationInstructions, previousSummary: effectivePreviousSummary, - thinkingLevel, - streamFn, }) : buildStructuredFallbackSummary(effectivePreviousSummary, summarizationInstructions); summaryWithoutPreservedTurns = historySummary; if (preparation.isSplitTurn && turnPrefixMessages.length > 0) { const prefixSummary = await summarizeViaLLM({ + ...llmSummaryParams, messages: turnPrefixMessages, - model, - apiKey, - headers, - signal, - reserveTokens, maxChunkTokens, - contextWindow: contextWindowTokens, customInstructions: composeSplitTurnInstructions( TURN_PREFIX_INSTRUCTIONS, currentInstructions, ), - summarizationInstructions, previousSummary: undefined, - thinkingLevel, - streamFn, }); splitTurnSectionLocal = `**Turn Context (split turn):**\n\n${prefixSummary}`; summaryWithoutPreservedTurns = historySummary.trim() @@ -1477,30 +1287,11 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { : `${structuredInstructions}\n\n${qualityFeedbackInstruction}`; } - // Cap the main history body first, then append split-turn context, preserved - // turns, diagnostics, and workspace rules so they survive truncation. - const workspaceContext = await readWorkspaceContextForSummary( - runtime?.postCompactionSections, - runtime?.workspaceDir, - ); - const suffix = assembleSuffix({ + // Cap history before suffixes so diagnostics and workspace rules survive. + return await finalizeSummary(lastHistorySummary || summary, { splitTurnSection: lastSplitTurnSection, preservedTurnsSection: preservedTurnsSectionLocal, - toolFailureSection, - fileOpsSummary, - workspaceContext, }); - const bodyToCap = lastHistorySummary || summary; - summary = capCompactionSummaryPreservingSuffix(bodyToCap, suffix); - - return { - compaction: { - summary, - firstKeptEntryId: preparation.firstKeptEntryId, - tokensBefore: preparation.tokensBefore, - details: { readFiles, modifiedFiles }, - }, - }; } catch (error) { // Caller cancellation is terminal, not a safeguard failure. Preserve the // original abort so the runner can classify it without a false data-loss warning. diff --git a/src/agents/compaction-planning-worker.test.ts b/src/agents/compaction-planning-worker.test.ts index 519b4e2d6dc4..e22b963ccf31 100644 --- a/src/agents/compaction-planning-worker.test.ts +++ b/src/agents/compaction-planning-worker.test.ts @@ -57,11 +57,23 @@ describe("compaction planning worker", () => { ).toBe("/repo/dist/agents/compaction-planning.worker.js"); }); - it("rejects invalid worker input", () => { - expect(runCompactionPlanningWorkerInput({ kind: "summaryChunks" })).toEqual({ - status: "failed", - error: "invalid compaction planning worker input", - }); + it("rejects invalid and retired worker input", () => { + for (const input of [ + { kind: "summaryChunks" }, + { + kind: "historyPrune", + messagesToSummarize: [], + turnPrefixMessages: [], + tokensBefore: 0, + contextWindowTokens: 1, + maxHistoryShare: 0.5, + }, + ]) { + expect(runCompactionPlanningWorkerInput(input)).toEqual({ + status: "failed", + error: "invalid compaction planning worker input", + }); + } }); it("plans summary chunks in the packaged worker", () => { @@ -167,6 +179,17 @@ describe("compaction planning worker", () => { expect(value.chunkIndexes.length).toBeGreaterThan(1); }); + it.each([ + { kind: "oversizedFallback", messages: [makeMessage(1)], contextWindow: 1200 }, + { kind: "stageSplit", messages: [makeMessage(1)], maxChunkTokens: 1200 }, + { kind: "adaptiveChunkRatio", messages: [makeMessage(1)], contextWindow: 1200 }, + ])("plans $kind for worker input", (input) => { + expect(runCompactionPlanningWorkerInput(input)).toMatchObject({ + status: "ok", + value: { kind: input.kind }, + }); + }); + it("preserves original user identity while worker fallback omits an oversized tool batch", async () => { const displacedUser = makeMessage(2, "keep the latest real user request"); const messages: AgentMessage[] = [ diff --git a/src/agents/compaction-planning-worker.ts b/src/agents/compaction-planning-worker.ts index 02631d5eb3a2..b80b28ae5ed0 100644 --- a/src/agents/compaction-planning-worker.ts +++ b/src/agents/compaction-planning-worker.ts @@ -163,62 +163,6 @@ function runCompactionPlanningWorker(params: { }); } -function shouldFallbackToMainThread(error: unknown): boolean { - return error instanceof CompactionPlanningWorkerError && error.code === "unavailable"; -} - -function shouldUsePlanningWorker(messageCount: number): boolean { - return messageCount >= COMPACTION_PLANNING_WORKER_MIN_MESSAGES; -} - -function indexSelectedMessages( - indexByMessage: ReadonlyMap, - selected: AgentMessage[], -): number[] { - return selected.map((message) => { - const index = indexByMessage.get(message); - if (index === undefined) { - throw new CompactionPlanningWorkerError( - "compaction planning result contains an unknown message", - "failed", - ); - } - return index; - }); -} - -function indexMessageChunks(source: AgentMessage[], chunks: AgentMessage[][]): number[][] { - const indexByMessage = new Map(source.map((message, index) => [message, index])); - return chunks.map((chunk) => indexSelectedMessages(indexByMessage, chunk)); -} - -function indexOversizedFallbackPlan( - source: AgentMessage[], - plan: OversizedFallbackPlan, -): Extract { - return { - kind: "oversizedFallback", - smallMessageIndexes: indexSelectedMessages( - new Map(source.map((message, index) => [message, index])), - plan.smallMessages, - ), - oversizedNotes: plan.oversizedNotes, - }; -} - -function indexStageSplitPlan( - source: AgentMessage[], - plan: StageSplitPlan, -): Extract { - return plan.mode === "split" - ? { - kind: "stageSplit", - mode: "split", - chunkIndexes: indexMessageChunks(source, plan.chunks), - } - : { kind: "stageSplit", mode: "single" }; -} - function restoreIndexedMessages(source: AgentMessage[], indexes: number[]): AgentMessage[] { return indexes.map((index) => { const message = source.at(index); @@ -232,27 +176,41 @@ function restoreIndexedMessages(source: AgentMessage[], indexes: number[]): Agen }); } -async function runWithUnavailableFallback(params: { - input: CompactionPlanningWorkerInput; +async function runCompactionPlan(params: { + input: TInput; signal?: AbortSignal; - fallback: () => T; - isExpected: (value: CompactionPlanningWorkerValue) => value is T; -}): Promise { + fallback: (messages: AgentMessage[]) => TResult; + restore: ( + value: Extract, + messages: AgentMessage[], + ) => TResult; +}): Promise { + const messages = sanitizeCompactionMessages(params.input.messages); + if (messages.length < COMPACTION_PLANNING_WORKER_MIN_MESSAGES) { + return params.fallback(params.input.messages); + } + try { const value = await runCompactionPlanningWorker({ - input: params.input, + input: { + ...params.input, + messages: projectCompactionMessagesForPlanning(messages), + }, signal: params.signal, }); - if (params.isExpected(value)) { - return value; + if (value.kind !== params.input.kind) { + throw new CompactionPlanningWorkerError( + "unexpected compaction planning worker result", + "failed", + ); } - throw new CompactionPlanningWorkerError( - "unexpected compaction planning worker result", - "failed", + return params.restore( + value as Extract, + messages, ); } catch (error) { - if (shouldFallbackToMainThread(error)) { - return params.fallback(); + if (error instanceof CompactionPlanningWorkerError && error.code === "unavailable") { + return params.fallback(messages); } throw error; } @@ -264,31 +222,17 @@ export async function buildSummaryChunksWithWorker(params: { maxChunkTokens: number; signal?: AbortSignal; }): Promise { - const messages = sanitizeCompactionMessages(params.messages); - if (!shouldUsePlanningWorker(messages.length)) { - return buildSummaryChunks(params); - } - const planningMessages = projectCompactionMessagesForPlanning(messages); - const value = await runWithUnavailableFallback({ + return runCompactionPlan({ input: { kind: "summaryChunks", - messages: planningMessages, + messages: params.messages, maxChunkTokens: params.maxChunkTokens, }, signal: params.signal, - fallback: () => ({ - kind: "summaryChunks" as const, - chunkIndexes: indexMessageChunks( - messages, - buildSummaryChunks({ messages, maxChunkTokens: params.maxChunkTokens }), - ), - }), - isExpected: ( - valueCandidate, - ): valueCandidate is Extract => - valueCandidate.kind === "summaryChunks", + fallback: (messages) => buildSummaryChunks({ messages, maxChunkTokens: params.maxChunkTokens }), + restore: (value, messages) => + value.chunkIndexes.map((indexes) => restoreIndexedMessages(messages, indexes)), }); - return value.chunkIndexes.map((indexes) => restoreIndexedMessages(messages, indexes)); } /** Builds an oversized-message fallback plan, using the worker when worthwhile. */ @@ -297,32 +241,20 @@ export async function buildOversizedFallbackPlanWithWorker(params: { contextWindow: number; signal?: AbortSignal; }): Promise { - const messages = sanitizeCompactionMessages(params.messages); - if (!shouldUsePlanningWorker(messages.length)) { - return buildOversizedFallbackPlan(params); - } - const planningMessages = projectCompactionMessagesForPlanning(messages); - const value = await runWithUnavailableFallback({ + return runCompactionPlan({ input: { kind: "oversizedFallback", - messages: planningMessages, + messages: params.messages, contextWindow: params.contextWindow, }, signal: params.signal, - fallback: () => - indexOversizedFallbackPlan( - messages, - buildOversizedFallbackPlan({ messages, contextWindow: params.contextWindow }), - ), - isExpected: ( - valueEntry, - ): valueEntry is Extract => - valueEntry.kind === "oversizedFallback", + fallback: (messages) => + buildOversizedFallbackPlan({ messages, contextWindow: params.contextWindow }), + restore: (value, messages) => ({ + smallMessages: restoreIndexedMessages(messages, value.smallMessageIndexes), + oversizedNotes: value.oversizedNotes, + }), }); - return { - smallMessages: restoreIndexedMessages(messages, value.smallMessageIndexes), - oversizedNotes: value.oversizedNotes, - }; } /** Builds a staged summarization split plan with worker fallback. */ @@ -333,41 +265,30 @@ export async function buildStageSplitPlanWithWorker(params: { minMessagesForSplit?: number; signal?: AbortSignal; }): Promise { - const messages = sanitizeCompactionMessages(params.messages); - if (!shouldUsePlanningWorker(messages.length)) { - return buildStageSplitPlan(params); - } - const planningMessages = projectCompactionMessagesForPlanning(messages); - const value = await runWithUnavailableFallback({ + return runCompactionPlan({ input: { kind: "stageSplit", - messages: planningMessages, + messages: params.messages, maxChunkTokens: params.maxChunkTokens, parts: params.parts, minMessagesForSplit: params.minMessagesForSplit, }, signal: params.signal, - fallback: () => - indexStageSplitPlan( + fallback: (messages) => + buildStageSplitPlan({ messages, - buildStageSplitPlan({ - messages, - maxChunkTokens: params.maxChunkTokens, - parts: params.parts, - minMessagesForSplit: params.minMessagesForSplit, - }), - ), - isExpected: ( - valueResult, - ): valueResult is Extract => - valueResult.kind === "stageSplit", + maxChunkTokens: params.maxChunkTokens, + parts: params.parts, + minMessagesForSplit: params.minMessagesForSplit, + }), + restore: (value, messages) => + value.mode === "split" + ? { + mode: "split", + chunks: value.chunkIndexes.map((indexes) => restoreIndexedMessages(messages, indexes)), + } + : { mode: "single" }, }); - return value.mode === "split" - ? { - mode: "split", - chunks: value.chunkIndexes.map((indexes) => restoreIndexedMessages(messages, indexes)), - } - : { mode: "single" }; } /** @@ -394,28 +315,16 @@ export async function computeAdaptiveChunkRatioWithWorker(params: { contextWindow: number; signal?: AbortSignal; }): Promise { - const messages = sanitizeCompactionMessages(params.messages); - if (!shouldUsePlanningWorker(messages.length)) { - return computeAdaptiveChunkRatio(params.messages, params.contextWindow); - } - const planningMessages = projectCompactionMessagesForPlanning(messages); - const value = await runWithUnavailableFallback({ + return runCompactionPlan({ input: { kind: "adaptiveChunkRatio", - messages: planningMessages, + messages: params.messages, contextWindow: params.contextWindow, }, signal: params.signal, - fallback: () => ({ - kind: "adaptiveChunkRatio" as const, - ratio: computeAdaptiveChunkRatio(params.messages, params.contextWindow), - }), - isExpected: ( - valueLocal, - ): valueLocal is Extract => - valueLocal.kind === "adaptiveChunkRatio", + fallback: () => computeAdaptiveChunkRatio(params.messages, params.contextWindow), + restore: (value) => value.ratio, }); - return value.ratio; } const compactionPlanningWorkerTesting = { diff --git a/src/agents/compaction-planning.worker.ts b/src/agents/compaction-planning.worker.ts index 147596965c61..e3867d5cc0d0 100644 --- a/src/agents/compaction-planning.worker.ts +++ b/src/agents/compaction-planning.worker.ts @@ -3,12 +3,10 @@ */ import { parentPort, workerData } from "node:worker_threads"; import { - buildHistoryPrunePlan, buildOversizedFallbackPlan, buildStageSplitPlan, buildSummaryChunks, computeAdaptiveChunkRatio, - type HistoryPrunePlan, } from "./compaction-planning.js"; import type { AgentMessage } from "./runtime/index.js"; @@ -31,15 +29,6 @@ export type CompactionPlanningWorkerInput = parts?: number; minMessagesForSplit?: number; } - | { - kind: "historyPrune"; - messagesToSummarize: AgentMessage[]; - turnPrefixMessages: AgentMessage[]; - tokensBefore: number; - contextWindowTokens: number; - maxHistoryShare: number; - parts?: number; - } | { kind: "adaptiveChunkRatio"; messages: AgentMessage[]; @@ -66,9 +55,6 @@ export type CompactionPlanningWorkerValue = mode: "split"; chunkIndexes: number[][]; } - | ({ - kind: "historyPrune"; - } & HistoryPrunePlan) | { kind: "adaptiveChunkRatio"; ratio: number; @@ -85,57 +71,72 @@ export type CompactionPlanningWorkerResult = error: string; }; -function isFiniteNumber(value: unknown): value is number { - return typeof value === "number" && Number.isFinite(value); -} - -function isMessageArray(value: unknown): value is AgentMessage[] { - return Array.isArray(value); -} - function isWorkerInput(value: unknown): value is CompactionPlanningWorkerInput { if (!value || typeof value !== "object" || !("kind" in value)) { return false; } const input = value as Record; + if (!Array.isArray(input.messages)) { + return false; + } switch (input.kind) { case "summaryChunks": - return isMessageArray(input.messages) && isFiniteNumber(input.maxChunkTokens); - case "oversizedFallback": - return isMessageArray(input.messages) && isFiniteNumber(input.contextWindow); case "stageSplit": - return isMessageArray(input.messages) && isFiniteNumber(input.maxChunkTokens); - case "historyPrune": - return ( - isMessageArray(input.messagesToSummarize) && - isMessageArray(input.turnPrefixMessages) && - isFiniteNumber(input.tokensBefore) && - isFiniteNumber(input.contextWindowTokens) && - isFiniteNumber(input.maxHistoryShare) - ); + return typeof input.maxChunkTokens === "number" && Number.isFinite(input.maxChunkTokens); + case "oversizedFallback": case "adaptiveChunkRatio": - return isMessageArray(input.messages) && isFiniteNumber(input.contextWindow); + return typeof input.contextWindow === "number" && Number.isFinite(input.contextWindow); default: return false; } } -function indexSelectedMessages( - indexByMessage: ReadonlyMap, - selected: AgentMessage[], -): number[] { - return selected.map((message) => { - const index = indexByMessage.get(message); - if (index === undefined) { - throw new Error("Compaction planning result contains an unknown message"); - } - return index; - }); +function createMessageIndexer(source: AgentMessage[]): (selected: AgentMessage[]) => number[] { + const indexByMessage = new Map(source.map((message, index) => [message, index])); + return (selected) => + selected.map((message) => { + const index = indexByMessage.get(message); + if (index === undefined) { + throw new Error("Compaction planning result contains an unknown message"); + } + return index; + }); } -function indexMessageChunks(source: AgentMessage[], chunks: AgentMessage[][]): number[][] { - const indexByMessage = new Map(source.map((message, index) => [message, index])); - return chunks.map((chunk) => indexSelectedMessages(indexByMessage, chunk)); +function planCompactionWorkerInput( + input: CompactionPlanningWorkerInput, +): CompactionPlanningWorkerValue { + switch (input.kind) { + case "summaryChunks": + return { + kind: input.kind, + chunkIndexes: buildSummaryChunks(input).map(createMessageIndexer(input.messages)), + }; + case "oversizedFallback": { + const plan = buildOversizedFallbackPlan(input); + return { + kind: input.kind, + smallMessageIndexes: createMessageIndexer(input.messages)(plan.smallMessages), + oversizedNotes: plan.oversizedNotes, + }; + } + case "stageSplit": { + const plan = buildStageSplitPlan(input); + return plan.mode === "split" + ? { + kind: input.kind, + mode: "split", + chunkIndexes: plan.chunks.map(createMessageIndexer(input.messages)), + } + : { kind: input.kind, mode: "single" }; + } + case "adaptiveChunkRatio": + return { + kind: input.kind, + ratio: computeAdaptiveChunkRatio(input.messages, input.contextWindow), + }; + } + throw new Error("unsupported compaction planning worker input"); } /** Run one compaction planning request and return a serializable result. */ @@ -148,65 +149,7 @@ export function runCompactionPlanningWorkerInput(input: unknown): CompactionPlan } try { - switch (input.kind) { - case "summaryChunks": { - const chunks = buildSummaryChunks(input); - return { - status: "ok", - value: { - kind: "summaryChunks", - chunkIndexes: indexMessageChunks(input.messages, chunks), - }, - }; - } - case "oversizedFallback": { - const plan = buildOversizedFallbackPlan(input); - const indexByMessage = new Map(input.messages.map((message, index) => [message, index])); - return { - status: "ok", - value: { - kind: "oversizedFallback", - smallMessageIndexes: indexSelectedMessages(indexByMessage, plan.smallMessages), - oversizedNotes: plan.oversizedNotes, - }, - }; - } - case "stageSplit": { - const plan = buildStageSplitPlan(input); - return { - status: "ok", - value: - plan.mode === "split" - ? { - kind: "stageSplit", - mode: "split", - chunkIndexes: indexMessageChunks(input.messages, plan.chunks), - } - : { kind: "stageSplit", mode: "single" }, - }; - } - case "historyPrune": - return { - status: "ok", - value: { - kind: "historyPrune", - ...buildHistoryPrunePlan(input), - }, - }; - case "adaptiveChunkRatio": - return { - status: "ok", - value: { - kind: "adaptiveChunkRatio", - ratio: computeAdaptiveChunkRatio(input.messages, input.contextWindow), - }, - }; - } - - return { - status: "failed", - error: "unsupported compaction planning worker input", - }; + return { status: "ok", value: planCompactionWorkerInput(input) }; } catch (error) { return { status: "failed", diff --git a/src/agents/embedded-agent-runner/context-engine-maintenance.ts b/src/agents/embedded-agent-runner/context-engine-maintenance.ts index 88e0464b2fa1..262510a48c34 100644 --- a/src/agents/embedded-agent-runner/context-engine-maintenance.ts +++ b/src/agents/embedded-agent-runner/context-engine-maintenance.ts @@ -40,24 +40,35 @@ import { rewriteTranscriptEntriesInSessionManager } from "./transcript-rewrite.j import { resolveRuntimeTranscriptReadTarget } from "./transcript-runtime-state.js"; const TURN_MAINTENANCE_TASK_KIND = "context_engine_turn_maintenance"; -const TURN_MAINTENANCE_TASK_LABEL = "Context engine turn maintenance"; -const TURN_MAINTENANCE_TASK_TASK = "Deferred context-engine maintenance after turn."; const TURN_MAINTENANCE_LANE_PREFIX = "context-engine-turn-maintenance:"; const TURN_MAINTENANCE_LONG_WAIT_MS = 10_000; const DEFERRED_TURN_MAINTENANCE_ABORT_STATE_KEY = Symbol.for( "openclaw.contextEngineTurnMaintenanceAbortState", ); -type DeferredTurnMaintenanceScheduleParams = { - contextEngine: ContextEngine; +type SessionManagerRewriteLock = (operation: () => Promise | T) => Promise; + +type ContextEngineMaintenanceParams = { + contextEngine?: ContextEngine; sessionId: string; - sessionKey: string; + sessionKey?: string; sessionTarget?: ContextEngineSessionTarget; sessionFile: string; + reason: "bootstrap" | "compaction" | "turn"; sessionManager?: Parameters[0]["sessionManager"]; + withSessionManagerRewriteLock?: SessionManagerRewriteLock; runtimeContext?: ContextEngineRuntimeContext; runtimeSettings?: ContextEngineRuntimeSettings; agentId?: string; + executionMode?: "foreground" | "background"; + onDeferredMaintenance?: (promise: Promise) => void; + onDeferredMaintenanceFailure?: (error: unknown) => void; config?: OpenClawConfig; + disposeDeferredContextEngineAfterMaintenance?: boolean; +}; + +type DeferredTurnMaintenanceScheduleParams = ContextEngineMaintenanceParams & { + contextEngine: ContextEngine; + sessionKey: string; disposeContextEngineAfterMaintenance?: boolean; onScheduleFailure?: (error: unknown) => void; }; @@ -70,55 +81,24 @@ type DeferredTurnMaintenanceRunState = { const activeDeferredTurnMaintenanceRuns = new Map(); -type SessionManagerRewriteLock = (operation: () => Promise | T) => Promise; - type DeferredTurnMaintenanceSignal = "SIGINT" | "SIGTERM"; type DeferredTurnMaintenanceProcessLike = Pick & Partial> & { [DEFERRED_TURN_MAINTENANCE_ABORT_STATE_KEY]?: DeferredTurnMaintenanceAbortState; }; type DeferredTurnMaintenanceAbortState = { - registered: boolean; controllers: Set; cleanupHandlers: Map void>; }; -function resolveDeferredTurnMaintenanceAbortState( - processLike: DeferredTurnMaintenanceProcessLike, -): DeferredTurnMaintenanceAbortState { - const existing = processLike[DEFERRED_TURN_MAINTENANCE_ABORT_STATE_KEY]; - if (existing) { - return existing; - } - const created: DeferredTurnMaintenanceAbortState = { - registered: false, - controllers: new Set(), - cleanupHandlers: new Map void>(), - }; - processLike[DEFERRED_TURN_MAINTENANCE_ABORT_STATE_KEY] = created; - return created; -} - function unregisterDeferredTurnMaintenanceAbortSignalHandlers( processLike: DeferredTurnMaintenanceProcessLike, state: DeferredTurnMaintenanceAbortState, ): void { - if (!state.registered) { - return; - } for (const [signal, handler] of state.cleanupHandlers) { processLike.off(signal, handler); } state.cleanupHandlers.clear(); - state.registered = false; -} - -function normalizeSessionKey(sessionKey?: string): string | undefined { - return normalizeOptionalString(sessionKey) || undefined; -} - -function resolveDeferredTurnMaintenanceLane(sessionKey: string): string { - return `${TURN_MAINTENANCE_LANE_PREFIX}${sessionKey}`; } async function disposeDeferredMaintenanceContextEngine( @@ -136,20 +116,16 @@ async function disposeDeferredMaintenanceContextEngine( function createDeferredTurnMaintenanceAbortSignal(params?: { processLike?: DeferredTurnMaintenanceProcessLike; }): { - abortSignal?: AbortSignal; + abortSignal: AbortSignal; dispose: () => void; } { - if (typeof AbortController === "undefined") { - return { abortSignal: undefined, dispose: () => {} }; - } - const processLike = (params?.processLike ?? process) as DeferredTurnMaintenanceProcessLike; - const state = resolveDeferredTurnMaintenanceAbortState(processLike); + const state = (processLike[DEFERRED_TURN_MAINTENANCE_ABORT_STATE_KEY] ??= { + controllers: new Set(), + cleanupHandlers: new Map void>(), + }); const handleTerminationSignal = (signalName: DeferredTurnMaintenanceSignal) => { - const shouldReraise = - typeof processLike.listenerCount === "function" - ? processLike.listenerCount(signalName) === 1 - : false; + const shouldReraise = processLike.listenerCount?.(signalName) === 1; for (const activeController of state.controllers) { if (!activeController.signal.aborted) { activeController.abort( @@ -167,34 +143,24 @@ function createDeferredTurnMaintenanceAbortSignal(params?: { } } }; - if (!state.registered) { - state.registered = true; - const onSigint = () => handleTerminationSignal("SIGINT"); - const onSigterm = () => handleTerminationSignal("SIGTERM"); - state.cleanupHandlers.set("SIGINT", onSigint); - state.cleanupHandlers.set("SIGTERM", onSigterm); - processLike.on("SIGINT", onSigint); - processLike.on("SIGTERM", onSigterm); + if (state.cleanupHandlers.size === 0) { + for (const signal of ["SIGINT", "SIGTERM"] as const) { + const handler = () => handleTerminationSignal(signal); + state.cleanupHandlers.set(signal, handler); + processLike.on(signal, handler); + } } const controller = new AbortController(); state.controllers.add(controller); - let disposed = false; - - const cleanup = () => { - if (disposed) { - return; - } - disposed = true; - state.controllers.delete(controller); - if (state.controllers.size === 0) { - unregisterDeferredTurnMaintenanceAbortSignalHandlers(processLike, state); - } - }; - return { abortSignal: controller.signal, - dispose: cleanup, + dispose: () => { + state.controllers.delete(controller); + if (state.controllers.size === 0) { + unregisterDeferredTurnMaintenanceAbortSignalHandlers(processLike, state); + } + }, }; } @@ -220,28 +186,13 @@ if (process.env.VITEST || process.env.NODE_ENV === "test") { } export async function waitForDeferredTurnMaintenanceForSession(sessionKey?: string): Promise { - const normalizedSessionKey = normalizeSessionKey(sessionKey); + const normalizedSessionKey = normalizeOptionalString(sessionKey); if (!normalizedSessionKey) { return; } await activeDeferredTurnMaintenanceRuns.get(normalizedSessionKey)?.promise; } -function markDeferredTurnMaintenanceTaskScheduleFailure(params: { - sessionKey: string; - taskId: string; - error: unknown; -}): void { - const errorMessage = formatErrorMessage(params.error); - log.warn(`failed to schedule deferred context engine maintenance: ${errorMessage}`); - cancelTaskByIdForOwner({ - taskId: params.taskId, - callerOwnerKey: params.sessionKey, - endedAt: Date.now(), - terminalSummary: `Deferred maintenance could not be scheduled: ${errorMessage}`, - }); -} - function buildTurnMaintenanceTaskDescriptor(params: { sessionKey: string; runId?: string; @@ -259,8 +210,8 @@ function buildTurnMaintenanceTaskDescriptor(params: { ownerKey: params.sessionKey, scopeKind: "session", runId, - label: TURN_MAINTENANCE_TASK_LABEL, - task: TURN_MAINTENANCE_TASK_TASK, + label: "Context engine turn maintenance", + task: "Deferred context-engine maintenance after turn.", notifyPolicy: params.notifyPolicy ?? "silent", // Fast maintenance stays silent and must not create a one-task flow. // Long-running and failed workers promote it to pending before notifying. @@ -269,37 +220,17 @@ function buildTurnMaintenanceTaskDescriptor(params: { }); } -function promoteTurnMaintenanceTaskVisibility(params: { - sessionKey: string; - runId: string; - notifyPolicy: "done_only" | "state_changes"; -}) { - return buildTurnMaintenanceTaskDescriptor({ - sessionKey: params.sessionKey, - runId: params.runId, - notifyPolicy: params.notifyPolicy, - deliveryStatus: "pending", - }); -} - /** * Attach runtime-owned transcript rewrite helpers to an existing * context-engine runtime context payload. */ -function buildContextEngineMaintenanceRuntimeContext(params: { - sessionId: string; - sessionKey?: string; - sessionTarget?: ContextEngineSessionTarget; - sessionFile: string; - sessionManager?: Parameters[0]["sessionManager"]; - withSessionManagerRewriteLock?: SessionManagerRewriteLock; - runtimeContext?: ContextEngineRuntimeContext; - agentId?: string; - allowDeferredCompactionExecution?: boolean; - config?: OpenClawConfig; - purpose?: string; - contextEnginePluginId?: string; -}): ContextEngineRuntimeContext { +function buildContextEngineMaintenanceRuntimeContext( + params: Omit & { + allowDeferredCompactionExecution?: boolean; + purpose?: string; + contextEnginePluginId?: string; + }, +): ContextEngineRuntimeContext { return { ...params.runtimeContext, ...resolveContextEngineCapabilities({ @@ -347,21 +278,12 @@ function buildContextEngineMaintenanceRuntimeContext(params: { }; } -async function executeContextEngineMaintenance(params: { - contextEngine: ContextEngine; - sessionId: string; - sessionKey?: string; - sessionTarget?: ContextEngineSessionTarget; - sessionFile: string; - reason: "bootstrap" | "compaction" | "turn"; - sessionManager?: Parameters[0]["sessionManager"]; - withSessionManagerRewriteLock?: SessionManagerRewriteLock; - runtimeContext?: ContextEngineRuntimeContext; - runtimeSettings?: ContextEngineRuntimeSettings; - agentId?: string; - executionMode: "foreground" | "background"; - config?: OpenClawConfig; -}): Promise { +async function executeContextEngineMaintenance( + params: ContextEngineMaintenanceParams & { + contextEngine: ContextEngine; + executionMode: "foreground" | "background"; + }, +): Promise { if (typeof params.contextEngine.maintain !== "function") { return undefined; } @@ -372,17 +294,11 @@ async function executeContextEngineMaintenance(params: { sessionFile: params.sessionFile, runtimeSettings: params.runtimeSettings, runtimeContext: buildContextEngineMaintenanceRuntimeContext({ - sessionId: params.sessionId, - sessionKey: params.sessionKey, - sessionTarget: params.sessionTarget, - sessionFile: params.sessionFile, + ...params, sessionManager: params.executionMode === "background" ? undefined : params.sessionManager, withSessionManagerRewriteLock: params.executionMode === "background" ? undefined : params.withSessionManagerRewriteLock, - runtimeContext: params.runtimeContext, - agentId: params.agentId, allowDeferredCompactionExecution: params.executionMode === "background", - config: params.config, purpose: `context-engine.${params.reason}.maintenance`, contextEnginePluginId: resolveContextEngineOwnerPluginId(params.contextEngine), }), @@ -397,46 +313,27 @@ async function executeContextEngineMaintenance(params: { return result; } -async function runDeferredTurnMaintenanceWorker(params: { - contextEngine: ContextEngine; - sessionId: string; - sessionKey: string; - sessionTarget?: ContextEngineSessionTarget; - sessionFile: string; - sessionManager?: Parameters[0]["sessionManager"]; - runtimeContext?: ContextEngineRuntimeContext; - runtimeSettings?: ContextEngineRuntimeSettings; - agentId?: string; - runId: string; - config?: OpenClawConfig; - disposeContextEngineAfterMaintenance?: boolean; -}): Promise { +async function runDeferredTurnMaintenanceWorker( + params: DeferredTurnMaintenanceScheduleParams & { + runId: string; + }, +): Promise { let surfacedUserNotice = false; - let longRunningTimer: ReturnType | null = null; + let longRunningTimer: ReturnType | undefined; const shutdownAbort = createDeferredTurnMaintenanceAbortSignal(); - const surfaceMaintenanceUpdate = (summary: string, eventSummary: string) => { - promoteTurnMaintenanceTaskVisibility({ + const taskRun = { runId: params.runId, runtime: "acp" as const, sessionKey: params.sessionKey }; + const makeTaskVisible = (notifyPolicy: "done_only" | "state_changes") => + buildTurnMaintenanceTaskDescriptor({ sessionKey: params.sessionKey, runId: params.runId, - notifyPolicy: "state_changes", + notifyPolicy, + deliveryStatus: "pending", }); - surfacedUserNotice = true; - recordTaskRunProgressByRunId({ - runId: params.runId, - runtime: "acp", - sessionKey: params.sessionKey, - lastEventAt: Date.now(), - progressSummary: summary, - eventSummary, - }); - }; try { const runningAt = Date.now(); startTaskRunByRunId({ - runId: params.runId, - runtime: "acp", - sessionKey: params.sessionKey, + ...taskRun, startedAt: runningAt, lastEventAt: runningAt, progressSummary: "Running deferred maintenance.", @@ -444,39 +341,27 @@ async function runDeferredTurnMaintenanceWorker(params: { }); longRunningTimer = setTimeout(() => { try { - surfaceMaintenanceUpdate( - "Deferred maintenance is still running.", - "Deferred maintenance is still running.", - ); + makeTaskVisible("state_changes"); + surfacedUserNotice = true; + const summary = "Deferred maintenance is still running."; + recordTaskRunProgressByRunId({ + ...taskRun, + lastEventAt: Date.now(), + progressSummary: summary, + eventSummary: summary, + }); } catch (error) { log.warn(`failed to surface deferred maintenance progress: ${String(error)}`); } }, TURN_MAINTENANCE_LONG_WAIT_MS); const result = await executeContextEngineMaintenance({ - contextEngine: params.contextEngine, - sessionId: params.sessionId, - sessionKey: params.sessionKey, - sessionTarget: params.sessionTarget, - sessionFile: params.sessionFile, - reason: "turn", - sessionManager: params.sessionManager, - runtimeContext: params.runtimeContext, - runtimeSettings: params.runtimeSettings, - agentId: params.agentId, - config: params.config, + ...params, executionMode: "background", }); - if (longRunningTimer) { - clearTimeout(longRunningTimer); - longRunningTimer = null; - } - const endedAt = Date.now(); completeTaskRunByRunId({ - runId: params.runId, - runtime: "acp", - sessionKey: params.sessionKey, + ...taskRun, endedAt, lastEventAt: endedAt, progressSummary: result?.changed @@ -487,11 +372,7 @@ async function runDeferredTurnMaintenanceWorker(params: { : "No transcript changes were needed.", }); } catch (err) { - if (shutdownAbort.abortSignal?.aborted) { - if (longRunningTimer) { - clearTimeout(longRunningTimer); - longRunningTimer = null; - } + if (shutdownAbort.abortSignal.aborted) { const task = findTaskByRunIdForOwner({ runId: params.runId, callerOwnerKey: params.sessionKey, @@ -506,23 +387,13 @@ async function runDeferredTurnMaintenanceWorker(params: { } return; } - if (longRunningTimer) { - clearTimeout(longRunningTimer); - longRunningTimer = null; - } const endedAt = Date.now(); const reason = formatErrorMessage(err); if (!surfacedUserNotice) { - promoteTurnMaintenanceTaskVisibility({ - sessionKey: params.sessionKey, - runId: params.runId, - notifyPolicy: "done_only", - }); + makeTaskVisible("done_only"); } failTaskRunByRunId({ - runId: params.runId, - runtime: "acp", - sessionKey: params.sessionKey, + ...taskRun, endedAt, lastEventAt: endedAt, error: reason, @@ -531,6 +402,9 @@ async function runDeferredTurnMaintenanceWorker(params: { }); log.warn(`deferred context engine maintenance failed: ${reason}`); } finally { + if (longRunningTimer) { + clearTimeout(longRunningTimer); + } shutdownAbort.dispose(); if (params.disposeContextEngineAfterMaintenance) { await disposeDeferredMaintenanceContextEngine(params.contextEngine); @@ -541,7 +415,7 @@ async function runDeferredTurnMaintenanceWorker(params: { function scheduleDeferredTurnMaintenance( params: DeferredTurnMaintenanceScheduleParams, ): Promise | undefined { - const sessionKey = normalizeSessionKey(params.sessionKey); + const sessionKey = normalizeOptionalString(params.sessionKey); if (!sessionKey) { return undefined; } @@ -592,37 +466,31 @@ function scheduleDeferredTurnMaintenance( log.warn("[context-engine] failed to create deferred turn maintenance task", { sessionKey }); return undefined; } + const lane = `${TURN_MAINTENANCE_LANE_PREFIX}${sessionKey}`; log.info( `[context-engine] deferred turn maintenance ${reusableTask ? "resuming" : "queued"} ` + - `taskId=${task.taskId} sessionKey=${sessionKey} lane=${resolveDeferredTurnMaintenanceLane(sessionKey)}`, + `taskId=${task.taskId} sessionKey=${sessionKey} lane=${lane}`, ); + const cancelFailedTask = (error: unknown) => { + const errorMessage = formatErrorMessage(error); + log.warn(`failed to schedule deferred context engine maintenance: ${errorMessage}`); + cancelTaskByIdForOwner({ + taskId: task.taskId, + callerOwnerKey: sessionKey, + endedAt: Date.now(), + terminalSummary: `Deferred maintenance could not be scheduled: ${errorMessage}`, + }); + }; const schedulerAbort = createDeferredTurnMaintenanceAbortSignal(); let runPromise: Promise; try { - runPromise = enqueueCommandInLane(resolveDeferredTurnMaintenanceLane(sessionKey), async () => - runDeferredTurnMaintenanceWorker({ - contextEngine: params.contextEngine, - sessionId: params.sessionId, - sessionKey, - sessionTarget: params.sessionTarget, - sessionFile: params.sessionFile, - sessionManager: params.sessionManager, - runtimeContext: params.runtimeContext, - runtimeSettings: params.runtimeSettings, - agentId: params.agentId, - config: params.config, - runId: task.runId!, - disposeContextEngineAfterMaintenance: params.disposeContextEngineAfterMaintenance, - }), + runPromise = enqueueCommandInLane(lane, () => + runDeferredTurnMaintenanceWorker({ ...params, sessionKey, runId: task.runId! }), ); } catch (err) { schedulerAbort.dispose(); - markDeferredTurnMaintenanceTaskScheduleFailure({ - sessionKey, - taskId: task.taskId, - error: err, - }); + cancelFailedTask(err); return undefined; } const cleanupDeferredTurnMaintenance = async () => { @@ -631,7 +499,7 @@ function scheduleDeferredTurnMaintenance( if (current !== state) { return; } - const shutdownTriggered = schedulerAbort.abortSignal?.aborted === true; + const shutdownTriggered = schedulerAbort.abortSignal.aborted; const rerunParams = current.rerunRequested && !shutdownTriggered ? current.latestParams : undefined; const discardedRerunParams = @@ -646,15 +514,11 @@ function scheduleDeferredTurnMaintenance( const trackedPromise = runPromise .catch((err: unknown) => { params.onScheduleFailure?.(err); - markDeferredTurnMaintenanceTaskScheduleFailure({ - sessionKey, - taskId: task.taskId, - error: err, - }); + cancelFailedTask(err); }) - .then(cleanupDeferredTurnMaintenance, async (err: unknown) => { + .then(cleanupDeferredTurnMaintenance, async (error: unknown) => { await cleanupDeferredTurnMaintenance(); - throw err; + throw error; }); const state: DeferredTurnMaintenanceRunState = { promise: trackedPromise, @@ -669,25 +533,11 @@ function scheduleDeferredTurnMaintenance( /** * Run optional context-engine transcript maintenance and normalize the result. */ -export async function runContextEngineMaintenance(params: { - contextEngine?: ContextEngine; - sessionId: string; - sessionKey?: string; - sessionTarget?: ContextEngineSessionTarget; - sessionFile: string; - reason: "bootstrap" | "compaction" | "turn"; - sessionManager?: Parameters[0]["sessionManager"]; - withSessionManagerRewriteLock?: SessionManagerRewriteLock; - runtimeContext?: ContextEngineRuntimeContext; - runtimeSettings?: ContextEngineRuntimeSettings; - agentId?: string; - executionMode?: "foreground" | "background"; - onDeferredMaintenance?: (promise: Promise) => void; - onDeferredMaintenanceFailure?: (error: unknown) => void; - config?: OpenClawConfig; - disposeDeferredContextEngineAfterMaintenance?: boolean; -}): Promise { - if (typeof params.contextEngine?.maintain !== "function") { +export async function runContextEngineMaintenance( + params: ContextEngineMaintenanceParams, +): Promise { + const contextEngine = params.contextEngine; + if (typeof contextEngine?.maintain !== "function") { return undefined; } @@ -695,21 +545,14 @@ export async function runContextEngineMaintenance(params: { const shouldDefer = params.reason === "turn" && executionMode !== "background" && - params.contextEngine.info.turnMaintenanceMode === "background"; + contextEngine.info.turnMaintenanceMode === "background"; if (shouldDefer) { try { const deferred = scheduleDeferredTurnMaintenance({ - contextEngine: params.contextEngine, - sessionId: params.sessionId, + ...params, + contextEngine, sessionKey: params.sessionKey ?? params.sessionId, - sessionTarget: params.sessionTarget, - sessionFile: params.sessionFile, - sessionManager: params.sessionManager, - runtimeContext: params.runtimeContext, - runtimeSettings: params.runtimeSettings, - agentId: params.agentId, - config: params.config, disposeContextEngineAfterMaintenance: params.disposeDeferredContextEngineAfterMaintenance, onScheduleFailure: params.onDeferredMaintenanceFailure, }); @@ -723,21 +566,7 @@ export async function runContextEngineMaintenance(params: { } try { - return await executeContextEngineMaintenance({ - contextEngine: params.contextEngine, - sessionId: params.sessionId, - sessionKey: params.sessionKey, - sessionTarget: params.sessionTarget, - sessionFile: params.sessionFile, - reason: params.reason, - sessionManager: params.sessionManager, - withSessionManagerRewriteLock: params.withSessionManagerRewriteLock, - runtimeContext: params.runtimeContext, - runtimeSettings: params.runtimeSettings, - agentId: params.agentId, - executionMode, - config: params.config, - }); + return await executeContextEngineMaintenance({ ...params, contextEngine, executionMode }); } catch (err) { log.warn(`context engine maintain failed (${params.reason}): ${String(err)}`); return undefined; diff --git a/src/agents/embedded-agent-runner/run/compaction-runtime.ts b/src/agents/embedded-agent-runner/run/compaction-runtime.ts index 571d7932dc33..cc3289924367 100644 --- a/src/agents/embedded-agent-runner/run/compaction-runtime.ts +++ b/src/agents/embedded-agent-runner/run/compaction-runtime.ts @@ -1,11 +1,177 @@ import type { resolveContextEngine } from "../../../context-engine/registry.js"; -import { resolveCompactionSuccessorTranscript } from "../../../context-engine/types.js"; +import type { buildContextEngineRuntimeSettings } from "../../../context-engine/runtime-settings.js"; +import { + resolveCompactionSuccessorTranscript, + type ContextEngineSessionTarget, +} from "../../../context-engine/types.js"; +import { resolveProcessToolScopeKey } from "../../agent-tools.js"; +import { listActiveProcessSessionReferences } from "../../bash-process-references.js"; +import { buildEmbeddedCompactionRuntimeContext } from "../compaction-runtime-context.js"; +import { + compactContextEngineWithSafetyTimeout, + resolveCompactionTimeoutMs, +} from "../compaction-safety-timeout.js"; +import { resolveContextEngineCapabilities } from "../context-engine-capabilities.js"; import { log } from "../logger.js"; +import type { EmbeddedRunContextRecoveryState } from "./context-recovery-state.js"; import type { PreparedEmbeddedRunInput } from "./execution-context.js"; +import type { RunEmbeddedAgentParams } from "./params.js"; +import { buildContextEngineCompactionSessionTarget } from "./session-bootstrap.js"; import type { createEmbeddedRunSessionPromptState } from "./session-prompt-state.js"; +import type { EmbeddedRunAttemptResult } from "./types.js"; type ContextEngine = Awaited>; type SessionPromptState = ReturnType; +type CompactionResult = Awaited>; + +export type EmbeddedRunCompactionRecoveryInput = { + runParams: RunEmbeddedAgentParams; + state: EmbeddedRunContextRecoveryState; + contextEngine: ContextEngine; + contextTokenBudget?: number; + genericCompactionRecoveryAllowed: boolean; + attempt: EmbeddedRunAttemptResult; + runtimeAuthPlan: Parameters[0]["runtimeAuthPlan"]; + resolvedSessionKey: string; + sessionAgentId: string; + agentDir: string; + workspaceDir: string; + provider: string; + modelId: string; + harnessRuntime: string; + thinkLevel: Parameters[0]["thinkLevel"]; + authProfileId?: string; + authProfileIdSource: "auto" | "user"; + resolveContextEnginePluginId: () => string | undefined; + buildRuntimeSettings: (settings: { + tokenBudget?: number | null; + degradedReason?: string | null; + }) => ReturnType; + onCompactionHookMessages: (payload: { + phase: "before" | "after"; + messages: string[]; + }) => Promise; + runOwnsCompactionBeforeHook: (reason: string) => Promise; + runOwnsCompactionAfterHook: ( + reason: string, + result: CompactionResult, + previousSessionId?: string, + ) => Promise; + adoptCompactionTranscript: (result: CompactionResult) => Promise; + getActiveSession: () => { + id: string; + file: string; + target?: ContextEngineSessionTarget; + }; + armPostCompactionGuard: () => void; +}; + +/** Preserve one prepared owner snapshot for both timeout and overflow recovery. */ +export async function compactEmbeddedRunForRecovery( + input: EmbeddedRunCompactionRecoveryInput, + recovery: { + tokenBudget: number; + trigger: "overflow" | "timeout_recovery"; + diagId: string; + attempt: number; + maxAttempts: number; + currentTokenCount?: number; + }, +) { + const { runParams } = input; + const activeSession = input.getActiveSession(); + const runtimeContext = { + ...buildEmbeddedCompactionRuntimeContext({ + sessionKey: runParams.sessionKey, + messageChannel: runParams.messageChannel, + messageProvider: runParams.messageProvider, + clientCaps: runParams.clientCaps, + chatType: runParams.chatType, + agentAccountId: runParams.agentAccountId, + currentChannelId: runParams.currentChannelId, + currentThreadTs: runParams.currentThreadTs, + currentMessageId: runParams.currentMessageId, + authProfileId: input.authProfileId, + authProfileIdSource: input.authProfileIdSource, + runtimeAuthPlan: input.runtimeAuthPlan, + workspaceDir: input.workspaceDir, + agentDir: input.agentDir, + config: runParams.config, + toolOverrides: runParams.toolOverrides, + skillsSnapshot: runParams.skillsSnapshot, + senderId: runParams.senderId, + provider: input.provider, + modelId: input.modelId, + harnessRuntime: input.harnessRuntime, + modelSelectionLocked: runParams.modelSelectionLocked, + modelFallbacksOverride: runParams.modelFallbacksOverride, + thinkLevel: input.thinkLevel, + reasoningLevel: runParams.reasoningLevel, + bashElevated: runParams.bashElevated, + extraSystemPrompt: runParams.extraSystemPrompt, + sourceReplyDeliveryMode: runParams.sourceReplyDeliveryMode, + ownerNumbers: runParams.ownerNumbers, + activeProcessSessions: listActiveProcessSessionReferences({ + scopeKey: resolveProcessToolScopeKey({ + sessionKey: runParams.sandboxSessionKey?.trim() || runParams.sessionKey, + sessionId: activeSession.id, + agentId: input.sessionAgentId, + }), + }), + }), + ...resolveContextEngineCapabilities({ + config: runParams.config, + sessionKey: runParams.sessionKey, + agentId: input.sessionAgentId, + contextEnginePluginId: input.resolveContextEnginePluginId(), + purpose: + recovery.trigger === "overflow" + ? "context-engine.overflow-compaction" + : "context-engine.timeout-compaction", + }), + onCompactionHookMessages: input.onCompactionHookMessages, + ...(input.attempt.promptCache ? { promptCache: input.attempt.promptCache } : {}), + runId: runParams.runId, + trigger: recovery.trigger, + ...(recovery.currentTokenCount !== undefined + ? { currentTokenCount: recovery.currentTokenCount } + : {}), + diagId: recovery.diagId, + attempt: recovery.attempt, + maxAttempts: recovery.maxAttempts, + }; + const runtimeSettings = input.buildRuntimeSettings({ + tokenBudget: recovery.tokenBudget, + ...(recovery.trigger === "overflow" ? { degradedReason: "context_overflow" } : {}), + }); + const result = await compactContextEngineWithSafetyTimeout( + input.contextEngine, + { + sessionId: activeSession.id, + sessionKey: input.resolvedSessionKey, + agentId: input.sessionAgentId, + sessionTarget: buildContextEngineCompactionSessionTarget({ + agentId: input.sessionAgentId, + config: runParams.config, + sessionFile: activeSession.file, + sessionId: activeSession.id, + sessionKey: input.resolvedSessionKey, + sessionTarget: activeSession.target, + }), + tokenBudget: recovery.tokenBudget, + ...(recovery.currentTokenCount !== undefined + ? { currentTokenCount: recovery.currentTokenCount } + : {}), + force: true, + compactionTarget: "budget", + runtimeContext, + runtimeSettings, + }, + resolveCompactionTimeoutMs(runParams.config), + runParams.abortSignal, + ); + return { result, runtimeContext, runtimeSettings }; +} export function createEmbeddedRunCompactionRuntime(input: { runParams: PreparedEmbeddedRunInput["runParams"]; diff --git a/src/agents/embedded-agent-runner/run/overflow-context-recovery.ts b/src/agents/embedded-agent-runner/run/overflow-context-recovery.ts index 782502ba89a8..eafc52fbca22 100644 --- a/src/agents/embedded-agent-runner/run/overflow-context-recovery.ts +++ b/src/agents/embedded-agent-runner/run/overflow-context-recovery.ts @@ -1,23 +1,14 @@ import { isContextOverflow } from "@openclaw/ai/internal/runtime"; import { truncateUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; -import { buildContextEngineRuntimeSettings } from "../../../context-engine/runtime-settings.js"; -import type { ContextEngine, ContextEngineSessionTarget } from "../../../context-engine/types.js"; +import type { ContextEngine } from "../../../context-engine/types.js"; import { formatErrorMessage } from "../../../infra/errors.js"; import type { AssistantMessage } from "../../../llm/types.js"; import { projectAgentRunAttemptTerminal } from "../../agent-run-terminal-outcome.js"; -import { resolveProcessToolScopeKey } from "../../agent-tools.js"; -import { listActiveProcessSessionReferences } from "../../bash-process-references.js"; import { extractObservedOverflowTokenCount, isCompactionFailureError, isLikelyContextOverflowError, } from "../../embedded-agent-helpers.js"; -import { buildEmbeddedCompactionRuntimeContext } from "../compaction-runtime-context.js"; -import { - compactContextEngineWithSafetyTimeout, - resolveCompactionTimeoutMs, -} from "../compaction-safety-timeout.js"; -import { resolveContextEngineCapabilities } from "../context-engine-capabilities.js"; import { runContextEngineMaintenance } from "../context-engine-maintenance.js"; import { log } from "../logger.js"; import { @@ -30,26 +21,20 @@ import { sessionLikelyHasOversizedToolResults, truncateOversizedToolResultsInActiveTarget, } from "../tool-result-truncation.js"; -import type { EmbeddedRunContextRecoveryState } from "./context-recovery-state.js"; -import { createCompactionDiagId } from "./helpers.js"; -import type { RunEmbeddedAgentParams } from "./params.js"; import { - buildContextEngineCompactionSessionTarget, + compactEmbeddedRunForRecovery, + type EmbeddedRunCompactionRecoveryInput, +} from "./compaction-runtime.js"; +import { createCompactionDiagId } from "./helpers.js"; +import { isNoRealConversationCompactionNoop, resetNoRealConversationTokenSnapshot, } from "./session-bootstrap.js"; -import type { EmbeddedRunAttemptResult } from "./types.js"; const MAX_OVERFLOW_COMPACTION_ATTEMPTS = 3; type CompactResult = Awaited>; -type ActiveSession = { - id: string; - file: string; - target?: ContextEngineSessionTarget; -}; - type EmbeddedRunOverflowRecoveryOutcome = | { action: "none" } | { action: "retry" } @@ -60,52 +45,19 @@ type EmbeddedRunOverflowRecoveryOutcome = userText: string; }; -export async function recoverEmbeddedRunOverflow(input: { - runParams: RunEmbeddedAgentParams; - state: EmbeddedRunContextRecoveryState; - contextEngine: ContextEngine; - contextTokenBudget?: number; - genericCompactionRecoveryAllowed: boolean; - aborted: boolean; - signalOwnedInterruption: boolean; - promptError: unknown; - assistantErrorText?: string; - assistantOverflowCandidate?: AssistantMessage; - attempt: EmbeddedRunAttemptResult; - toolResultPromptProjectionState: ToolResultPromptProjectionState; - attemptCompactionCount: number; - runtimeAuthPlan: Parameters[0]["runtimeAuthPlan"]; - resolvedSessionKey: string; - sessionAgentId: string; - agentDir: string; - workspaceDir: string; - provider: string; - modelId: string; - harnessRuntime: string; - thinkLevel: Parameters[0]["thinkLevel"]; - authProfileId?: string; - authProfileIdSource: "auto" | "user"; - resolveContextEnginePluginId: () => string | undefined; - buildRuntimeSettings: (settings: { - tokenBudget?: number | null; - degradedReason?: string | null; - }) => ReturnType; - onCompactionHookMessages: (payload: { - phase: "before" | "after"; - messages: string[]; - }) => Promise; - runOwnsCompactionBeforeHook: (reason: string) => Promise; - runOwnsCompactionAfterHook: ( - reason: string, - result: CompactResult, - previousSessionId?: string, - ) => Promise; - adoptCompactionTranscript: (result: CompactResult) => Promise; - getActiveSession: () => ActiveSession; - prepareCurrentTranscriptRetry: () => void; - prepareCompactedTranscriptRetry: () => Promise; - armPostCompactionGuard: () => void; -}): Promise { +export async function recoverEmbeddedRunOverflow( + input: EmbeddedRunCompactionRecoveryInput & { + aborted: boolean; + signalOwnedInterruption: boolean; + promptError: unknown; + assistantErrorText?: string; + assistantOverflowCandidate?: AssistantMessage; + toolResultPromptProjectionState: ToolResultPromptProjectionState; + attemptCompactionCount: number; + prepareCurrentTranscriptRetry: () => void; + prepareCompactedTranscriptRetry: () => Promise; + }, +): Promise { const contextOverflowError = !input.aborted && !input.signalOwnedInterruption ? (() => { @@ -213,94 +165,15 @@ export async function recoverEmbeddedRunOverflow(input: { let previousSessionId: string | undefined; await input.runOwnsCompactionBeforeHook("overflow recovery"); try { - const sessionBeforeCompaction = input.getActiveSession(); - const overflowCompactionRuntimeContext = { - ...buildEmbeddedCompactionRuntimeContext({ - sessionKey: runParams.sessionKey, - messageChannel: runParams.messageChannel, - messageProvider: runParams.messageProvider, - clientCaps: runParams.clientCaps, - chatType: runParams.chatType, - agentAccountId: runParams.agentAccountId, - currentChannelId: runParams.currentChannelId, - currentThreadTs: runParams.currentThreadTs, - currentMessageId: runParams.currentMessageId, - authProfileId: input.authProfileId, - authProfileIdSource: input.authProfileIdSource, - runtimeAuthPlan: input.runtimeAuthPlan, - workspaceDir: input.workspaceDir, - agentDir: input.agentDir, - config: runParams.config, - toolOverrides: runParams.toolOverrides, - skillsSnapshot: runParams.skillsSnapshot, - senderId: runParams.senderId, - provider: input.provider, - modelId: input.modelId, - harnessRuntime: input.harnessRuntime, - modelSelectionLocked: runParams.modelSelectionLocked, - modelFallbacksOverride: runParams.modelFallbacksOverride, - thinkLevel: input.thinkLevel, - reasoningLevel: runParams.reasoningLevel, - bashElevated: runParams.bashElevated, - extraSystemPrompt: runParams.extraSystemPrompt, - sourceReplyDeliveryMode: runParams.sourceReplyDeliveryMode, - ownerNumbers: runParams.ownerNumbers, - activeProcessSessions: listActiveProcessSessionReferences({ - scopeKey: resolveProcessToolScopeKey({ - sessionKey: runParams.sandboxSessionKey?.trim() || runParams.sessionKey, - sessionId: sessionBeforeCompaction.id, - agentId: input.sessionAgentId, - }), - }), - }), - ...resolveContextEngineCapabilities({ - config: runParams.config, - sessionKey: runParams.sessionKey, - agentId: input.sessionAgentId, - contextEnginePluginId: input.resolveContextEnginePluginId(), - purpose: "context-engine.overflow-compaction", - }), - onCompactionHookMessages: input.onCompactionHookMessages, - ...(input.attempt.promptCache ? { promptCache: input.attempt.promptCache } : {}), - runId: runParams.runId, + const compaction = await compactEmbeddedRunForRecovery(input, { + tokenBudget: input.contextTokenBudget, trigger: "overflow", - ...(overflowTokenCountForCompaction !== undefined - ? { currentTokenCount: overflowTokenCountForCompaction } - : {}), diagId: overflowDiagId, attempt: input.state.overflowCompactionAttempts, maxAttempts: MAX_OVERFLOW_COMPACTION_ATTEMPTS, - }; - const overflowCompactionRuntimeSettings = input.buildRuntimeSettings({ - tokenBudget: input.contextTokenBudget, - degradedReason: "context_overflow", + currentTokenCount: overflowTokenCountForCompaction, }); - compactResult = await compactContextEngineWithSafetyTimeout( - input.contextEngine, - { - sessionId: sessionBeforeCompaction.id, - sessionKey: input.resolvedSessionKey, - agentId: input.sessionAgentId, - sessionTarget: buildContextEngineCompactionSessionTarget({ - agentId: input.sessionAgentId, - config: runParams.config, - sessionFile: sessionBeforeCompaction.file, - sessionId: sessionBeforeCompaction.id, - sessionKey: input.resolvedSessionKey, - sessionTarget: sessionBeforeCompaction.target, - }), - tokenBudget: input.contextTokenBudget, - ...(overflowTokenCountForCompaction !== undefined - ? { currentTokenCount: overflowTokenCountForCompaction } - : {}), - force: true, - compactionTarget: "budget", - runtimeContext: overflowCompactionRuntimeContext, - runtimeSettings: overflowCompactionRuntimeSettings, - }, - resolveCompactionTimeoutMs(runParams.config), - runParams.abortSignal, - ); + compactResult = compaction.result; if (compactResult.ok && compactResult.compacted) { previousSessionId = await input.adoptCompactionTranscript(compactResult); const sessionAfterCompaction = input.getActiveSession(); @@ -311,8 +184,8 @@ export async function recoverEmbeddedRunOverflow(input: { sessionTarget: sessionAfterCompaction.target, sessionFile: sessionAfterCompaction.file, reason: "compaction", - runtimeContext: overflowCompactionRuntimeContext, - runtimeSettings: overflowCompactionRuntimeSettings, + runtimeContext: compaction.runtimeContext, + runtimeSettings: compaction.runtimeSettings, config: runParams.config, agentId: input.sessionAgentId, }); diff --git a/src/agents/embedded-agent-runner/run/timeout-context-recovery.ts b/src/agents/embedded-agent-runner/run/timeout-context-recovery.ts index 88f42a021a0c..7c9b52acedf3 100644 --- a/src/agents/embedded-agent-runner/run/timeout-context-recovery.ts +++ b/src/agents/embedded-agent-runner/run/timeout-context-recovery.ts @@ -1,74 +1,24 @@ -import { buildContextEngineRuntimeSettings } from "../../../context-engine/runtime-settings.js"; -import type { ContextEngine, ContextEngineSessionTarget } from "../../../context-engine/types.js"; -import { resolveProcessToolScopeKey } from "../../agent-tools.js"; -import { listActiveProcessSessionReferences } from "../../bash-process-references.js"; import { deriveContextPromptTokens, normalizeUsage } from "../../usage.js"; import { runPostCompactionSideEffects } from "../compaction-hooks.js"; -import { buildEmbeddedCompactionRuntimeContext } from "../compaction-runtime-context.js"; -import { - compactContextEngineWithSafetyTimeout, - resolveCompactionTimeoutMs, -} from "../compaction-safety-timeout.js"; -import { resolveContextEngineCapabilities } from "../context-engine-capabilities.js"; import { log } from "../logger.js"; -import type { EmbeddedRunContextRecoveryState } from "./context-recovery-state.js"; +import { + compactEmbeddedRunForRecovery, + type EmbeddedRunCompactionRecoveryInput, +} from "./compaction-runtime.js"; import { createCompactionDiagId } from "./helpers.js"; -import type { RunEmbeddedAgentParams } from "./params.js"; -import { buildContextEngineCompactionSessionTarget } from "./session-bootstrap.js"; -import type { EmbeddedRunAttemptResult } from "./types.js"; const MAX_TIMEOUT_COMPACTION_ATTEMPTS = 2; -type CompactResult = Awaited>; - -type ActiveSession = { - id: string; - file: string; - target?: ContextEngineSessionTarget; -}; - -export async function recoverEmbeddedRunTimeout(input: { - runParams: RunEmbeddedAgentParams; - state: EmbeddedRunContextRecoveryState; - contextEngine: ContextEngine; - contextTokenBudget?: number; - genericCompactionRecoveryAllowed: boolean; - timedOut: boolean; - signalOwnedInterruption: boolean; - timedOutDuringCompaction: boolean; - timedOutDuringToolExecution: boolean; - timedOutByRunBudget: boolean; - lastRunPromptUsage?: ReturnType; - attempt: EmbeddedRunAttemptResult; - runtimeAuthPlan: Parameters[0]["runtimeAuthPlan"]; - resolvedSessionKey: string; - sessionAgentId: string; - agentDir: string; - workspaceDir: string; - provider: string; - modelId: string; - harnessRuntime: string; - thinkLevel: Parameters[0]["thinkLevel"]; - authProfileId?: string; - authProfileIdSource: "auto" | "user"; - resolveContextEnginePluginId: () => string | undefined; - buildRuntimeSettings: (settings: { - tokenBudget?: number | null; - }) => ReturnType; - onCompactionHookMessages: (payload: { - phase: "before" | "after"; - messages: string[]; - }) => Promise; - runOwnsCompactionBeforeHook: (reason: string) => Promise; - runOwnsCompactionAfterHook: ( - reason: string, - result: CompactResult, - previousSessionId?: string, - ) => Promise; - adoptCompactionTranscript: (result: CompactResult) => Promise; - getActiveSession: () => ActiveSession; - armPostCompactionGuard: () => void; -}): Promise { +export async function recoverEmbeddedRunTimeout( + input: EmbeddedRunCompactionRecoveryInput & { + timedOut: boolean; + signalOwnedInterruption: boolean; + timedOutDuringCompaction: boolean; + timedOutDuringToolExecution: boolean; + timedOutByRunBudget: boolean; + lastRunPromptUsage?: ReturnType; + }, +): Promise { if ( !input.genericCompactionRecoveryAllowed || input.contextTokenBudget === undefined || @@ -106,88 +56,16 @@ export async function recoverEmbeddedRunTimeout(input: { `[timeout-compaction] LLM timed out with high prompt token usage (${Math.round(tokenUsedRatio * 100)}%); ` + `attempting compaction before retry (attempt ${input.state.timeoutCompactionAttempts}/${MAX_TIMEOUT_COMPACTION_ATTEMPTS}) diagId=${timeoutDiagId}`, ); - let timeoutCompactResult: CompactResult; + let timeoutCompactResult: Awaited>; await input.runOwnsCompactionBeforeHook("timeout recovery"); try { - const activeSession = input.getActiveSession(); - const runParams = input.runParams; - const timeoutCompactionRuntimeContext = { - ...buildEmbeddedCompactionRuntimeContext({ - sessionKey: runParams.sessionKey, - messageChannel: runParams.messageChannel, - messageProvider: runParams.messageProvider, - clientCaps: runParams.clientCaps, - chatType: runParams.chatType, - agentAccountId: runParams.agentAccountId, - currentChannelId: runParams.currentChannelId, - currentThreadTs: runParams.currentThreadTs, - currentMessageId: runParams.currentMessageId, - authProfileId: input.authProfileId, - authProfileIdSource: input.authProfileIdSource, - runtimeAuthPlan: input.runtimeAuthPlan, - workspaceDir: input.workspaceDir, - agentDir: input.agentDir, - config: runParams.config, - toolOverrides: runParams.toolOverrides, - skillsSnapshot: runParams.skillsSnapshot, - senderId: runParams.senderId, - provider: input.provider, - modelId: input.modelId, - harnessRuntime: input.harnessRuntime, - modelSelectionLocked: runParams.modelSelectionLocked, - modelFallbacksOverride: runParams.modelFallbacksOverride, - thinkLevel: input.thinkLevel, - reasoningLevel: runParams.reasoningLevel, - bashElevated: runParams.bashElevated, - extraSystemPrompt: runParams.extraSystemPrompt, - sourceReplyDeliveryMode: runParams.sourceReplyDeliveryMode, - ownerNumbers: runParams.ownerNumbers, - activeProcessSessions: listActiveProcessSessionReferences({ - scopeKey: resolveProcessToolScopeKey({ - sessionKey: runParams.sandboxSessionKey?.trim() || runParams.sessionKey, - sessionId: activeSession.id, - agentId: input.sessionAgentId, - }), - }), - }), - ...resolveContextEngineCapabilities({ - config: runParams.config, - sessionKey: runParams.sessionKey, - agentId: input.sessionAgentId, - contextEnginePluginId: input.resolveContextEnginePluginId(), - purpose: "context-engine.timeout-compaction", - }), - onCompactionHookMessages: input.onCompactionHookMessages, - ...(input.attempt.promptCache ? { promptCache: input.attempt.promptCache } : {}), - runId: runParams.runId, + ({ result: timeoutCompactResult } = await compactEmbeddedRunForRecovery(input, { + tokenBudget: input.contextTokenBudget, trigger: "timeout_recovery", diagId: timeoutDiagId, attempt: input.state.timeoutCompactionAttempts, maxAttempts: MAX_TIMEOUT_COMPACTION_ATTEMPTS, - }; - timeoutCompactResult = await compactContextEngineWithSafetyTimeout( - input.contextEngine, - { - sessionId: activeSession.id, - sessionKey: input.resolvedSessionKey, - agentId: input.sessionAgentId, - sessionTarget: buildContextEngineCompactionSessionTarget({ - agentId: input.sessionAgentId, - config: runParams.config, - sessionFile: activeSession.file, - sessionId: activeSession.id, - sessionKey: input.resolvedSessionKey, - sessionTarget: activeSession.target, - }), - tokenBudget: input.contextTokenBudget, - force: true, - compactionTarget: "budget", - runtimeContext: timeoutCompactionRuntimeContext, - runtimeSettings: input.buildRuntimeSettings({ tokenBudget: input.contextTokenBudget }), - }, - resolveCompactionTimeoutMs(runParams.config), - runParams.abortSignal, - ); + })); } catch (compactErr) { log.warn( `[timeout-compaction] contextEngine.compact() threw during timeout recovery for ${input.provider}/${input.modelId}: ${String(compactErr)}`, diff --git a/src/agents/embedded-agent-runner/tool-result-char-estimator.ts b/src/agents/embedded-agent-runner/tool-result-char-estimator.ts index f57c28ba1fbd..b6ea979ca8d2 100644 --- a/src/agents/embedded-agent-runner/tool-result-char-estimator.ts +++ b/src/agents/embedded-agent-runner/tool-result-char-estimator.ts @@ -207,10 +207,3 @@ export function estimateMessageCharsCached( cache.set(msg, estimated); return estimated; } - -export function invalidateMessageCharsCacheEntry( - cache: MessageCharEstimateCache, - msg: AgentMessage, -): void { - cache.delete(msg); -} diff --git a/src/agents/embedded-agent-runner/tool-result-context-guard.ts b/src/agents/embedded-agent-runner/tool-result-context-guard.ts index 1fdad4fe9285..864f35c06f18 100644 --- a/src/agents/embedded-agent-runner/tool-result-context-guard.ts +++ b/src/agents/embedded-agent-runner/tool-result-context-guard.ts @@ -1,4 +1,3 @@ -import { truncateUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; /** * Installs context guards for oversized tool-result histories. */ @@ -19,13 +18,9 @@ import { createMessageCharEstimateCache, estimateMessageCharsCached, getToolResultText, - invalidateMessageCharsCacheEntry, isToolResultMessage, } from "./tool-result-char-estimator.js"; -import { - estimateToolResultTextChars, - sliceToolResultTextToBudget, -} from "./tool-result-text-budget.js"; +import { truncateToolResultText } from "./tool-result-truncation.js"; const SINGLE_TOOL_RESULT_CONTEXT_SHARE = 0.5; const TRANSCRIPT_PROMPT_TEXT_KEY = "__openclawTranscriptPromptText"; @@ -135,46 +130,6 @@ function stripTranscriptPromptMarkers(messages: AgentMessage[]): AgentMessage[] return changed ? stripped : messages; } -function truncateTextToBudget(text: string, maxChars: number): string { - const budgetOptions = { minimumRawWeight: TOOL_RESULT_CHARS_PER_TOKEN_ESTIMATE }; - if (estimateToolResultTextChars(text, budgetOptions) <= maxChars) { - return text; - } - - if (maxChars <= 0) { - return formatContextLimitTruncationNotice(text.length); - } - - let prefix = sliceToolResultTextToBudget(text, maxChars, budgetOptions); - for (let i = 0; i < 4; i += 1) { - const suffix = formatContextLimitTruncationNotice(Math.max(1, text.length - prefix.length)); - prefix = sliceToolResultTextToBudget( - text, - Math.max(0, maxChars - estimateToolResultTextChars(suffix, budgetOptions)), - budgetOptions, - ); - } - - const newline = prefix.lastIndexOf("\n"); - if (newline > prefix.length * 0.7) { - prefix = truncateUtf16Safe(prefix, newline); - } - - for (let i = 0; i < 4; i += 1) { - const suffix = formatContextLimitTruncationNotice(text.length - prefix.length); - const nextPrefix = sliceToolResultTextToBudget( - prefix, - Math.max(0, maxChars - estimateToolResultTextChars(suffix, budgetOptions)), - budgetOptions, - ); - if (nextPrefix.length === prefix.length) { - return prefix + suffix; - } - prefix = nextPrefix; - } - return prefix + formatContextLimitTruncationNotice(text.length - prefix.length); -} - function replaceToolResultText(msg: AgentMessage, text: string): AgentMessage { const content = (msg as { content?: unknown }).content; const replacementContent = @@ -219,69 +174,27 @@ function truncateToolResultToChars( return replaceToolResultText(msg, formatContextLimitTruncationNotice(rawText.length)); } - const truncatedText = truncateTextToBudget(rawText, maxChars); + const truncatedText = truncateToolResultText(rawText, maxChars, { + minKeepChars: 0, + minimumRawWeight: TOOL_RESULT_CHARS_PER_TOKEN_ESTIMATE, + preserveImportantTail: false, + }); return replaceToolResultText(msg, truncatedText); } -function cloneMessagesForGuard(messages: AgentMessage[]): AgentMessage[] { - return messages.map( - (msg) => ({ ...(msg as unknown as Record) }) as unknown as AgentMessage, - ); -} - -function toolResultsNeedTruncation(params: { +function enforceToolResultLimit(params: { messages: AgentMessage[]; maxSingleToolResultChars: number; -}): boolean { +}): AgentMessage[] { const { messages, maxSingleToolResultChars } = params; const estimateCache = createMessageCharEstimateCache(); - for (const message of messages) { - if (!isToolResultMessage(message)) { - continue; - } - if (estimateMessageCharsCached(message, estimateCache) > maxSingleToolResultChars) { - return true; - } - } - return false; -} - -function applyMessageMutationInPlace( - target: AgentMessage, - source: AgentMessage, - cache?: MessageCharEstimateCache, -): void { - if (target === source) { - return; - } - - const targetRecord = target as unknown as Record; - const sourceRecord = source as unknown as Record; - for (const key of Object.keys(targetRecord)) { - if (!(key in sourceRecord)) { - delete targetRecord[key]; - } - } - Object.assign(targetRecord, sourceRecord); - if (cache) { - invalidateMessageCharsCacheEntry(cache, target); - } -} - -function enforceToolResultLimitInPlace(params: { - messages: AgentMessage[]; - maxSingleToolResultChars: number; -}): void { - const { messages, maxSingleToolResultChars } = params; - const estimateCache = createMessageCharEstimateCache(); - - for (const message of messages) { - if (!isToolResultMessage(message)) { - continue; - } - const truncated = truncateToolResultToChars(message, maxSingleToolResultChars, estimateCache); - applyMessageMutationInPlace(message, truncated, estimateCache); - } + let changed = false; + const guarded = messages.map((message) => { + const next = truncateToolResultToChars(message, maxSingleToolResultChars, estimateCache); + changed ||= next !== message; + return next; + }); + return changed ? guarded : messages; } function hasNewToolResultAfterFence(params: { @@ -487,18 +400,10 @@ export function installToolResultContextGuard(params: { : messages; const sourceMessages = Array.isArray(transformed) ? transformed : messages; - const contextMessages = toolResultsNeedTruncation({ + const contextMessages = enforceToolResultLimit({ messages: sourceMessages, maxSingleToolResultChars, - }) - ? cloneMessagesForGuard(sourceMessages) - : sourceMessages; - if (contextMessages !== sourceMessages) { - enforceToolResultLimitInPlace({ - messages: contextMessages, - maxSingleToolResultChars, - }); - } + }); if (params.midTurnPrecheck?.enabled) { const prePromptMessageCount = Math.max( 0, diff --git a/src/agents/embedded-agent-runner/tool-result-truncation.ts b/src/agents/embedded-agent-runner/tool-result-truncation.ts index 4dbed8e2acd5..e414cbfb38a1 100644 --- a/src/agents/embedded-agent-runner/tool-result-truncation.ts +++ b/src/agents/embedded-agent-runner/tool-result-truncation.ts @@ -3,6 +3,7 @@ */ import { createHash } from "node:crypto"; import { existsSync } from "node:fs"; +import { isRecord } from "@openclaw/normalization-core/record-coerce"; import { normalizeLowercaseStringOrEmpty } from "@openclaw/normalization-core/string-coerce"; import { sliceUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; import { createDedupeCache } from "../../infra/dedupe.js"; @@ -56,6 +57,8 @@ export const toolResultWarningDedupe = { type ToolResultTruncationOptions = { suffix?: string | ((truncatedChars: number) => string); minKeepChars?: number; + minimumRawWeight?: number; + preserveImportantTail?: boolean; }; const DEFAULT_SUFFIX = (truncatedChars: number) => @@ -82,11 +85,10 @@ function logToolResultSessionTruncation(params: { `aggregateBudgetChars=${params.aggregateBudgetChars} ` + `oversized=${params.oversizedReplacementCount} aggregate=${params.aggregateReplacementCount}) ` + `sessionKey=${sessionLogKey}`; - if (params.aggregateReplacementCount <= 0) { - log.info(message); - return; - } - if (toolResultWarningDedupe.sessionRecovery.check(sessionLogKey)) { + if ( + params.aggregateReplacementCount <= 0 || + toolResultWarningDedupe.sessionRecovery.check(sessionLogKey) + ) { log.info(message); return; } @@ -95,32 +97,25 @@ function logToolResultSessionTruncation(params: { ); } -async function openRuntimeTranscriptSessionManager(scope: RuntimeTranscriptScope): Promise<{ - sessionManager: SessionManager; - target: Awaited>; -}> { - const target = await resolveRuntimeTranscriptReadTarget(scope); - return { sessionManager: SessionManager.open(target), target }; -} - function resolveSuffixFactory( suffix: ToolResultTruncationOptions["suffix"], ): (truncatedChars: number) => string { - if (typeof suffix === "function") { - return suffix; - } - if (typeof suffix === "string") { - return () => suffix; - } - return DEFAULT_SUFFIX; + return typeof suffix === "function" + ? suffix + : typeof suffix === "string" + ? () => suffix + : DEFAULT_SUFFIX; } function resolveEffectiveMinKeepChars(params: { maxChars: number; minKeepChars: number; suffixFactory: (truncatedChars: number) => string; + minimumRawWeight?: number; }): number { - const suffixFloor = estimateToolResultTextChars(params.suffixFactory(1)); + const suffixFloor = estimateToolResultTextChars(params.suffixFactory(1), { + minimumRawWeight: params.minimumRawWeight, + }); return Math.max(0, Math.min(params.minKeepChars, Math.max(0, params.maxChars - suffixFloor))); } @@ -129,25 +124,31 @@ function appendBoundedTruncationSuffix(params: { originalTextLength: number; maxChars: number; suffixFactory: (truncatedChars: number) => string; + minimumRawWeight?: number; }): string { let keptText = params.keptText; + const budgetOptions = { minimumRawWeight: params.minimumRawWeight }; while (true) { const suffix = params.suffixFactory(Math.max(1, params.originalTextLength - keptText.length)); - const suffixChars = estimateToolResultTextChars(suffix); + const suffixChars = estimateToolResultTextChars(suffix, budgetOptions); if (suffixChars >= params.maxChars) { const fullOmissionSuffix = params.suffixFactory(Math.max(1, params.originalTextLength)); - return sliceToolResultTextToBudget(fullOmissionSuffix, params.maxChars); + return sliceToolResultTextToBudget(fullOmissionSuffix, params.maxChars, budgetOptions); } - const nextKeptText = sliceToolResultTextToBudget(keptText, params.maxChars - suffixChars); + const nextKeptText = sliceToolResultTextToBudget( + keptText, + params.maxChars - suffixChars, + budgetOptions, + ); const finalText = nextKeptText + suffix; if ( nextKeptText.length === keptText.length && - estimateToolResultTextChars(finalText) <= params.maxChars + estimateToolResultTextChars(finalText, budgetOptions) <= params.maxChars ) { return finalText; } if (nextKeptText.length === 0 && keptText.length === 0) { - return sliceToolResultTextToBudget(finalText, params.maxChars); + return sliceToolResultTextToBudget(finalText, params.maxChars, budgetOptions); } keptText = nextKeptText; } @@ -183,38 +184,48 @@ function hasImportantTail(text: string): boolean { * This ensures error messages and summaries at the end of tool output * aren't lost during truncation. */ -function truncateToolResultText( +export function truncateToolResultText( text: string, maxChars: number, options: ToolResultTruncationOptions = {}, ): string { const suffixFactory = resolveSuffixFactory(options.suffix); + const budgetOptions = { minimumRawWeight: options.minimumRawWeight }; const minKeepChars = resolveEffectiveMinKeepChars({ maxChars, minKeepChars: options.minKeepChars ?? MIN_KEEP_CHARS, suffixFactory, + minimumRawWeight: options.minimumRawWeight, }); - if (estimateToolResultTextChars(text) <= maxChars) { + if (estimateToolResultTextChars(text, budgetOptions) <= maxChars) { return text; } - const initialKeptText = sliceToolResultTextToBudget(text, maxChars); + const initialKeptText = sliceToolResultTextToBudget(text, maxChars, budgetOptions); const defaultSuffix = suffixFactory(Math.max(1, text.length - initialKeptText.length)); - const budget = Math.max(minKeepChars, maxChars - estimateToolResultTextChars(defaultSuffix)); + const budget = Math.max( + minKeepChars, + maxChars - estimateToolResultTextChars(defaultSuffix, budgetOptions), + ); // If tail looks important, split budget between head and tail - if (hasImportantTail(text) && budget > minKeepChars * 2) { + if ( + options.preserveImportantTail !== false && + hasImportantTail(text) && + budget > minKeepChars * 2 + ) { const tailBudget = Math.min(Math.floor(budget * 0.3), 4_000); - const headBudget = budget - tailBudget - estimateToolResultTextChars(MIDDLE_OMISSION_MARKER); + const headBudget = + budget - tailBudget - estimateToolResultTextChars(MIDDLE_OMISSION_MARKER, budgetOptions); if (headBudget > minKeepChars) { // Find clean cut points at newline boundaries - let headText = sliceToolResultTextToBudget(text, headBudget); + let headText = sliceToolResultTextToBudget(text, headBudget, budgetOptions); const headNewline = headText.lastIndexOf("\n"); if (headNewline > headText.length * 0.8) { headText = sliceUtf16Safe(headText, 0, headNewline); } - let tailText = sliceToolResultTextTailToBudget(text, tailBudget); + let tailText = sliceToolResultTextTailToBudget(text, tailBudget, budgetOptions); const tailNewline = tailText.indexOf("\n"); if (tailNewline !== -1 && tailNewline < tailText.length * 0.2) { tailText = sliceUtf16Safe(tailText, tailNewline + 1); @@ -226,13 +237,14 @@ function truncateToolResultText( originalTextLength: text.length, maxChars, suffixFactory, + minimumRawWeight: options.minimumRawWeight, }); } } } // Default: keep the beginning - let keptText = sliceToolResultTextToBudget(text, budget); + let keptText = sliceToolResultTextToBudget(text, budget, budgetOptions); const lastNewline = keptText.lastIndexOf("\n"); if (lastNewline > keptText.length * 0.8) { keptText = sliceUtf16Safe(keptText, 0, lastNewline); @@ -242,6 +254,7 @@ function truncateToolResultText( originalTextLength: text.length, maxChars, suffixFactory, + minimumRawWeight: options.minimumRawWeight, }); } @@ -291,23 +304,17 @@ export function resolveLiveToolResultAggregateMaxChars(params: { * Get the total token-budget character estimate for text blocks in a tool result message. */ function getToolResultTextBudget(msg: AgentMessage): number { - if (!msg || (msg as { role?: string }).role !== "toolResult") { + if (!msg || msg.role !== "toolResult") { return 0; } const content = (msg as { content?: unknown }).content; - if (!Array.isArray(content)) { - return 0; - } - let totalLength = 0; - for (const block of content) { - if (isToolResultTextBlock(block)) { - const text = block.text; - if (typeof text === "string") { - totalLength += estimateToolResultTextChars(text); - } - } - } - return totalLength; + return Array.isArray(content) + ? content.reduce( + (total, block) => + total + (isToolResultTextBlock(block) ? estimateToolResultTextChars(block.text) : 0), + 0, + ) + : 0; } /** @@ -421,50 +428,25 @@ type ToolResultSpillDetails = { function getToolResultSpillDetails(message: AgentMessage): ToolResultSpillDetails | undefined { const details = (message as { details?: unknown }).details; - if (!details || typeof details !== "object" || Array.isArray(details)) { + if (!isRecord(details)) { return undefined; } - const nested = (details as { spill?: unknown }).spill; - const nestedSpill = - nested && typeof nested === "object" && !Array.isArray(nested) - ? (nested as Record) - : undefined; + const nestedSpill = isRecord(details.spill) ? details.spill : undefined; // web_fetch owns the nested contract. Exec tools still own the flat spill fields. - const path = nestedSpill?.path ?? (details as { fullOutputPath?: unknown }).fullOutputPath; + const path = nestedSpill?.path ?? details.fullOutputPath; if (typeof path !== "string" || path.length === 0) { return undefined; } - const truncated = - nestedSpill?.truncated === true || - (details as { spillTruncated?: unknown }).spillTruncated === true; - const chars = nestedSpill?.chars ?? (details as { spilledChars?: unknown }).spilledChars; + const chars = nestedSpill?.chars ?? details.spilledChars; return { path, - truncated, + truncated: nestedSpill?.truncated === true || details.spillTruncated === true, ...(typeof chars === "number" && Number.isFinite(chars) ? { chars: Math.max(0, Math.floor(chars)) } : {}), }; } -function toolResultTextContainsFullOutputFooter( - message: AgentMessage, - fullOutputPath: string, -): boolean { - const content = (message as { content?: unknown }).content; - if (!Array.isArray(content)) { - return false; - } - const footer = formatFullOutputFooter(fullOutputPath); - const escapedFooter = JSON.stringify(footer).slice(1, -1); - return content.some((block: unknown) => { - if (!isToolResultTextBlock(block)) { - return false; - } - return block.text.includes(footer) || block.text.includes(escapedFooter); - }); -} - type AggregateElisionMarkers = { full: string; compact: string; @@ -478,9 +460,19 @@ function resolveAggregateElisionMarkers( if (!spill) { return undefined; } + const content = (message as { content?: unknown }).content; + const footer = formatFullOutputFooter(spill.path); + const escapedFooter = JSON.stringify(footer).slice(1, -1); // Details alone are not model-visible. Only preserve paths that already // appeared in the original footer, so elision discloses nothing new. - if (!toolResultTextContainsFullOutputFooter(message, spill.path)) { + if ( + !Array.isArray(content) || + !content.some( + (block) => + isToolResultTextBlock(block) && + (block.text.includes(footer) || block.text.includes(escapedFooter)), + ) + ) { return undefined; } // Aggregate elision is a rare recovery path, not a request hot path; one @@ -490,20 +482,16 @@ function resolveAggregateElisionMarkers( } // The path was already disclosed in the original tool footer; preserving it // here adds no new disclosure and only keeps recovery possible. - if (spill.truncated) { - const count = spill.chars === undefined ? "capped content" : `first ${spill.chars} chars`; - return { - full: `[tool result elided: partial output preserved at ${spill.path} (${count}); read it if the output is needed]`, - compact: `[partial: ${spill.path}]`, - truncationSuffix: (truncatedChars) => - `[... ${Math.max(1, Math.floor(truncatedChars))} chars truncated; partial output at ${spill.path}]`, - }; - } + const kind = spill.truncated ? "partial" : "full"; + const count = spill.truncated + ? ` (${spill.chars === undefined ? "capped content" : `first ${spill.chars} chars`})` + : ""; + const output = `${kind} output`; return { - full: `[tool result elided: full output preserved at ${spill.path}; read it if the output is needed]`, - compact: `[read ${spill.path}]`, + full: `[tool result elided: ${output} preserved at ${spill.path}${count}; read it if the output is needed]`, + compact: spill.truncated ? `[partial: ${spill.path}]` : `[read ${spill.path}]`, truncationSuffix: (truncatedChars) => - `[... ${Math.max(1, Math.floor(truncatedChars))} chars truncated; full output at ${spill.path}]`, + `[... ${Math.max(1, Math.floor(truncatedChars))} chars truncated; ${output} at ${spill.path}]`, }; } @@ -514,14 +502,10 @@ function formatAggregateElisionText( if (remainingTextBudget <= 0) { return ""; } - if (spillMarkers?.full && estimateToolResultTextChars(spillMarkers.full) <= remainingTextBudget) { - return spillMarkers.full; - } - if ( - spillMarkers?.compact && - estimateToolResultTextChars(spillMarkers.compact) <= remainingTextBudget - ) { - return spillMarkers.compact; + for (const marker of [spillMarkers?.full, spillMarkers?.compact]) { + if (marker && estimateToolResultTextChars(marker) <= remainingTextBudget) { + return marker; + } } return sliceToolResultTextToBudget(AGGREGATE_ELISION_MARKER, remainingTextBudget); } @@ -598,29 +582,6 @@ export function truncateOversizedToolResultsInMessages( minKeepChars: RECOVERY_MIN_KEEP_CHARS, protectTrailingToolResults: Boolean(projectionState), }); - if (projectionState) { - for (const [index] of messages.entries()) { - const projectionKey = projectionKeys[index]; - if (projectionKey) { - projectionState.frozen.add(projectionKey); - } - } - } - if (plan.replacements.length === 0) { - const projectedMessages = branch.map((entry) => entry.message); - const hasProjectedChanges = projectedMessages.some( - (message, index) => message !== messages[index], - ); - return { - messages: hasProjectedChanges ? projectedMessages : messages, - truncatedCount: 0, - aggregateTruncatedCount: 0, - aggregatePressureEngaged: plan.aggregatePressureExceeded, - aggregateBudgetChars, - }; - } - - const replacementIds = new Set(plan.replacements.map((replacement) => replacement.entryId)); const replacedBranch = applyToolResultReplacementsToBranch(branch, plan.replacements); if (projectionState) { for (const [index, originalMessage] of messages.entries()) { @@ -628,15 +589,20 @@ export function truncateOversizedToolResultsInMessages( const projectionKey = projectionKeys[index]; if (projectionKey) { projectionState.frozen.add(projectionKey); - if (projectedMessage && projectedMessage !== originalMessage) { + if ( + plan.replacements.length > 0 && + projectedMessage && + projectedMessage !== originalMessage + ) { projectionState.replacements.set(projectionKey, projectedMessage); } } } } + const output = replacedBranch.map((entry) => entry.message as AgentMessage); return { - messages: replacedBranch.map((entry) => entry.message as AgentMessage), - truncatedCount: replacementIds.size, + messages: output.some((message, index) => message !== messages[index]) ? output : messages, + truncatedCount: new Set(plan.replacements.map((replacement) => replacement.entryId)).size, aggregateTruncatedCount: plan.aggregateReplacementCount, aggregatePressureEngaged: plan.aggregatePressureExceeded, aggregateBudgetChars, @@ -703,12 +669,11 @@ function getToolResultProjectionKeys( const baseKeyCounts = new Map(); for (const baseKey of baseKeys) { if (baseKey) { - baseKeyCounts.set(baseKey, (baseKeyCounts.get(baseKey) ?? 0) + 1); - } - } - for (const [baseKey, count] of baseKeyCounts) { - if (count > 1) { - projectionState.ambiguousBaseKeys.add(baseKey); + const count = (baseKeyCounts.get(baseKey) ?? 0) + 1; + baseKeyCounts.set(baseKey, count); + if (count > 1) { + projectionState.ambiguousBaseKeys.add(baseKey); + } } } const occurrences = new Map(); @@ -825,18 +790,13 @@ function seedRecoveryBranchFromFrozenProjection(params: { function getToolResultTextBlocks(message: AgentMessage): string[] { const content = (message as { content?: unknown }).content; - if (!Array.isArray(content)) { - return []; - } - return content.flatMap((block) => - block && typeof block === "object" && (block as { type?: unknown }).type === "text" - ? [ - typeof (block as { text?: unknown }).text === "string" - ? (block as { text: string }).text - : "", - ] - : [], - ); + return Array.isArray(content) + ? content.flatMap((block) => + isRecord(block) && block.type === "text" + ? [typeof block.text === "string" ? block.text : ""] + : [], + ) + : []; } function buildAggregateToolResultReplacements(params: { @@ -851,34 +811,23 @@ function buildAggregateToolResultReplacements(params: { ? getTrailingToolResultEntryIds(params.branch) : new Set(); const candidates = params.branch - .map((entry, index) => ({ entry, index })) - .filter( - ( - item, - ): item is { - entry: { - id: string; - type: string; - message: AgentMessage; - aggregateEligible?: boolean; - deferAggregateRecovery?: boolean; - }; - index: number; - } => - item.entry.type === "message" && - Boolean(item.entry.message) && - (item.entry.message as { role?: string }).role === "toolResult", - ) - .map((item) => ({ - index: item.index, - entryId: item.entry.id, - message: item.entry.message, - spillSourceMessage: params.spillSourceBranch?.[item.index]?.message ?? item.entry.message, - textLength: getToolResultTextBudget(item.entry.message), - aggregateEligible: item.entry.aggregateEligible !== false, - deferredByFreshProjection: item.entry.deferAggregateRecovery === true, - protectedByTrailingBatch: protectedEntryIds.has(item.entry.id), - })) + .flatMap((entry, index) => { + const message = entry.message; + return entry.type === "message" && message?.role === "toolResult" + ? [ + { + index, + entryId: entry.id, + message, + spillSourceMessage: params.spillSourceBranch?.[index]?.message ?? message, + textLength: getToolResultTextBudget(message), + aggregateEligible: entry.aggregateEligible !== false, + deferredByFreshProjection: entry.deferAggregateRecovery === true, + protectedByTrailingBatch: protectedEntryIds.has(entry.id), + }, + ] + : []; + }) .filter((item) => item.textLength > 0); if (candidates.length < 2) { @@ -898,15 +847,10 @@ function buildAggregateToolResultReplacements(params: { } let remainingReduction = totalChars - params.aggregateBudgetChars; - const replacements: Array<{ entryId: string; message: AgentMessage }> = []; + const replacements = new Map(); const aggregateRecoveryCandidates = candidates .filter((item) => !item.deferredByFreshProjection && !item.protectedByTrailingBatch) - .toSorted((a, b) => { - if (a.index !== b.index) { - return a.index - b.index; - } - return b.textLength - a.textLength; - }); + .toSorted((a, b) => a.index - b.index); const recoveryCandidates = [ ...aggregateRecoveryCandidates.filter((item) => item.aggregateEligible), // Start from frozen projections before touching deferred fresh results. Reusing their @@ -945,41 +889,31 @@ function buildAggregateToolResultReplacements(params: { continue; } - replacements.push({ entryId: candidate.entryId, message: truncatedMessage }); + replacements.set(candidate.entryId, { entryId: candidate.entryId, message: truncatedMessage }); remainingReduction -= actualReduction; } - if (remainingReduction > 0) { - for (const candidate of recoveryCandidates) { - if (remainingReduction <= 0) { - break; - } - const existingReplacement = replacements.find( - (replacement) => replacement.entryId === candidate.entryId, - ); - const baseMessage = existingReplacement?.message ?? candidate.message; - const baseTextLength = getToolResultTextBudget(baseMessage); - const targetTextChars = Math.max(0, baseTextLength - remainingReduction); - const spillMarkers = resolveAggregateElisionMarkers(candidate.spillSourceMessage); - const emptyMessage = clearToolResultText(candidate.message, targetTextChars, spillMarkers); - const actualReduction = Math.max(0, baseTextLength - getToolResultTextBudget(emptyMessage)); - if (actualReduction <= 0 && !spillMarkers) { - continue; - } - const replacement = { entryId: candidate.entryId, message: emptyMessage }; - const existingIndex = replacements.findIndex( - (existing) => existing.entryId === candidate.entryId, - ); - if (existingIndex >= 0) { - replacements[existingIndex] = replacement; - } else { - replacements.push(replacement); - } - remainingReduction -= actualReduction; + for (const candidate of recoveryCandidates) { + if (remainingReduction <= 0) { + break; } + const baseMessage = replacements.get(candidate.entryId)?.message ?? candidate.message; + const baseTextLength = getToolResultTextBudget(baseMessage); + const spillMarkers = resolveAggregateElisionMarkers(candidate.spillSourceMessage); + const emptyMessage = clearToolResultText( + candidate.message, + Math.max(0, baseTextLength - remainingReduction), + spillMarkers, + ); + const actualReduction = Math.max(0, baseTextLength - getToolResultTextBudget(emptyMessage)); + if (actualReduction <= 0 && !spillMarkers) { + continue; + } + replacements.set(candidate.entryId, { entryId: candidate.entryId, message: emptyMessage }); + remainingReduction -= actualReduction; } - return { replacements, pressureExceeded: true }; + return { replacements: [...replacements.values()], pressureExceeded: true }; } function getTrailingToolResultEntryIds(branch: ToolResultBranchEntry[]): Set { @@ -1089,20 +1023,19 @@ function calculateReplacementReduction( return 0; } const branchById = new Map(branch.map((entry) => [entry.id, entry])); - let reduction = 0; - - for (const replacement of replacements) { + return replacements.reduce((reduction, replacement) => { const entry = branchById.get(replacement.entryId); if (!entry?.message) { - continue; + return reduction; } - reduction += Math.max( - 0, - getToolResultTextBudget(entry.message) - getToolResultTextBudget(replacement.message), + return ( + reduction + + Math.max( + 0, + getToolResultTextBudget(entry.message) - getToolResultTextBudget(replacement.message), + ) ); - } - - return reduction; + }, 0); } function applyToolResultReplacementsToBranch( @@ -1112,18 +1045,10 @@ function applyToolResultReplacementsToBranch( if (replacements.length === 0) { return branch; } - const replacementsById = new Map( - replacements.map((replacement) => [replacement.entryId, replacement]), - ); + const replacementsById = new Map(replacements.map(({ entryId, message }) => [entryId, message])); return branch.map((entry) => { - const replacement = replacementsById.get(entry.id); - if (!replacement || entry.type !== "message") { - return entry; - } - return { - ...entry, - message: replacement.message, - }; + const message = replacementsById.get(entry.id); + return message && entry.type === "message" ? { ...entry, message } : entry; }); } @@ -1404,7 +1329,8 @@ export async function truncateOversizedToolResultsInActiveTarget(params: { projectionState?: ToolResultPromptProjectionState; }): Promise<{ truncated: boolean; truncatedCount: number; reason?: string }> { try { - const { sessionManager, target } = await openRuntimeTranscriptSessionManager(params.scope); + const target = await resolveRuntimeTranscriptReadTarget(params.scope); + const sessionManager = SessionManager.open(target); return truncateOversizedToolResultsInExistingSessionManager({ sessionManager, contextWindowTokens: params.contextWindowTokens, diff --git a/src/agents/harness/tool-result-middleware.ts b/src/agents/harness/tool-result-middleware.ts index e6d9dc69f72d..26b51d765013 100644 --- a/src/agents/harness/tool-result-middleware.ts +++ b/src/agents/harness/tool-result-middleware.ts @@ -56,10 +56,7 @@ function isValidMiddlewareContentBlock(value: unknown): boolean { function hasValidMiddlewareDetailsShape( value: unknown, - state: { keys: number; seen: WeakSet } = { - keys: 0, - seen: new WeakSet(), - }, + state: { keys: number; seen: WeakSet } = { keys: 0, seen: new WeakSet() }, depth = 0, ): boolean { if (value === undefined || value === null) { @@ -71,35 +68,16 @@ function hasValidMiddlewareDetailsShape( if (typeof value === "string" || typeof value === "number" || typeof value === "boolean") { return true; } - if (typeof value !== "object") { - return false; - } - if (state.seen.has(value)) { + if (typeof value !== "object" || state.seen.has(value)) { return false; } state.seen.add(value); - if (Array.isArray(value)) { - state.keys += value.length; - if (state.keys > MAX_MIDDLEWARE_DETAILS_KEYS) { - return false; - } - for (const entry of value) { - if (!hasValidMiddlewareDetailsShape(entry, state, depth + 1)) { - return false; - } - } - return true; - } - for (const entry of Object.values(value)) { - state.keys += 1; - if (state.keys > MAX_MIDDLEWARE_DETAILS_KEYS) { - return false; - } - if (!hasValidMiddlewareDetailsShape(entry, state, depth + 1)) { - return false; - } - } - return true; + const entries = Array.isArray(value) ? value : Object.values(value); + state.keys += entries.length; + return ( + state.keys <= MAX_MIDDLEWARE_DETAILS_KEYS && + entries.every((entry) => hasValidMiddlewareDetailsShape(entry, state, depth + 1)) + ); } function isValidMiddlewareDetails(value: unknown): boolean { @@ -125,10 +103,6 @@ function isValidMiddlewareToolResult(value: unknown): value is OpenClawAgentTool ); } -function createMiddlewareContentCoerceState(): MiddlewareContentCoerceState { - return { depth: 0, seen: new Set() }; -} - function descendMiddlewareContentCoerceState( value: unknown, state: MiddlewareContentCoerceState, @@ -136,18 +110,15 @@ function descendMiddlewareContentCoerceState( if (state.depth >= MAX_MIDDLEWARE_CONTENT_DEPTH) { return undefined; } - if (value !== null && typeof value === "object") { - if (state.seen.has(value)) { - return undefined; - } - const seen = new Set(state.seen); - seen.add(value); - return { depth: state.depth + 1, seen }; + if (value === null || typeof value !== "object") { + return { depth: state.depth + 1, seen: state.seen }; } - return { depth: state.depth + 1, seen: state.seen }; + return state.seen.has(value) + ? undefined + : { depth: state.depth + 1, seen: new Set([...state.seen, value]) }; } -function stringifyMiddlewareTextPayload(value: unknown): string | undefined { +function serializeMiddlewareValue(value: unknown): string | undefined { const seen = new WeakSet(); try { return JSON.stringify(value, (_key, val) => { @@ -172,7 +143,7 @@ function stringifyMiddlewareTextPayload(value: unknown): string | undefined { function coerceMiddlewareText( value: unknown, - state: MiddlewareContentCoerceState = createMiddlewareContentCoerceState(), + state: MiddlewareContentCoerceState, options: MiddlewareToolResultCoerceOptions = {}, ): string | undefined { if (typeof value === "string") { @@ -194,18 +165,13 @@ function coerceMiddlewareText( return text; } } - const content = value.content; - if (Array.isArray(content)) { - const chunks = coerceMiddlewareContentArray(content, nextState, options) - .filter( - (block): block is Extract => - block.type === "text", - ) - .map((block) => block.text) - .filter((text) => text.length > 0); - return chunks.length > 0 ? chunks.join("\n") : undefined; + if (Array.isArray(value.content)) { + const text = coerceMiddlewareContentArray(value.content, nextState, options) + .flatMap((block) => (block.type === "text" && block.text ? [block.text] : [])) + .join("\n"); + return text || undefined; } - return stringifyMiddlewareTextPayload(value); + return serializeMiddlewareValue(value); } function appendMiddlewareContentBlock( @@ -231,10 +197,9 @@ function appendMiddlewareContentBlock( return; } const remainingChars = MAX_MIDDLEWARE_TEXT_CHARS - previous.text.length - 1; - if (remainingChars <= 0) { - return; + if (remainingChars > 0) { + previous.text = `${previous.text}\n${truncateUtf16Safe(block.text, remainingChars)}`; } - previous.text = `${previous.text}\n${truncateUtf16Safe(block.text, remainingChars)}`; } function coerceMiddlewareContentArray( @@ -243,31 +208,16 @@ function coerceMiddlewareContentArray( options: MiddlewareToolResultCoerceOptions = {}, ): MiddlewareContentBlock[] { const blocks: MiddlewareContentBlock[] = []; - let inspectedBlocks = 0; - for (const entry of content) { - inspectedBlocks += 1; - if ( - inspectedBlocks > MAX_MIDDLEWARE_CONTENT_BLOCKS || - blocks.length >= MAX_MIDDLEWARE_CONTENT_BLOCKS - ) { + for (const entry of content.slice(0, MAX_MIDDLEWARE_CONTENT_BLOCKS)) { + if (blocks.length >= MAX_MIDDLEWARE_CONTENT_BLOCKS) { break; } - const coercedBlocks = coerceMiddlewareContentBlocks(entry, state, options); - if (coercedBlocks.length > 0) { - for (const block of coercedBlocks) { - appendMiddlewareContentBlock(blocks, block); - if (blocks.length >= MAX_MIDDLEWARE_CONTENT_BLOCKS) { - break; - } - } - continue; - } - const text = coerceMiddlewareText(entry, state, options); - if (text) { - appendMiddlewareContentBlock(blocks, { - type: "text", - text: truncateUtf16Safe(text, MAX_MIDDLEWARE_TEXT_CHARS), - }); + const coerced = coerceMiddlewareContentBlocks(entry, state, options); + const text = coerced.length === 0 ? coerceMiddlewareText(entry, state, options) : undefined; + for (const block of text + ? [{ type: "text" as const, text: truncateUtf16Safe(text, MAX_MIDDLEWARE_TEXT_CHARS) }] + : coerced) { + appendMiddlewareContentBlock(blocks, block); } } return blocks; @@ -275,7 +225,7 @@ function coerceMiddlewareContentArray( function coerceMiddlewareContentBlocks( value: unknown, - state: MiddlewareContentCoerceState = createMiddlewareContentCoerceState(), + state: MiddlewareContentCoerceState, options: MiddlewareToolResultCoerceOptions = {}, ): MiddlewareContentBlock[] { if (isValidMiddlewareContentBlock(value)) { @@ -327,19 +277,14 @@ function coerceMiddlewareToolResult( if (!isRecord(value) || !Array.isArray(value.content)) { return undefined; } + const state: MiddlewareContentCoerceState = { depth: 0, seen: new Set() }; const content: OpenClawAgentToolResult["content"] = []; - const state = createMiddlewareContentCoerceState(); - let inspectedBlocks = 0; - for (const block of value.content) { - inspectedBlocks += 1; - if (inspectedBlocks > MAX_MIDDLEWARE_CONTENT_BLOCKS) { - break; - } + for (const block of value.content.slice(0, MAX_MIDDLEWARE_CONTENT_BLOCKS)) { for (const coerced of coerceMiddlewareContentBlocks(block, state, options)) { - content.push(coerced); if (content.length >= MAX_MIDDLEWARE_CONTENT_BLOCKS) { break; } + content.push(coerced); } if (content.length >= MAX_MIDDLEWARE_CONTENT_BLOCKS) { break; @@ -374,31 +319,14 @@ function coerceMiddlewareToolResult( * cannot be represented at all (top-level function/symbol/undefined). */ function sanitizeMiddlewareDetailsValue(value: unknown): unknown { - const seen = new WeakSet(); - try { - const serialized = JSON.stringify(value, (_key, val) => { - if (typeof val === "bigint") { - return val.toString(); - } - if (val !== null && typeof val === "object") { - if (seen.has(val)) { - return undefined; - } - seen.add(val); - } - return val; - }); - if (serialized === undefined) { - return null; - } - const serializedBytes = Buffer.byteLength(serialized, "utf8"); - if (serializedBytes > MAX_MIDDLEWARE_DETAILS_BYTES) { - return { truncated: true, originalSizeBytes: serializedBytes }; - } - return JSON.parse(serialized); - } catch { + const serialized = serializeMiddlewareValue(value); + if (serialized === undefined) { return null; } + const bytes = Buffer.byteLength(serialized, "utf8"); + return bytes > MAX_MIDDLEWARE_DETAILS_BYTES + ? { truncated: true, originalSizeBytes: bytes } + : JSON.parse(serialized); } /** @@ -417,13 +345,9 @@ function sanitizeToolResultForMiddleware(result: OpenClawAgentToolResult): OpenC if (coerced) { return coerced; } - if (result.details === undefined || result.details === null) { - return result; - } - if (isValidMiddlewareDetails(result.details)) { - return result; - } - return { ...result, details: sanitizeMiddlewareDetailsValue(result.details) }; + return result.details == null || isValidMiddlewareDetails(result.details) + ? result + : { ...result, details: sanitizeMiddlewareDetailsValue(result.details) }; } function buildMiddlewareFailureResult(): OpenClawAgentToolResult { diff --git a/src/agents/session-tool-result-guard.ts b/src/agents/session-tool-result-guard.ts index d49a8a90a009..f298f6518876 100644 --- a/src/agents/session-tool-result-guard.ts +++ b/src/agents/session-tool-result-guard.ts @@ -4,6 +4,7 @@ * Caps large tool results, repairs missing results, applies redaction, and emits transcript update events. */ import { resolveIntegerOption } from "@openclaw/normalization-core/number-coercion"; +import { isRecord } from "@openclaw/normalization-core/record-coerce"; import { normalizeOptionalString } from "@openclaw/normalization-core/string-coerce"; import { sliceUtf16Safe, truncateUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; import { @@ -200,15 +201,11 @@ function redactPersistedDetailString( )} original chars omitted]`; } -function isSensitivePersistedDetailKey(key: string | undefined): boolean { - return Boolean(key && isSensitiveFieldKey(key)); -} - function selectPersistedDetailRedactionKey( key: string, inheritedKey: string | undefined, ): string | undefined { - return isSensitivePersistedDetailKey(key) ? key : inheritedKey; + return isSensitiveFieldKey(key) ? key : inheritedKey; } function redactedOriginalDetailKeys( @@ -287,6 +284,26 @@ function redactPersistedSummaryField( ); } +function copyPersistedSummaryFields(params: { + target: Record; + source: Record; + keys: readonly string[]; + maxChars: number; + redactionConfig?: ToolResultDetailRedactionConfig; +}): void { + for (const key of params.keys) { + const value = params.source[key]; + if (value !== undefined) { + params.target[key] = redactPersistedSummaryField( + key, + value, + params.maxChars, + params.redactionConfig, + ); + } + } +} + function sanitizePersistedSessionDetail( value: unknown, redactionConfig?: ToolResultDetailRedactionConfig, @@ -296,24 +313,25 @@ function sanitizePersistedSessionDetail( } const src = value as Record; const out: Record = {}; - for (const key of [ - "sessionId", - "status", - "pid", - "startedAt", - "endedAt", - "runtimeMs", - "cwd", - "name", - "truncated", - "exitCode", - "exitSignal", - ]) { - const field = src[key]; - if (field !== undefined) { - out[key] = redactPersistedSummaryField(key, field, 500, redactionConfig); - } - } + copyPersistedSummaryFields({ + target: out, + source: src, + keys: [ + "sessionId", + "status", + "pid", + "startedAt", + "endedAt", + "runtimeMs", + "cwd", + "name", + "truncated", + "exitCode", + "exitSignal", + ], + maxChars: 500, + redactionConfig, + }); if (typeof src.command === "string") { out.command = redactPersistedDetailString(src.command, 500, redactionConfig); } @@ -357,28 +375,24 @@ function buildPersistedDetailsFallback( } if (src) { fallback.originalDetailKeys = redactedOriginalDetailKeys(src, redactionConfig); - for (const key of [ - "status", - "sessionId", - "pid", - "exitCode", - "exitSignal", - "truncated", - "spill", - "fullOutputPath", - "spilledChars", - "spillTruncated", - ]) { - const field = src[key]; - if (field !== undefined) { - fallback[key] = redactPersistedSummaryField( - key, - field, - MAX_PERSISTED_DETAIL_FALLBACK_STRING_CHARS, - redactionConfig, - ); - } - } + copyPersistedSummaryFields({ + target: fallback, + source: src, + keys: [ + "status", + "sessionId", + "pid", + "exitCode", + "exitSignal", + "truncated", + "spill", + "fullOutputPath", + "spilledChars", + "spillTruncated", + ], + maxChars: MAX_PERSISTED_DETAIL_FALLBACK_STRING_CHARS, + redactionConfig, + }); copyPersistedResultStateFields( fallback, src, @@ -390,21 +404,23 @@ function buildPersistedDetailsFallback( } function enforcePersistedDetailsByteCap( - value: Record, - src: Record | undefined, + value: unknown, + originalDetails: unknown, originalSize: BoundedJsonUtf8Bytes, redactionConfig?: ToolResultDetailRedactionConfig, -): Record { +): unknown { const sanitizedBytes = jsonUtf8BytesOrInfinity(value); if (sanitizedBytes <= MAX_PERSISTED_TOOL_RESULT_DETAILS_BYTES) { return value; } - const fallback = buildPersistedDetailsFallback( - src, - originalSize, - sanitizedBytes, - redactionConfig, - ); + const fallback = isRecord(originalDetails) + ? buildPersistedDetailsFallback(originalDetails, originalSize, sanitizedBytes, redactionConfig) + : { + persistedDetailsTruncated: true, + finalDetailsTruncated: true, + ...originalDetailsSizeFields(originalSize), + sanitizedDetailsBytes: sanitizedBytes, + }; if (jsonUtf8BytesOrInfinity(fallback) <= MAX_PERSISTED_TOOL_RESULT_DETAILS_BYTES) { return fallback; } @@ -416,32 +432,6 @@ function enforcePersistedDetailsByteCap( }; } -function enforceRedactedPersistedDetailsByteCap( - redacted: unknown, - originalDetails: unknown, - originalSize: BoundedJsonUtf8Bytes, - redactionConfig?: ToolResultDetailRedactionConfig, -): unknown { - const redactedBytes = jsonUtf8BytesOrInfinity(redacted); - if (redactedBytes <= MAX_PERSISTED_TOOL_RESULT_DETAILS_BYTES) { - return redacted; - } - if (originalDetails && typeof originalDetails === "object" && !Array.isArray(originalDetails)) { - return buildPersistedDetailsFallback( - originalDetails as Record, - originalSize, - redactedBytes, - redactionConfig, - ); - } - return { - persistedDetailsTruncated: true, - finalDetailsTruncated: true, - ...originalDetailsSizeFields(originalSize), - sanitizedDetailsBytes: redactedBytes, - }; -} - function sanitizeToolResultDetailsForPersistence( details: unknown, redactionConfig?: ToolResultDetailRedactionConfig, @@ -453,7 +443,7 @@ function sanitizeToolResultDetailsForPersistence( // need to be fully stringified just to learn they exceed the persistence cap. const originalSize = boundedJsonUtf8Bytes(details, MAX_PERSISTED_TOOL_RESULT_DETAILS_BYTES); if (originalSize.complete && originalSize.bytes <= MAX_PERSISTED_TOOL_RESULT_DETAILS_BYTES) { - return enforceRedactedPersistedDetailsByteCap( + return enforcePersistedDetailsByteCap( redactPersistedDetailValue(details, 0, undefined, redactionConfig), details, originalSize, @@ -478,37 +468,33 @@ function sanitizeToolResultDetailsForPersistence( ...originalDetailsSizeFields(originalSize), originalDetailKeys: redactedOriginalDetailKeys(src, redactionConfig), }; - for (const key of [ - "status", - "sessionId", - "pid", - "startedAt", - "endedAt", - "cwd", - "name", - "exitCode", - "exitSignal", - "retryInMs", - "total", - "totalLines", - "totalChars", - "truncated", - "spill", - "fullOutputPath", - "spilledChars", - "spillTruncated", - "truncation", - ]) { - const field = src[key]; - if (field !== undefined) { - out[key] = redactPersistedSummaryField( - key, - field, - MAX_PERSISTED_DETAIL_STRING_CHARS, - redactionConfig, - ); - } - } + copyPersistedSummaryFields({ + target: out, + source: src, + keys: [ + "status", + "sessionId", + "pid", + "startedAt", + "endedAt", + "cwd", + "name", + "exitCode", + "exitSignal", + "retryInMs", + "total", + "totalLines", + "totalChars", + "truncated", + "spill", + "fullOutputPath", + "spilledChars", + "spillTruncated", + "truncation", + ], + maxChars: MAX_PERSISTED_DETAIL_STRING_CHARS, + redactionConfig, + }); copyPersistedResultStateFields(out, src, MAX_PERSISTED_DETAIL_STRING_CHARS, redactionConfig); if (typeof src.tail === "string") { out.tail = redactPersistedDetailString( @@ -528,29 +514,18 @@ function sanitizeToolResultDetailsForPersistence( return enforcePersistedDetailsByteCap(out, src, originalSize, redactionConfig); } -function capToolResultDetails( - msg: AgentMessage, - redactionConfig?: ToolResultDetailRedactionConfig, -): AgentMessage { - if ((msg as { role?: string }).role !== "toolResult") { - return msg; - } - const details = (msg as { details?: unknown }).details; - const sanitizedDetails = sanitizeToolResultDetailsForPersistence(details, redactionConfig); - if (sanitizedDetails === details) { - return msg; - } - const next = { ...msg } as AgentMessage & { details?: unknown }; - next.details = sanitizedDetails; - return next; -} - function capToolResultForPersistence( msg: AgentMessage, maxChars: number, redactionConfig?: ToolResultDetailRedactionConfig, ): AgentMessage { - return capToolResultDetails(capToolResultSize(msg, maxChars), redactionConfig); + const capped = capToolResultSize(msg, maxChars); + if (capped.role !== "toolResult") { + return capped; + } + const details = (capped as { details?: unknown }).details; + const sanitizedDetails = sanitizeToolResultDetailsForPersistence(details, redactionConfig); + return sanitizedDetails === details ? capped : { ...capped, details: sanitizedDetails }; } function normalizePersistedToolResultName( diff --git a/src/context-engine/registry.ts b/src/context-engine/registry.ts index b68e2e1dbd40..a1b205360982 100644 --- a/src/context-engine/registry.ts +++ b/src/context-engine/registry.ts @@ -11,15 +11,12 @@ import { recordPersistedContextEngineQuarantine, } from "./quarantine-health.js"; import type { - AssembleResult, BootstrapResult, - CompactResult, ContextEngine, + ContextEngineInfo, ContextEngineMaintenanceResult, IngestBatchResult, IngestResult, - SubagentSpawnPreparation, - ContextEngineInfo, } from "./types.js"; /** @@ -58,10 +55,12 @@ type RegisterContextEngineForOwnerOptions = { lifecycle?: ContextEngineRegistrationLifecycle; }; -const HOST_PARAM_METHODS = +type GuardedContextEngineMethodName = Exclude; +const GUARDED_CONTEXT_ENGINE_METHODS = new Set( "bootstrap maintain ingest ingestBatch afterTurn assemble compact prepareSubagentSpawn onSubagentEnded".split( " ", - ); + ), +); export const CONTEXT_ENGINE_HOST_PARAMS = new Set( "sessionKey prompt runtimeSettings sessionTarget runtimeContext".split(" "), ); @@ -71,7 +70,7 @@ function wrapContextEngineWithHostParamProjection(engine: ContextEngine): Contex const engineRecord = engine as unknown as Record; const wrappedRecord: Record = {}; Object.defineProperty(wrappedRecord, "info", { get: () => engine.info }); - for (const methodName of HOST_PARAM_METHODS) { + for (const methodName of GUARDED_CONTEXT_ENGINE_METHODS) { const method = engineRecord[methodName]; if (typeof method !== "function") { continue; @@ -168,10 +167,6 @@ const contextEngineRegistryState = resolveGlobalSingleton ({ ...quarantine, failedAt: new Date(failedAt) }), + ); const seenEngineIds = new Set(quarantines.map((entry) => entry.engineId)); for (const entry of listPersistedContextEngineQuarantines()) { if (seenEngineIds.has(entry.engineId)) { @@ -251,13 +236,12 @@ export function listContextEngineQuarantines(): ContextEngineRuntimeQuarantine[] } function clearContextEngineRuntimeQuarantine(engineId?: string): void { - const quarantinedEngines = getContextEngineRegistryState().quarantinedEngines; + const quarantinedEngines = contextEngineRegistryState.quarantinedEngines; if (engineId === undefined) { quarantinedEngines.clear(); - clearPersistedContextEngineQuarantineForProcess(undefined, process.pid); - return; + } else { + quarantinedEngines.delete(engineId); } - quarantinedEngines.delete(engineId); clearPersistedContextEngineQuarantineForProcess(engineId, process.pid); } @@ -272,7 +256,7 @@ export function registerContextEngineForOwner( ): ContextEngineRegistrationResult { const normalizedOwner = requireContextEngineOwner(owner); const lifecycle = opts?.lifecycle ?? "runtime"; - const registry = getContextEngineRegistryState().engines; + const registry = contextEngineRegistryState.engines; const existing = registry.get(id); if ( id === defaultSlotIdForKey("contextEngine") && @@ -301,19 +285,19 @@ export function registerContextEngineForOwner( /** Returns registration metadata so callers can distinguish discovery snapshots from runtime entries. */ export function getContextEngineRegistration(id: string): ContextEngineRegistration | undefined { - return getContextEngineRegistryState().engines.get(id); + return contextEngineRegistryState.engines.get(id); } /** * List all registered engine ids. */ function listContextEngineIds(): string[] { - return [...getContextEngineRegistryState().engines.keys()]; + return [...contextEngineRegistryState.engines.keys()]; } export function clearContextEnginesForOwner(owner: string): void { const normalizedOwner = requireContextEngineOwner(owner); - const registry = getContextEngineRegistryState().engines; + const registry = contextEngineRegistryState.engines; for (const [id, entry] of registry.entries()) { if (entry.owner === normalizedOwner) { registry.delete(id); @@ -328,10 +312,7 @@ export function clearContextEnginesForOwner(owner: string): void { export function resolveContextEngineOwnerPluginId( engine: ContextEngine | undefined | null, ): string | undefined { - if (!engine) { - return undefined; - } - const owner = resolveEffectiveContextEngineMetadata(engine)?.owner; + const owner = engine && resolveEffectiveContextEngineMetadata(engine)?.owner; if (!owner?.startsWith("plugin:")) { return undefined; } @@ -370,89 +351,46 @@ function describeResolvedContextEngineContractError( if (!info || typeof info !== "object") { issues.push("missing info"); } else { - const infoRecord = info as Record; // Engines own their internal info.id; it is metadata, not a handle into the // registry. The registered id (plugin slot id) and the engine's own id are // allowed to differ, so we only require that info.id is a non-empty string // for display/logging purposes and do not enforce equality with engineId. - const infoId = typeof infoRecord.id === "string" ? infoRecord.id.trim() : ""; - if (!infoId) { - issues.push("missing info.id"); - } - if (typeof infoRecord.name !== "string" || !infoRecord.name.trim()) { - issues.push("missing info.name"); + const infoRecord = info as Record; + for (const field of ["id", "name"]) { + const value = infoRecord[field]; + if (typeof value !== "string" || !value.trim()) { + issues.push(`missing info.${field}`); + } } } - if (typeof candidate.ingest !== "function") { - issues.push("missing ingest()"); - } - if (typeof candidate.assemble !== "function") { - issues.push("missing assemble()"); - } - if (typeof candidate.compact !== "function") { - issues.push("missing compact()"); + for (const method of ["ingest", "assemble", "compact"]) { + if (typeof candidate[method] !== "function") { + issues.push(`missing ${method}()`); + } } - if (issues.length === 0) { - return null; - } - - return `Context engine "${engineId}" factory returned an invalid ContextEngine: ${issues.join(", ")}.`; + return issues.length === 0 + ? null + : `Context engine "${engineId}" factory returned an invalid ContextEngine: ${issues.join(", ")}.`; } -type GuardedContextEngineMethodName = - | "bootstrap" - | "maintain" - | "ingest" - | "ingestBatch" - | "afterTurn" - | "assemble" - | "compact" - | "prepareSubagentSpawn" - | "onSubagentEnded"; - -const GUARDED_CONTEXT_ENGINE_METHODS = new Set([ - "bootstrap", - "maintain", - "ingest", - "ingestBatch", - "afterTurn", - "assemble", - "compact", - "prepareSubagentSpawn", - "onSubagentEnded", -] satisfies GuardedContextEngineMethodName[]); - -function contextEngineFallbackResult( - methodName: GuardedContextEngineMethodName, -): BootstrapResult | ContextEngineMaintenanceResult | IngestResult | IngestBatchResult | void { - switch (methodName) { - case "bootstrap": - return { - bootstrapped: false, - reason: "context engine downgraded to legacy", - }; - case "maintain": - return { - changed: false, - bytesFreed: 0, - rewrittenEntries: 0, - reason: "context engine downgraded to legacy", - }; - case "ingest": - return { ingested: false }; - case "ingestBatch": - return { ingestedCount: 0 }; - case "afterTurn": - case "prepareSubagentSpawn": - case "onSubagentEnded": - return undefined; - case "assemble": - case "compact": - throw new Error(`No legacy fallback result for ${methodName}`); - } -} +const CONTEXT_ENGINE_FALLBACK_RESULTS = { + bootstrap: { bootstrapped: false, reason: "context engine downgraded to legacy" }, + maintain: { + changed: false, + bytesFreed: 0, + rewrittenEntries: 0, + reason: "context engine downgraded to legacy", + }, + ingest: { ingested: false }, + ingestBatch: { ingestedCount: 0 }, +} as const satisfies { + bootstrap: BootstrapResult; + maintain: ContextEngineMaintenanceResult; + ingest: IngestResult; + ingestBatch: IngestBatchResult; +}; function contextEngineAbortSignal(methodParams: unknown): AbortSignal | undefined { if (!methodParams || typeof methodParams !== "object") { @@ -465,22 +403,7 @@ function contextEngineAbortSignal(methodParams: unknown): AbortSignal | undefine return undefined; } -function contextEngineAbortError(methodParams: unknown): Error | undefined { - const signal = contextEngineAbortSignal(methodParams); - if (!signal?.aborted) { - return undefined; - } - const reason = signal.reason; - if (reason instanceof Error) { - return reason; - } - return createAbortError( - typeof reason === "string" && reason ? reason : "Context engine operation aborted.", - ); -} - -function isContextEngineAbortRejection(error: unknown, methodParams: unknown): boolean { - const signal = contextEngineAbortSignal(methodParams); +function isContextEngineAbortRejection(error: unknown, signal: AbortSignal | undefined): boolean { if (!signal?.aborted) { return false; } @@ -488,13 +411,7 @@ function isContextEngineAbortRejection(error: unknown, methodParams: unknown): b return true; } if (error instanceof Error) { - const message = error.message.toLowerCase(); - return ( - error.name === "AbortError" || - message.includes("abort") || - message.includes("cancelled") || - message.includes("canceled") - ); + return error.name === "AbortError" || /abort|cancelled|canceled/iu.test(error.message); } return typeof error === "string" && /abort|cancelled|canceled/iu.test(error); } @@ -503,32 +420,22 @@ async function invokeFallbackContextEngineMethod(params: { getFallbackEngine: () => Promise; methodName: GuardedContextEngineMethodName; methodParams: unknown; -}): Promise< - | AssembleResult - | BootstrapResult - | CompactResult - | ContextEngineMaintenanceResult - | IngestBatchResult - | IngestResult - | SubagentSpawnPreparation - | void -> { +}): Promise { const fallbackEngine = await params.getFallbackEngine(); const fallbackMethod = fallbackEngine[params.methodName] as | ((methodParams: unknown) => unknown) | undefined; if (typeof fallbackMethod === "function") { - return (await fallbackMethod.call(fallbackEngine, params.methodParams)) as - | AssembleResult - | BootstrapResult - | CompactResult - | ContextEngineMaintenanceResult - | IngestBatchResult - | IngestResult - | SubagentSpawnPreparation - | void; + return await fallbackMethod.call(fallbackEngine, params.methodParams); } - return contextEngineFallbackResult(params.methodName); + if (params.methodName === "assemble" || params.methodName === "compact") { + throw new Error(`No legacy fallback result for ${params.methodName}`); + } + const fallbackResult = + CONTEXT_ENGINE_FALLBACK_RESULTS[ + params.methodName as keyof typeof CONTEXT_ENGINE_FALLBACK_RESULTS + ]; + return fallbackResult ? { ...fallbackResult } : undefined; } function wrapContextEngineWithRuntimeQuarantine(params: { @@ -550,17 +457,14 @@ function wrapContextEngineWithRuntimeQuarantine(params: { }); return fallbackEnginePromise; }; - const fallbackInfo = (): ContextEngineInfo => { - return ( - resolvedFallbackEngine?.info ?? { - id: params.defaultEngineId, - name: - params.defaultEngineId === "legacy" - ? "Legacy Context Engine" - : `${params.defaultEngineId} Context Engine`, - } - ); - }; + const fallbackInfo = (): ContextEngineInfo => + resolvedFallbackEngine?.info ?? { + id: params.defaultEngineId, + name: + params.defaultEngineId === "legacy" + ? "Legacy Context Engine" + : `${params.defaultEngineId} Context Engine`, + }; const isQuarantined = () => Boolean(getContextEngineQuarantine(params.engineId)); const proxy = new Proxy(params.engine, { @@ -575,23 +479,26 @@ function wrapContextEngineWithRuntimeQuarantine(params: { const methodName = property as GuardedContextEngineMethodName; return async (methodParams: unknown) => { - const aborted = contextEngineAbortError(methodParams); - if (aborted) { - throw aborted; + const abortSignal = contextEngineAbortSignal(methodParams); + if (abortSignal?.aborted) { + const reason = abortSignal.reason; + throw reason instanceof Error + ? reason + : createAbortError( + typeof reason === "string" && reason ? reason : "Context engine operation aborted.", + ); } + const invokeFallback = () => + invokeFallbackContextEngineMethod({ getFallbackEngine, methodName, methodParams }); if (isQuarantined()) { // Runtime failures downgrade future guarded calls for this process. - return await invokeFallbackContextEngineMethod({ - getFallbackEngine, - methodName, - methodParams, - }); + return await invokeFallback(); } try { return await (value as (methodParams: unknown) => unknown).call(target, methodParams); } catch (error) { - if (isContextEngineAbortRejection(error, methodParams)) { + if (isContextEngineAbortRejection(error, abortSignal)) { // Abort is caller intent, not engine instability; never quarantine for it. throw error; } @@ -605,15 +512,9 @@ function wrapContextEngineWithRuntimeQuarantine(params: { if (methodName === "compact" || methodName === "prepareSubagentSpawn") { throw error; } - try { - return await invokeFallbackContextEngineMethod({ - getFallbackEngine, - methodName, - methodParams, - }); - } catch { + return await invokeFallback().catch(() => { throw error; - } + }); } }; }, @@ -656,13 +557,10 @@ export async function resolveContextEngine( config?: OpenClawConfig, options?: ResolveContextEngineOptions, ): Promise { + const defaultEngineId = defaultSlotIdForKey("contextEngine"); const slotValue = config?.plugins?.slots?.contextEngine; const engineId = - typeof slotValue === "string" && slotValue.trim() - ? slotValue.trim() - : defaultSlotIdForKey("contextEngine"); - - const defaultEngineId = defaultSlotIdForKey("contextEngine"); + typeof slotValue === "string" && slotValue.trim() ? slotValue.trim() : defaultEngineId; const isDefaultEngine = engineId === defaultEngineId; const factoryCtx: ContextEngineFactoryContext = { @@ -677,7 +575,7 @@ export async function resolveContextEngine( return resolveDefaultContextEngine(defaultEngineId, factoryCtx); } - const entry = getContextEngineRegistryState().engines.get(engineId); + const entry = contextEngineRegistryState.engines.get(engineId); if (!entry) { if (isDefaultEngine) { throw new Error( @@ -702,47 +600,23 @@ export async function resolveContextEngine( } let engine: ContextEngine; + let operation: "factory" | "contract-validation" = "factory"; try { engine = await entry.factory(factoryCtx); - } catch (factoryError) { - if (isDefaultEngine) { - throw factoryError; - } - recordContextEngineQuarantine({ - engineId, - owner: entry.owner, - operation: "factory", - error: factoryError, - defaultEngineId, - }); - return resolveDefaultContextEngine(defaultEngineId, factoryCtx); - } - - let contractError: string | null; - try { - contractError = describeResolvedContextEngineContractError(engineId, engine); - } catch (validationError) { - if (isDefaultEngine) { - throw validationError; - } - recordContextEngineQuarantine({ - engineId, - owner: entry.owner, - operation: "contract-validation", - error: validationError, - defaultEngineId, - }); - return resolveDefaultContextEngine(defaultEngineId, factoryCtx); - } - if (contractError) { - if (isDefaultEngine) { + operation = "contract-validation"; + const contractError = describeResolvedContextEngineContractError(engineId, engine); + if (contractError) { throw new Error(contractError); } + } catch (error) { + if (isDefaultEngine) { + throw error; + } recordContextEngineQuarantine({ engineId, owner: entry.owner, - operation: "contract-validation", - error: contractError, + operation, + error, defaultEngineId, }); return resolveDefaultContextEngine(defaultEngineId, factoryCtx); @@ -766,7 +640,7 @@ async function resolveDefaultContextEngine( defaultEngineId: string, factoryCtx: ContextEngineFactoryContext, ): Promise { - const defaultEntry = getContextEngineRegistryState().engines.get(defaultEngineId); + const defaultEntry = contextEngineRegistryState.engines.get(defaultEngineId); if (!defaultEntry) { throw new Error( `[context-engine] fallback failed: default engine "${defaultEngineId}" is not registered. ` +