diff --git a/src/agent/Agent.ts b/src/agent/Agent.ts index fe3979f0..9fdb407c 100644 --- a/src/agent/Agent.ts +++ b/src/agent/Agent.ts @@ -7,7 +7,7 @@ import type { BaseCheckpointSaver, CheckpointTuple } from "@langchain/langgraph" import { MemorySaver } from "@langchain/langgraph"; import { type ReactAgent, createAgent, summarizationMiddleware } from "langchain"; import { createSubAgentMiddleware } from "deepagents/browser"; -import { TFile } from "obsidian"; +import { Notice, TFile } from "obsidian"; import { ProviderEndpointError } from "../providers/errors"; import type { ChatModelConfig } from "../providers/index"; @@ -27,6 +27,7 @@ import { import { toBase64, toBase64DataUri } from "../utils/attachments"; import { extractTextFromPdf } from "../utils/pdfExtractor"; import { Logger } from "../utils/logging"; +import { countMessageTokens, estimateSystemPromptTokens } from "../utils/tokenEstimator"; import { normalizeMessages } from "./messageNormalization"; import { type ThreadSnapshot, type ThreadStore, createSnapshot } from "./memory/ThreadStore"; import { @@ -613,7 +614,7 @@ export class Agent { systemPrompt, checkpointer: this.checkpointer, middleware: [ - ...this.buildSummarizationMiddleware(selectedModel, resolvedSummarization), + ...this.buildSummarizationMiddleware(selectedModel, systemPrompt, resolvedSummarization), ...this.buildSubAgentMiddleware(selectedModel, subAgents, tools), ] as never, }); @@ -632,7 +633,11 @@ export class Agent { }; } - private buildSummarizationMiddleware(selectedModel: SelectedModel, summarizationModel?: SelectedModel) { + private buildSummarizationMiddleware( + selectedModel: SelectedModel, + systemPrompt: string, + summarizationModel?: SelectedModel, + ) { const model = summarizationModel?.instance ?? selectedModel.instance; const contextWindow = selectedModel.options?.contextWindow; const triggerTokens = getSummarizationTriggerTokens(contextWindow); @@ -643,7 +648,14 @@ export class Agent { return []; } - const trimTokensToSummarize = getTrimTokensToSummarize(triggerTokens); + // The summarizer reads what it replaces, bounded by its OWN window: the + // dedicated summary model's when one is set, else the chat model's. + const summarizerContextWindow = summarizationModel ? summarizationModel.options?.contextWindow : contextWindow; + const trimTokensToSummarize = getTrimTokensToSummarize(summarizerContextWindow, triggerTokens); + // The middleware only counts messages; the system prompt rides along on + // every request outside them. Lowering the trigger by its size makes it + // fire at the same total the context meter shows (#545). + const messageTriggerTokens = Math.max(triggerTokens - estimateSystemPromptTokens(systemPrompt), 1); return [ // Guarded so a failed summary aborts the trim instead of replacing the @@ -651,12 +663,27 @@ export class Agent { guardSummarizationFailure( summarizationMiddleware({ model, - trigger: { tokens: triggerTokens, messages: SUMMARY_KEEP_MESSAGE_COUNT + 2 }, + trigger: { tokens: messageTriggerTokens, messages: SUMMARY_KEEP_MESSAGE_COUNT + 2 }, keep: { messages: SUMMARY_KEEP_MESSAGE_COUNT }, summaryPrefix: SUMMARY_PREFIX, summaryPrompt: SUMMARY_PROMPT, + tokenCounter: countMessageTokens, trimTokensToSummarize, }), + { + tokenCounter: countMessageTokens, + trimTokensToSummarize, + onHistoryDropped: ({ droppedTokens }) => { + new Notice( + `Smart Second Brain: about ${formatThousands(droppedTokens)} tokens of the oldest chat history did not fit the summarization model's context window and were left out of the summary.`, + ); + }, + onRepeatedFailure: (failures) => { + new Notice( + `Smart Second Brain: summarizing this chat's history failed ${failures} times in a row, so the full history is still sent and may exceed the model's context limit. Check the agent's summarization model; the developer console has the error.`, + ); + }, + }, ), ] as const; } @@ -1949,3 +1976,7 @@ ${userMessage}`; return undefined; } } + +function formatThousands(tokens: number): string { + return tokens >= 1000 ? `${Math.round(tokens / 1000)}k` : String(tokens); +} diff --git a/src/agent/summarization.ts b/src/agent/summarization.ts index 44a8ae83..6166fa73 100644 --- a/src/agent/summarization.ts +++ b/src/agent/summarization.ts @@ -1,3 +1,4 @@ +import type { BaseMessage } from "@langchain/core/messages"; import { Logger } from "../utils/logging"; const SUMMARY_TRIGGER_RATIO = 0.8; @@ -5,6 +6,11 @@ export const SUMMARY_KEEP_MESSAGE_COUNT = 12; const SUMMARY_MIN_TRIGGER_TOKENS = 12_000; const SUMMARY_MIN_TRIM_TOKENS = 6_000; const SUMMARY_MAX_TRIM_TOKENS = 24_000; +// Share of the summarization model's own window the history may fill; the rest +// holds the prompt template and the summary it writes back. +const SUMMARY_INPUT_RATIO = 0.85; +const SUMMARY_MIN_INPUT_TOKENS = 1_000; +export const SUMMARY_FAILURE_NOTICE_THRESHOLD = 3; export const SUMMARY_PREFIX = "Previous conversation summary"; export const SUMMARY_PROMPT = `Summarize the earlier conversation so the assistant can continue the thread accurately. @@ -27,7 +33,24 @@ export function getSummarizationTriggerTokens(contextWindow: number | undefined) return Math.max(Math.floor(contextWindow * SUMMARY_TRIGGER_RATIO), SUMMARY_MIN_TRIGGER_TOKENS); } -export function getTrimTokensToSummarize(triggerTokens: number): number { +/** + * How much of the history being replaced the summarization model gets to read. + * + * LangChain trims the to-be-summarized messages to this budget (keeping the + * newest) and then removes ALL of them, so anything over it is lost without a + * trace (#545). The budget is therefore the summarization model's own window, + * minus room for the prompt and the summary, so it reads everything it + * replaces. Only when that window is unknown (a dedicated summary model with no + * configured context window) does it fall back to a conservative fixed cap. + */ +export function getTrimTokensToSummarize(summarizerContextWindow: number | undefined, triggerTokens: number): number { + if (summarizerContextWindow && summarizerContextWindow > 0) { + const promptTokens = Math.ceil(SUMMARY_PROMPT.length / 4); + return Math.max( + Math.floor(summarizerContextWindow * SUMMARY_INPUT_RATIO) - promptTokens, + SUMMARY_MIN_INPUT_TOKENS, + ); + } return Math.max( Math.min(Math.floor(triggerTokens / 2), SUMMARY_MAX_TRIM_TOKENS), Math.min(SUMMARY_MIN_TRIM_TOKENS, triggerTokens), @@ -40,18 +63,38 @@ export function shouldSummarizeForEstimatedTokens(estimatedTokens: number, conte } /** - * LangChain's `summarizationMiddleware` swallows summary-model failures: - * `createSummary` catches the error and returns `"Error generating summary: - * ${e}"` (or "…: Invalid response format") as if it were the summary, and the - * middleware then commits the trim — permanently replacing the trimmed turns - * with an error stub in the model's context (#435, observed in a real thread - * where a mid-stream network error became the entire "memory" of 195 - * messages). The summary message content is always - * `${SUMMARY_PREFIX}\n\n`, so a failed one is - * identified by its exact prefix. + * Stub bodies LangChain's `createSummary` returns in place of a summary: + * - `"Error generating summary: ${e}"` (or "…: Invalid response format") — it + * catches summary-model failures and returns this as if it were the summary, + * and the middleware then commits the trim, permanently replacing the trimmed + * turns with an error stub (#435, observed in a real thread where a mid-stream + * network error became the entire "memory" of 195 messages); + * - `"Previous conversation was too long to summarize."` — the trim left nothing; + * - an empty string — the model answered with no text. + * The summary message content is always `${SUMMARY_PREFIX}\n\n`, so a failed one is identified by what follows the prefix. */ export function isFailedSummaryContent(content: unknown): boolean { - return typeof content === "string" && content.startsWith(`${SUMMARY_PREFIX}\n\nError generating summary:`); + if (typeof content !== "string") return false; + const head = `${SUMMARY_PREFIX}\n\n`; + if (!content.startsWith(head)) return false; + const body = content.slice(head.length).trim(); + return ( + body.length === 0 || + body.startsWith("Error generating summary:") || + body === "Previous conversation was too long to summarize." + ); +} + +function isSummaryMessage(msg: unknown): msg is { additional_kwargs: Record; content?: unknown } { + const message = msg as { additional_kwargs?: Record } | null; + return message?.additional_kwargs?.lc_source === "summarization"; +} + +function getUpdateMessages(update: unknown): unknown[] | null { + if (!update || typeof update !== "object") return null; + const messages = (update as { messages?: unknown }).messages; + return Array.isArray(messages) ? messages : null; } /** @@ -60,13 +103,38 @@ export function isFailedSummaryContent(content: unknown): boolean { * summary `HumanMessage`, tagged `lc_source: "summarization"`. */ export function isFailedSummaryUpdate(update: unknown): boolean { - if (!update || typeof update !== "object") return false; - const messages = (update as { messages?: unknown }).messages; - if (!Array.isArray(messages)) return false; - return messages.some((msg) => { - const message = msg as { additional_kwargs?: Record; content?: unknown } | null; - return message?.additional_kwargs?.lc_source === "summarization" && isFailedSummaryContent(message.content); - }); + const messages = getUpdateMessages(update); + if (!messages) return false; + return messages.some((msg) => isSummaryMessage(msg) && isFailedSummaryContent(msg.content)); +} + +/** The messages a successful summarization update replaced with its summary. */ +function getSummarizedMessages(state: unknown, updateMessages: unknown[]): BaseMessage[] { + const stateMessages = (state as { messages?: unknown } | null)?.messages; + if (!Array.isArray(stateMessages)) return []; + // The update is [RemoveMessage(all), summary, ...preserved]; everything in + // the old state that is not preserved was summarized. + const preservedIds = new Set( + updateMessages + .filter((msg) => !isSummaryMessage(msg)) + .map((msg) => (msg as { id?: unknown }).id) + .filter((id) => typeof id === "string"), + ); + return (stateMessages as BaseMessage[]).filter((msg) => !msg.id || !preservedIds.has(msg.id)); +} + +export interface SummarizationGuardOptions { + /** The middleware's `tokenCounter`, used to measure what was summarized. */ + tokenCounter?: (messages: BaseMessage[]) => number | Promise; + /** The middleware's `trimTokensToSummarize`. */ + trimTokensToSummarize?: number; + /** + * Called when the history being replaced was larger than the summarizer's + * budget, so its oldest part was dropped instead of summarized. + */ + onHistoryDropped?: (info: { summarizedTokens: number; budgetTokens: number; droppedTokens: number }) => void; + /** Called once when summarization has failed this many turns in a row. */ + onRepeatedFailure?: (consecutiveFailures: number) => void; } /** @@ -75,20 +143,72 @@ export function isFailedSummaryUpdate(update: unknown): boolean { * untrimmed history (one turn over the trigger ratio is harmless) and * summarization simply retries on the next turn. Success and no-op updates * pass through untouched. + * + * Failures are counted across turns, per thread (the runnable is cached per + * agent config and shared by every chat, so a single counter would mix + * chats); after `SUMMARY_FAILURE_NOTICE_THRESHOLD` in a row in one thread + * `onRepeatedFailure` fires once, since by then that thread's history keeps + * growing toward the provider's limit. A + * successful summary is measured against the trim budget and reports any + * history that was dropped unsummarized. */ -export function guardSummarizationFailure(middleware: T): T { +export function guardSummarizationFailure( + middleware: T, + options: SummarizationGuardOptions = {}, +): T { const inner = middleware.beforeModel; if (typeof inner !== "function") return middleware; + const failuresByThread = new Map(); middleware.beforeModel = async (...args: unknown[]) => { const update = await inner.apply(middleware, args); + const threadId = getThreadId(args[1]); if (isFailedSummaryUpdate(update)) { + const consecutiveFailures = (failuresByThread.get(threadId) ?? 0) + 1; + failuresByThread.set(threadId, consecutiveFailures); + const summary = getUpdateMessages(update)?.find(isSummaryMessage)?.content; Logger.warn( - "[Agent] Summary generation failed — keeping the full history this turn; summarization retries next turn.", + `[Agent] Summary generation failed (${consecutiveFailures} in a row) — keeping the full history this turn; summarization retries next turn.`, + summary, ); + if (consecutiveFailures === SUMMARY_FAILURE_NOTICE_THRESHOLD) { + options.onRepeatedFailure?.(consecutiveFailures); + } return undefined; } + + const updateMessages = getUpdateMessages(update); + if (updateMessages?.some(isSummaryMessage)) { + failuresByThread.delete(threadId); + await reportDroppedHistory(args[0], updateMessages, options); + } return update; }; return middleware; } + +/** The run's thread id from the middleware runtime (`configurable.thread_id`). */ +function getThreadId(runtime: unknown): string { + const threadId = (runtime as { configurable?: { thread_id?: unknown } } | null)?.configurable?.thread_id; + return typeof threadId === "string" ? threadId : ""; +} + +async function reportDroppedHistory( + state: unknown, + updateMessages: unknown[], + options: SummarizationGuardOptions, +): Promise { + const { tokenCounter, trimTokensToSummarize: budgetTokens } = options; + if (!tokenCounter || budgetTokens === undefined) return; + try { + const summarizedTokens = await tokenCounter(getSummarizedMessages(state, updateMessages)); + if (summarizedTokens <= budgetTokens) return; + const droppedTokens = summarizedTokens - budgetTokens; + Logger.warn( + `[Agent] Summarized ~${budgetTokens} of ~${summarizedTokens} tokens of history; the oldest ~${droppedTokens} did not fit the summarization model's budget and were dropped.`, + ); + options.onHistoryDropped?.({ summarizedTokens, budgetTokens, droppedTokens }); + } catch (err) { + Logger.debug("[Agent] Could not measure summarized history", err); + } +} diff --git a/src/utils/tokenEstimator.ts b/src/utils/tokenEstimator.ts index f11ac893..381c1317 100644 --- a/src/utils/tokenEstimator.ts +++ b/src/utils/tokenEstimator.ts @@ -1,19 +1,22 @@ import type { BaseMessage } from "@langchain/core/messages"; -import { AIMessage, HumanMessage } from "@langchain/core/messages"; +import { AIMessage, HumanMessage, ToolMessage } from "@langchain/core/messages"; -const TOKENS_PER_WORD = 1.3; // Conservative estimate for English text -const PER_MESSAGE_OVERHEAD = 80; // Separators, metadata, framing +// Same ratio as LangChain's `countTokensApproximately`. `countMessageTokens` is +// also the summarization middleware's `tokenCounter`, so the meter, the +// pre-send prediction and the actual trigger all measure with one ruler (#545). +const CHARS_PER_TOKEN = 4; +const PER_MESSAGE_OVERHEAD = 4; // Role marker and separators const SYSTEM_PROMPT_OVERHEAD = 120; const ATTACHMENT_OVERHEAD = 200; const VISIBLE_NOTES_OVERHEAD = 100; const SELECTION_OVERHEAD = 50; const GRAPH_NOTES_OVERHEAD = 100; -const TOOL_CALL_OVERHEAD = 60; +const TOOL_CALL_OVERHEAD = 10; const DRAFT_OVERHEAD = 20; /** * Estimated context usage breakdown. - * All token counts are approximations based on word counts. + * All token counts are approximations based on character counts. */ export interface UsageEstimate { /** Estimated tokens used by the entire conversation (messages + draft) */ @@ -77,8 +80,9 @@ export interface LiveContextOptions { * Estimates context usage for a conversation thread including the current draft. * * Estimation logic: - * - Word count approximation: ~1.3 tokens per word (conservative for English) - * - Per-message overhead: 80 tokens (for separators, metadata, etc.) + * - Character approximation: ~4 characters per token (LangChain's ratio) + * - Per-message overhead: 4 tokens (role marker, separators) + * - Tool results (ToolMessage bodies) count like any other message text * - Attachments: +200 tokens per attachment (conservative approximation) * - Visible notes context: +100 tokens total if present * - Selection context: +50 tokens if present @@ -186,15 +190,6 @@ export function buildUsageEstimate(totalTokens: number, contextWindow: number | }; } -/** - * Simple word count approximation. - * Splits on whitespace and filters out empty strings. - */ -function countWords(text: string): number { - if (!text) return 0; - return text.split(/\s+/).filter((word) => word.length > 0).length; -} - function stringifyForTokenEstimate(value: unknown): string { if (value === undefined || value === null) return ""; if (typeof value === "string") return value; @@ -207,11 +202,11 @@ function stringifyForTokenEstimate(value: unknown): string { function estimateTextTokens(text: string, fixedOverhead = 0): number { if (!text.trim()) return 0; - const wordCount = countWords(text); - return Math.ceil(wordCount * TOKENS_PER_WORD) + fixedOverhead; + return Math.ceil(text.length / CHARS_PER_TOKEN) + fixedOverhead; } -function estimateSystemPromptTokens(systemPrompt: string | undefined): number { +/** Estimated tokens of the assembled system prompt, as the meter counts it. */ +export function estimateSystemPromptTokens(systemPrompt: string | undefined): number { if (!systemPrompt?.trim()) return 0; return estimateTextTokens(systemPrompt, SYSTEM_PROMPT_OVERHEAD); } @@ -227,6 +222,16 @@ function estimateAdditionalTextBlockTokens(additionalTextBlocks: string[] | unde return totalTokens; } +/** + * Estimated tokens of a message list, using the meter's estimator. Passed to the + * summarization middleware as its `tokenCounter` so the trigger fires where the + * meter says it will. + */ +export function countMessageTokens(messages: BaseMessage[]): number { + const { humanTokens, assistantTokens, toolTokens } = estimateBaseMessagePayloadTokens(messages); + return humanTokens + assistantTokens + toolTokens; +} + function estimateBaseMessagePayloadTokens(messages: BaseMessage[]): { humanTokens: number; assistantTokens: number; @@ -278,6 +283,11 @@ function estimateBaseMessagePayloadTokens(messages: BaseMessage[]): { toolTokens += estimateTextTokens(payload, TOOL_CALL_OVERHEAD); } } + continue; + } + + if (ToolMessage.isInstance(message)) { + toolTokens += estimateTextTokens(message.text || "") + PER_MESSAGE_OVERHEAD; } } diff --git a/test/agent/Agent.test.ts b/test/agent/Agent.test.ts index aeabb365..6b65f894 100644 --- a/test/agent/Agent.test.ts +++ b/test/agent/Agent.test.ts @@ -14,6 +14,8 @@ vi.mock("langchain", () => ({ import type { BaseChatModel } from "@langchain/core/language_models/chat_models"; import { Agent, sanitizeRunnableName } from "../../src/agent/Agent"; +import { getTrimTokensToSummarize } from "../../src/agent/summarization"; +import { countMessageTokens, estimateSystemPromptTokens } from "../../src/utils/tokenEstimator"; function makeRegistry() { return { @@ -32,6 +34,10 @@ const BASE_RESOLVE_PARAMS = { subAgents: [] as const, } as const; +// The middleware counts messages only, so its trigger is the meter's trigger +// minus the system prompt the meter adds on top. +const SYSTEM_PROMPT_TOKENS = estimateSystemPromptTokens(BASE_RESOLVE_PARAMS.systemPrompt); + describe("Agent summarization middleware", () => { beforeEach(() => { createAgentMock.mockClear(); @@ -56,8 +62,10 @@ describe("Agent summarization middleware", () => { expect(summarizationMiddlewareMock).toHaveBeenCalledTimes(1); expect((summarizationMiddlewareMock.mock.calls[0] ?? [])[0]).toMatchObject({ - trigger: { tokens: 80_000, messages: 14 }, + trigger: { tokens: 80_000 - SYSTEM_PROMPT_TOKENS, messages: 14 }, keep: { messages: 12 }, + tokenCounter: countMessageTokens, + trimTokensToSummarize: getTrimTokensToSummarize(100_000, 80_000), }); expect(registry.createChatInstance).toHaveBeenCalledWith("openai", "gpt-4o-mini", { contextWindow: 100_000 }); expect(createAgentMock).toHaveBeenCalledWith( @@ -82,9 +90,30 @@ describe("Agent summarization middleware", () => { expect(summarizationMiddlewareMock).toHaveBeenCalledTimes(1); expect(registry.createChatInstance).toHaveBeenCalledTimes(1); expect((summarizationMiddlewareMock.mock.calls[0] ?? [])[0]).toMatchObject({ - trigger: { tokens: 40_000, messages: 14 }, + trigger: { tokens: 40_000 - SYSTEM_PROMPT_TOKENS, messages: 14 }, keep: { messages: 12 }, + tokenCounter: countMessageTokens, + // The chat model summarizes, so its own window bounds what it reads. + trimTokensToSummarize: getTrimTokensToSummarize(50_000, 40_000), + }); + }); + + it("bounds the summarizer's input by the dedicated summary model's own window", async () => { + const registry = makeRegistry(); + const agent = new Agent({ registry: registry as never }); + + await agent.resolveRun({ + ...BASE_RESOLVE_PARAMS, + cacheKey: "test-key-4", + provider: "openai", + chatModel: "gpt-4o", + options: { contextWindow: 200_000 }, + summarizationModel: { provider: "openai", chatModel: "small", options: { contextWindow: 32_000 } }, }); + + const config = (summarizationMiddlewareMock.mock.calls[0] ?? [])[0] as { trimTokensToSummarize: number }; + expect(config.trimTokensToSummarize).toBe(getTrimTokensToSummarize(32_000, 160_000)); + expect(config.trimTokensToSummarize).toBeLessThan(32_000); }); it("skips summarization middleware when the context window is unknown", async () => { diff --git a/test/agent/summarizationFailureGuard.test.ts b/test/agent/summarizationFailureGuard.test.ts index abc70f6f..af76f191 100644 --- a/test/agent/summarizationFailureGuard.test.ts +++ b/test/agent/summarizationFailureGuard.test.ts @@ -1,12 +1,15 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { AIMessage, HumanMessage } from "@langchain/core/messages"; import { summarizationMiddleware } from "langchain"; import { + getTrimTokensToSummarize, guardSummarizationFailure, isFailedSummaryContent, isFailedSummaryUpdate, + SUMMARY_FAILURE_NOTICE_THRESHOLD, SUMMARY_PREFIX, } from "../../src/agent/summarization"; +import { countMessageTokens } from "../../src/utils/tokenEstimator"; /* -------------------------------------------------------------------------- * Pure helpers @@ -24,6 +27,14 @@ describe("isFailedSummaryContent", () => { ); }); + it("matches an empty summary and the too-long-to-summarize stub (#545)", () => { + expect(isFailedSummaryContent(`${SUMMARY_PREFIX}\n\n`)).toBe(true); + expect(isFailedSummaryContent(`${SUMMARY_PREFIX}\n\n \n`)).toBe(true); + expect(isFailedSummaryContent(`${SUMMARY_PREFIX}\n\nPrevious conversation was too long to summarize.`)).toBe( + true, + ); + }); + it("does not match real summaries, even ones mentioning errors", () => { expect(isFailedSummaryContent(`${SUMMARY_PREFIX}\n\nUser debugged an "Error generating summary" bug.`)).toBe( false, @@ -65,6 +76,24 @@ describe("isFailedSummaryUpdate", () => { }); }); +describe("getTrimTokensToSummarize", () => { + it("lets the summarizer read nearly its whole window, so it sees everything it replaces (#545)", () => { + // At a 128k window the trigger is ~102k; the old cap read only 24k of it. + const budget = getTrimTokensToSummarize(128_000, 102_400); + expect(budget).toBeGreaterThan(102_400 - 12 * 1_000); + expect(budget).toBeLessThan(128_000); + }); + + it("is bounded by a smaller dedicated summary model's window", () => { + expect(getTrimTokensToSummarize(16_000, 160_000)).toBeLessThan(16_000); + }); + + it("keeps the conservative cap when the summarizer's window is unknown", () => { + expect(getTrimTokensToSummarize(undefined, 102_400)).toBe(24_000); + expect(getTrimTokensToSummarize(undefined, 12_000)).toBe(6_000); + }); +}); + /* -------------------------------------------------------------------------- * Guard against the real middleware * ------------------------------------------------------------------------*/ @@ -81,16 +110,31 @@ function longConversation() { ]; } -function makeMiddleware(model: { invoke: (...args: unknown[]) => Promise }) { +function makeMiddleware( + model: { invoke: (...args: unknown[]) => Promise }, + extra: { trimTokensToSummarize?: number } = {}, +) { return summarizationMiddleware({ model: model as never, trigger: { tokens: 10 }, keep: { messages: 2 }, summaryPrefix: SUMMARY_PREFIX, summaryPrompt: "Summarize:\n{messages}", + tokenCounter: countMessageTokens, + ...extra, }); } +type BeforeModel = (s: unknown, r: unknown) => Promise; + +function failingModel() { + return { + invoke: async () => { + throw new Error("boom"); + }, + }; +} + const runtime = { context: {} } as never; describe("guardSummarizationFailure (real summarizationMiddleware)", () => { @@ -163,3 +207,128 @@ describe("guardSummarizationFailure (real summarizationMiddleware)", () => { expect(bare.beforeModel).toBeUndefined(); }); }); + +describe("guardSummarizationFailure — repeated failures and dropped history", () => { + it("aborts the trim for an empty summary", async () => { + const middleware = guardSummarizationFailure( + makeMiddleware({ invoke: async () => new AIMessage({ content: " " }) }), + ); + const update = await (middleware.beforeModel as BeforeModel)({ messages: longConversation() }, runtime); + expect(update).toBeUndefined(); + }); + + it("aborts the trim when nothing survived trimming (too-long stub)", async () => { + const invoke = vi.fn(); + const middleware = guardSummarizationFailure(makeMiddleware({ invoke }, { trimTokensToSummarize: 0 })); + const update = await (middleware.beforeModel as BeforeModel)({ messages: longConversation() }, runtime); + expect(invoke).not.toHaveBeenCalled(); + expect(update).toBeUndefined(); + }); + + it("reports repeated consecutive failures exactly once per streak", async () => { + const onRepeatedFailure = vi.fn(); + let fail = true; + const middleware = guardSummarizationFailure( + makeMiddleware({ + invoke: async () => { + if (fail) throw new Error("boom"); + return new AIMessage({ content: "A proper summary." }); + }, + }), + { onRepeatedFailure }, + ); + const run = () => (middleware.beforeModel as BeforeModel)({ messages: longConversation() }, runtime); + + for (let i = 1; i < SUMMARY_FAILURE_NOTICE_THRESHOLD; i++) await run(); + expect(onRepeatedFailure).not.toHaveBeenCalled(); + await run(); + expect(onRepeatedFailure).toHaveBeenCalledTimes(1); + expect(onRepeatedFailure).toHaveBeenCalledWith(SUMMARY_FAILURE_NOTICE_THRESHOLD); + await run(); + expect(onRepeatedFailure).toHaveBeenCalledTimes(1); + + // A success resets the streak; a fresh streak reports again. + fail = false; + await run(); + fail = true; + for (let i = 0; i < SUMMARY_FAILURE_NOTICE_THRESHOLD; i++) await run(); + expect(onRepeatedFailure).toHaveBeenCalledTimes(2); + }); + + it("counts failure streaks per thread, not across chats sharing the runnable", async () => { + const onRepeatedFailure = vi.fn(); + const failingThreads = new Set(["a", "b"]); + const middleware = guardSummarizationFailure( + makeMiddleware({ + invoke: async (_prompt: unknown, config: unknown) => { + const threadId = (config as { configurable?: { thread_id?: string } }).configurable?.thread_id; + if (threadId && failingThreads.has(threadId)) throw new Error("boom"); + return new AIMessage({ content: "A proper summary." }); + }, + }), + { onRepeatedFailure }, + ); + const run = (threadId: string) => + (middleware.beforeModel as BeforeModel)( + { messages: longConversation() }, + { context: {}, configurable: { thread_id: threadId } }, + ); + + // Two failures in "a", then one in "b": neither thread reached the threshold. + for (let i = 1; i < SUMMARY_FAILURE_NOTICE_THRESHOLD; i++) await run("a"); + await run("b"); + expect(onRepeatedFailure).not.toHaveBeenCalled(); + + // A success in "c" must not reset "a"'s streak. + await run("c"); + await run("a"); + expect(onRepeatedFailure).toHaveBeenCalledTimes(1); + }); + + it("does not report a failure streak for failures interrupted by no-ops", async () => { + const onRepeatedFailure = vi.fn(); + const middleware = guardSummarizationFailure(makeMiddleware(failingModel()), { onRepeatedFailure }); + await (middleware.beforeModel as BeforeModel)({ messages: longConversation() }, runtime); + // A below-trigger turn neither fails nor succeeds; it must not reset or count. + await (middleware.beforeModel as BeforeModel)({ messages: [new HumanMessage("hi")] }, runtime); + expect(onRepeatedFailure).not.toHaveBeenCalled(); + }); + + it("reports history that exceeded the summarizer's budget", async () => { + const onHistoryDropped = vi.fn(); + const trimTokensToSummarize = 20; + const middleware = guardSummarizationFailure( + makeMiddleware( + { invoke: async () => new AIMessage({ content: "Summary of what fit." }) }, + { trimTokensToSummarize }, + ), + { tokenCounter: countMessageTokens, trimTokensToSummarize, onHistoryDropped }, + ); + const messages = longConversation(); + const update = await (middleware.beforeModel as BeforeModel)({ messages }, runtime); + + expect(update).toBeDefined(); + expect(onHistoryDropped).toHaveBeenCalledTimes(1); + // Everything but the 2 kept messages was summarized. + const summarizedTokens = countMessageTokens(messages.slice(0, -2)); + expect(onHistoryDropped).toHaveBeenCalledWith({ + summarizedTokens, + budgetTokens: trimTokensToSummarize, + droppedTokens: summarizedTokens - trimTokensToSummarize, + }); + }); + + it("stays silent when the summarizer read everything it replaced", async () => { + const onHistoryDropped = vi.fn(); + const trimTokensToSummarize = 100_000; + const middleware = guardSummarizationFailure( + makeMiddleware( + { invoke: async () => new AIMessage({ content: "Full summary." }) }, + { trimTokensToSummarize }, + ), + { tokenCounter: countMessageTokens, trimTokensToSummarize, onHistoryDropped }, + ); + await (middleware.beforeModel as BeforeModel)({ messages: longConversation() }, runtime); + expect(onHistoryDropped).not.toHaveBeenCalled(); + }); +}); diff --git a/test/utils/tokenEstimator.test.ts b/test/utils/tokenEstimator.test.ts index d72b2d53..52e85694 100644 --- a/test/utils/tokenEstimator.test.ts +++ b/test/utils/tokenEstimator.test.ts @@ -1,9 +1,11 @@ import { describe, it, expect, vi } from "vitest"; -import { AIMessage, HumanMessage } from "@langchain/core/messages"; +import { AIMessage, HumanMessage, ToolMessage } from "@langchain/core/messages"; +import { countTokensApproximately } from "langchain"; vi.mock("obsidian", () => import("../__mocks__/obsidian")); import { + countMessageTokens, estimateContextUsage, estimateContextUsageBreakdown, formatContextUsage, @@ -40,26 +42,26 @@ describe("tokenEstimator", () => { const messages = createConversation("hello world", "goodbye world"); const result = estimateContextUsage(messages, "", 128000); - // "hello world" = 2 words * 1.3 = 2.6 ≈ 3 tokens - // "goodbye world" = 2 words * 1.3 = 2.6 ≈ 3 tokens - // Per-message overhead: 80 + 80 = 160 - // Total: 3 + 3 + 160 = 166 - expect(result.estimatedUsedTokens).toBe(166); - expect(result.usagePercent).toBeCloseTo(0.13, 2); + // "hello world" = 11 chars / 4 ≈ 3 tokens + // "goodbye world" = 13 chars / 4 ≈ 4 tokens + // Per-message overhead: 4 + 4 = 8 + // Total: 3 + 4 + 8 = 15 + expect(result.estimatedUsedTokens).toBe(15); + expect(result.usagePercent).toBeCloseTo(0.0117, 4); }); it("should include draft text in estimation", () => { const messages = createConversation("hello", "hi"); const result = estimateContextUsage(messages, "what is this", 128000); - // hello = 1 word * 1.3 ≈ 2 tokens (rounded up) - // hi = 1 word * 1.3 ≈ 2 tokens (rounded up) - // Message overheads: 160 - // Draft "what is this" = 3 words * 1.3 = 3.9 ≈ 4 tokens + // hello = 5 chars / 4 ≈ 2 tokens (rounded up) + // hi = 2 chars / 4 ≈ 1 token (rounded up) + // Message overheads: 8 + // Draft "what is this" = 12 chars / 4 = 3 tokens // Draft overhead: 20 - // Total: 2 + 2 + 160 + 4 + 20 = 188 - expect(result.estimatedUsedTokens).toBe(188); - expect(result.usagePercent).toBeCloseTo(0.147, 2); + // Total: 2 + 1 + 8 + 3 + 20 = 34 + expect(result.estimatedUsedTokens).toBe(34); + expect(result.usagePercent).toBeCloseTo(0.0266, 4); }); it("should ignore empty draft text", () => { @@ -75,8 +77,8 @@ describe("tokenEstimator", () => { const result = estimateContextUsage(messages, "", 128000); // 2 attachments * 200 tokens each = 400 - // Plus other content: ~164 (hello + hi + message overheads) - // Expected: ~564 + // Plus other content: 11 (hello + hi + message overheads) + // Expected: 411 expect(result.estimatedUsedTokens).toBeGreaterThan(400); expect(result.estimatedUsedTokens).toBeLessThan(600); }); @@ -181,7 +183,7 @@ describe("tokenEstimator", () => { const result = estimateContextUsage(messages, "", 128000); // Should accumulate tokens from all three pairs - expect(result.estimatedUsedTokens).toBeGreaterThan(300); // At least 3 * ~100 + expect(result.estimatedUsedTokens).toBeGreaterThan(30); expect(result.usagePercent).toBeGreaterThan(0); }); @@ -190,7 +192,7 @@ describe("tokenEstimator", () => { const messages = createConversation(longText, "short response"); const result = estimateContextUsage(messages, "", 128000); - // long text (1000 words) * 1.3 ≈ 1300 tokens + overhead + // long text (4999 chars) / 4 ≈ 1250 tokens + overhead expect(result.estimatedUsedTokens).toBeGreaterThan(1000); }); @@ -198,8 +200,8 @@ describe("tokenEstimator", () => { const messages = createConversation("", ""); const result = estimateContextUsage(messages, "", 128000); - // Only message overhead: 80 + 80 = 160 - expect(result.estimatedUsedTokens).toBe(160); + // Only message overhead: 4 + 4 = 8 + expect(result.estimatedUsedTokens).toBe(8); }); it("should not double-count when messages and draft are both empty", () => { @@ -262,6 +264,55 @@ describe("tokenEstimator", () => { expect(withToolCall.estimatedUsedTokens).toBeGreaterThan(withoutToolCall.estimatedUsedTokens); }); + + it("counts tool result bodies under tool tokens (#545)", () => { + const toolOutput = "x".repeat(4000); + const breakdown = estimateContextUsageBreakdown( + [ + new HumanMessage("read it"), + new AIMessage({ + content: "", + tool_calls: [{ id: "tc-1", name: "read_content", args: { path: "note.md" } }], + }), + new ToolMessage({ content: toolOutput, tool_call_id: "tc-1" }), + ], + "", + ); + + // 4000 chars / 4 = 1000, plus the call payload and framing. + expect(breakdown.toolTokens).toBeGreaterThanOrEqual(1000); + expect(breakdown.toolTokens).toBeLessThan(1100); + }); + }); + + describe("countMessageTokens", () => { + it("matches the meter's message payload, so meter and summarization trigger agree", () => { + const messages = [ + new HumanMessage("hello there"), + new AIMessage({ + content: "on it", + tool_calls: [{ id: "tc-1", name: "read_content", args: { path: "a.md" } }], + }), + new ToolMessage({ content: "note body ".repeat(200), tool_call_id: "tc-1" }), + ]; + const breakdown = estimateContextUsageBreakdown(messages, ""); + + expect(countMessageTokens(messages)).toBe( + breakdown.humanTokens + breakdown.assistantTokens + breakdown.toolTokens, + ); + }); + + it("stays within a few percent of LangChain's chars/4 counter on text-heavy history", () => { + const messages = [ + new HumanMessage("Summarize my notes about the project roadmap. ".repeat(40)), + new AIMessage("Here is what the roadmap says about the next quarter. ".repeat(60)), + new ToolMessage({ content: "A long note body with plenty of text. ".repeat(300), tool_call_id: "t" }), + ]; + const ours = countMessageTokens(messages); + const langchain = countTokensApproximately(messages); + + expect(Math.abs(ours - langchain) / langchain).toBeLessThan(0.02); + }); }); describe("formatContextUsage", () => {