diff --git a/CHANGELOG.md b/CHANGELOG.md index 1df12b35..3acaabea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,8 +6,13 @@ ## [Unreleased] +### 🐛 修复 + +- **Prompt cache 命中率回归防护** ([#120](https://github.com/TencentCloud/TencentDB-Agent-Memory/issues/120)):将 `` 注入块清理逻辑抽取为 `src/utils/injected-memory.ts`,并在 `before_message_write` 中复用,避免 `showInjected=true` 时动态 L1 召回内容被持久化进历史消息。新增单测覆盖字符串消息、多 part 消息和历史膨胀估算。 + ### ✨ 新功能 +- **Cache-Aware Context Engine D 方案核心能力** ([#120](https://github.com/TencentCloud/TencentDB-Agent-Memory/issues/120)):新增 `recall.mode="tool-only"` 默认路径、确定性 Session Snapshot、前置工具结果卸载、`tdai_offload_read`、Cache Epoch 纯合约与 Task Snapshot/Delta 纯合约。默认不再每轮自动 prepend 动态 L1 记忆;大工具结果会先写入 `refs/*.md`,prompt 只保留摘要、`result_ref` 和恢复说明。 - **时区可配置** ([#75](https://github.com/Tencent/TencentDB-Agent-Memory/issues/75) / [#87](https://github.com/Tencent/TencentDB-Agent-Memory/issues/87)):新增顶层 `timezone` 配置项,支持 IANA 时区名(`Asia/Shanghai`、`Europe/Berlin`)和 UTC 偏移串(`+08:00`、`-05:30`)。默认 `"system"`(跟随进程系统时区),升级零感。 - **暴露给 LLM 的时间戳**统一为带显式 offset 的 ISO 8601(如 `2026-04-07T11:04:45+08:00`),修复 #87 报告的 UTC/本地时区混用导致 LLM 误算时间差的问题。 - **L1 / L2 prompt 顶部**自动插入时区声明,指引 LLM 按正确时区推算"昨天"、"上周"等相对时间。 diff --git a/index.ts b/index.ts index 868a7701..d1875e7d 100644 --- a/index.ts +++ b/index.ts @@ -44,6 +44,10 @@ import { decideHookPolicy, } from "./src/utils/ensure-hook-policy.js"; import { resolveOpenClawStateDir } from "./src/utils/openclaw-state-dir.js"; +import { + measureMessageContentChars, + stripInjectedRelevantMemoriesFromContent, +} from "./src/utils/injected-memory.js"; const TAG = "[memory-tdai]"; @@ -613,41 +617,25 @@ export default function register(api: OpenClawPluginApi) { } // Strip from user messages before they are persisted to - // the session JSONL. The current-turn LLM already saw the full prompt - // (effectivePrompt lives in memory), but we don't want recall artifacts - // polluting the historical transcript for future replays. - api.logger.debug?.(`${TAG} Registering before_message_write hook (strip )`); + // the session JSONL unless the operator explicitly opts into persistence. + // The current-turn LLM already saw the full prompt (effectivePrompt lives in + // memory), but D-plan defaults keep dynamic recall out of replay history. + api.logger.debug?.(`${TAG} Registering before_message_write hook (persistInjected=${cfg.recall.persistInjected})`); api.on("before_message_write", (event) => { + if (cfg.recall.persistInjected) return; const msg = event.message as { role?: string; content?: unknown }; const contentType = typeof msg.content === "string" ? "string" : Array.isArray(msg.content) ? "parts" : typeof msg.content; api.logger.debug?.(`${TAG} [before_message_write] role=${msg.role}, contentType=${contentType}`); if (msg.role !== "user") return; - // UserMessage.content: string | (TextContent | ImageContent)[] - const STRIP_RE = /[\s\S]*?<\/relevant-memories>\s*/g; - - if (typeof msg.content === "string") { - if (!msg.content.includes("")) return; - const cleaned = msg.content.replace(STRIP_RE, "").trim(); - if (cleaned === msg.content) return; - api.logger.debug?.(`${TAG} [before_message_write] Stripped: ${msg.content.length} → ${cleaned.length} chars`); - return { message: { ...event.message, content: cleaned } as typeof event.message }; - } + const cleanedContent = stripInjectedRelevantMemoriesFromContent(msg.content); + if (cleanedContent === msg.content) return; - if (Array.isArray(msg.content)) { - let totalStripped = 0; - const cleanedParts = (msg.content as Array>).map((part) => { - if (part.type !== "text" || typeof part.text !== "string") return part; - if (!(part.text as string).includes("")) return part; - const cleaned = (part.text as string).replace(STRIP_RE, "").trim(); - totalStripped += (part.text as string).length - cleaned.length; - return { ...part, text: cleaned }; - }); - if (totalStripped === 0) return; - api.logger.debug?.(`${TAG} [before_message_write] Stripped from parts: removed ${totalStripped} chars`); - return { message: { ...event.message, content: cleanedParts } as unknown as typeof event.message }; - } + const beforeLen = measureMessageContentChars(msg.content); + const afterLen = measureMessageContentChars(cleanedContent); + api.logger.debug?.(`${TAG} [before_message_write] Stripped injected memories: ${beforeLen} → ${afterLen} chars`); + return { message: { ...event.message, content: cleanedContent } as typeof event.message }; }); // After agent end: auto-capture + L0 record + L1/L2/L3 schedule diff --git a/openclaw.plugin.json b/openclaw.plugin.json index 0a6d7e9a..703956f5 100644 --- a/openclaw.plugin.json +++ b/openclaw.plugin.json @@ -7,7 +7,7 @@ "onStartup": true }, "contracts": { - "tools": ["tdai_memory_search", "tdai_conversation_search"] + "tools": ["tdai_memory_search", "tdai_conversation_search", "tdai_offload_read"] }, "configSchema": { "type": "object", @@ -74,12 +74,20 @@ "description": "记忆召回设置", "properties": { "enabled": { "type": "boolean", "default": true, "description": "是否启用自动召回" }, + "mode": { "type": "string", "enum": ["tool-only", "auto"], "default": "tool-only", "description": "召回模式:tool-only 默认不向每轮 prompt 注入动态 L1 记忆,由 Agent 按需调用工具;auto 可启用旧式自动召回" }, + "showInjected": { "type": "boolean", "default": false, "description": "mode=auto 时是否把动态召回的 注入当前 prompt" }, + "persistInjected": { "type": "boolean", "default": false, "description": "是否允许动态召回内容写入历史消息。默认 false,避免污染后续重放和 prompt cache 前缀" }, "maxResults": { "type": "number", "default": 5, "description": "召回最大结果数" }, "maxCharsPerMemory": { "type": "number", "default": 0, "description": "单条 L1 记忆注入的最大字符数;填 0 表示不限制" }, "maxTotalRecallChars": { "type": "number", "default": 0, "description": "本轮 auto-recall 注入的 L1 记忆总字符预算;填 0 表示不限制" }, "scoreThreshold": { "type": "number", "default": 0.3, "description": "最低分数阈值" }, "strategy": { "type": "string", "enum": ["embedding", "keyword", "hybrid"], "default": "hybrid", "description": "搜索策略:keyword(关键词)、embedding(向量)、hybrid(混合RRF融合,推荐)" }, - "timeoutMs": { "type": "number", "default": 5000, "description": "召回整体超时(毫秒),超时后跳过记忆注入并打印警告日志" } + "timeoutMs": { "type": "number", "default": 5000, "description": "召回整体超时(毫秒),超时后跳过记忆注入并打印警告日志" }, + "sessionSnapshotMaxTokens": { "type": "number", "default": 1600, "description": "稳定 session snapshot 的 token 预算" }, + "stableMemoryMaxItems": { "type": "number", "default": 8, "description": "稳定快照中最多纳入的稳定记忆条数" }, + "sceneSummaryMaxItems": { "type": "number", "default": 6, "description": "稳定快照中最多纳入的场景摘要条数" }, + "dynamicRecallMaxTokens": { "type": "number", "default": 900, "description": "mode=auto 时动态召回片段预算" }, + "maxSearchCallsPerTurn": { "type": "number", "default": 3, "description": "tdai_memory_search 与 tdai_conversation_search 每轮合计调用预算" } } }, "embedding": { @@ -161,6 +169,14 @@ "mildOffloadRatio": { "type": "number", "default": 0.5, "description": "温和压缩触发比例(占 context window)" }, "aggressiveCompressRatio": { "type": "number", "default": 0.85, "description": "激进压缩触发比例" }, "mmdMaxTokenRatio": { "type": "number", "default": 0.2, "description": "MMD 注入 token 预算比例" }, + "inlineToolResultMaxTokens": { "type": "number", "default": 1200, "description": "单个工具结果超过该 token 估算值时,前置写入 refs 并在 prompt 中替换为 result_ref 摘要" }, + "summaryMaxTokens": { "type": "number", "default": 350, "description": "工具结果前置卸载摘要 token 预算" }, + "previewMaxChars": { "type": "number", "default": 1200, "description": "用于生成确定性摘要的预览字符数" }, + "readChunkMaxTokens": { "type": "number", "default": 1600, "description": "tdai_offload_read 单次返回 token 预算" }, + "epochTriggerRatio": { "type": "number", "default": 0.8, "description": "Cache Epoch 触发比例(占 context window)" }, + "epochTargetRatio": { "type": "number", "default": 0.55, "description": "Cache Epoch 批量压缩后的目标比例" }, + "epochMinimumTurns": { "type": "number", "default": 4, "description": "允许切换 Cache Epoch 前的最小轮数" }, + "epochMinimumIntervalTurns": { "type": "number", "default": 2, "description": "可批量压缩的最小连续旧轮数" }, "backendUrl": { "type": "string", "description": "后端服务 URL(如 https://offload-api.example.com),配置后 L1/L1.5/L2/L4 走后端" }, "backendApiKey": { "type": "string", "description": "后端 API 认证 token" }, "backendTimeoutMs": { "type": "number", "default": 10000, "description": "后端调用超时(毫秒)" } diff --git a/src/config.cache-aware.test.ts b/src/config.cache-aware.test.ts new file mode 100644 index 00000000..98a52f2f --- /dev/null +++ b/src/config.cache-aware.test.ts @@ -0,0 +1,33 @@ +import { describe, expect, it } from "vitest"; +import { parseConfig } from "./config.js"; + +describe("cache-aware context defaults", () => { + it("defaults recall to tool-only with injected dynamic memories hidden and unpersisted", () => { + const cfg = parseConfig({}); + + expect(cfg.recall.mode).toBe("tool-only"); + expect(cfg.recall.showInjected).toBe(false); + expect(cfg.recall.persistInjected).toBe(false); + expect(cfg.recall.maxSearchCallsPerTurn).toBe(3); + }); + + it("parses front-offload and cache-epoch controls", () => { + const cfg = parseConfig({ + recall: { mode: "auto", showInjected: true, persistInjected: false }, + offload: { + inlineToolResultMaxTokens: 32, + summaryMaxTokens: 12, + previewMaxChars: 80, + epochTriggerRatio: 0.7, + }, + }); + + expect(cfg.recall.mode).toBe("auto"); + expect(cfg.recall.showInjected).toBe(true); + expect(cfg.recall.persistInjected).toBe(false); + expect(cfg.offload.inlineToolResultMaxTokens).toBe(32); + expect(cfg.offload.summaryMaxTokens).toBe(12); + expect(cfg.offload.previewMaxChars).toBe(80); + expect(cfg.offload.epochTriggerRatio).toBe(0.7); + }); +}); diff --git a/src/config.ts b/src/config.ts index ef614415..4bc43238 100644 --- a/src/config.ts +++ b/src/config.ts @@ -81,6 +81,12 @@ export interface PipelineTriggerConfig { export interface RecallConfig { /** Enable auto-recall (default: true) */ enabled: boolean; + /** Recall mode. "tool-only" keeps dynamic L1 recall out of the prompt by default. */ + mode: "auto" | "tool-only"; + /** Show dynamic injected memories in the prompt when mode="auto" (default: false) */ + showInjected: boolean; + /** Persist dynamic injected memories into conversation history (default: false) */ + persistInjected: boolean; /** Max results to return (default: 5) */ maxResults: number; /** Max characters injected for a single recalled L1 memory. 0 disables the per-memory limit. */ @@ -93,6 +99,16 @@ export interface RecallConfig { strategy: "embedding" | "keyword" | "hybrid"; /** Overall recall timeout in milliseconds (default: 5000). When exceeded, recall is skipped with a warning. */ timeoutMs: number; + /** Stable session snapshot token budget (default: 1600) */ + sessionSnapshotMaxTokens: number; + /** Max stable memories to include in the frozen snapshot (default: 8) */ + stableMemoryMaxItems: number; + /** Max scene summary items to include in the frozen snapshot (default: 6) */ + sceneSummaryMaxItems: number; + /** Dynamic recall token budget for mode="auto" (default: 900) */ + dynamicRecallMaxTokens: number; + /** Combined active-search call budget per user turn (default: 3) */ + maxSearchCallsPerTurn: number; } /** Embedding service configuration for vector search. */ @@ -260,6 +276,22 @@ export interface OffloadConfig { aggressiveCompressRatio: number; /** MMD injection token budget ratio (default: 0.2) */ mmdMaxTokenRatio: number; + /** Inline tool-result budget before front offload replaces the prompt payload (default: 1200) */ + inlineToolResultMaxTokens: number; + /** Summary token budget for offloaded tool result stubs (default: 350) */ + summaryMaxTokens: number; + /** Preview character budget used to build deterministic tool result summaries (default: 1200) */ + previewMaxChars: number; + /** Max tokens returned by tdai_offload_read in one call (default: 1600) */ + readChunkMaxTokens: number; + /** Cache epoch transition trigger ratio of context window (default: 0.8) */ + epochTriggerRatio: number; + /** Cache epoch target ratio after batch compaction (default: 0.55) */ + epochTargetRatio: number; + /** Minimum turns before an epoch can transition (default: 4) */ + epochMinimumTurns: number; + /** Minimum continuous old interval turns to compact (default: 2) */ + epochMinimumIntervalTurns: number; /** Backend service URL. When set, L1/L1.5/L2/L4 LLM calls go through the backend. */ backendUrl?: string; /** Backend API authentication token */ @@ -489,6 +521,14 @@ export function parseConfig(raw: Record | undefined): MemoryTda mildOffloadRatio: num(offloadGroup, "mildOffloadRatio") ?? 0.5, aggressiveCompressRatio: num(offloadGroup, "aggressiveCompressRatio") ?? 0.85, mmdMaxTokenRatio: num(offloadGroup, "mmdMaxTokenRatio") ?? 0.2, + inlineToolResultMaxTokens: num(offloadGroup, "inlineToolResultMaxTokens") ?? 1200, + summaryMaxTokens: num(offloadGroup, "summaryMaxTokens") ?? 350, + previewMaxChars: num(offloadGroup, "previewMaxChars") ?? 1200, + readChunkMaxTokens: num(offloadGroup, "readChunkMaxTokens") ?? 1600, + epochTriggerRatio: num(offloadGroup, "epochTriggerRatio") ?? 0.8, + epochTargetRatio: num(offloadGroup, "epochTargetRatio") ?? 0.55, + epochMinimumTurns: num(offloadGroup, "epochMinimumTurns") ?? 4, + epochMinimumIntervalTurns: num(offloadGroup, "epochMinimumIntervalTurns") ?? 2, backendUrl: optStr(offloadGroup, "backendUrl"), backendApiKey: optStr(offloadGroup, "backendApiKey"), backendTimeoutMs: num(offloadGroup, "backendTimeoutMs") ?? 120000, @@ -529,12 +569,20 @@ export function parseConfig(raw: Record | undefined): MemoryTda }, recall: { enabled: bool(recallGroup, "enabled") ?? true, + mode: validateRecallMode(str(recallGroup, "mode")) ?? "tool-only", + showInjected: bool(recallGroup, "showInjected") ?? false, + persistInjected: bool(recallGroup, "persistInjected") ?? false, maxResults: num(recallGroup, "maxResults") ?? 5, maxCharsPerMemory: num(recallGroup, "maxCharsPerMemory") ?? 0, maxTotalRecallChars: num(recallGroup, "maxTotalRecallChars") ?? 0, scoreThreshold: num(recallGroup, "scoreThreshold") ?? 0.3, strategy: validateStrategy(str(recallGroup, "strategy")) ?? "hybrid", timeoutMs: num(recallGroup, "timeoutMs") ?? 5000, + sessionSnapshotMaxTokens: num(recallGroup, "sessionSnapshotMaxTokens") ?? 1600, + stableMemoryMaxItems: num(recallGroup, "stableMemoryMaxItems") ?? 8, + sceneSummaryMaxItems: num(recallGroup, "sceneSummaryMaxItems") ?? 6, + dynamicRecallMaxTokens: num(recallGroup, "dynamicRecallMaxTokens") ?? 900, + maxSearchCallsPerTurn: num(recallGroup, "maxSearchCallsPerTurn") ?? 3, }, embedding: { enabled: embeddingEnabled, @@ -634,6 +682,7 @@ function strArray(src: Record, key: string): string[] | undefin } const VALID_STRATEGIES: RecallConfig["strategy"][] = ["embedding", "keyword", "hybrid"]; +const VALID_RECALL_MODES: RecallConfig["mode"][] = ["auto", "tool-only"]; /** * Validate recall strategy against whitelist. @@ -646,6 +695,13 @@ function validateStrategy(value: string | undefined): RecallConfig["strategy"] | : undefined; } +function validateRecallMode(value: string | undefined): RecallConfig["mode"] | undefined { + if (!value) return undefined; + return VALID_RECALL_MODES.includes(value as RecallConfig["mode"]) + ? (value as RecallConfig["mode"]) + : undefined; +} + /** * Normalize a cleanup time string. * diff --git a/src/core/hooks/auto-recall.tool-only.test.ts b/src/core/hooks/auto-recall.tool-only.test.ts new file mode 100644 index 00000000..5b6c706b --- /dev/null +++ b/src/core/hooks/auto-recall.tool-only.test.ts @@ -0,0 +1,42 @@ +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { parseConfig } from "../../config.js"; +import { performAutoRecall } from "./auto-recall.js"; + +const tempDirs: string[] = []; + +afterEach(async () => { + for (const dir of tempDirs.splice(0)) { + await rm(dir, { recursive: true, force: true }); + } +}); + +describe("performAutoRecall tool-only mode", () => { + it("does not search or prepend dynamic memories by default", async () => { + const pluginDataDir = await mkdtemp(join(tmpdir(), "tdai-recall-")); + tempDirs.push(pluginDataDir); + const cfg = parseConfig({}); + const searchL1Fts = vi.fn(async () => []); + const vectorStore = { + isFtsAvailable: () => true, + searchL1Fts, + getCapabilities: () => ({ nativeHybridSearch: false }), + } as any; + + const result = await performAutoRecall({ + userText: "帮我回忆项目规划", + actorId: "user", + sessionKey: "agent:test:session", + cfg, + pluginDataDir, + vectorStore, + }); + + expect(searchL1Fts).not.toHaveBeenCalled(); + expect(result?.prependContext).toBeUndefined(); + expect(result?.appendSystemContext).toContain(" +function buildMemoryToolsGuide(maxSearchCallsPerTurn: number): string { + const maxCalls = Number.isFinite(maxSearchCallsPerTurn) && maxSearchCallsPerTurn > 0 + ? Math.floor(maxSearchCallsPerTurn) + : 3; + return ` ## 记忆工具调用指南 当上方注入的记忆片段不足以回答用户问题时,可主动调用以下工具获取更多信息: @@ -42,10 +47,11 @@ const MEMORY_TOOLS_GUIDE = ` - **read_file**(Scene Navigation 中的路径):当已定位到相关情境,且需要该场景的完整画像、事件经过或阶段结论时使用。 ### ⚠️ 调用次数限制 -每轮对话中,tdai_memory_search 和 tdai_conversation_search **合计最多调用 3 次**。 -- 首次搜索无结果时,可换关键词或换工具重试,但总调用次数不要超过 3 次。 -- 若 3 次搜索后仍无结果,说明该信息不在记忆中,请直接根据已有信息回复用户,不要继续搜索。 -` +每轮对话中,tdai_memory_search 和 tdai_conversation_search **合计最多调用 ${maxCalls} 次**。 +- 首次搜索无结果时,可换关键词或换工具重试,但总调用次数不要超过 ${maxCalls} 次。 +- 若 ${maxCalls} 次搜索后仍无结果,说明该信息不在记忆中,请直接根据已有信息回复用户,不要继续搜索。 +`; +} /** A single recalled L1 memory with its search score and type. */ export interface RecalledMemory { @@ -112,7 +118,9 @@ async function performAutoRecallInner(params: { const { userText, cfg, pluginDataDir, logger, vectorStore, embeddingService } = params; const tRecallStart = performance.now(); - // Search relevant memories (L1 layer) — skip only when userText is empty/undefined + // Search relevant memories (L1 layer) only in explicit auto mode. The + // zero-config D-plan path is tool-only: stable context is appended to the + // system prompt and dynamic recall is retrieved through tools on demand. const tSearchStart = performance.now(); let memoryLines: string[] = []; let effectiveStrategy = "skipped"; @@ -120,12 +128,17 @@ async function performAutoRecallInner(params: { let searchTiming: SearchTiming = { ftsMs: 0, embeddingMs: 0, ftsHits: 0, embeddingHits: 0 }; if (!userText || userText.length === 0) { logger?.debug?.(`${TAG} User text empty/undefined, skipping memory search (persona/scene still injected)`); + } else if (cfg.recall.mode !== "auto") { + logger?.debug?.(`${TAG} recall.mode=${cfg.recall.mode}, skipping automatic L1 search; use memory tools on demand`); } else { effectiveStrategy = cfg.recall.strategy ?? "hybrid"; const searchResult = await searchMemories(userText, pluginDataDir, cfg, logger, effectiveStrategy as "keyword" | "embedding" | "hybrid", vectorStore, embeddingService); memoryLines = searchResult.lines; searchTiming = searchResult.timing; - memoryLines = applyRecallBudget(memoryLines, cfg.recall, logger); + memoryLines = applyRecallBudget(memoryLines, { + ...cfg.recall, + maxTotalRecallChars: cfg.recall.maxTotalRecallChars || cfg.recall.dynamicRecallMaxTokens * 4, + }, logger); // Extract structured RecalledMemory from formatted lines for metric reporting recalledL1Memories = memoryLines.map((line) => { @@ -161,28 +174,16 @@ async function performAutoRecallInner(params: { try { const sceneIndex = await readSceneIndex(pluginDataDir); if (sceneIndex.length > 0) { - sceneNavigation = generateSceneNavigation(sceneIndex, pluginDataDir); - logger?.debug?.(`${TAG} Scene navigation generated: ${sceneIndex.length} scenes`); + const sceneLimit = cfg.recall.sceneSummaryMaxItems > 0 ? cfg.recall.sceneSummaryMaxItems : sceneIndex.length; + const boundedSceneIndex = sceneIndex.slice(0, sceneLimit); + sceneNavigation = generateSceneNavigation(boundedSceneIndex, pluginDataDir); + logger?.debug?.(`${TAG} Scene navigation generated: ${boundedSceneIndex.length}/${sceneIndex.length} scenes`); } } catch { logger?.debug?.(`${TAG} No scene index found`); } const tSceneEnd = performance.now(); - if (memoryLines.length === 0 && !personaContent && !sceneNavigation) { - const totalMs = performance.now() - tRecallStart; - logger?.info( - `${TAG} ⏱ Recall timing: total=${totalMs.toFixed(0)}ms, ` + - `search=${(tSearchEnd - tSearchStart).toFixed(0)}ms(strategy=${effectiveStrategy},hits=${memoryLines.length},` + - `fts=${searchTiming.ftsMs.toFixed(0)}ms/${searchTiming.ftsHits}hits,` + - `vec=${searchTiming.embeddingMs.toFixed(0)}ms/${searchTiming.embeddingHits}hits), ` + - `persona=${(tPersonaEnd - tPersonaStart).toFixed(0)}ms, ` + - `scene=${(tSceneEnd - tSceneStart).toFixed(0)}ms — no context to inject`, - ); - logger?.debug?.(`${TAG} No memories/persona/scenes to inject`); - return undefined; - } - // Split recall context into stable and dynamic parts to optimize prompt caching. // // appendSystemContext (system prompt end — stable, cacheable): @@ -194,16 +195,16 @@ async function performAutoRecallInner(params: { // L1 relevant memories — different every turn, moved out of system prompt // so it doesn't bust the system prompt cache. const stableParts: string[] = []; - if (personaContent) { - stableParts.push(`\n${personaContent}\n`); - } - if (sceneNavigation) { - stableParts.push(`\n${sceneNavigation}\n`); - } + const sessionSnapshot = buildSessionSnapshot({ + persona: personaContent, + sceneNavigation, + maxTokens: cfg.recall.sessionSnapshotMaxTokens, + }); + stableParts.push(sessionSnapshot.text); // Dynamic part: L1 relevant memories (changes every turn) → prependContext (user prompt) let prependContext: string | undefined; - if (memoryLines.length > 0) { + if (cfg.recall.mode === "auto" && cfg.recall.showInjected && memoryLines.length > 0) { prependContext = `\n以下是当前对话召回的相关记忆,不代表当前任务进程,仅作为参考:\n\n${memoryLines.join(RECALL_LINE_SEPARATOR)}\n`; } @@ -211,9 +212,7 @@ async function performAutoRecallInner(params: { // Append memory tools usage guide to the stable part so the agent knows // how to actively retrieve deeper context when the injected snippets // are not enough. This is static content and benefits from caching. - if (stableParts.length > 0 || prependContext) { - stableParts.push(MEMORY_TOOLS_GUIDE); - } + stableParts.push(buildMemoryToolsGuide(cfg.recall.maxSearchCallsPerTurn)); const appendSystemContext = stableParts.length > 0 ? stableParts.join("\n\n") : undefined; @@ -630,7 +629,10 @@ async function searchHybrid( // Sort by combined RRF score and take top results const sorted = [...mergedMap.entries()] - .sort((a, b) => b[1].rrfScore - a[1].rrfScore) + .sort((a, b) => { + const scoreDiff = b[1].rrfScore - a[1].rrfScore; + return scoreDiff !== 0 ? scoreDiff : a[0].localeCompare(b[0]); + }) .slice(0, maxResults); if (sorted.length > 0) { diff --git a/src/core/session/session-snapshot.test.ts b/src/core/session/session-snapshot.test.ts new file mode 100644 index 00000000..973644e7 --- /dev/null +++ b/src/core/session/session-snapshot.test.ts @@ -0,0 +1,28 @@ +import { describe, expect, it } from "vitest"; +import { buildSessionSnapshot } from "./session-snapshot.js"; + +describe("buildSessionSnapshot", () => { + it("is deterministic and excludes dynamic or volatile fields", () => { + const input = { + persona: " 用户偏好中文回答。\n", + sceneNavigation: "- 项目A\n- 项目B", + stableMemories: [ + { id: "m2", content: "第二条", type: "instruction" }, + { id: "m1", content: "第一条", type: "persona", score: 0.98, createdAt: "2026-01-01T00:00:00Z" }, + ], + maxTokens: 200, + }; + + const first = buildSessionSnapshot(input); + const second = buildSessionSnapshot({ ...input, now: "2026-07-06T00:00:00Z" }); + + expect(second).toEqual(first); + expect(first.text).toContain("\n${boundedBody}\n`; + return { text, hash, estimatedTokens: estimateTokens(text) }; +} + +function normalizeSnapshotInput(input: SessionSnapshotInput): Required> { + const stableMemories = [...(input.stableMemories ?? [])] + .filter((m) => m.id && m.content) + .map((m) => ({ + id: String(m.id), + type: m.type ? String(m.type) : "memory", + content: normalizeText(m.content), + sceneName: m.sceneName ? normalizeText(m.sceneName) : undefined, + })) + .sort((a, b) => a.id.localeCompare(b.id)); + + return { + persona: normalizeText(input.persona ?? ""), + sceneNavigation: normalizeText(input.sceneNavigation ?? ""), + stableMemories, + }; +} + +function renderSnapshotBody(input: Required>): string { + const parts: string[] = []; + if (input.persona) { + parts.push(["## persona", input.persona].join("\n")); + } + if (input.sceneNavigation) { + parts.push(["## scene_navigation", input.sceneNavigation].join("\n")); + } + if (input.stableMemories.length > 0) { + parts.push([ + "## stable_memories", + ...input.stableMemories.map((m) => { + const scene = m.sceneName ? ` scene="${escapeAttr(m.sceneName)}"` : ""; + return `- id="${escapeAttr(m.id)}" type="${escapeAttr(m.type ?? "memory")}"${scene}: ${m.content}`; + }), + ].join("\n")); + } + if (parts.length === 0) { + parts.push("## stable_context\n(empty)"); + } + return parts.join("\n\n"); +} + +export function estimateTokens(text: string): number { + let cjk = 0; + for (let i = 0; i < text.length; i++) { + const code = text.charCodeAt(i); + if ((code >= 0x4e00 && code <= 0x9fff) || (code >= 0x3400 && code <= 0x4dbf) || (code >= 0xf900 && code <= 0xfaff)) cjk++; + } + return Math.ceil(cjk * 1.5 + (text.length - cjk) / 4); +} + +function truncateToEstimatedTokens(text: string, maxTokens: number): string { + if (estimateTokens(text) <= maxTokens) return text; + const chars = Array.from(text); + let lo = 0; + let hi = chars.length; + while (lo < hi) { + const mid = Math.ceil((lo + hi) / 2); + if (estimateTokens(chars.slice(0, mid).join("")) <= maxTokens) lo = mid; + else hi = mid - 1; + } + return `${chars.slice(0, Math.max(0, lo - 20)).join("").trimEnd()}\n[truncated]`; +} + +function normalizeText(text: string): string { + return text.replace(/\r\n/g, "\n").replace(/[ \t]+$/gm, "").trim(); +} + +function escapeAttr(text: string): string { + return text.replace(/&/g, "&").replace(/"/g, """).replace(/ 0 ? Math.floor(value) : undefined; +} diff --git a/src/offload/epoch-manager.test.ts b/src/offload/epoch-manager.test.ts new file mode 100644 index 00000000..5597bb46 --- /dev/null +++ b/src/offload/epoch-manager.test.ts @@ -0,0 +1,30 @@ +import { describe, expect, it } from "vitest"; +import { createInitialCacheEpoch, shouldTransitionCacheEpoch, transitionCacheEpoch } from "./epoch-manager.js"; + +describe("cache epoch manager", () => { + it("keeps append-only epochs below threshold", () => { + const epoch = createInitialCacheEpoch({ snapshotHash: "h1", startedAtTurn: 1 }); + + expect(shouldTransitionCacheEpoch(epoch, { + totalTokens: 100, + contextWindow: 1_000, + currentTurn: 3, + triggerRatio: 0.8, + minimumTurns: 4, + })).toBe(false); + }); + + it("creates a new frozen epoch when threshold and minimum turns are reached", () => { + const epoch = createInitialCacheEpoch({ snapshotHash: "h1", startedAtTurn: 1 }); + const next = transitionCacheEpoch(epoch, { + snapshotHash: "h2", + compactedRange: { startTurn: 1, endTurn: 8 }, + currentTurn: 9, + }); + + expect(next.id).toBe(epoch.id + 1); + expect(next.previousSnapshotHash).toBe("h1"); + expect(next.snapshotHash).toBe("h2"); + expect(next.compactedRanges).toEqual([{ startTurn: 1, endTurn: 8 }]); + }); +}); diff --git a/src/offload/epoch-manager.ts b/src/offload/epoch-manager.ts new file mode 100644 index 00000000..44986113 --- /dev/null +++ b/src/offload/epoch-manager.ts @@ -0,0 +1,49 @@ +export interface CacheEpoch { + id: number; + snapshotHash: string; + previousSnapshotHash?: string; + startedAtTurn: number; + compactedRanges: Array<{ startTurn: number; endTurn: number }>; +} + +export function createInitialCacheEpoch(input: { snapshotHash: string; startedAtTurn?: number }): CacheEpoch { + return { + id: 1, + snapshotHash: input.snapshotHash, + startedAtTurn: input.startedAtTurn ?? 1, + compactedRanges: [], + }; +} + +export function shouldTransitionCacheEpoch( + epoch: CacheEpoch, + input: { + totalTokens: number; + contextWindow: number; + currentTurn: number; + triggerRatio: number; + minimumTurns: number; + }, +): boolean { + if (input.contextWindow <= 0) return false; + const turnsInEpoch = input.currentTurn - epoch.startedAtTurn + 1; + if (turnsInEpoch < input.minimumTurns) return false; + return input.totalTokens / input.contextWindow >= input.triggerRatio; +} + +export function transitionCacheEpoch( + epoch: CacheEpoch, + input: { + snapshotHash: string; + compactedRange: { startTurn: number; endTurn: number }; + currentTurn: number; + }, +): CacheEpoch { + return { + id: epoch.id + 1, + previousSnapshotHash: epoch.snapshotHash, + snapshotHash: input.snapshotHash, + startedAtTurn: input.currentTurn, + compactedRanges: [...epoch.compactedRanges, input.compactedRange], + }; +} diff --git a/src/offload/hooks/after-tool-call.ts b/src/offload/hooks/after-tool-call.ts index 64e38395..4804906c 100644 --- a/src/offload/hooks/after-tool-call.ts +++ b/src/offload/hooks/after-tool-call.ts @@ -7,7 +7,8 @@ import { nowChinaISO } from "../time-utils.js"; import { buildTiktokenContextSnapshot, type ContextSnapshot } from "../context-token-tracker.js"; import { traceOffloadDecision, traceMessagesSnapshot } from "../opik-tracer.js"; import { PLUGIN_DEFAULTS } from "../types.js"; -import { readOffloadEntries, markOffloadStatus, readMmd } from "../storage.js"; +import { readOffloadEntries, markOffloadStatus, writeRefMd } from "../storage.js"; +import { normalizeToolResultForPrompt } from "../tool-result-normalizer.js"; import { createL3TokenCounter } from "../l3-token-counter.js"; import { normalizeToolCallIdForLookup, @@ -28,7 +29,7 @@ import { isTokenOverflowError, dumpMessagesSnapshot, } from "./llm-input-l3.js"; -import { MMD_MESSAGE_MARKER, findActiveMmdInsertionPoint, findHistoryMmdInsertionPoint } from "../mmd-injector.js"; +import { maybeUpdateMmdInMessages, findHistoryMmdInsertionPoint } from "../mmd-injector.js"; import type { OffloadStateManager } from "../state-manager.js"; import type { PluginConfig, PluginLogger, ToolPair } from "../types.js"; import type { BackendClient } from "../backend-client.js"; @@ -173,13 +174,55 @@ export function createAfterToolCallHandler( return; } + const rawResult = event.result; + const timestamp = nowChinaISO(); + let promptResult = rawResult; + let frontOffload: { resultRef?: string; contentHash?: string; originalTokens?: number } = {}; + try { + const normalized = await normalizeToolResultForPrompt({ + toolName: event.toolName, + toolCallId, + timestamp, + result: rawResult, + maxTokens: pluginConfig?.inlineToolResultMaxTokens ?? PLUGIN_DEFAULTS.inlineToolResultMaxTokens, + summaryMaxTokens: pluginConfig?.summaryMaxTokens ?? PLUGIN_DEFAULTS.summaryMaxTokens, + previewMaxChars: pluginConfig?.previewMaxChars ?? PLUGIN_DEFAULTS.previewMaxChars, + writeRef: async (content) => writeRefMd( + stateManager.ctx, + timestamp, + event.toolName, + content, + `${toolCallId}-${normalizedHashPrefix(content)}`, + ), + }); + if (normalized.offloaded) { + promptResult = normalized.promptResult; + frontOffload = { + resultRef: normalized.resultRef, + contentHash: normalized.contentHash, + originalTokens: normalized.originalTokens, + }; + event.result = promptResult; + replacePromptFacingToolResult(event.messages, toolCallId, promptResult); + logger.debug?.( + `[context-offload] front-offload: ${event.toolName} (${toolCallId}) ` + + `tokens=${normalized.originalTokens}, ref=${normalized.resultRef}`, + ); + } + } catch (err) { + logger.warn?.(`[context-offload] front-offload failed for ${event.toolName} (${toolCallId}): ${String(err)}`); + } + const pair: ToolPair = { toolName: event.toolName, toolCallId, params: resolvedParams, - result: event.result, + result: rawResult, + result_ref: frontOffload.resultRef, + content_hash: frontOffload.contentHash, + original_tokens: frontOffload.originalTokens, error: event.error, - timestamp: nowChinaISO(), + timestamp, durationMs: event.durationMs, }; stateManager.addToolPair(pair); @@ -191,61 +234,10 @@ export function createAfterToolCallHandler( if (turn) stateManager.cachedLatestTurnMessages = turn; } - // In-loop active MMD injection / update. - // Only inject after L1.5 has settled (task boundary determined, activeMmdFile set). - // This also picks up L2 MMD content updates (L2 runs async and may patch the MMD - // file between tool calls). if (event.messages && Array.isArray(event.messages)) { try { - const l15Settled = stateManager.l15Settled; - const activeMmdFile = stateManager.getActiveMmdFile(); - if (!l15Settled) { - logger.debug?.(`[context-offload] after_tool_call MMD: SKIP (L1.5 not settled yet)`); - } else if (!activeMmdFile) { - logger.debug?.(`[context-offload] after_tool_call MMD: SKIP (no active MMD file)`); - } else { - const mmdContent = await readMmd(stateManager.ctx, activeMmdFile); - if (mmdContent) { - let taskGoal = ""; - const metaMatch = mmdContent.match(/^%%\{\s*(.*?)\s*\}%%/); - if (metaMatch) { - try { const meta = JSON.parse(`{${metaMatch[1]}}`); taskGoal = meta.taskGoal || ""; } catch { /* */ } - } - const mmdText = [ - ``, - `【当前活跃任务的mermaid流程图】这是你最近正在执行的任务的阶段性记录(此条下方的tool use未被汇总,进程可能有延迟,仅供参考)。`, - taskGoal ? `**任务目标:** ${taskGoal}` : "", - `**任务文件:** ${activeMmdFile}`, - "```mermaid", mmdContent, "```", - `标记为 "doing" 的节点是近期焦点(注:可能有延迟,下方的tool use未被统计,仅供参考),"done" 的已完成。请参考此保持方向感,避免重复已完成的工作。`, - ``, - ].filter((line) => line !== "").join("\n"); - - const existingIdx = event.messages.findIndex((m: any) => m._mmdContextMessage === "active"); - const newMsg = { role: "user", content: [{ type: "text", text: mmdText }], _mmdContextMessage: "active" }; - if (existingIdx >= 0) { - // Check if content changed (L1.5 switched file or L2 updated content) - const oldContent = Array.isArray(event.messages[existingIdx].content) - ? event.messages[existingIdx].content.map((c: any) => c.text ?? "").join("") - : (event.messages[existingIdx].content ?? ""); - const contentChanged = !oldContent.includes(activeMmdFile) || oldContent !== mmdText; - if (contentChanged) { - event.messages[existingIdx] = newMsg; - logger.debug?.(`[context-offload] after_tool_call MMD: UPDATED at [${existingIdx}], file=${activeMmdFile}, contentChanged=true`); - _dumpMessagesAfterMmd(event.messages, "UPDATED", logger); - } else { - logger.debug?.(`[context-offload] after_tool_call MMD: unchanged, skip update`); - } - } else { - const insertIdx = findActiveMmdInsertionPoint(event.messages); - event.messages.splice(insertIdx, 0, newMsg); - logger.debug?.(`[context-offload] after_tool_call MMD: INJECTED at [${insertIdx}], file=${activeMmdFile}, msgs=${event.messages.length}`); - _dumpMessagesAfterMmd(event.messages, "INJECTED", logger); - } - } else { - logger.debug?.(`[context-offload] after_tool_call MMD: file=${activeMmdFile} content is null`); - } - } + const changed = await maybeUpdateMmdInMessages(event.messages, stateManager, logger, getContextWindow, pluginConfig); + if (changed) _dumpMessagesAfterMmd(event.messages, "APPENDED", logger); } catch (err) { logger.warn(`[context-offload] after_tool_call MMD error: ${err}`); } @@ -586,6 +578,75 @@ function _extractText(msg: any): string { return ""; } +function normalizedHashPrefix(content: string): string { + let hash = 0; + for (let i = 0; i < content.length; i++) { + hash = ((hash << 5) - hash + content.charCodeAt(i)) | 0; + } + return Math.abs(hash).toString(16); +} + +function replacePromptFacingToolResult(messages: any, toolCallId: string, promptResult: unknown): boolean { + if (!Array.isArray(messages)) return false; + const target = normalizeToolCallIdForLookup(toolCallId); + for (const msg of messages) { + const id = extractPromptToolResultId(msg); + if (id && normalizeToolCallIdForLookup(id) === target) { + replaceMessageContent(msg, promptResult); + stripLargePromptFields(msg); + return true; + } + } + return false; +} + +function replaceMessageContent(msg: any, promptResult: unknown): void { + const text = typeof promptResult === "string" ? promptResult : JSON.stringify(promptResult, null, 2); + if (msg.type === "message" && msg.message) { + msg.message.content = [{ type: "text", text }]; + } else { + msg.content = [{ type: "text", text }]; + } +} + +function extractPromptToolResultId(msg: any): string | null { + const direct = extractToolCallId(msg) + ?? msg?.id + ?? msg?.tool_use_id + ?? msg?.message?.id + ?? msg?.message?.tool_use_id; + if (typeof direct === "string" && direct) return direct; + + const content = msg?.content ?? msg?.message?.content; + if (Array.isArray(content)) { + for (const block of content) { + const id = block?.toolCallId ?? block?.tool_call_id ?? block?.tool_use_id ?? block?.id; + if (typeof id === "string" && id) return id; + } + } + return null; +} + +function stripLargePromptFields(msg: any): void { + const preserve = new Set([ + "role", "type", "name", "id", "toolCallId", "tool_call_id", "tool_use_id", + "content", "message", "status", "_offloaded", "_mmdContextMessage", + "_mmdInjection", "_contextOffloadProcessed", "_cachedTokens", "_tokenCount", + ]); + const stripObj = (obj: any) => { + if (!obj || typeof obj !== "object") return; + for (const key of Object.keys(obj)) { + if (preserve.has(key)) continue; + const value = obj[key]; + if (value == null) continue; + const serialized = typeof value === "string" ? value : JSON.stringify(value); + if (serialized && serialized.length > 500) delete obj[key]; + } + }; + stripObj(msg); + if (msg?.message && typeof msg.message === "object") stripObj(msg.message); +} + /** Dump all messages after MMD injection for diagnostics (debug-level only). */ function _dumpMessagesAfterMmd(messages: any[], action: string, logger: PluginLogger): void { const mmdCount = messages.filter((m: any) => m._mmdContextMessage || m._mmdInjection).length; diff --git a/src/offload/index.test.ts b/src/offload/index.test.ts new file mode 100644 index 00000000..87222f74 --- /dev/null +++ b/src/offload/index.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, it } from "vitest"; +import { sliceRefContent } from "./index.js"; + +describe("sliceRefContent", () => { + it("returns the query hit from an oversized single-line ref payload", () => { + const before = Array.from({ length: 400 }, (_, i) => `TDAI_OFFLOAD_SENTINEL ${i} ${"X".repeat(80)}`).join("\\n"); + const hit = `TDAI_OFFLOAD_SENTINEL 400 ${"X".repeat(80)}`; + const after = Array.from({ length: 49 }, (_, i) => `TDAI_OFFLOAD_SENTINEL ${401 + i} ${"X".repeat(80)}`).join("\\n"); + const raw = `# Tool Result\n\n~~~json\n{"aggregated":"${before}\\n${hit}\\n${after}"}\n~~~`; + + const sliced = sliceRefContent(raw, { + query: "TDAI_OFFLOAD_SENTINEL 400", + maxTokens: 1200, + }); + + expect(sliced).toContain("TDAI_OFFLOAD_SENTINEL 400"); + expect(sliced).not.toBe("[truncated: max_tokens=1200]"); + }); +}); diff --git a/src/offload/index.ts b/src/offload/index.ts index 35c18bb4..4fbd0f20 100644 --- a/src/offload/index.ts +++ b/src/offload/index.ts @@ -7,6 +7,7 @@ * This module is the merged equivalent of the standalone context-offload-plugin's index.js, * adapted to co-exist with the memory-tencentdb plugin. */ +import { join } from "node:path"; import { OffloadStateManager } from "./state-manager.js"; import { createAfterToolCallHandler } from "./hooks/after-tool-call.js"; import { createBeforePromptBuildHandler } from "./hooks/before-prompt-build.js"; @@ -20,6 +21,8 @@ import { readOffloadEntries, markOffloadStatus, DEFAULT_DATA_ROOT, + readRefMd, + readRefMdFromDataDir, } from "./storage.js"; import { buildTiktokenContextSnapshot, configureTokenTracker, tiktokenCount, jsonReplacer } from "./context-token-tracker.js"; import { fastEstimateMessages } from "./fast-token-estimate.js"; @@ -96,6 +99,74 @@ let _sharedSessions: SessionRegistry | null = null; // ─── Helpers ───────────────────────────────────────────────────────────────── +function isSafeRefPath(refPath: string): boolean { + return /^refs\/[^/\\]+\.md$/u.test(refPath) && + !refPath.includes("..") && + !refPath.includes(":") && + !refPath.startsWith("/") && + !refPath.startsWith("\\"); +} + +export function sliceRefContent( + raw: string, + options: { query?: string; startLine?: number; endLine?: number; maxTokens: number }, +): string { + if (options.query && options.query.trim()) { + const q = options.query.trim().toLowerCase(); + const rawLower = raw.toLowerCase(); + const matchIndex = rawLower.indexOf(q); + if (matchIndex < 0) return "(empty)"; + + const maxTokens = Math.max(80, Math.floor(options.maxTokens)); + const budgetChars = Math.max(400, maxTokens * 4); + let startOffset = Math.max(0, matchIndex - Math.floor(budgetChars / 2)); + let endOffset = Math.min(raw.length, startOffset + budgetChars); + if (endOffset - startOffset < budgetChars) { + startOffset = Math.max(0, endOffset - budgetChars); + } + let excerpt = raw.slice(startOffset, endOffset); + + while (tiktokenCount(excerpt) > maxTokens && excerpt.length > q.length) { + const excess = Math.max(100, Math.ceil(excerpt.length * 0.1)); + const matchInExcerpt = Math.max(0, matchIndex - startOffset); + const trimStart = matchInExcerpt > excerpt.length / 2; + if (trimStart) { + startOffset = Math.min(matchIndex, startOffset + excess); + } else { + endOffset = Math.max(matchIndex + options.query.trim().length, endOffset - excess); + } + excerpt = raw.slice(startOffset, endOffset); + } + + const prefix = startOffset > 0 ? "[truncated before]\n" : ""; + const suffix = endOffset < raw.length ? "\n[truncated after]" : ""; + return `${prefix}${excerpt}${suffix}`; + } + + const lines = raw.split(/\r?\n/u).map((line, idx) => ({ lineNo: idx + 1, text: line })); + const start = Math.max(1, Math.floor(options.startLine ?? 1)); + const end = Math.max(start, Math.floor(options.endLine ?? lines.length)); + const selectedLines = lines.filter((line) => line.lineNo >= start && line.lineNo <= end); + const maxTokens = Math.max(80, Math.floor(options.maxTokens)); + const output: string[] = []; + for (const item of selectedLines) { + const next = `${item.lineNo}: ${item.text}`; + const candidate = [...output, next].join("\n"); + if (tiktokenCount(candidate) > maxTokens) { + output.push(`[truncated: max_tokens=${maxTokens}]`); + break; + } + output.push(next); + } + return output.join("\n") || "(empty)"; +} + +function resolveReadMaxTokens(requested: number | undefined, configuredLimit: number): number { + const hardLimit = Math.max(80, Math.floor(configuredLimit)); + if (requested == null || !Number.isFinite(requested)) return hardLimit; + return Math.min(hardLimit, Math.max(80, Math.floor(requested))); +} + function parseCreateSkillCommand( prompt: string, ): { mmdName: string | null; skillFocus: string | null } | null { @@ -298,6 +369,14 @@ export function registerOffload(api: any, offloadConfig: OffloadConfig): void { mildOffloadRatio: offloadConfig.mildOffloadRatio, aggressiveCompressRatio: offloadConfig.aggressiveCompressRatio, mmdMaxTokenRatio: offloadConfig.mmdMaxTokenRatio, + inlineToolResultMaxTokens: offloadConfig.inlineToolResultMaxTokens, + summaryMaxTokens: offloadConfig.summaryMaxTokens, + previewMaxChars: offloadConfig.previewMaxChars, + readChunkMaxTokens: offloadConfig.readChunkMaxTokens, + epochTriggerRatio: offloadConfig.epochTriggerRatio, + epochTargetRatio: offloadConfig.epochTargetRatio, + epochMinimumTurns: offloadConfig.epochMinimumTurns, + epochMinimumIntervalTurns: offloadConfig.epochMinimumIntervalTurns, }; // Fix 4: Configure token tracker encoding to match plugin config (default: o200k_base) @@ -439,6 +518,10 @@ export function registerOffload(api: any, offloadConfig: OffloadConfig): void { const refByToolCallId = new Map(); for (const p of pairs) { try { + if (p.result_ref) { + refByToolCallId.set(p.toolCallId, p.result_ref); + continue; + } const resultStr = typeof p.result === "string" ? sanitizeText(p.result) : sanitizeText(JSON.stringify(p.result, null, 2)); @@ -887,6 +970,50 @@ export function registerOffload(api: any, offloadConfig: OffloadConfig): void { return entry.manager; }; + if (typeof api.registerTool === "function") { + api.registerTool( + { + name: "tdai_offload_read", + label: "Offload Read", + description: + "Read full tool results previously replaced by context offload. Use result_ref from an offloaded tool result stub. " + + "Only refs/... paths are allowed. Prefer query/start_line/end_line/max_tokens to read the smallest useful chunk.", + parameters: { + type: "object", + properties: { + result_ref: { type: "string", description: "Relative ref path, must start with refs/ and end with .md" }, + query: { type: "string", description: "Optional substring filter. Returns matching lines with line numbers." }, + start_line: { type: "number", description: "Optional 1-based inclusive start line" }, + end_line: { type: "number", description: "Optional 1-based inclusive end line" }, + max_tokens: { type: "number", description: "Optional max token budget for returned content" }, + }, + required: ["result_ref"], + }, + async execute(_toolCallId: string, params: Record, ctx?: any) { + const refPath = String(params?.result_ref ?? ""); + if (!isSafeRefPath(refPath)) { + return "tdai_offload_read: invalid result_ref. Only refs/*.md relative paths are allowed."; + } + const mgr = ctx?.sessionKey ? await _resolveSession(ctx.sessionKey, ctx?.sessionId) : _lastActiveMgr; + const raw = mgr + ? await readRefMd(mgr.ctx, refPath) + : await readRefMdFromDataDir(join(dataRoot, "main"), refPath); + if (raw == null) return `tdai_offload_read: result_ref not found: ${refPath}`; + return sliceRefContent(raw, { + query: typeof params?.query === "string" ? params.query : undefined, + startLine: typeof params?.start_line === "number" ? params.start_line : undefined, + endLine: typeof params?.end_line === "number" ? params.end_line : undefined, + maxTokens: resolveReadMaxTokens( + typeof params?.max_tokens === "number" ? params.max_tokens : undefined, + pCfg.readChunkMaxTokens ?? PLUGIN_DEFAULTS.readChunkMaxTokens, + ), + }); + }, + }, + { name: "tdai_offload_read" }, + ); + } + // L2 Scheduler — uses module-level state (_l2Running, _l2PollHandle, _l2FirstNotifyAt) // Clean up any lingering poll timer from previous registerOffload() call if (_l2PollHandle !== null) { clearTimeout(_l2PollHandle); _l2PollHandle = null; } diff --git a/src/offload/l3-helpers.ts b/src/offload/l3-helpers.ts index d9d08f90..75222472 100644 --- a/src/offload/l3-helpers.ts +++ b/src/offload/l3-helpers.ts @@ -225,7 +225,8 @@ export function replaceWithSummary(msg: any, entry: OffloadEntry): { originalLen const summaryContent = [ `[Offloaded Tool Result | node: ${entry.node_id ?? "N/A"}]`, `Summary: ${entry.summary}`, - `result_ref: ${entry.result_ref} (read this file for full tool call and raw result)`, + `result_ref: ${entry.result_ref}`, + `instruction: call tdai_offload_read({ result_ref: "${entry.result_ref}" }) for full tool call and raw result`, ].join("\n"); // Measure original content length diff --git a/src/offload/mmd-injector.ts b/src/offload/mmd-injector.ts index dbf7b54c..bd010d02 100644 --- a/src/offload/mmd-injector.ts +++ b/src/offload/mmd-injector.ts @@ -9,12 +9,14 @@ * The marker property `_mmdContextMessage` is used to locate the message for * replacement. L3 compression must skip messages carrying this marker. */ -import { readMmd, listMmds } from "./storage.js"; +import { readMmd, listMmds, writeRefMd } from "./storage.js"; import { PLUGIN_DEFAULTS, type PluginConfig, type PluginLogger } from "./types.js"; import { createL3TokenCounter } from "./l3-token-counter.js"; import { traceOffloadDecision } from "./opik-tracer.js"; import { isToolResultMessage, isAssistantMessageWithToolUse } from "./l3-helpers.js"; import type { OffloadStateManager } from "./state-manager.js"; +import { buildTaskSnapshot, buildTaskDeltaMessage } from "./task-snapshot.js"; +import { nowChinaISO } from "./time-utils.js"; /** Marker property on the injected message object. */ export const MMD_MESSAGE_MARKER = "_mmdContextMessage"; @@ -53,9 +55,7 @@ export async function injectMmdIntoMessages( `[context-offload] mmd-injector inject: injectionReady=${injReady}, activeMmdFile=${actFile ?? "null"}, msgs=${messages.length}`, ); if (!injReady) { - removeMmdMessages(messages); - stateManager.lastMmdInjectedTokens = 0; - return { mmdTokens: 0 }; + return { mmdTokens: stateManager.lastMmdInjectedTokens }; } const contextWindow = @@ -70,7 +70,6 @@ export async function injectMmdIntoMessages( logger.debug?.( `[context-offload] mmd-injector inject: activeMmdText=${activeMmdText ? `${activeMmdText.length} chars` : "null"}, contextWindow=${contextWindow}`, ); - removeMmdMessages(messages); let totalMmdTokens = 0; @@ -80,8 +79,7 @@ export async function injectMmdIntoMessages( content: [{ type: "text", text: activeMmdText }], [MMD_MESSAGE_MARKER]: "active", }; - const insertIdx = findActiveMmdInsertionPoint(messages); - messages.splice(insertIdx, 0, activeMsg); + messages.push(activeMsg); totalMmdTokens += countTokens(activeMmdText); } @@ -89,7 +87,7 @@ export async function injectMmdIntoMessages( const activeMmd = stateManager.getActiveMmdFile(); logger.debug?.( - `[context-offload] mmd-injector: injected active MMD into messages (${totalMmdTokens} tokens, file=${activeMmd})`, + `[context-offload] mmd-injector: appended active MMD snapshot/delta (${totalMmdTokens} tokens, file=${activeMmd})`, ); // Summary after active MMD injection (was full dump, now aggregated) @@ -109,7 +107,7 @@ export async function injectMmdIntoMessages( mmdMaxTokenRatio, }, output: { - result: `MMD 注入 messages:${totalMmdTokens} tokens (active only)`, + result: `MMD append-only snapshot/delta:${totalMmdTokens} tokens (active only)`, mmdTokens: totalMmdTokens, hasActive: !!activeMmdText, hasHistory: false, @@ -325,10 +323,9 @@ async function buildActiveMmdBlock( try { const mmdContent = await readMmd(stateManager.ctx, activeMmdFile); if (!mmdContent) return null; - stateManager.setInjectedMmdVersion( - activeMmdFile, - computeFingerprint(mmdContent), - ); + const fingerprint = computeFingerprint(mmdContent); + const previousFingerprint = stateManager.getInjectedMmdVersion(activeMmdFile); + if (previousFingerprint === fingerprint) return null; const metaMatch = mmdContent.match(/^%%\{\s*(.*?)\s*\}%%/); let taskGoal = ""; if (metaMatch) { @@ -339,28 +336,30 @@ async function buildActiveMmdBlock( /* ignore */ } } - const nodePattern = /\b(\d+-N\d+|N\d+)\b/g; - const nodeIds: string[] = []; - let match: RegExpExecArray | null; - while ((match = nodePattern.exec(mmdContent)) !== null) { - if (!nodeIds.includes(match[1])) nodeIds.push(match[1]); - } - return [ - ``, - `【当前活跃任务的mermaid流程图】这是你最近正在执行的任务的阶段性记录(此条下方的tool use未被汇总,进程可能有延迟,仅供参考)。`, - taskGoal ? `**任务目标:** ${taskGoal}` : "", - `**任务文件:** ${activeMmdFile}`, - nodeIds.length > 0 - ? `**节点索引:** 可通过 node_id 在 offload.{sessionid}.jsonl 中查找对应的工具调用记录。如需查看某个节点对应的原始工具调用与完整结果,请在 offload.{sessionid}.jsonl 中找到对应条目的 result_ref 并读取该文件。` - : "", - "```mermaid", + const refPath = await writeRefMd( + stateManager.ctx, + nowChinaISO(), + "task-mermaid", mmdContent, - "```", - `标记为 "doing" 的节点是近期焦点(注:可能有延迟,下方的tool use未被统计,仅供参考),"done" 的已完成。请参考此保持方向感,避免重复已完成的工作。`, - ``, - ] - .filter((line) => line !== "") - .join("\n"); + `${activeMmdFile}-${fingerprint}`, + ); + stateManager.setInjectedMmdVersion(activeMmdFile, fingerprint); + + if (!previousFingerprint) { + return buildTaskSnapshot({ + taskGoal, + mmdFile: activeMmdFile, + mermaid: mmdContent, + resultRef: refPath, + }).text; + } + + return buildTaskDeltaMessage({ + taskGoal, + mmdFile: activeMmdFile, + changedNodeIds: extractNodeIds(mmdContent), + resultRef: refPath, + }).content[0].text; } catch (err) { logger.error( `[context-offload] mmd-injector: Error building active MMD block: ${err}`, @@ -372,3 +371,16 @@ async function buildActiveMmdBlock( function computeFingerprint(content: string): string { return `${content.length}:${content.slice(0, 64)}`; } + +function extractNodeIds(content: string): string[] { + const nodePattern = /\b(\d+-N\d+|N\d+|[A-Za-z]\w*)\b/g; + const nodeIds: string[] = []; + let match: RegExpExecArray | null; + while ((match = nodePattern.exec(content)) !== null) { + const id = match[1]; + if (!["flowchart", "graph", "subgraph", "end", "classDef"].includes(id) && !nodeIds.includes(id)) { + nodeIds.push(id); + } + } + return nodeIds.slice(0, 40); +} diff --git a/src/offload/storage.test.ts b/src/offload/storage.test.ts new file mode 100644 index 00000000..24d86f7e --- /dev/null +++ b/src/offload/storage.test.ts @@ -0,0 +1,25 @@ +import { mkdtemp, rm } from "node:fs/promises"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; +import { describe, expect, it } from "vitest"; +import { + createStorageContext, + ensureDirs, + readRefMdFromDataDir, + writeRefMd, +} from "./storage.js"; + +describe("readRefMdFromDataDir", () => { + it("reads a safe ref from an agent data directory without a session manager", async () => { + const dataRoot = await mkdtemp(join(tmpdir(), "tdai-offload-")); + try { + const ctx = createStorageContext(dataRoot, "main", "session-1"); + await ensureDirs(ctx); + const ref = await writeRefMd(ctx, "2026-07-06T00:00:00.000Z", "exec", "TDAI_OFFLOAD_SENTINEL 400"); + + await expect(readRefMdFromDataDir(ctx.dataDir, ref)).resolves.toContain("TDAI_OFFLOAD_SENTINEL 400"); + } finally { + await rm(dataRoot, { recursive: true, force: true }); + } + }); +}); diff --git a/src/offload/storage.ts b/src/offload/storage.ts index 222a66b7..0619b1b5 100644 --- a/src/offload/storage.ts +++ b/src/offload/storage.ts @@ -534,8 +534,10 @@ export async function writeRefMd( timestamp: string, toolName: string, content: string, + suffix?: string, ): Promise { - const filename = `${isoToFilename(timestamp)}.md`; + const safeSuffix = suffix ? `-${sanitizePath(suffix).slice(0, 80)}` : ""; + const filename = `${isoToFilename(timestamp)}${safeSuffix}.md`; const filePath = join(ctx.refsDir, filename); const safeContent = (content ?? "").replace(UNSAFE_CHAR_RE, ""); const header = `# Tool Result: ${toolName}\n\n**Timestamp:** ${timestamp}\n\n---\n\n`; @@ -548,7 +550,15 @@ export async function readRefMd( ctx: StorageContext, refPath: string, ): Promise { - const filePath = join(ctx.dataDir, refPath); + return readRefMdFromDataDir(ctx.dataDir, refPath); +} + +/** Read a ref MD file from an agent data directory by relative path */ +export async function readRefMdFromDataDir( + dataDir: string, + refPath: string, +): Promise { + const filePath = join(dataDir, refPath); if (!existsSync(filePath)) return null; return readFile(filePath, "utf-8"); } diff --git a/src/offload/task-snapshot.test.ts b/src/offload/task-snapshot.test.ts new file mode 100644 index 00000000..00c366a0 --- /dev/null +++ b/src/offload/task-snapshot.test.ts @@ -0,0 +1,31 @@ +import { describe, expect, it } from "vitest"; +import { appendTaskDeltaMessage, buildTaskDeltaMessage, buildTaskSnapshot } from "./task-snapshot.js"; + +describe("task snapshot and delta", () => { + it("keeps full Mermaid out of delta messages and appends without mutating old messages", () => { + const fullMermaid = "flowchart TD\nA[done] --> B[doing]\nB --> C[todo]"; + const snapshot = buildTaskSnapshot({ + taskGoal: "完成 D 方案", + mmdFile: "task.mmd", + mermaid: fullMermaid, + resultRef: "refs/task-full.md", + }); + const delta = buildTaskDeltaMessage({ + taskGoal: "完成 D 方案", + mmdFile: "task.mmd", + changedNodeIds: ["B", "C"], + resultRef: "refs/task-full.md", + maxChars: 200, + }); + const before = [{ role: "user", content: "hello" }]; + const copy = JSON.stringify(before); + const after = appendTaskDeltaMessage(before, delta); + + expect(snapshot.text).toContain("refs/task-full.md"); + expect(snapshot.text).toContain("hash="); + expect(delta.content[0].text).not.toContain("flowchart TD"); + expect(delta.content[0].text).toContain("changed_nodes"); + expect(JSON.stringify(before)).toBe(copy); + expect(after).toHaveLength(2); + }); +}); diff --git a/src/offload/task-snapshot.ts b/src/offload/task-snapshot.ts new file mode 100644 index 00000000..42791c92 --- /dev/null +++ b/src/offload/task-snapshot.ts @@ -0,0 +1,58 @@ +import { createHash } from "node:crypto"; + +export interface TaskSnapshotInput { + taskGoal: string; + mmdFile: string; + mermaid: string; + resultRef: string; +} + +export interface TaskSnapshot { + text: string; + hash: string; +} + +export function buildTaskSnapshot(input: TaskSnapshotInput): TaskSnapshot { + const hash = sha256(input.mermaid).slice(0, 16); + const text = [ + ``, + `task_goal: ${input.taskGoal}`, + `mmd_file: ${input.mmdFile}`, + `full_mermaid_ref: ${input.resultRef}`, + `mermaid_hash: ${hash}`, + ``, + ].join("\n"); + return { text, hash }; +} + +export function buildTaskDeltaMessage(input: { + taskGoal: string; + mmdFile: string; + changedNodeIds: string[]; + resultRef: string; + maxChars?: number; +}): { role: "user"; content: Array<{ type: "text"; text: string }>; _mmdContextMessage: "delta" } { + const changedNodes = [...new Set(input.changedNodeIds)].sort(); + const text = [ + ``, + `task_goal: ${input.taskGoal}`, + `mmd_file: ${input.mmdFile}`, + `full_mermaid_ref: ${input.resultRef}`, + `changed_nodes: ${changedNodes.join(", ") || "(none)"}`, + ``, + ].join("\n"); + const maxChars = input.maxChars ?? 1200; + return { + role: "user", + content: [{ type: "text", text: text.length > maxChars ? `${text.slice(0, maxChars).trimEnd()}\n[truncated]` : text }], + _mmdContextMessage: "delta", + }; +} + +export function appendTaskDeltaMessage>(messages: T[], delta: T): T[] { + return [...messages, delta]; +} + +function sha256(text: string): string { + return createHash("sha256").update(text).digest("hex"); +} diff --git a/src/offload/tool-result-normalizer.test.ts b/src/offload/tool-result-normalizer.test.ts new file mode 100644 index 00000000..3aca0d25 --- /dev/null +++ b/src/offload/tool-result-normalizer.test.ts @@ -0,0 +1,45 @@ +import { describe, expect, it } from "vitest"; +import { normalizeToolResultForPrompt, stableSerialize } from "./tool-result-normalizer.js"; + +describe("normalizeToolResultForPrompt", () => { + it("keeps small tool results inline", async () => { + const result = await normalizeToolResultForPrompt({ + toolName: "small", + toolCallId: "tc-small", + timestamp: "2026-07-06T00:00:00.000Z", + result: { ok: true }, + maxTokens: 100, + writeRef: async () => "refs/unused.md", + }); + + expect(result.offloaded).toBe(false); + expect(result.promptResult).toEqual({ ok: true }); + }); + + it("offloads large tool results with deterministic serialization and content hash", async () => { + const refs: string[] = []; + const result = await normalizeToolResultForPrompt({ + toolName: "read_file", + toolCallId: "tc-large", + timestamp: "2026-07-06T00:00:00.000Z", + result: { b: "x".repeat(1000), a: 1 }, + maxTokens: 20, + summaryMaxTokens: 20, + previewMaxChars: 60, + writeRef: async (content) => { + refs.push(content); + return "refs/tc-large.md"; + }, + }); + + expect(result.offloaded).toBe(true); + expect(refs).toHaveLength(1); + expect(refs[0]).toBe(stableSerialize({ a: 1, b: "x".repeat(1000) })); + expect(result.promptResult).toMatchObject({ + _tdai_offloaded: true, + result_ref: "refs/tc-large.md", + tool_call_id: "tc-large", + }); + expect(JSON.stringify(result.promptResult)).toContain(result.contentHash); + }); +}); diff --git a/src/offload/tool-result-normalizer.ts b/src/offload/tool-result-normalizer.ts new file mode 100644 index 00000000..a74d0a05 --- /dev/null +++ b/src/offload/tool-result-normalizer.ts @@ -0,0 +1,74 @@ +import { createHash } from "node:crypto"; +import { estimateTokens } from "../core/session/session-snapshot.js"; + +export interface NormalizeToolResultInput { + toolName: string; + toolCallId: string; + timestamp: string; + result: unknown; + maxTokens: number; + summaryMaxTokens?: number; + previewMaxChars?: number; + writeRef: (content: string) => Promise; +} + +export interface NormalizeToolResultOutput { + offloaded: boolean; + promptResult: unknown; + resultRef?: string; + contentHash?: string; + originalTokens: number; +} + +export async function normalizeToolResultForPrompt(input: NormalizeToolResultInput): Promise { + const serialized = stableSerialize(input.result); + const originalTokens = estimateTokens(serialized); + if (originalTokens <= Math.max(1, input.maxTokens)) { + return { offloaded: false, promptResult: input.result, originalTokens }; + } + + const contentHash = sha256(serialized); + const resultRef = await input.writeRef(serialized); + const summary = buildDeterministicSummary(serialized, input.previewMaxChars ?? 1200, input.summaryMaxTokens ?? 350); + return { + offloaded: true, + resultRef, + contentHash, + originalTokens, + promptResult: { + _tdai_offloaded: true, + tool_name: input.toolName, + tool_call_id: input.toolCallId, + result_ref: resultRef, + content_hash: contentHash, + original_tokens: originalTokens, + summary, + instruction: "需要完整工具结果时调用 tdai_offload_read({ result_ref }),可用 start_line/end_line/query/max_tokens 精确读取。", + }, + }; +} + +export function stableSerialize(value: unknown): string { + return JSON.stringify(sortForJson(value), null, 2); +} + +function sortForJson(value: unknown): unknown { + if (Array.isArray(value)) return value.map(sortForJson); + if (value && typeof value === "object") { + const entries = Object.entries(value as Record) + .sort(([a], [b]) => a.localeCompare(b)) + .map(([k, v]) => [k, sortForJson(v)]); + return Object.fromEntries(entries); + } + return value; +} + +function buildDeterministicSummary(serialized: string, previewMaxChars: number, summaryMaxTokens: number): string { + const maxChars = Math.max(80, Math.min(previewMaxChars, summaryMaxTokens * 4)); + const preview = Array.from(serialized).slice(0, maxChars).join("").trimEnd(); + return preview.length < serialized.length ? `${preview}\n[truncated]` : preview; +} + +function sha256(text: string): string { + return createHash("sha256").update(text).digest("hex"); +} diff --git a/src/offload/types.ts b/src/offload/types.ts index fce1f319..86c16ddd 100644 --- a/src/offload/types.ts +++ b/src/offload/types.ts @@ -35,6 +35,9 @@ export interface ToolPair { toolCallId: string; params: Record | string; result: unknown; + result_ref?: string; + content_hash?: string; + original_tokens?: number; error?: string; timestamp: string; durationMs?: number; @@ -185,6 +188,22 @@ export interface PluginConfig { emergencyTargetRatio?: number; /** Max ratio of total tokens that injected MMDs may occupy. Default: 0.2 */ mmdMaxTokenRatio?: number; + /** Inline tool-result budget before front offload. Default: 1200 */ + inlineToolResultMaxTokens?: number; + /** Summary token budget for front offload stubs. Default: 350 */ + summaryMaxTokens?: number; + /** Preview character budget for deterministic summaries. Default: 1200 */ + previewMaxChars?: number; + /** Max tokens returned by tdai_offload_read. Default: 1600 */ + readChunkMaxTokens?: number; + /** Cache epoch transition trigger ratio. Default: 0.8 */ + epochTriggerRatio?: number; + /** Cache epoch target ratio after compaction. Default: 0.55 */ + epochTargetRatio?: number; + /** Minimum turns before epoch transition. Default: 4 */ + epochMinimumTurns?: number; + /** Minimum compactable continuous interval. Default: 2 */ + epochMinimumIntervalTurns?: number; /** * L3 token counting: `tiktoken` uses js-tiktoken (exact BPE for chosen encoding); * `heuristic` uses 中文/1.7 + 其余/4. Default: tiktoken. @@ -243,6 +262,14 @@ export const PLUGIN_DEFAULTS = { /** Emergency target: delete until tokens <= contextWindow * 0.6 */ emergencyTargetRatio: 0.6, mmdMaxTokenRatio: 0.2, + inlineToolResultMaxTokens: 1200, + summaryMaxTokens: 350, + previewMaxChars: 1200, + readChunkMaxTokens: 1600, + epochTriggerRatio: 0.8, + epochTargetRatio: 0.55, + epochMinimumTurns: 4, + epochMinimumIntervalTurns: 2, l3TokenCountMode: "tiktoken" as const, l3TiktokenEncoding: "cl100k_base" as const, defaultSystemOverheadRatio: 0.12, diff --git a/src/utils/injected-memory.test.ts b/src/utils/injected-memory.test.ts new file mode 100644 index 00000000..e0dfe653 --- /dev/null +++ b/src/utils/injected-memory.test.ts @@ -0,0 +1,139 @@ +import { describe, expect, test } from "vitest"; +import { + commonPrefixChars, + estimatePrefixReuse, + estimateInjectedHistoryChars, + stripInjectedRelevantMemoriesFromContent, + stripInjectedRelevantMemoriesFromText, +} from "./injected-memory.js"; + +describe("injected memory cleanup", () => { + test("removes relevant memories from a persisted string user message", () => { + const content = [ + "", + "memory line 1", + "memory line 2", + "", + "What should I do next?", + ].join("\n"); + + expect(stripInjectedRelevantMemoriesFromText(content)).toBe("What should I do next?"); + }); + + test("removes injected text parts and preserves non-text parts", () => { + const content = [ + { + type: "text", + text: "\nold memory\n\n", + }, + { + type: "image_url", + image_url: { url: "data:image/png;base64,abc" }, + }, + { + type: "text", + text: "Continue the task", + }, + ]; + + expect(stripInjectedRelevantMemoriesFromContent(content)).toEqual([ + { + type: "image_url", + image_url: { url: "data:image/png;base64,abc" }, + }, + { + type: "text", + text: "Continue the task", + }, + ]); + }); + + test("reports the showInjected history bloat avoided by stripping", () => { + const turns = [ + { + role: "user", + content: "\n".concat("a".repeat(500), "\n\nturn 1"), + }, + { + role: "assistant", + content: "ok", + }, + { + role: "user", + content: "\n".concat("b".repeat(700), "\n\nturn 2"), + }, + ]; + + expect(estimateInjectedHistoryChars(turns)).toEqual({ + beforeChars: 1298, + afterChars: 14, + removedChars: 1284, + removedBlocks: 2, + }); + }); + + test("improves reusable prompt prefix when history is tail-truncated", () => { + const stableSystem = "system:".concat("s".repeat(1100)); + const historyBudgetChars = 320; + const injectedA = "\n".concat("a".repeat(600), "\n\nfirst turn"); + const injectedB = "\n".concat("b".repeat(700), "\n\nsecond turn"); + const injectedC = "\n".concat("c".repeat(500), "\n\nthird turn"); + + const buildPrompt = ( + history: Array<{ role: "user" | "assistant"; content: string }>, + currentUser: string, + ) => { + const historyText = history.map((m) => `${m.role}:${m.content}`).join("\n"); + const tailHistory = historyText.length > historyBudgetChars + ? historyText.slice(historyText.length - historyBudgetChars) + : historyText; + return `${stableSystem}\n${tailHistory}\nuser:${currentUser}`; + }; + + const pollutedTurn2 = buildPrompt( + [ + { role: "user", content: injectedA }, + { role: "assistant", content: "assistant one" }, + ], + injectedB, + ); + const pollutedTurn3 = buildPrompt( + [ + { role: "user", content: injectedA }, + { role: "assistant", content: "assistant one" }, + { role: "user", content: injectedB }, + { role: "assistant", content: "assistant two" }, + ], + injectedC, + ); + + const cleanTurn2 = buildPrompt( + [ + { role: "user", content: stripInjectedRelevantMemoriesFromText(injectedA) }, + { role: "assistant", content: "assistant one" }, + ], + injectedB, + ); + const cleanTurn3 = buildPrompt( + [ + { role: "user", content: stripInjectedRelevantMemoriesFromText(injectedA) }, + { role: "assistant", content: "assistant one" }, + { role: "user", content: stripInjectedRelevantMemoriesFromText(injectedB) }, + { role: "assistant", content: "assistant two" }, + ], + injectedC, + ); + + const pollutedPrefix = commonPrefixChars(pollutedTurn2, pollutedTurn3); + const cleanPrefix = commonPrefixChars(cleanTurn2, cleanTurn3); + const pollutedReuse = estimatePrefixReuse(pollutedTurn2, pollutedTurn3); + const cleanReuse = estimatePrefixReuse(cleanTurn2, cleanTurn3); + + expect(pollutedPrefix).toBe(1108); + expect(cleanPrefix).toBe(1153); + expect(cleanPrefix - pollutedPrefix).toBe(45); + expect(pollutedReuse.reuseRatio).toBeCloseTo(0.5579, 4); + expect(cleanReuse.reuseRatio).toBeCloseTo(0.6604, 4); + expect(cleanReuse.reuseRatio).toBeGreaterThan(pollutedReuse.reuseRatio); + }); +}); diff --git a/src/utils/injected-memory.ts b/src/utils/injected-memory.ts new file mode 100644 index 00000000..8baf5025 --- /dev/null +++ b/src/utils/injected-memory.ts @@ -0,0 +1,129 @@ +const RELEVANT_MEMORIES_RE = /]*>[\s\S]*?<\/relevant-memories>\s*/gi; + +type MessageContent = string | Array> | unknown; + +export interface InjectedHistoryEstimate { + beforeChars: number; + afterChars: number; + removedChars: number; + removedBlocks: number; +} + +export interface PrefixReuseEstimate { + prefixChars: number; + comparableChars: number; + reuseRatio: number; +} + +export function stripInjectedRelevantMemoriesFromText(text: string): string { + return text.replace(RELEVANT_MEMORIES_RE, "").trim(); +} + +export function stripInjectedRelevantMemoriesFromContent(content: MessageContent): MessageContent { + if (typeof content === "string") { + if (!hasInjectedRelevantMemories(content)) return content; + return stripInjectedRelevantMemoriesFromText(content); + } + + if (!Array.isArray(content)) return content; + + let changed = false; + const cleanedParts: Array> = []; + for (const part of content) { + if (!part || typeof part !== "object") { + cleanedParts.push(part); + continue; + } + if (part.type !== "text" || typeof part.text !== "string") { + cleanedParts.push(part); + continue; + } + if (!hasInjectedRelevantMemories(part.text)) { + cleanedParts.push(part); + continue; + } + + changed = true; + const cleanedText = stripInjectedRelevantMemoriesFromText(part.text); + if (cleanedText.length > 0) { + cleanedParts.push({ ...part, text: cleanedText }); + } + } + + return changed ? cleanedParts : content; +} + +export function hasInjectedRelevantMemories(text: string): boolean { + RELEVANT_MEMORIES_RE.lastIndex = 0; + return RELEVANT_MEMORIES_RE.test(text); +} + +export function countInjectedRelevantMemoryBlocks(text: string): number { + RELEVANT_MEMORIES_RE.lastIndex = 0; + return [...text.matchAll(RELEVANT_MEMORIES_RE)].length; +} + +export function estimateInjectedHistoryChars(messages: Array<{ role?: unknown; content?: unknown }>): InjectedHistoryEstimate { + let beforeChars = 0; + let afterChars = 0; + let removedBlocks = 0; + + for (const message of messages) { + const before = contentToText(message.content); + if (!before) continue; + beforeChars += before.length; + + if (message.role === "user") { + removedBlocks += countInjectedRelevantMemoryBlocks(before); + const stripped = stripInjectedRelevantMemoriesFromContent(message.content); + afterChars += contentToText(stripped).length; + } else { + afterChars += before.length; + } + } + + return { + beforeChars, + afterChars, + removedChars: beforeChars - afterChars, + removedBlocks, + }; +} + +export function commonPrefixChars(a: string, b: string): number { + const max = Math.min(a.length, b.length); + let i = 0; + while (i < max && a.charCodeAt(i) === b.charCodeAt(i)) i++; + return i; +} + +export function estimatePrefixReuse(previousPrompt: string, currentPrompt: string): PrefixReuseEstimate { + const comparableChars = Math.max(1, Math.min(previousPrompt.length, currentPrompt.length)); + const prefixChars = commonPrefixChars(previousPrompt, currentPrompt); + return { + prefixChars, + comparableChars, + reuseRatio: prefixChars / comparableChars, + }; +} + +export function measureMessageContentChars(content: unknown): number { + return contentToText(content).length; +} + +function contentToText(content: unknown): string { + if (typeof content === "string") return content; + if (!Array.isArray(content)) return ""; + const parts: string[] = []; + for (const part of content) { + if ( + part && + typeof part === "object" && + (part as Record).type === "text" && + typeof (part as Record).text === "string" + ) { + parts.push((part as Record).text); + } + } + return parts.join("\n"); +}