diff --git a/packages/opencode/src/memory/injector.ts b/packages/opencode/src/memory/injector.ts index 7ba77db15baa..f0e03bc70271 100644 --- a/packages/opencode/src/memory/injector.ts +++ b/packages/opencode/src/memory/injector.ts @@ -18,6 +18,14 @@ function relevanceWeight(entry: Memory.Info): number { return entry.relevanceScore * recencyWeight * Math.log2(entry.accessCount + 2) } +const DEFAULT_BUDGET_ALLOCATION = { + agent: 0.30, + project: 0.25, + feedback: 0.20, + user: 0.15, + reference: 0.10, +} as const + export namespace MemoryInjector { export async function load(agent?: string): Promise { const config = await Config.get() @@ -57,18 +65,21 @@ export namespace MemoryInjector { const generalIds = new Set(entries.map((e) => e.id)) agentEntries = agentEntries.filter((e) => !generalIds.has(e.id)) - // Build sections within token budget + // Build sections with proportional budget, redistributing unused portions + const allocation = DEFAULT_BUDGET_ALLOCATION const sections: string[] = [] - let tokenBudget = maxTokens - const injectedIds: string[] = [] + const allIncludedIds: string[] = [] + + // Only reserve agent budget when agent entries exist + const agentBudget = agentEntries.length > 0 ? Math.floor(maxTokens * allocation.agent) : 0 + const generalPool = maxTokens - agentBudget // Agent-specific section first (highest priority) if (agentEntries.length > 0) { - const agentSection = buildSection("Agent-Specific Knowledge", agentEntries, tokenBudget) + const agentSection = buildSection("Agent-Specific Knowledge", agentEntries, agentBudget) if (agentSection.text) { sections.push(agentSection.text) - tokenBudget -= agentSection.tokens - injectedIds.push(...agentSection.includedIds) + allIncludedIds.push(...agentSection.includedIds) } } @@ -78,27 +89,51 @@ export namespace MemoryInjector { const feedbackEntries = entries.filter((e) => e.type === "feedback") const referenceEntries = entries.filter((e) => e.type === "reference") - for (const [title, group] of [ - ["Project Knowledge", projectEntries], - ["User Preferences", userEntries], - ["Feedback & Patterns", feedbackEntries], - ["Reference", referenceEntries], - ] as const) { - if (group.length === 0 || tokenBudget <= 0) continue - const section = buildSection(title, group, tokenBudget) - if (section.text) { - sections.push(section.text) - tokenBudget -= section.tokens - injectedIds.push(...section.includedIds) + // Compute active sections and redistribute budget proportionally + const generalSections = [ + ["Project Knowledge", projectEntries, allocation.project] as const, + ["User Preferences", userEntries, allocation.user] as const, + ["Feedback & Patterns", feedbackEntries, allocation.feedback] as const, + ["Reference", referenceEntries, allocation.reference] as const, + ] + const activeSections = generalSections.filter(([, group]) => group.length > 0) + const activeWeight = activeSections.reduce((sum, [, , weight]) => sum + weight, 0) + + // Pass 1: Build sections with initial proportional budgets + const sectionResults = activeSections.map(([title, group, weight]) => { + const budget = activeWeight > 0 ? Math.floor(generalPool * (weight / activeWeight)) : 0 + return { title, group, weight, budget, result: buildSection(title, group, budget) } + }) + + // Pass 2: Redistribute budget from sections that couldn't fit any entries + const failedBudget = sectionResults + .filter((s) => !s.result.text) + .reduce((sum, s) => sum + s.budget, 0) + + if (failedBudget > 0) { + const viableWeight = sectionResults + .filter((s) => s.result.text) + .reduce((sum, s) => sum + s.weight, 0) + for (const s of sectionResults) { + if (!s.result.text || viableWeight <= 0) continue + const extra = Math.floor(failedBudget * (s.weight / viableWeight)) + s.result = buildSection(s.title, s.group, s.budget + extra) + } + } + + for (const s of sectionResults) { + if (s.result.text) { + sections.push(s.result.text) + allIncludedIds.push(...s.result.includedIds) } } if (sections.length === 0) return undefined // Increment access counts only for entries actually injected (within token budget) - if (injectedIds.length > 0) { + if (allIncludedIds.length > 0) { try { - await MemoryStore.runPromise((svc) => svc.incrementAccessBatch(injectedIds)) + await MemoryStore.runPromise((svc) => svc.incrementAccessBatch(allIncludedIds)) } catch { // Non-critical: access count tracking failure should not block injection } diff --git a/packages/opencode/test/memory/injector.test.ts b/packages/opencode/test/memory/injector.test.ts index 74a963555002..de919ddaa4be 100644 --- a/packages/opencode/test/memory/injector.test.ts +++ b/packages/opencode/test/memory/injector.test.ts @@ -278,6 +278,121 @@ describe("MemoryInjector.load", () => { }) }) + test("all sections get entries when budget is sufficient", async () => { + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const ts = Date.now() + Database.use((d) => { + seed(d, { id: `proj-${ts}`, projectPath: tmp.path, topic: "proj-entry", type: "project", content: "Project content" }) + seed(d, { id: `user-${ts}`, projectPath: tmp.path, topic: "user-entry", type: "user", content: "User content" }) + seed(d, { id: `fb-${ts}`, projectPath: tmp.path, topic: "fb-entry", type: "feedback", content: "Feedback content" }) + seed(d, { id: `ref-${ts}`, projectPath: tmp.path, topic: "ref-entry", type: "reference", content: "Reference content" }) + }) + + const result = await MemoryInjector.load() + expect(result).toBeDefined() + expect(result).toContain("proj-entry") + expect(result).toContain("user-entry") + expect(result).toContain("fb-entry") + expect(result).toContain("ref-entry") + }, + }) + }) + + test("large agent section does not starve other sections", async () => { + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const ts = Date.now() + Database.use((d) => { + // Fill agent section with many large entries + for (let i = 0; i < 20; i++) { + seed(d, { + id: `agent-bulk-${i}-${ts}`, + projectPath: tmp.path, + topic: `agent-entry-${i}`, + type: "project", + content: "A".repeat(400), + agent: "code-reviewer", + }) + } + // Add general entries that should still appear + seed(d, { id: `gen-proj-${ts}`, projectPath: tmp.path, topic: "general-project", type: "project", content: "General project entry" }) + seed(d, { id: `gen-fb-${ts}`, projectPath: tmp.path, topic: "general-feedback", type: "feedback", content: "General feedback entry" }) + }) + + const result = await MemoryInjector.load("code-reviewer") + expect(result).toBeDefined() + // General sections must not be starved even with large agent section + expect(result).toContain("general-project") + expect(result).toContain("general-feedback") + }, + }) + }) + + test("empty sections produce no output headers", async () => { + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const ts = Date.now() + // Only seed project and feedback — user and reference are empty + Database.use((d) => { + seed(d, { id: `only-proj-${ts}`, projectPath: tmp.path, topic: "only-project", type: "project", content: "Only project entry" }) + seed(d, { id: `only-fb-${ts}`, projectPath: tmp.path, topic: "only-feedback", type: "feedback", content: "Only feedback entry" }) + }) + + const result = await MemoryInjector.load() + expect(result).toBeDefined() + // Present sections appear + expect(result).toContain("only-project") + expect(result).toContain("only-feedback") + // Empty sections produce no header + expect(result).not.toContain("## User Preferences") + expect(result).not.toContain("## Reference") + }, + }) + }) + + test("budget from empty sections is redistributed to active sections", async () => { + await using tmp = await tmpdir({ + git: true, + config: { memory: { max_memory_tokens: 200 } } as any, + }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const ts = Date.now() + // Only seed "project" entries — user/feedback/reference are empty + Database.use((d) => { + for (let i = 0; i < 10; i++) { + seed(d, { + id: `redist-${i}-${ts}`, + projectPath: tmp.path, + topic: `redist-entry-${i}`, + type: "project", + content: `Redistributed budget content item ${i}`, + }) + } + }) + + const result = await MemoryInjector.load() + expect(result).toBeDefined() + // With redistribution, project section gets the full general pool (no agent entries) + // so it should fit more than 1 entry even with a tight 200-token budget + const entryCount = (result!.match(/redist-entry-\d+/g) || []).length + expect(entryCount).toBeGreaterThan(1) + // Empty sections should not appear + expect(result).not.toContain("## User Preferences") + expect(result).not.toContain("## Feedback & Patterns") + expect(result).not.toContain("## Reference") + }, + }) + }) + test("includes description in entry output when present", async () => { await using tmp = await tmpdir({ git: true }) await Instance.provide({