From 026b2ebfdeb31b60613f47399afc533924cefe78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B0=8F=E9=B3=84=E9=B1=BC?= Date: Wed, 28 Oct 2026 17:43:57 +0800 Subject: [PATCH] fix(memory): enforce recalled context budgets (MEMMY-815) --- .../local-api-contracts/src/memory-runtime.ts | 1 + Memory/src/contracts/memory-runtime.ts | 1 + .../src/service/retrieval/injection-budget.ts | 62 +++++++++++ .../service/retrieval/retrieval-service.ts | 78 +++++++------- Memory/src/types.ts | 2 + .../service/retrieval/context-budget.test.ts | 102 ++++++++++++++++++ .../retrieval/injected-context.test.ts | 6 +- .../retrieval/query-and-filter.test.ts | 11 ++ 8 files changed, 219 insertions(+), 44 deletions(-) create mode 100644 Memory/src/service/retrieval/injection-budget.ts create mode 100644 Memory/tests/service/retrieval/context-budget.test.ts diff --git a/App/backend/local-api-contracts/src/memory-runtime.ts b/App/backend/local-api-contracts/src/memory-runtime.ts index 3c909be58..9be76e2e9 100644 --- a/App/backend/local-api-contracts/src/memory-runtime.ts +++ b/App/backend/local-api-contracts/src/memory-runtime.ts @@ -425,6 +425,7 @@ export const StartTurnInputSchema = RuntimeRequestFieldsSchema.extend({ query: NonEmptyStringSchema, turnId: NonEmptyStringSchema.optional(), contextHints: UnknownRecordSchema.optional(), + /** Estimated recalled-context token limit; fixed task instructions are exempt. Default: 1800. */ contextBudget: z.number().int().nonnegative().optional() }); export type StartTurnInput = z.infer; diff --git a/Memory/src/contracts/memory-runtime.ts b/Memory/src/contracts/memory-runtime.ts index 1ee3d88ca..9bdd3344a 100644 --- a/Memory/src/contracts/memory-runtime.ts +++ b/Memory/src/contracts/memory-runtime.ts @@ -423,6 +423,7 @@ export const StartTurnInputSchema = RuntimeRequestFieldsSchema.extend({ query: NonEmptyStringSchema, turnId: NonEmptyStringSchema.optional(), contextHints: UnknownRecordSchema.optional(), + /** Estimated recalled-context token limit; fixed task instructions are exempt. Default: 1800. */ contextBudget: z.number().int().nonnegative().optional() }); export type StartTurnInput = z.infer; diff --git a/Memory/src/service/retrieval/injection-budget.ts b/Memory/src/service/retrieval/injection-budget.ts new file mode 100644 index 000000000..d2fbb6cfb --- /dev/null +++ b/Memory/src/service/retrieval/injection-budget.ts @@ -0,0 +1,62 @@ +import type { InjectedContext, MemoryKind, RecallMemoryLayer } from "../../types.js"; + +export interface InjectionPacket { + injectedContext: InjectedContext; + sourceMemoryIds: string[]; +} + +export interface BudgetSource { + id: string; + kind: MemoryKind; + memoryLayer: RecallMemoryLayer; + tokenEstimate?: number; +} + +/** The renderer is the authority for both content and sources, including wrappers. */ +export function budgetInjection( + candidates: T[], + budget: number, + render: (selected: T[]) => InjectionPacket, + sources: (candidate: T) => BudgetSource[], + fixedContextChars = 0 +): InjectionPacket & { droppedDueToBudget: Array } { + const limit = Number.isFinite(budget) ? Math.max(0, Math.floor(budget)) : 0; + const measure = (items: T[]): InjectionPacket => { + const packet = render(items); + return { ...packet, injectedContext: { ...packet.injectedContext, + tokenEstimate: Math.ceil(packet.injectedContext.markdown.length / 4) } }; + }; + // Fixed task instructions are independent of recalled memories. Charge only + // the additional rendered text; keep tokenEstimate truthful for the full packet. + const fits = (packet: InjectionPacket) => Math.ceil(Math.max(0, + packet.injectedContext.markdown.length - fixedContextChars) / 4) <= limit; + const full = measure(candidates); + if ((limit > 0 || candidates.length === 0) && fits(full)) { + return { ...full, droppedDueToBudget: [] }; + } + let result: InjectionPacket = { + injectedContext: { markdown: "", sections: [], tokenEstimate: 0 }, sourceMemoryIds: [] + }; + const selected: T[] = []; + const rejected: T[] = []; + const empty = measure([]); + if (fits(empty)) result = empty; + for (const candidate of candidates) { + if (limit === 0) { rejected.push(candidate); continue; } + const trial = measure([...selected, candidate]); + if (fits(trial)) { + selected.push(candidate); + result = trial; + } else { + rejected.push(candidate); + } + } + const retained = new Set(result.sourceMemoryIds); + const dropped = new Map(); + for (const candidate of rejected) { + for (const source of sources(candidate)) { + if (!retained.has(source.id)) dropped.set(source.id, { ...source, reason: "token_budget" }); + } + } + return { ...result, droppedDueToBudget: [...dropped.values()] }; +} diff --git a/Memory/src/service/retrieval/retrieval-service.ts b/Memory/src/service/retrieval/retrieval-service.ts index 2a0e9624b..56bc2315c 100644 --- a/Memory/src/service/retrieval/retrieval-service.ts +++ b/Memory/src/service/retrieval/retrieval-service.ts @@ -1,3 +1,4 @@ +import { budgetInjection } from "./injection-budget.js"; import type { RetrievalQueryExtract, RetrievalResult @@ -652,50 +653,42 @@ export function buildInjectedContext( ) : rendered; - void budget; - const sections: InjectedContext["sections"] = memories.map((section) => section.section); - const renderedSections: RenderedInjectedSection[] = [...memories]; - const sourceMemoryIds: string[] = memories.flatMap((section) => section.section.memoryIds); - const droppedDueToBudget: Array<{ - id: string; - kind: MemoryKind; - memoryLayer: RecallMemoryLayer; - reason: "token_budget"; - tokenEstimate?: number; - }> = []; - let used = sections.reduce((sum, section) => sum + (section.tokenEstimate ?? 0), 0); - const guidance = decisionGuidanceSection( - contextMemoriesForInjectedSources(contextMemories, sourceMemoryIds) - ); - const avoidance = failureAvoidanceSection( - contextMemoriesForInjectedSources(contextMemories, sourceMemoryIds) - ); - if (guidance) { - const estimate = guidance.tokenEstimate ?? 0; - sections.push(guidance); - sourceMemoryIds.push(...guidance.memoryIds); - used += estimate; - } - if (avoidance) { - sections.push(avoidance); - sourceMemoryIds.push(...avoidance.memoryIds); - used += avoidance.tokenEstimate ?? 0; - } - - const markdown = renderInjectedMarkdown(renderedSections, guidance, avoidance, retrievalMode, options); + const fixedContextChars = renderInjectedMarkdown([], undefined, undefined, retrievalMode, options).length; + return budgetInjection(memories, budget, (selected) => { + const sections: InjectedContext["sections"] = selected.map((item) => item.section); + const sourceMemoryIds = selected.flatMap((item) => item.section.memoryIds); + const sourceMemories = contextMemoriesForInjectedSources(contextMemories, sourceMemoryIds); + const guidance = decisionGuidanceSection(sourceMemories); + const avoidance = failureAvoidanceSection(sourceMemories); + if (guidance) { + sections.push(guidance); + sourceMemoryIds.push(...guidance.memoryIds); + } + if (avoidance) { + sections.push(avoidance); + sourceMemoryIds.push(...avoidance.memoryIds); + } + return { + injectedContext: { + markdown: renderInjectedMarkdown(selected, guidance, avoidance, retrievalMode, options), + sections + }, + sourceMemoryIds: uniq(sourceMemoryIds) + }; + }, (item) => item.section.memoryIds.map((id) => ({ + id, kind: item.section.kind, memoryLayer: item.section.memoryLayer, + tokenEstimate: item.section.tokenEstimate + })), fixedContextChars); +} - return { - injectedContext: { - markdown, - sections, - tokenEstimate: used - }, - sourceMemoryIds: uniq(sourceMemoryIds), - droppedDueToBudget - }; +function buildTimeFilteredInjectedContext(memories: MemoryRow[], timeZone: string, budget: number) { + return budgetInjection(memories, budget, + (selected) => renderTimeFilteredInjectedContext(selected, timeZone), + (memory) => [{ id: memory.id, kind: kindFromMemory(memory), memoryLayer: memory.memoryLayer, + tokenEstimate: renderTimeFilteredInjectedContext([memory], timeZone).injectedContext.tokenEstimate }]); } -function buildTimeFilteredInjectedContext( +function renderTimeFilteredInjectedContext( memories: MemoryRow[], timeZone: string ): { @@ -2039,7 +2032,8 @@ export class RetrievalService { const contextPacket = timeFilter ? buildTimeFilteredInjectedContext( memories.filter((memory) => hits.some((hit) => hit.id === memory.id)), - timeZone + timeZone, + request.contextBudget ?? 1800 ) : buildInjectedContext( hits, diff --git a/Memory/src/types.ts b/Memory/src/types.ts index 7ef5d0a1a..6886bd1f9 100644 --- a/Memory/src/types.ts +++ b/Memory/src/types.ts @@ -351,6 +351,7 @@ export interface TurnStartRequest extends RequestEnvelope { turnId?: string; layers?: MemoryLayer[]; contextHints?: Record; + /** Estimated recalled-context token limit; fixed task instructions are exempt. Default: 1800. */ contextBudget?: number; } @@ -498,6 +499,7 @@ export interface MemorySearchRequest extends RequestEnvelope { layers?: MemoryLayer[]; tags?: string[]; limit?: number; + /** Estimated recalled-context token limit; fixed task instructions are exempt. Default: 1800. */ contextBudget?: number; includeInjectedContext?: boolean; verbose?: boolean; diff --git a/Memory/tests/service/retrieval/context-budget.test.ts b/Memory/tests/service/retrieval/context-budget.test.ts new file mode 100644 index 000000000..933880d83 --- /dev/null +++ b/Memory/tests/service/retrieval/context-budget.test.ts @@ -0,0 +1,102 @@ +import { describe, expect, it } from "vitest"; +import { buildInjectedContext } from "../../../src/service/retrieval/retrieval-service.js"; +import type { MemoryRow, RecallHit } from "../../../src/types.js"; + +function skill(id: string, guide: string, procedure: Record = {}) { + const hit: RecallHit = { id, kind: "skill", memoryLayer: "Skill", status: "activated", title: id, + snippet: guide, score: 1, tags: [], source: "skill" }; + const memory: MemoryRow = { id, timeline: "2026-10-28T00:00:00Z", userId: "test", memoryType: "SkillMemory", + status: "activated", visibility: "private", memoryKey: id, memoryValue: guide, tags: [], info: {}, + properties: { internal_info: { memory_layer: "Skill", memory_kind: "skill", + skill: { name: id, invocation_guide: guide, procedure_json: procedure } } }, + memoryLayer: "Skill", version: 1, createdAt: "2026-10-28T00:00:00Z", updatedAt: "2026-10-28T00:00:00Z" }; + return { hit, memory }; +} + +describe("Injected context budget", () => { + const tuning = { skillInjectionMode: "full" as const }; + const packet = (rows: ReturnType[], budget: number, query?: string) => buildInjectedContext( + rows.map(row => row.hit), budget, rows.map(row => row.memory), "turn_start", undefined, query, tuning); + + it("keeps sufficient-budget output and counts the final markdown", () => { + const rows = [skill("one", "Short complete guide"), skill("two", "Second guide")]; + const ample = packet(rows, 10000); + const exact = Math.ceil(ample.injectedContext.markdown.length / 4); + const bounded = packet(rows, exact); + expect(bounded.injectedContext.markdown).toBe(ample.injectedContext.markdown); + expect(bounded.injectedContext.tokenEstimate).toBe(exact); + expect(bounded.droppedDueToBudget).toEqual([]); + expect(bounded.sourceMemoryIds).toEqual(["one", "two"]); + }); + + it.each([0, 1, 5])("returns no oversized content at budget %s", budget => { + const result = packet([skill("long", "procedure ".repeat(1000))], budget); + expect(result.injectedContext.markdown).toBe(""); + expect(result.injectedContext.sections).toEqual([]); + expect(result.sourceMemoryIds).toEqual([]); + expect(result.droppedDueToBudget).toEqual([expect.objectContaining({ id: "long", reason: "token_budget" })]); + }); + + it("skips a long full skill and keeps a later short one without mutating hits", () => { + const rows = [skill("long", "procedure ".repeat(1500)), skill("short", "Useful small guide")]; + const original = structuredClone(rows); + const result = packet(rows, 800); + expect(result.sourceMemoryIds).toEqual(["short"]); + expect(result.injectedContext.markdown).not.toContain("procedure "); + expect(result.injectedContext.tokenEstimate).toBeLessThanOrEqual(800); + expect(result.droppedDueToBudget.map(item => item.id)).toEqual(["long"]); + expect(rows).toEqual(original); + }); + + it("includes attached decision guidance in admission and removes its source on rejection", () => { + const rows = [skill("guided", "Short guide", { decisionGuidance: { + preference: ["Always preserve the user's verified report requirements before writing"], + antiPattern: ["Do not invent successful tests"] + } })]; + const full = packet(rows, 10000); + expect(full.injectedContext.markdown).toContain("Decision guidance"); + const result = packet(rows, Math.ceil(full.injectedContext.markdown.length / 4) - 1); + expect(result.sourceMemoryIds).toEqual([]); + expect(result.injectedContext.markdown).not.toContain("Decision guidance"); + expect(result.droppedDueToBudget.map(item => item.id)).toEqual(["guided"]); + }); + + it("budgets failure-avoidance reminders together with their source", () => { + const row = skill("avoidance", "Verify the backup before migration"); + row.hit.kind = "policy"; row.hit.memoryLayer = "L2"; row.hit.source = "search"; + row.memory.memoryLayer = "L2"; + row.memory.properties.internal_info = { memory_layer: "L2", memory_kind: "policy", policy: { + title: "Migration guardrail", trigger: "Before migrating a database", confidence: 0.9, status: "active", experience_type: "failure_avoidance", + evidence_polarity: "negative", procedure: "Verify the backup before migration", + decision_guidance: { preference: ["Verify backup integrity"], anti_pattern: ["Never delete the only backup"] } + } }; + const full = packet([row], 10000); + expect(full.injectedContext.markdown).toContain("Failure avoidance"); + const bounded = packet([row], Math.ceil(full.injectedContext.markdown.length / 4) - 1); + expect(bounded.injectedContext.markdown).not.toContain("Failure avoidance"); + expect(bounded.sourceMemoryIds).toEqual([]); + expect(bounded.droppedDueToBudget).toEqual([expect.objectContaining({ id: "avoidance", kind: "policy" })]); + }); + + it("does not report a shared retained source as dropped", () => { + const rows = [skill("large", "very long guide ".repeat(1000)), skill("small", "Small guide")]; + rows[0]!.hit.memberMemoryIds = ["shared", "large-only"]; + rows[1]!.hit.memberMemoryIds = ["shared"]; + const bounded = packet(rows, 800); + expect(bounded.sourceMemoryIds).toEqual(["shared"]); + expect(bounded.droppedDueToBudget.map(item => item.id)).toEqual(["large-only"]); + }); + + it("preserves fixed task instructions at zero budget but excludes recalled sources", () => { + const query = "You need to fix a bug in the example repository.\n\n## Issue Description\nA lookup filter subquery returns too many columns because annotations leak into the select list projection."; + const fixed = packet([], 10000, query); + expect(fixed.injectedContext.markdown).toContain("Repository repair task protocol"); + const bounded = packet([skill("long", "guide ".repeat(2000))], 0, query); + expect(bounded.injectedContext.markdown).toBe(fixed.injectedContext.markdown); + expect(bounded.sourceMemoryIds).toEqual([]); + expect(bounded.droppedDueToBudget.map(item => item.id)).toEqual(["long"]); + const allowed = packet([skill("short", "Short guide")], 800, query); + expect(allowed.sourceMemoryIds).toEqual(["short"]); + expect(Math.ceil((allowed.injectedContext.markdown.length - fixed.injectedContext.markdown.length) / 4)).toBeLessThanOrEqual(800); + }); +}); diff --git a/Memory/tests/service/retrieval/injected-context.test.ts b/Memory/tests/service/retrieval/injected-context.test.ts index 033e7843e..5087ba651 100644 --- a/Memory/tests/service/retrieval/injected-context.test.ts +++ b/Memory/tests/service/retrieval/injected-context.test.ts @@ -473,8 +473,10 @@ describe("MemoryService / retrieval / injected context", () => { }); expect(prepared.hits.length).toBeGreaterThan(1); - expect(prepared.sourceMemoryIds.length).toBeGreaterThanOrEqual(prepared.hits.length); - expect(prepared.droppedDueToBudget).toEqual([]); + expect(prepared.sourceMemoryIds).toEqual([]); + expect(prepared.injectedContext.markdown).toBe(""); + expect(prepared.droppedDueToBudget.length).toBeGreaterThan(0); + expect(prepared.droppedDueToBudget.every(item => item.reason === "token_budget")).toBe(true); expect(db.db.prepare( `SELECT turn_id, json_extract(request_json, '$.retrievalMode') AS retrieval_mode diff --git a/Memory/tests/service/retrieval/query-and-filter.test.ts b/Memory/tests/service/retrieval/query-and-filter.test.ts index 88f404e1f..fb7f9aa95 100644 --- a/Memory/tests/service/retrieval/query-and-filter.test.ts +++ b/Memory/tests/service/retrieval/query-and-filter.test.ts @@ -594,6 +594,17 @@ describe("MemoryService / retrieval / query and filtering", () => { stats: { llmFilter: { durationMs?: number } }; }).stats; expect(timeFilterStats.llmFilter.durationMs).toBeUndefined(); + const bounded = await service.search({ + namespace: { source: "codex", profileId: "jiang", userId: "user-time-filter" }, + timeZone: "Asia/Shanghai", query: "我今天做了什么，总结一下", limit: 100, contextBudget: 40 + }); + expect(bounded.hits.map(hit => hit.id)).toEqual(recall.hits.map(hit => hit.id)); + expect(bounded.injectedContext.tokenEstimate).toBeLessThanOrEqual(40); + expect(bounded.sourceMemoryIds.length).toBeGreaterThan(0); + expect(bounded.sourceMemoryIds.length).toBeLessThan(recall.sourceMemoryIds.length); + expect(bounded.droppedDueToBudget.length + bounded.sourceMemoryIds.length).toBe(20); + expect(bounded.droppedDueToBudget.every(item => item.reason === "token_budget")).toBe(true); + db.close(); });