From 565180d9fe6cc1a98de2b86781926514789ce2b5 Mon Sep 17 00:00:00 2001 From: zhanba Date: Mon, 31 Aug 2026 11:54:46 +0800 Subject: [PATCH 1/5] fix(livekit): bound memory background task steps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Allow forked background tasks to define task-specific step limits so memory workflows stop before consuming the generic 50-step budget. 🤖 Generated with [Pochi](https://getpochi.com) | [Task](https://app.getpochi.com/share/p-756b25d05481454296b52ae4fcf65ad7) Co-Authored-By: Pochi --- packages/common/src/base/index.ts | 2 ++ .../__tests__/fork-agent.test.ts | 4 ++++ .../livekit/src/background-task/fork-agent.ts | 4 ++++ .../memory/__tests__/adaptors.test.ts | 7 ++++++ .../src/background-task/memory/auto-memory.ts | 4 ++++ .../src/background-task/memory/task-memory.ts | 2 ++ .../__tests__/task-executor.test.ts | 22 +++++++++++++++++++ .../task-executor/task-executor.ts | 3 ++- packages/livekit/src/chat/live-chat-kit.ts | 1 + 9 files changed, 48 insertions(+), 1 deletion(-) diff --git a/packages/common/src/base/index.ts b/packages/common/src/base/index.ts index 057cc57c13..b601a5214e 100644 --- a/packages/common/src/base/index.ts +++ b/packages/common/src/base/index.ts @@ -108,6 +108,8 @@ export interface BackgroundTaskState { tools?: readonly ToolSpecInput[]; parentTaskId?: string; useCase?: ForkAgentUseCase; + /** Maximum number of steps this background task may run. */ + maxSteps?: number; /** Step-start count inherited from the parent, excluded from the max-step guard. */ baselineStepCount?: number; } diff --git a/packages/livekit/src/background-task/__tests__/fork-agent.test.ts b/packages/livekit/src/background-task/__tests__/fork-agent.test.ts index 6a381ffed7..131da6ef60 100644 --- a/packages/livekit/src/background-task/__tests__/fork-agent.test.ts +++ b/packages/livekit/src/background-task/__tests__/fork-agent.test.ts @@ -28,6 +28,7 @@ describe("createForkAgent", () => { parentMessages, parentCwd: "/repo", directive: "extract memory", + maxSteps: 3, }).initMessages; expect(result).toHaveLength(3); @@ -50,6 +51,7 @@ describe("createForkAgent", () => { parentCwd: "/repo", directive: "extract memory", tools: ["readFile(/memory/**)", "writeToFile(/memory/**)"], + maxSteps: 3, }); expect(agent).toMatchObject({ @@ -62,6 +64,7 @@ describe("createForkAgent", () => { "writeToFile(/memory/**)", "attemptCompletion", ], + maxSteps: 3, baselineStepCount: 0, }); expect(agent).not.toHaveProperty("taskId"); @@ -97,6 +100,7 @@ describe("createForkAgent", () => { parentMessages, parentCwd: "/repo", directive: "extract memory", + maxSteps: 3, }); expect(agent.baselineStepCount).toBe(3); diff --git a/packages/livekit/src/background-task/fork-agent.ts b/packages/livekit/src/background-task/fork-agent.ts index 1f9177f491..6ba3a895b4 100644 --- a/packages/livekit/src/background-task/fork-agent.ts +++ b/packages/livekit/src/background-task/fork-agent.ts @@ -22,6 +22,7 @@ type ForkAgentInput = { parentCwd: string | undefined; directive: string; tools?: readonly ToolSpecInput[]; + maxSteps: number; }; export type ForkAgent = { @@ -31,6 +32,7 @@ export type ForkAgent = { initTitle: string | undefined; parentTaskId: string | undefined; tools: readonly ToolSpecInput[] | undefined; + maxSteps: number; baselineStepCount: number; }; @@ -94,6 +96,7 @@ export function createForkAgent( parentMessages: input.parentMessages.length, initMessages: initMessages.length, tools: input.tools?.length, + maxSteps: input.maxSteps, baselineStepCount, }, "Creating fork agent", @@ -106,6 +109,7 @@ export function createForkAgent( initTitle: input.initTitle, parentTaskId: input.parentTaskId, tools, + maxSteps: input.maxSteps, baselineStepCount, }; } diff --git a/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts b/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts index 10b945ff28..644de43a40 100644 --- a/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts +++ b/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts @@ -50,6 +50,7 @@ describe("task-memory adaptor", () => { expect(await stateStore.read(task.id)).toMatchObject({ parentTaskId: "parent", useCase: "task-memory", + maxSteps: 3, }); }); @@ -252,6 +253,7 @@ describe("auto-memory adaptor", () => { expect(await stateStore.read(task.id)).toMatchObject({ parentTaskId: "parent", useCase: "auto-memory", + maxSteps: 5, tools: [ "readFile(/repo/.pochi/memory/**)", "readFile(/repo/.pochi/transcripts/**)", @@ -478,6 +480,10 @@ describe("auto-memory adaptor", () => { status: "pending-model", title: "[Auto Memory Dream]", }); + expect(await stateStore.read(dreamTask?.id ?? "")).toMatchObject({ + useCase: "auto-memory-dream", + maxSteps: 20, + }); expect(manager.beginDreamRun).toHaveBeenCalledTimes(1); store.updateTaskStatus(dreamTask?.id ?? "", "completed"); @@ -547,6 +553,7 @@ function createTestBackgroundTask({ parentTaskId: agent.parentTaskId, tools: agent.tools, useCase: agent.label, + maxSteps: agent.maxSteps, baselineStepCount: agent.baselineStepCount, }); store.commit( diff --git a/packages/livekit/src/background-task/memory/auto-memory.ts b/packages/livekit/src/background-task/memory/auto-memory.ts index b6075fd231..2d9cfe43eb 100644 --- a/packages/livekit/src/background-task/memory/auto-memory.ts +++ b/packages/livekit/src/background-task/memory/auto-memory.ts @@ -31,6 +31,8 @@ const MemoryReadToolNames = [ "searchFiles", ] as const; const MemoryAgentWriteToolNames = ["writeToFile", "applyDiff"] as const; +const AutoMemoryMaxSteps = 5; +const AutoMemoryDreamMaxSteps = 20; const MaxSessionTranscriptChars = 24_000; const MaxPartChars = 4_000; @@ -76,6 +78,7 @@ async function startAutoMemoryExtraction({ previousMessageCount, }), tools: buildMemoryTools(context), + maxSteps: AutoMemoryMaxSteps, }); const handle = await startForkAgent(agent); @@ -148,6 +151,7 @@ async function startAutoMemoryDream({ sessions, }), tools: buildMemoryTools(run.context), + maxSteps: AutoMemoryDreamMaxSteps, }); const handle = await startForkAgent(agent); diff --git a/packages/livekit/src/background-task/memory/task-memory.ts b/packages/livekit/src/background-task/memory/task-memory.ts index 61106b567e..24bed21b1c 100644 --- a/packages/livekit/src/background-task/memory/task-memory.ts +++ b/packages/livekit/src/background-task/memory/task-memory.ts @@ -43,6 +43,7 @@ const TaskMemoryAllowedTools: readonly ToolSpecInput[] = [ "readFile", `writeToFile(${TaskMemoryFileUri})`, ]; +const TaskMemoryMaxSteps = 3; const TaskMemoryStoreFilePath = new URL(TaskMemoryFileUri).pathname; @@ -124,6 +125,7 @@ async function startTaskMemoryExtraction({ parentCwd, directive: prompts.taskMemory.buildExtractionDirective(existingMemory), tools: TaskMemoryAllowedTools, + maxSteps: TaskMemoryMaxSteps, }); const handle = await startForkAgent(agent); diff --git a/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts b/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts index 0dd53729fd..9a44e0a546 100644 --- a/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts +++ b/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts @@ -123,6 +123,28 @@ describe("TaskExecutor", () => { await executor.dispose(); }); + it("uses the background task's configured max steps", async () => { + const store = new FakeLiveKitStore([ + makeTask({ id: "task", status: "pending-model" }), + ]); + store.setMessages("task", [ + makeAssistantMessage([{ type: "step-start" }, { type: "step-start" }]), + ]); + const adaptor = makeAdaptor({ executeToolCall: vi.fn() }); + const executor = makeExecutor(store, adaptor, { maxSteps: 1 }); + + await executor.drain(); + + expect(store.readTask("task")).toMatchObject({ + status: "failed", + error: { + kind: "InternalError", + message: "The task failed to complete, max step count reached.", + }, + }); + await executor.dispose(); + }); + it("does not start duplicate running tasks for the same active task", async () => { const store = new FakeLiveKitStore([ makeTask({ id: "task", status: "pending-tool" }), diff --git a/packages/livekit/src/background-task/task-executor/task-executor.ts b/packages/livekit/src/background-task/task-executor/task-executor.ts index 4aba34fe98..4e16caff2f 100644 --- a/packages/livekit/src/background-task/task-executor/task-executor.ts +++ b/packages/livekit/src/background-task/task-executor/task-executor.ts @@ -604,8 +604,9 @@ class RunningTask { 0, stepCount - (this.taskState.baselineStepCount ?? 0), ); + const maxSteps = this.taskState.maxSteps ?? TaskExecutorMaxStep; - if (effectiveStepCount > TaskExecutorMaxStep) { + if (effectiveStepCount > maxSteps) { throw new Error("The task failed to complete, max step count reached."); } } diff --git a/packages/livekit/src/chat/live-chat-kit.ts b/packages/livekit/src/chat/live-chat-kit.ts index 12635a8623..09165c4a37 100644 --- a/packages/livekit/src/chat/live-chat-kit.ts +++ b/packages/livekit/src/chat/live-chat-kit.ts @@ -177,6 +177,7 @@ async function createBackgroundTaskFromForkAgent({ parentTaskId: agent.parentTaskId, tools: agent.tools, useCase: agent.label, + maxSteps: agent.maxSteps, baselineStepCount: agent.baselineStepCount, }); From 9e40991e8d6e9fe05908ca51d16ae5bdf2d68364 Mon Sep 17 00:00:00 2001 From: zhanba Date: Tue, 1 Sep 2026 20:08:20 +0800 Subject: [PATCH 2/5] fix(livekit): extract task memory once per compaction cycle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task memory extracted on warm-up, again on a token-increment rule, and once more synchronously right before compaction. That burned background forks and shortened the verbatim tail compaction is supposed to preserve, since the tail is whatever sits between the extraction boundary and the compaction point. Extract once when the context reaches 80% of the auto-compact threshold, mark the cycle done on success, and reset the flag when compaction finishes. Only failed attempts retry, capped at two per cycle, so a stale boundary cannot push compaction into its degraded attach path. Compaction no longer triggers an extraction itself and only waits up to 5s for an in-flight one to settle. 🤖 Generated with [Pochi](https://getpochi.com) | [Task](https://app.getpochi.com/share/p-8a63e58579ce42a1a32c7848ee4f84e9) Co-Authored-By: Pochi --- packages/common/src/base/constants.ts | 10 +- packages/common/src/base/memory.ts | 5 +- .../memory/__tests__/adaptors.test.ts | 140 ++++++++++++++++++ .../src/background-task/memory/task-memory.ts | 77 +++++----- .../__tests__/live-chat-kit-memory.test.ts | 6 +- .../livekit/src/chat/auto-compact-policy.ts | 13 ++ packages/livekit/src/chat/live-chat-kit.ts | 30 +++- .../src/lib/hooks/use-task-memory-state.ts | 4 +- 8 files changed, 228 insertions(+), 57 deletions(-) diff --git a/packages/common/src/base/constants.ts b/packages/common/src/base/constants.ts index 2b875e5250..325761834c 100644 --- a/packages/common/src/base/constants.ts +++ b/packages/common/src/base/constants.ts @@ -24,11 +24,13 @@ export const PochiClientHeader = "x-pochi-client"; export const PochiRequestUseCaseHeader = "x-pochi-request-use-case"; /** - * Task Memory thresholds — background extraction of session notes. + * Task Memory extracts once per compaction cycle. Everything after the + * extraction boundary survives compaction verbatim, so a second extraction in + * the same cycle would only shrink that tail; only failures are retried. */ -export const TaskMemoryInitTokenThreshold = 10_000; -export const TaskMemoryUpdateTokenIncrement = 5_000; -export const TaskMemoryUpdateToolCallThreshold = 3; +export const TaskMemoryExtractionThresholdRatio = 0.8; +export const MaxTaskMemoryExtractionAttemptsPerCycle = 2; +export const TaskMemoryFallbackCompactThreshold = 120_000; /** * Timeout (ms) for any single git operation. diff --git a/packages/common/src/base/memory.ts b/packages/common/src/base/memory.ts index 8e3d406563..f6053d0cb7 100644 --- a/packages/common/src/base/memory.ts +++ b/packages/common/src/base/memory.ts @@ -1,9 +1,6 @@ import type { AutoMemoryContext } from "./prompts/auto-memory"; export interface TaskMemoryState { - initialized: boolean; - lastExtractionTokens: number; - lastExtractionToolCalls: number; /** * UUID of the last message incorporated into memory.md by the most recent * successful extraction. Compaction uses this as the boundary to know @@ -17,6 +14,8 @@ export interface TaskMemoryState { * fork agent writes memory.md successfully. */ pendingExtractionMessageId?: string; + extractionAttemptsSinceCompact: number; + extractedSinceCompact?: boolean; isExtracting: boolean; extractionCount: number; activeTaskId?: string; diff --git a/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts b/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts index 644de43a40..ad1d68bb14 100644 --- a/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts +++ b/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts @@ -12,6 +12,9 @@ import { AutoMemoryAdaptor, type AutoMemoryManager } from "../auto-memory"; import { TaskMemoryAdaptor } from "../task-memory"; import { describe, expect, it, vi } from "vitest"; +/** Pinned so extraction fires at 16k tokens. */ +const TestCompactThreshold = 20_000; + describe("task-memory adaptor", () => { it("starts extraction from stream-finish usage before the main task completes", async () => { const store = new FakeStore([ @@ -32,6 +35,7 @@ describe("task-memory adaptor", () => { backgroundTask, parentTaskId: "parent", parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, }); await expect( @@ -79,6 +83,7 @@ describe("task-memory adaptor", () => { }, parentTaskId: "parent", parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, }); await expect( @@ -136,6 +141,7 @@ describe("task-memory adaptor", () => { }, parentTaskId: "parent", parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, }); await adaptor.update({ @@ -181,6 +187,7 @@ describe("task-memory adaptor", () => { }, parentTaskId: "parent", parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, }); await adaptor.update({ @@ -212,6 +219,137 @@ describe("task-memory adaptor", () => { activeTaskId: undefined, }); }); + + it("waits until the context approaches the auto-compact threshold", async () => { + const store = new FakeStore([ + makeTask({ + id: "parent", + status: "pending-tool", + background: false, + title: "Build shared runner", + }), + ]); + const adaptor = new TaskMemoryAdaptor({ + store: store as unknown as LiveKitStore, + backgroundTask: createTestBackgroundTask({ + store: store as unknown as LiveKitStore, + stateStore: new BackgroundTaskStateStore(), + }), + parentTaskId: "parent", + parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, + }); + + await expect( + adaptor.update({ + messages: makeParentMessages(), + contextWindowUsage: usage(10_000), + }), + ).resolves.toBe(false); + expect(store.backgroundTasks()).toHaveLength(0); + }); + + it("extracts only once per compaction cycle", async () => { + const store = new FakeStore([ + makeTask({ + id: "parent", + status: "pending-tool", + background: false, + title: "Build shared runner", + }), + ]); + let taskMemoryState: TaskMemoryState | undefined = { + extractionAttemptsSinceCompact: 1, + isExtracting: false, + extractionCount: 1, + extractedSinceCompact: true, + lastExtractionMessageId: "assistant-1", + }; + const adaptor = new TaskMemoryAdaptor({ + store: store as unknown as LiveKitStore, + backgroundTask: createTestBackgroundTask({ + store: store as unknown as LiveKitStore, + stateStore: new BackgroundTaskStateStore(), + }), + taskMemoryStateStore: { + get: () => taskMemoryState, + set: (state) => { + taskMemoryState = state; + }, + }, + parentTaskId: "parent", + parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, + }); + + await expect( + adaptor.update({ + messages: makeParentMessages(), + contextWindowUsage: usage(19_000), + }), + ).resolves.toBe(false); + expect(store.backgroundTasks()).toHaveLength(0); + + await adaptor.resetForNewCompactionCycle(); + await expect( + adaptor.update({ + messages: makeParentMessages(), + contextWindowUsage: usage(19_000), + }), + ).resolves.toBe(true); + }); + + it("retries a failed extraction once, then stops for the cycle", async () => { + const store = new FakeStore([ + makeTask({ + id: "parent", + status: "pending-tool", + background: false, + title: "Build shared runner", + }), + ]); + let taskMemoryState: TaskMemoryState | undefined = { + extractionAttemptsSinceCompact: 1, + isExtracting: false, + extractionCount: 0, + }; + const adaptor = new TaskMemoryAdaptor({ + store: store as unknown as LiveKitStore, + backgroundTask: createTestBackgroundTask({ + store: store as unknown as LiveKitStore, + stateStore: new BackgroundTaskStateStore(), + }), + taskMemoryStateStore: { + get: () => taskMemoryState, + set: (state) => { + taskMemoryState = state; + }, + }, + parentTaskId: "parent", + parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, + }); + + await expect( + adaptor.update({ + messages: makeParentMessages(), + contextWindowUsage: usage(17_000), + }), + ).resolves.toBe(true); + + taskMemoryState = { + ...(taskMemoryState as TaskMemoryState), + isExtracting: false, + activeTaskId: undefined, + }; + + await expect( + adaptor.update({ + messages: makeParentMessages(), + contextWindowUsage: usage(17_000), + }), + ).resolves.toBe(false); + }); }); describe("auto-memory adaptor", () => { @@ -708,6 +846,8 @@ function makeParentMessages(): Message[] { ] as Message[]; } + + function usage(tokens: number) { return { system: tokens, diff --git a/packages/livekit/src/background-task/memory/task-memory.ts b/packages/livekit/src/background-task/memory/task-memory.ts index 24bed21b1c..86a92e6a9a 100644 --- a/packages/livekit/src/background-task/memory/task-memory.ts +++ b/packages/livekit/src/background-task/memory/task-memory.ts @@ -25,16 +25,13 @@ const logger = getLogger("TaskMemory"); type ExtractionMetrics = { tokens: number; - toolCalls: number; trailingMessageId: string | undefined; }; type TaskMemoryExtractionResult = "pending" | "succeeded" | "failed"; const DefaultTaskMemoryState: TaskMemoryState = { - initialized: false, - lastExtractionTokens: 0, - lastExtractionToolCalls: 0, + extractionAttemptsSinceCompact: 0, isExtracting: false, extractionCount: 0, }; @@ -54,27 +51,32 @@ function getExtractionMetrics(data: { const last = data.messages.at(-1); return { tokens: computeTotalTokens(data.contextWindowUsage), - toolCalls: countToolCalls(data.messages), trailingMessageId: last?.id, }; } +export function resolveExtractionTrigger( + compactThreshold: number | undefined, +): number { + const base = + compactThreshold && compactThreshold > 0 + ? compactThreshold + : constants.TaskMemoryFallbackCompactThreshold; + return Math.round(base * constants.TaskMemoryExtractionThresholdRatio); +} + function shouldExtractTaskMemory( state: TaskMemoryState, metrics: ExtractionMetrics, + trigger: number, ): boolean { if (state.isExtracting) return false; - - if (!state.initialized) { - return metrics.tokens >= constants.TaskMemoryInitTokenThreshold; - } - - const tokenDelta = metrics.tokens - state.lastExtractionTokens; - const toolCallDelta = metrics.toolCalls - state.lastExtractionToolCalls; + if (metrics.tokens < trigger) return false; + if (state.extractedSinceCompact) return false; return ( - tokenDelta >= constants.TaskMemoryUpdateTokenIncrement && - toolCallDelta >= constants.TaskMemoryUpdateToolCallThreshold + state.extractionAttemptsSinceCompact < + constants.MaxTaskMemoryExtractionAttemptsPerCycle ); } @@ -84,10 +86,8 @@ function toExtractingState( ): TaskMemoryState { return { ...state, - initialized: true, isExtracting: true, - lastExtractionTokens: metrics.tokens, - lastExtractionToolCalls: metrics.toolCalls, + extractionAttemptsSinceCompact: state.extractionAttemptsSinceCompact + 1, pendingExtractionMessageId: metrics.trailingMessageId, }; } @@ -168,6 +168,7 @@ function resolveTaskMemoryExtractionState({ return { ...state, isExtracting: false, + extractedSinceCompact: succeeded ? true : state.extractedSinceCompact, extractionCount: succeeded ? state.extractionCount + 1 : state.extractionCount, @@ -215,19 +216,6 @@ function computeTotalTokens(usage?: ContextWindowUsage) { ); } -function countToolCalls(messages: UIMessage[]): number { - let count = 0; - for (const message of messages) { - if (message.role !== "assistant") continue; - for (const part of message.parts) { - if (isStaticToolUIPart(part)) { - count++; - } - } - } - return count; -} - type TaskMemoryAdaptorOptions = { store: LiveKitStore; backgroundTask: { @@ -238,6 +226,7 @@ type TaskMemoryAdaptorOptions = { parentTaskId: string; parentCwd: string | undefined | (() => string | undefined); isSubTask?: boolean; + getCompactThreshold?: () => number | undefined; }; export class TaskMemoryAdaptor { @@ -253,10 +242,11 @@ export class TaskMemoryAdaptor { return this.stateStore.get() ?? { ...DefaultTaskMemoryState }; } - resetTokenBaseline() { + resetForNewCompactionCycle() { return this.stateStore.set({ ...this.getState(), - lastExtractionTokens: 0, + extractionAttemptsSinceCompact: 0, + extractedSinceCompact: false, }); } @@ -268,16 +258,27 @@ export class TaskMemoryAdaptor { await this.settle(); const state = this.getState(); + const metrics = getExtractionMetrics(data); + const trigger = resolveExtractionTrigger( + this.options.getCompactThreshold?.(), + ); + if (!shouldExtractTaskMemory(state, metrics, trigger)) { + return false; + } + + return this.startExtraction(state, metrics, data.messages); + } + + private async startExtraction( + state: TaskMemoryState, + metrics: ExtractionMetrics, + messages: Message[], + ) { const task = this.options.store.query( makeTaskQuery(this.options.parentTaskId), ); if (!task) return false; - const metrics = getExtractionMetrics(data); - if (!shouldExtractTaskMemory(state, metrics)) { - return false; - } - try { const parentCwd = this.getParentCwd(); const memoryFile = this.options.store.query( @@ -290,7 +291,7 @@ export class TaskMemoryAdaptor { startForkAgent: (agent) => this.options.backgroundTask.startForkAgent(agent), parentTaskId: this.options.parentTaskId, - parentMessages: data.messages, + parentMessages: messages, parentCwd, parentTaskTitle: task.title ?? undefined, existingMemory: memoryFile?.content ?? undefined, diff --git a/packages/livekit/src/chat/__tests__/live-chat-kit-memory.test.ts b/packages/livekit/src/chat/__tests__/live-chat-kit-memory.test.ts index 31aece827b..1f838d89b5 100644 --- a/packages/livekit/src/chat/__tests__/live-chat-kit-memory.test.ts +++ b/packages/livekit/src/chat/__tests__/live-chat-kit-memory.test.ts @@ -394,7 +394,8 @@ describe("LiveChatKit memory lifecycle", () => { }); chatKit.chat.messages = [userMessage(), assistantMessage()]; - setLatestRequestSnapshot(chatKit, 20_000, 0); + // Above 80% of the 67k auto-compact threshold. + setLatestRequestSnapshot(chatKit, 60_000, 0); chatKit.chat.finish(assistantMessage()); await chatKit.drainBackgroundTasksAndSettleMemory(); @@ -496,7 +497,8 @@ describe("LiveChatKit memory lifecycle", () => { }); chatKit.chat.messages = [userMessage(), assistantMessage()]; - setLatestRequestSnapshot(chatKit, 20_000, 0); + // Above 80% of the 67k auto-compact threshold. + setLatestRequestSnapshot(chatKit, 60_000, 0); chatKit.chat.finish(assistantMessage()); await chatKit.drainBackgroundTasksAndSettleMemory(); diff --git a/packages/livekit/src/chat/auto-compact-policy.ts b/packages/livekit/src/chat/auto-compact-policy.ts index d39304851a..d291535dd5 100644 --- a/packages/livekit/src/chat/auto-compact-policy.ts +++ b/packages/livekit/src/chat/auto-compact-policy.ts @@ -28,6 +28,19 @@ function resolveContextWindow(llm: RequestData["llm"] | undefined): number { return declared || constants.DefaultContextWindow; } +export function resolveAutoCompactThreshold({ + llm, + effectiveContextWindow, +}: { + llm: RequestData["llm"] | undefined; + effectiveContextWindow?: number; +}): number { + return getAutoCompactThreshold( + resolveContextWindow(llm), + effectiveContextWindow, + ); +} + export function shouldAutoCompact({ messages, llm, diff --git a/packages/livekit/src/chat/live-chat-kit.ts b/packages/livekit/src/chat/live-chat-kit.ts index 09165c4a37..26a20f5fc1 100644 --- a/packages/livekit/src/chat/live-chat-kit.ts +++ b/packages/livekit/src/chat/live-chat-kit.ts @@ -49,6 +49,7 @@ import { toTaskError, toTaskGitInfo, toTaskStatus } from "../task"; import type { LiveKitStore, Message, Task } from "../types"; import { MaxConsecutiveAutoCompactFailures, + resolveAutoCompactThreshold, shouldAutoCompact, } from "./auto-compact-policy"; import { scheduleGenerateTitleJob } from "./background-job"; @@ -214,18 +215,20 @@ async function readRecentFilesForCompact( /** Polls until no extraction is in progress or `timeoutMs` elapses. */ async function settleTaskMemoryExtraction( - readTaskMemoryState: (() => TaskMemoryState | undefined) | undefined, + adaptor: TaskMemoryAdaptor | undefined, timeoutMs: number, ): Promise { - if (!readTaskMemoryState) return; - if (!readTaskMemoryState()?.isExtracting) return; + if (!adaptor?.getState().isExtracting) return; const start = Date.now(); while (Date.now() - start < timeoutMs) { - if (!readTaskMemoryState()?.isExtracting) return; + // Hosts without `waitForTaskDone` have no other chance to call `settle()`. + await adaptor.settle(); + if (!adaptor.getState().isExtracting) return; await new Promise((resolve) => setTimeout(resolve, TaskMemorySettlePollIntervalMs), ); } + logger.debug("Timed out waiting for the task-memory extraction to settle."); } async function runSideEffectSafely({ @@ -497,6 +500,7 @@ export class LiveChatKit< parentTaskId: taskId, parentCwd: defaultMemoryParentCwd, isSubTask, + getCompactThreshold: () => this.getAutoCompactThreshold(), }) : undefined; this.autoMemoryAdaptor = @@ -599,7 +603,7 @@ export class LiveChatKit< try { // Wait briefly so memory.md and boundary id are fresh. await settleTaskMemoryExtraction( - readEffectiveTaskMemoryState, + this.taskMemoryAdaptor, TaskMemorySettleTimeoutMs, ); const model = createModel({ llm: getters.getLLM() }); @@ -674,7 +678,7 @@ export class LiveChatKit< const { messages } = this.chat; // Wait briefly so memory.md and boundary id are fresh. await settleTaskMemoryExtraction( - readEffectiveTaskMemoryState, + this.taskMemoryAdaptor, TaskMemorySettleTimeoutMs, ); const model = createModel({ llm: getters.getLLM() }); @@ -1070,6 +1074,18 @@ export class LiveChatKit< this.backgroundTaskExecutor?.start(); } + private getAutoCompactThreshold(): number | undefined { + try { + return resolveAutoCompactThreshold({ + llm: this.getters.getLLM(), + effectiveContextWindow: this.getters.getEffectiveContextWindow?.(), + }); + } catch (error) { + logger.debug("Failed to resolve the auto-compact threshold", error); + return undefined; + } + } + private scheduleMemoryUpdate(data: { messages: Message[]; status?: string; @@ -1118,7 +1134,7 @@ export class LiveChatKit< onCompactFinish: ((success: boolean) => MaybePromise) | undefined, ) { if (success) { - await this.taskMemoryAdaptor?.resetTokenBaseline(); + await this.taskMemoryAdaptor?.resetForNewCompactionCycle(); } try { diff --git a/packages/vscode-webui/src/lib/hooks/use-task-memory-state.ts b/packages/vscode-webui/src/lib/hooks/use-task-memory-state.ts index af858030e2..9dfb71b8a3 100644 --- a/packages/vscode-webui/src/lib/hooks/use-task-memory-state.ts +++ b/packages/vscode-webui/src/lib/hooks/use-task-memory-state.ts @@ -4,9 +4,7 @@ import { threadSignal } from "@quilted/threads/signals"; import { useQuery } from "@tanstack/react-query"; const defaultTaskMemoryState: TaskMemoryState = { - initialized: false, - lastExtractionTokens: 0, - lastExtractionToolCalls: 0, + extractionAttemptsSinceCompact: 0, isExtracting: false, extractionCount: 0, }; From 98ed0573b454970d80af8cda35a7b140d12b45b9 Mon Sep 17 00:00:00 2001 From: zhanba Date: Tue, 1 Sep 2026 21:29:43 +0800 Subject: [PATCH 3/5] fix(livekit): throttle auto-memory extraction and detect its writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Auto-memory forked a full-context extraction agent at every completed turn, and judged success by attemptCompletion alone — so a systematically failing extraction never advanced the watermark and re-ran on every later turn. Require three new user turns before extracting, and treat a successful write by the fork as success while advancing the watermark even on failure. 🤖 Generated with [Pochi](https://getpochi.com) | [Task](https://app.getpochi.com/share/p-8a63e58579ce42a1a32c7848ee4f84e9) Co-Authored-By: Pochi --- .../memory/__tests__/adaptors.test.ts | 188 ++++++++++++++++-- .../src/background-task/memory/auto-memory.ts | 75 +++++-- .../__tests__/live-chat-kit-memory.test.ts | 12 +- 3 files changed, 242 insertions(+), 33 deletions(-) diff --git a/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts b/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts index ad1d68bb14..c06f8c9465 100644 --- a/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts +++ b/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts @@ -377,7 +377,7 @@ describe("auto-memory adaptor", () => { await expect( adaptor.update({ - messages: makeParentMessages(), + messages: makeAutoMemoryParentMessages(), status: "completed", }), ).resolves.toBe(true); @@ -440,7 +440,7 @@ describe("auto-memory adaptor", () => { await expect( adaptor.update({ - messages: makeParentMessages(), + messages: makeAutoMemoryParentMessages(), status: "completed", }), ).resolves.toBe(true); @@ -481,30 +481,21 @@ describe("auto-memory adaptor", () => { manager, }); + const messages = [ + ...makeAutoMemoryParentMessages(), + makeMemoryWriteMessage("write-memory"), + ]; + await expect( adaptor.update({ - messages: [ - { - id: "write-memory", - role: "assistant", - parts: [ - { - type: "tool-writeToFile", - toolCallId: "write", - state: "output-available", - input: { path: ".pochi/memory/index.md" }, - output: { success: true }, - }, - ], - }, - ] as Message[], + messages, status: "completed", }), ).resolves.toBe(false); expect(store.backgroundTasks()).toHaveLength(0); expect(autoMemoryState).toMatchObject({ - lastExtractionMessageCount: 1, + lastExtractionMessageCount: messages.length, isExtracting: false, }); expect(manager.beginDreamRun).toHaveBeenCalledTimes(1); @@ -596,7 +587,7 @@ describe("auto-memory adaptor", () => { }); await adaptor.update({ - messages: makeParentMessages(), + messages: makeAutoMemoryParentMessages(), status: "completed", }); const [extractionTask] = store.backgroundTasks(); @@ -634,8 +625,127 @@ describe("auto-memory adaptor", () => { success: true, }); }); + + it("waits for three new user turns before extracting", async () => { + const store = new FakeStore([ + makeTask({ id: "parent", status: "completed", background: false }), + ]); + const adaptor = makeAutoMemoryAdaptor({ store }); + + await expect( + adaptor.adaptor.update({ + messages: makeAutoMemoryParentMessages(2), + status: "completed", + }), + ).resolves.toBe(false); + expect(store.backgroundTasks()).toHaveLength(0); + + await expect( + adaptor.adaptor.update({ + messages: makeAutoMemoryParentMessages(3), + status: "completed", + }), + ).resolves.toBe(true); + expect(store.backgroundTasks()).toHaveLength(1); + }); + + it("treats a memory write by the extraction fork as success", async () => { + const store = new FakeStore([ + makeTask({ id: "parent", status: "completed", background: false }), + ]); + const { adaptor, getState } = makeAutoMemoryAdaptor({ store }); + const parentMessages = makeAutoMemoryParentMessages(); + + await adaptor.update({ messages: parentMessages, status: "completed" }); + + const [extractionTask] = store.backgroundTasks(); + store.setMessages(extractionTask.id, [ + ...parentMessages, + makeDirectiveMessage(), + makeMemoryWriteMessage("fork-write"), + ]); + // No attemptCompletion: the fork ran out of steps after writing memory. + store.updateTaskStatus(extractionTask.id, "failed"); + + await adaptor.settleAndMaybeContinue(); + + expect(getState()).toMatchObject({ + isExtracting: false, + extractionCount: 1, + lastExtractionMessageCount: parentMessages.length, + activeExtractionTaskId: undefined, + }); + }); + + it("advances the extraction mark after a failed extraction so it does not re-run", async () => { + const store = new FakeStore([ + makeTask({ id: "parent", status: "completed", background: false }), + ]); + const { adaptor, getState } = makeAutoMemoryAdaptor({ store }); + // The parent wrote memory itself in the first stretch, so the cloned prefix + // carries a successful write that must not be credited to the fork. + const earlierMessages = [ + ...makeAutoMemoryParentMessages(), + makeMemoryWriteMessage("parent-write"), + ]; + await adaptor.update({ messages: earlierMessages, status: "completed" }); + expect(store.backgroundTasks()).toHaveLength(0); + + const parentMessages = [ + ...earlierMessages, + ...makeAutoMemoryParentMessages(), + ]; + await adaptor.update({ messages: parentMessages, status: "completed" }); + + const [extractionTask] = store.backgroundTasks(); + store.setMessages(extractionTask.id, [ + ...parentMessages, + makeDirectiveMessage(), + ]); + store.updateTaskStatus(extractionTask.id, "failed"); + + await adaptor.settleAndMaybeContinue(); + + expect(getState()).toMatchObject({ + isExtracting: false, + extractionCount: 0, + lastExtractionMessageCount: parentMessages.length, + }); + + await expect( + adaptor.update({ messages: parentMessages, status: "completed" }), + ).resolves.toBe(false); + expect(store.backgroundTasks()).toHaveLength(1); + }); }); +function makeAutoMemoryAdaptor({ + store, + manager = makeAutoMemoryManager(), +}: { + store: FakeStore; + manager?: AutoMemoryManager; +}) { + let state: AutoMemoryTaskState | undefined; + const adaptor = new AutoMemoryAdaptor({ + store: store as unknown as LiveKitStore, + backgroundTask: createTestBackgroundTask({ + store: store as unknown as LiveKitStore, + stateStore: new BackgroundTaskStateStore(), + }), + autoMemoryStateStore: { + get: () => state, + set: (next) => { + state = next; + }, + }, + parentTaskId: "parent", + parentCwd: "/repo", + manager, + }); + return { adaptor, getState: () => state }; +} + const autoMemoryContext: AutoMemoryContext = { enabled: true, repoKey: "repo", @@ -846,6 +956,46 @@ function makeParentMessages(): Message[] { ] as Message[]; } +/** Auto-memory extraction only triggers once three new user turns exist. */ +function makeAutoMemoryParentMessages(userTurns = 3): Message[] { + return Array.from({ length: userTurns }, (_, index) => [ + { + id: `user-${index + 1}`, + role: "user", + parts: [{ type: "text", text: `please implement step ${index + 1}` }], + }, + { + id: `assistant-${index + 1}`, + role: "assistant", + parts: [{ type: "text", text: "done" }], + }, + ]).flat() as Message[]; +} + +function makeDirectiveMessage(): Message { + return { + id: "directive", + role: "user", + parts: [{ type: "text", text: "Extract durable long-term memories." }], + } as Message; +} + +function makeMemoryWriteMessage(id: string): Message { + return { + id, + role: "assistant", + parts: [ + { + type: "tool-writeToFile", + toolCallId: `write-${id}`, + state: "output-available", + input: { path: ".pochi/memory/index.md" }, + output: { success: true }, + }, + ], + } as Message; +} + function usage(tokens: number) { diff --git a/packages/livekit/src/background-task/memory/auto-memory.ts b/packages/livekit/src/background-task/memory/auto-memory.ts index 2d9cfe43eb..75f9e49582 100644 --- a/packages/livekit/src/background-task/memory/auto-memory.ts +++ b/packages/livekit/src/background-task/memory/auto-memory.ts @@ -11,7 +11,10 @@ import { import { type ToolSpecInput, ToolsByPermission } from "@getpochi/tools"; import { type UIMessage, getStaticToolName, isStaticToolUIPart } from "ai"; import { isPlainObject } from "remeda"; -import { makeTaskQuery } from "../../livestore/default-queries"; +import { + makeMessagesQuery, + makeTaskQuery, +} from "../../livestore/default-queries"; import type { LiveKitStore, Message } from "../../types"; import { type StartForkAgent, @@ -33,6 +36,7 @@ const MemoryReadToolNames = [ const MemoryAgentWriteToolNames = ["writeToFile", "applyDiff"] as const; const AutoMemoryMaxSteps = 5; const AutoMemoryDreamMaxSteps = 20; +const MinNewUserTurnsPerExtraction = 3; const MaxSessionTranscriptChars = 24_000; const MaxPartChars = 4_000; @@ -178,16 +182,25 @@ async function startAutoMemoryDream({ function resolveAutoMemoryExtractionState({ state, activeExtractionTask, + activeMessages, }: { state: AutoMemoryTaskState; activeExtractionTask: { status: string } | null | undefined; + activeMessages: readonly UIMessage[]; }): { nextState: AutoMemoryTaskState; success: boolean } | undefined { if (!state.isExtracting || !state.activeExtractionTaskId) return undefined; - if (activeExtractionTask && ActiveStatuses.has(activeExtractionTask.status)) { - return undefined; + + // The fork clones the parent conversation, so only messages past the cloned + // prefix belong to the extraction agent itself. + const wroteMemory = didExtractionWriteMemory( + activeMessages.slice(state.pendingExtractionMessageCount ?? 0), + ); + if (!wroteMemory) { + if (!activeExtractionTask) return undefined; + if (ActiveStatuses.has(activeExtractionTask.status)) return undefined; } - const success = activeExtractionTask?.status === "completed"; + const success = wroteMemory || activeExtractionTask?.status === "completed"; return { success, nextState: { @@ -196,10 +209,10 @@ function resolveAutoMemoryExtractionState({ extractionCount: success ? state.extractionCount + 1 : state.extractionCount, - lastExtractionMessageCount: success - ? (state.pendingExtractionMessageCount ?? - state.lastExtractionMessageCount) - : state.lastExtractionMessageCount, + // Advance even on failure: leaving the mark behind re-runs the same + // extraction on every following turn. + lastExtractionMessageCount: + state.pendingExtractionMessageCount ?? state.lastExtractionMessageCount, pendingExtractionMessageCount: undefined, activeExtractionTaskId: undefined, }, @@ -286,6 +299,39 @@ function didConversationWriteMemory( ); } +/** + * The extraction fork can only write inside the memory directory + * ({@link buildMemoryTools}), so any successful write means memory changed — + * regardless of whether the fork also reached attemptCompletion. + */ +function didExtractionWriteMemory(messages: readonly UIMessage[]): boolean { + return messages.some((message) => + message.parts.some((part) => { + if (!isStaticToolUIPart(part)) return false; + const toolName = getStaticToolName(part); + if (!MemoryAgentWriteToolNames.some((name) => name === toolName)) { + return false; + } + return isSuccessfulToolOutput(part); + }), + ); +} + +function countUserTurns(messages: readonly UIMessage[]): number { + return messages.filter(isUserTurn).length; +} + +function isUserTurn(message: UIMessage): boolean { + if (message.role !== "user") return false; + return message.parts.some( + (part) => + part.type === "text" && + part.text.trim().length > 0 && + !prompts.isSystemReminder(part.text) && + !prompts.isCompact(part.text), + ); +} + function serializeSessionTranscript(messages: readonly UIMessage[]): string { const chunks = messages.map((message, index) => { const parts = message.parts @@ -477,10 +523,10 @@ export class AutoMemoryAdaptor { } : undefined; - if ( - !state.isExtracting && - messageCount > state.lastExtractionMessageCount - ) { + const newUserTurns = countUserTurns( + data.messages.slice(state.lastExtractionMessageCount), + ); + if (!state.isExtracting && newUserTurns >= MinNewUserTurnsPerExtraction) { if ( didConversationWriteMemory( data.messages.slice(state.lastExtractionMessageCount), @@ -532,6 +578,11 @@ export class AutoMemoryAdaptor { makeTaskQuery(state.activeExtractionTaskId), ) : undefined, + activeMessages: state.activeExtractionTaskId + ? this.options.store + .query(makeMessagesQuery(state.activeExtractionTaskId)) + .map((row) => row.data as Message) + : [], }); if (extractionResolution) { await this.stateStore.set(extractionResolution.nextState); diff --git a/packages/livekit/src/chat/__tests__/live-chat-kit-memory.test.ts b/packages/livekit/src/chat/__tests__/live-chat-kit-memory.test.ts index 1f838d89b5..239acba3c4 100644 --- a/packages/livekit/src/chat/__tests__/live-chat-kit-memory.test.ts +++ b/packages/livekit/src/chat/__tests__/live-chat-kit-memory.test.ts @@ -393,7 +393,7 @@ describe("LiveChatKit memory lifecycle", () => { }, }); - chatKit.chat.messages = [userMessage(), assistantMessage()]; + chatKit.chat.messages = threeUserTurns(); // Above 80% of the 67k auto-compact threshold. setLatestRequestSnapshot(chatKit, 60_000, 0); chatKit.chat.finish(assistantMessage()); @@ -445,7 +445,7 @@ describe("LiveChatKit memory lifecycle", () => { }, }); - chatKit.chat.messages = [userMessage(), assistantMessage()]; + chatKit.chat.messages = threeUserTurns(); chatKit.chat.finish(assistantMessage()); await chatKit.drainBackgroundTasksAndSettleMemory(); @@ -869,6 +869,14 @@ function assistantMessage(): Message { } as unknown as Message; } +/** Auto-memory extraction only triggers once three new user turns exist. */ +function threeUserTurns(): Message[] { + return [1, 2, 3].flatMap((index) => [ + { ...userMessage(), id: `user-${index}` }, + { ...assistantMessage(), id: `assistant-${index}` }, + ]) as Message[]; +} + function assistantUserInputToolMessage( type: string, toolCallId: string, From 78f59126c4a8da3dc8d40cf073c314215566dc27 Mon Sep 17 00:00:00 2001 From: zhanba Date: Wed, 2 Sep 2026 21:54:31 +0800 Subject: [PATCH 4/5] fix(livekit): prevent stale task memory compaction --- .../memory/__tests__/adaptors.test.ts | 90 ++++++++++++++++- .../src/background-task/memory/task-memory.ts | 74 +++++++++++--- .../src/chat/__tests__/compact-task.test.ts | 99 ++++++++++++++++++- packages/livekit/src/chat/live-chat-kit.ts | 16 ++- packages/livekit/src/chat/llm/compact-task.ts | 29 ++++-- 5 files changed, 271 insertions(+), 37 deletions(-) diff --git a/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts b/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts index c06f8c9465..a6a934c500 100644 --- a/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts +++ b/packages/livekit/src/background-task/memory/__tests__/adaptors.test.ts @@ -162,7 +162,12 @@ describe("task-memory adaptor", () => { }); }); - it("marks extraction successful when the background task writes memory", async () => { + it.each([ + { output: { success: true }, extracted: true }, + { output: { error: "disk full" }, extracted: false }, + ])( + "publishes the extraction boundary only after a successful memory write ($extracted)", + async ({ output, extracted }) => { const store = new FakeStore([ makeTask({ id: "parent", @@ -206,18 +211,88 @@ describe("task-memory adaptor", () => { toolCallId: "write-1", state: "output-available", input: { path: TaskMemoryFileUri, content: "# Session Title" }, - output: { success: true }, + output: output as never, }, ], }, ] as Message[]); + store.updateTaskStatus(activeTaskId ?? "", "failed"); await expect(adaptor.settle()).resolves.toBe(true); expect(taskMemoryState).toMatchObject({ isExtracting: false, - extractionCount: 1, + extractionCount: extracted ? 1 : 0, + extractedSinceCompact: extracted ? true : undefined, + lastExtractionMessageId: extracted ? "assistant-1" : undefined, activeTaskId: undefined, }); + }, + ); + + it("discards an extraction that finishes after compaction starts", async () => { + const store = new FakeStore([ + makeTask({ + id: "parent", + status: "pending-tool", + background: false, + title: "Build shared runner", + }), + ]); + let taskMemoryState: TaskMemoryState | undefined; + const adaptor = new TaskMemoryAdaptor({ + store: store as unknown as LiveKitStore, + backgroundTask: createTestBackgroundTask({ + store: store as unknown as LiveKitStore, + stateStore: new BackgroundTaskStateStore(), + }), + taskMemoryStateStore: { + get: () => taskMemoryState, + set: (state) => { + taskMemoryState = state; + }, + }, + parentTaskId: "parent", + parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, + }); + + await adaptor.update({ + messages: makeParentMessages(), + contextWindowUsage: usage(20_000), + }); + const activeTaskId = taskMemoryState?.activeTaskId; + + await expect(adaptor.takeCompactionBoundaryMessageId()).resolves.toBe( + undefined, + ); + expect(taskMemoryState).toMatchObject({ + isExtracting: true, + pendingExtractionMessageId: undefined, + extractedSinceCompact: false, + }); + + store.setMessages(activeTaskId ?? "", [ + { + id: "write-memory", + role: "assistant", + parts: [ + { + type: "tool-writeToFile", + toolCallId: "write-1", + state: "output-available", + input: { path: TaskMemoryFileUri, content: "stale memory" }, + output: { success: true }, + }, + ], + }, + ] as Message[]); + + await adaptor.settle(); + expect(adaptor.getState()).toMatchObject({ + isExtracting: false, + extractedSinceCompact: false, + lastExtractionMessageId: undefined, + }); }); it("waits until the context approaches the auto-compact threshold", async () => { @@ -290,7 +365,14 @@ describe("task-memory adaptor", () => { ).resolves.toBe(false); expect(store.backgroundTasks()).toHaveLength(0); - await adaptor.resetForNewCompactionCycle(); + await expect(adaptor.takeCompactionBoundaryMessageId()).resolves.toBe( + "assistant-1", + ); + expect(taskMemoryState).toMatchObject({ + extractionAttemptsSinceCompact: 0, + extractedSinceCompact: false, + lastExtractionMessageId: undefined, + }); await expect( adaptor.update({ messages: makeParentMessages(), diff --git a/packages/livekit/src/background-task/memory/task-memory.ts b/packages/livekit/src/background-task/memory/task-memory.ts index 86a92e6a9a..e1e8e84df5 100644 --- a/packages/livekit/src/background-task/memory/task-memory.ts +++ b/packages/livekit/src/background-task/memory/task-memory.ts @@ -165,15 +165,18 @@ function resolveTaskMemoryExtractionState({ if (extractionResult === "pending") return undefined; const succeeded = extractionResult === "succeeded"; + // Compaction clears the pending boundary. The background task may still + // finish, but its result belongs to the previous cycle and must be ignored. + const usable = succeeded && state.pendingExtractionMessageId !== undefined; return { ...state, isExtracting: false, - extractedSinceCompact: succeeded ? true : state.extractedSinceCompact, + extractedSinceCompact: usable ? true : state.extractedSinceCompact, extractionCount: succeeded ? state.extractionCount + 1 : state.extractionCount, - lastExtractionMessageId: succeeded - ? (state.pendingExtractionMessageId ?? state.lastExtractionMessageId) + lastExtractionMessageId: usable + ? state.pendingExtractionMessageId : state.lastExtractionMessageId, pendingExtractionMessageId: undefined, activeTaskId: undefined, @@ -190,7 +193,14 @@ function getTaskMemoryExtractionResult( continue; } if (!part.input || typeof part.input !== "object") continue; - if ("path" in part.input && part.input.path === TaskMemoryFileUri) { + if ( + "path" in part.input && + part.input.path === TaskMemoryFileUri && + typeof part.output === "object" && + part.output !== null && + "success" in part.output && + part.output.success === true + ) { return "succeeded"; } } @@ -231,6 +241,8 @@ type TaskMemoryAdaptorOptions = { export class TaskMemoryAdaptor { private readonly stateStore: MemoryStateStore; + private state: TaskMemoryState | undefined; + private transitionQueue = Promise.resolve(); constructor(private readonly options: TaskMemoryAdaptorOptions) { this.stateStore = @@ -239,23 +251,39 @@ export class TaskMemoryAdaptor { } getState() { - return this.stateStore.get() ?? { ...DefaultTaskMemoryState }; + return this.state ?? this.stateStore.get() ?? { ...DefaultTaskMemoryState }; } - resetForNewCompactionCycle() { - return this.stateStore.set({ - ...this.getState(), - extractionAttemptsSinceCompact: 0, - extractedSinceCompact: false, + takeCompactionBoundaryMessageId() { + return this.enqueueTransition(async () => { + const state = this.getState(); + const boundaryMessageId = state.extractedSinceCompact + ? state.lastExtractionMessageId + : undefined; + await this.setState({ + ...state, + extractionAttemptsSinceCompact: 0, + extractedSinceCompact: false, + lastExtractionMessageId: undefined, + pendingExtractionMessageId: undefined, + }); + return boundaryMessageId; }); } - async update(data: { + update(data: { + messages: Message[]; + contextWindowUsage?: ContextWindowUsage; + }) { + return this.enqueueTransition(() => this.updateInner(data)); + } + + private async updateInner(data: { messages: Message[]; contextWindowUsage?: ContextWindowUsage; }) { if (this.options.isSubTask) return false; - await this.settle(); + await this.settleInner(); const state = this.getState(); const metrics = getExtractionMetrics(data); @@ -287,7 +315,7 @@ export class TaskMemoryAdaptor { const handle = await startTaskMemoryExtraction({ state, metrics, - setTaskMemoryState: (nextState) => this.stateStore.set(nextState), + setTaskMemoryState: (nextState) => this.setState(nextState), startForkAgent: (agent) => this.options.backgroundTask.startForkAgent(agent), parentTaskId: this.options.parentTaskId, @@ -305,6 +333,10 @@ export class TaskMemoryAdaptor { } async settle() { + return this.enqueueTransition(() => this.settleInner()); + } + + private async settleInner() { const state = this.getState(); if (!state.activeTaskId || !state.isExtracting) return false; @@ -317,10 +349,24 @@ export class TaskMemoryAdaptor { }); if (!nextState) return false; - await this.stateStore.set(nextState); + await this.setState(nextState); return true; } + private async setState(state: TaskMemoryState) { + await this.stateStore.set(state); + this.state = state; + } + + private enqueueTransition(run: () => T | PromiseLike): Promise { + const result = this.transitionQueue.then(run); + this.transitionQueue = result.then( + () => undefined, + () => undefined, + ); + return result; + } + private getParentCwd() { const { parentCwd } = this.options; return typeof parentCwd === "function" ? parentCwd() : parentCwd; diff --git a/packages/livekit/src/chat/__tests__/compact-task.test.ts b/packages/livekit/src/chat/__tests__/compact-task.test.ts index 01c68eba69..dc89639742 100644 --- a/packages/livekit/src/chat/__tests__/compact-task.test.ts +++ b/packages/livekit/src/chat/__tests__/compact-task.test.ts @@ -1,4 +1,14 @@ -import { describe, expect, it } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const generateTextMock = vi.hoisted(() => vi.fn()); + +vi.mock("ai", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + generateText: generateTextMock, + }; +}); import type { Message } from "../../types"; import { compactTask, @@ -156,6 +166,11 @@ describe("findInlineCompactAttachIndex", () => { }); describe("compactTask", () => { + beforeEach(() => { + generateTextMock.mockReset(); + generateTextMock.mockResolvedValue({ text: "fresh LLM summary" }); + }); + it("persists an inline compact block attached to a historical user message", async () => { const messages = [ userMsg("u0"), @@ -201,5 +216,87 @@ describe("compactTask", () => { text: expect.stringContaining(""), }, }); + expect(generateTextMock).not.toHaveBeenCalled(); + }); + + it("ignores old memory while the current extraction has no boundary", async () => { + const messages = [ + userMsg("u0"), + assistantMsg("a0"), + compactUserMsg("u1-compact"), + assistantMsg("a1"), + userMsg("u2", "latest request"), + ]; + const store = { + query: vi.fn(() => ({ content: "stale memory through u0" })), + commit: vi.fn(), + }; + + await compactTask({ + blobStore: {} as never, + taskId: "task-1", + storeId: "store-1", + model: {} as never, + messages, + inline: true, + store: store as never, + }); + + expect(store.query).not.toHaveBeenCalled(); + expect(generateTextMock).toHaveBeenCalledTimes(1); + expect(messages.at(-1)?.parts[0]).toMatchObject({ + type: "text", + text: expect.stringContaining("fresh LLM summary"), + }); + }); + + it("falls back non-inline when memory leaves a tail uncovered", async () => { + const messages = [ + userMsg("u0", "initial request"), + assistantMsg("a0"), + userMsg("u1"), + assistantMsg("a1", "latest assistant response"), + ]; + const store = { + query: vi.fn(() => ({ content: "stale memory" })), + commit: vi.fn(), + }; + + const summary = await compactTask({ + blobStore: {} as never, + taskId: "task-1", + storeId: "store-1", + model: {} as never, + messages, + taskMemoryBoundaryMessageId: "a0", + store: store as never, + }); + + expect(store.query).not.toHaveBeenCalled(); + expect(JSON.stringify(generateTextMock.mock.calls[0]?.[0]?.prompt)).toContain( + "latest assistant response", + ); + expect(summary).toContain("Previous conversation summary (4 messages)"); + }); + + it("uses non-inline memory when its boundary is the final message", async () => { + const messages = [userMsg("u0"), assistantMsg("a0")]; + const store = { + query: vi.fn(() => ({ content: "current memory" })), + commit: vi.fn(), + }; + + const summary = await compactTask({ + blobStore: {} as never, + taskId: "task-1", + storeId: "store-1", + model: {} as never, + messages, + taskMemoryBoundaryMessageId: "a0", + store: store as never, + }); + + expect(generateTextMock).not.toHaveBeenCalled(); + expect(summary).toContain("current memory"); }); }); diff --git a/packages/livekit/src/chat/live-chat-kit.ts b/packages/livekit/src/chat/live-chat-kit.ts index 26a20f5fc1..b65da62e7d 100644 --- a/packages/livekit/src/chat/live-chat-kit.ts +++ b/packages/livekit/src/chat/live-chat-kit.ts @@ -519,8 +519,6 @@ export class LiveChatKit< }) : undefined; - const readEffectiveTaskMemoryState = () => - this.taskMemoryAdaptor?.getState(); this.transport = new FlexibleChatTransport({ store, blobStore, @@ -607,6 +605,8 @@ export class LiveChatKit< TaskMemorySettleTimeoutMs, ); const model = createModel({ llm: getters.getLLM() }); + const taskMemoryBoundaryMessageId = + await this.taskMemoryAdaptor?.takeCompactionBoundaryMessageId(); if (isAutoCompact) { logger.info( `Auto-compact triggered (totalTokens=${ @@ -623,8 +623,7 @@ export class LiveChatKit< recentFiles: await readRecentFilesForCompact( getRecentFilesForCompact, ), - taskMemoryBoundaryMessageId: - readEffectiveTaskMemoryState()?.lastExtractionMessageId, + taskMemoryBoundaryMessageId, abortSignal, inline: true, store: this.store, @@ -682,6 +681,8 @@ export class LiveChatKit< TaskMemorySettleTimeoutMs, ); const model = createModel({ llm: getters.getLLM() }); + const taskMemoryBoundaryMessageId = + await this.taskMemoryAdaptor?.takeCompactionBoundaryMessageId(); const summary = await compactTask({ blobStore: this.blobStore, taskId: this.taskId, @@ -691,8 +692,7 @@ export class LiveChatKit< recentFiles: await readRecentFilesForCompact( getRecentFilesForCompact, ), - taskMemoryBoundaryMessageId: - readEffectiveTaskMemoryState()?.lastExtractionMessageId, + taskMemoryBoundaryMessageId, store: this.store, }); @@ -1133,10 +1133,6 @@ export class LiveChatKit< success: boolean, onCompactFinish: ((success: boolean) => MaybePromise) | undefined, ) { - if (success) { - await this.taskMemoryAdaptor?.resetForNewCompactionCycle(); - } - try { await onCompactFinish?.(success); } catch (notifyErr) { diff --git a/packages/livekit/src/chat/llm/compact-task.ts b/packages/livekit/src/chat/llm/compact-task.ts index 2bc45282d3..b1ffe04e2f 100644 --- a/packages/livekit/src/chat/llm/compact-task.ts +++ b/packages/livekit/src/chat/llm/compact-task.ts @@ -56,11 +56,21 @@ export async function compactTask({ const inlineAttachIndex = inline ? findInlineCompactAttachIndex(messages) : undefined; + const taskMemoryAttachIndex = inline + ? findVerbatimAttachIndex(messages, taskMemoryBoundaryMessageId) + : undefined; + const taskMemoryCoversAllMessages = + taskMemoryBoundaryMessageId === lastMessage.id; + const canUseTaskMemory = Boolean( + taskMemoryBoundaryMessageId && + (taskMemoryCoversAllMessages || taskMemoryAttachIndex !== undefined), + ); - // Prefer task memory if available + // Use memory only when it covers the whole transcript or an inline + // compaction can preserve the uncovered tail verbatim. let summaryText: string | undefined; let usedTaskMemory = false; - if (store) { + if (store && canUseTaskMemory) { const memoryFile = store.query( makeStoreFileQuery(TaskMemoryStoreFilePath), ); @@ -72,10 +82,13 @@ export async function compactTask({ // Fall back to LLM-generated summary if (!summaryText) { - const inputMessages = - inline && inlineAttachIndex !== undefined - ? messages.slice(0, inlineAttachIndex) - : messages.slice(0, -1); + let inputMessages = messages; + if (inline) { + inputMessages = + inlineAttachIndex !== undefined + ? messages.slice(0, inlineAttachIndex) + : messages.slice(0, -1); + } summaryText = await createSummary( blobStore, taskId, @@ -92,7 +105,7 @@ export async function compactTask({ if (inline) { // Preferred: attach at the boundary so trailing messages survive verbatim. const memoryAttachIndex = usedTaskMemory - ? findVerbatimAttachIndex(messages, taskMemoryBoundaryMessageId) + ? taskMemoryAttachIndex : undefined; const memoryAttachMessage = memoryAttachIndex !== undefined @@ -133,7 +146,7 @@ export async function compactTask({ // Non-inline: return the summary for callers seeding a fresh task. return prompts.inlineCompact( summaryText, - messages.length - 1, + messages.length, recentFileContext, ); } catch (err) { From 0b358465f9975edd2ab233a83b47e8f3fccfcaa1 Mon Sep 17 00:00:00 2001 From: zhanba Date: Wed, 2 Sep 2026 21:54:36 +0800 Subject: [PATCH 5/5] fix(livekit): preserve the final background task step --- .../__tests__/task-executor.test.ts | 122 +++++++++++++++++- .../task-executor/task-executor.ts | 53 ++++++-- 2 files changed, 161 insertions(+), 14 deletions(-) diff --git a/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts b/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts index 9a44e0a546..4fb5b8f52a 100644 --- a/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts +++ b/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts @@ -87,6 +87,12 @@ class MockLiveChatKit { markAsFailed(error: Error) { this.store.failTask(this.taskId, error.message); } + + markStartToolsExecution() {} + + markEndToolsExecution() { + this.store.setMessages(this.taskId, this.chat.messages); + } } type TestTask = { @@ -125,16 +131,63 @@ describe("TaskExecutor", () => { it("uses the background task's configured max steps", async () => { const store = new FakeLiveKitStore([ - makeTask({ id: "task", status: "pending-model" }), + makeTask({ id: "task", status: "pending-tool" }), ]); store.setMessages("task", [ - makeAssistantMessage([{ type: "step-start" }, { type: "step-start" }]), + makeAssistantMessage([ + { type: "step-start" }, + { type: "step-start" }, + makeToolPart("readFile", "read", { path: "a.ts" }), + ]), ]); - const adaptor = makeAdaptor({ executeToolCall: vi.fn() }); - const executor = makeExecutor(store, adaptor, { maxSteps: 1 }); + const executeToolCall = vi.fn(); + const adaptor = makeAdaptor({ executeToolCall }); + const executor = makeExecutor(store, adaptor, { + tools: ["readFile"], + maxSteps: 1, + }); + + await executor.drain(); + + expect(executeToolCall).not.toHaveBeenCalled(); + expect(mockState.instances[0].chat.sendMessageCalls).toBe(0); + expect(store.readTask("task")).toMatchObject({ + status: "failed", + error: { + kind: "InternalError", + message: "The task failed to complete, max step count reached.", + }, + }); + await executor.dispose(); + }); + + it("does not request another model step after processing the configured last step", async () => { + const store = new FakeLiveKitStore([ + makeTask({ id: "task", status: "pending-tool" }), + ]); + store.setMessages("task", [ + makeAssistantMessage([ + { type: "step-start" }, + makeToolPart("readFile", "read", { path: "a.ts" }), + ]), + ]); + const executeToolCall = vi.fn(async () => ({ content: "hello" })); + const adaptor = makeAdaptor({ executeToolCall }); + const executor = makeExecutor(store, adaptor, { + tools: ["readFile"], + maxSteps: 1, + }); await executor.drain(); + expect(executeToolCall).toHaveBeenCalledTimes(1); + expect(mockState.instances[0].chat.sendMessageCalls).toBe(0); + expect(getToolPart(store.readMessages("task").at(-1), "read")).toMatchObject( + { + state: "output-available", + output: { content: "hello" }, + }, + ); expect(store.readTask("task")).toMatchObject({ status: "failed", error: { @@ -145,6 +198,67 @@ describe("TaskExecutor", () => { await executor.dispose(); }); + it("does not overwrite a terminal state reached while the last tools run", async () => { + const store = new FakeLiveKitStore([ + makeTask({ id: "task", status: "pending-tool" }), + ]); + store.setMessages("task", [ + makeAssistantMessage([ + { type: "step-start" }, + makeToolPart("readFile", "read", { path: "a.ts" }), + ]), + ]); + const executeToolCall = vi.fn(async () => { + store.completeTask("task"); + return { content: "hello" }; + }); + const adaptor = makeAdaptor({ executeToolCall }); + const executor = makeExecutor(store, adaptor, { + tools: ["readFile"], + maxSteps: 1, + }); + + await executor.drain(); + + expect(executeToolCall).toHaveBeenCalledTimes(1); + expect(store.readTask("task")?.status).toBe("completed"); + expect(getToolPart(store.readMessages("task").at(-1), "read")).toMatchObject( + { + state: "output-available", + output: { content: "hello" }, + }, + ); + await executor.dispose(); + }); + + it("does not overwrite a completed response when observing the step limit", async () => { + const store = new FakeLiveKitStore([ + makeTask({ id: "task", status: "pending-model" }), + ]); + store.setMessages("task", [ + makeAssistantMessage([ + { type: "step-start" }, + { type: "step-start" }, + makeToolPart("attemptCompletion", "complete", { result: "done" }), + ]), + ]); + const ready = deferred(); + const adaptor: RunningTaskAdaptor = { + ...makeAdaptor({ executeToolCall: vi.fn() }), + waitUntilReady: () => ready.promise, + }; + const executor = makeExecutor(store, adaptor, { maxSteps: 1 }); + + executor.start(); + store.completeTask("task"); + ready.resolve(); + await executor.drain(); + + expect(mockState.instances[0].chat.sendMessageCalls).toBe(0); + expect(store.readTask("task")?.status).toBe("completed"); + await executor.dispose(); + }); + it("does not start duplicate running tasks for the same active task", async () => { const store = new FakeLiveKitStore([ makeTask({ id: "task", status: "pending-tool" }), diff --git a/packages/livekit/src/background-task/task-executor/task-executor.ts b/packages/livekit/src/background-task/task-executor/task-executor.ts index 4e16caff2f..5350878feb 100644 --- a/packages/livekit/src/background-task/task-executor/task-executor.ts +++ b/packages/livekit/src/background-task/task-executor/task-executor.ts @@ -89,6 +89,8 @@ type RunningTaskChat = { type RunningTaskChatKit = { chat: RunningTaskChat; task?: Task; + markStartToolsExecution: () => void; + markEndToolsExecution: () => void; markAsFailed: (error: Error) => MaybePromise; }; @@ -326,6 +328,13 @@ class RunningTask { if (stepResult === "finished") { return; } + const currentStatus = this.task?.status; + if ( + currentStatus === "completed" || + currentStatus === "pending-input" + ) { + return; + } if (stepResult === "retry") { this.retryCount += 1; @@ -338,6 +347,7 @@ class RunningTask { this.retryCount = 0; } + this.throwIfMaxStepReached(); await this.chat.sendMessage(); } } catch (error) { @@ -386,17 +396,16 @@ class RunningTask { } private async step(): Promise<"finished" | "next" | "retry"> { - this.throwIfMaxStepReached(); - const lastMessage = this.chat.messages.at(-1); if (!lastMessage) { throw new Error("No messages in the task chat."); } - return ( - (await this.processMessage(lastMessage)) ?? - (await this.processToolCalls(lastMessage)) - ); + const messageResult = await this.processMessage(lastMessage); + if (messageResult) return messageResult; + + this.throwIfMaxStepExceeded(); + return this.processToolCalls(lastMessage); } private async processMessage( @@ -484,7 +493,17 @@ class RunningTask { }); } - await this.toolCallQueue.start(); + const chatKit = this.chatKit; + if (!chatKit) { + throw new Error("Task chat is not initialized."); + } + + chatKit.markStartToolsExecution(); + try { + await this.toolCallQueue.start(); + } finally { + chatKit.markEndToolsExecution(); + } return "next"; } @@ -599,6 +618,22 @@ class RunningTask { } private throwIfMaxStepReached() { + const { effectiveStepCount, maxSteps } = this.getStepLimitState(); + + if (effectiveStepCount >= maxSteps) { + throw new Error("The task failed to complete, max step count reached."); + } + } + + private throwIfMaxStepExceeded() { + const { effectiveStepCount, maxSteps } = this.getStepLimitState(); + + if (effectiveStepCount > maxSteps) { + throw new Error("The task failed to complete, max step count reached."); + } + } + + private getStepLimitState() { const stepCount = countStepStarts(this.chat.messages); const effectiveStepCount = Math.max( 0, @@ -606,9 +641,7 @@ class RunningTask { ); const maxSteps = this.taskState.maxSteps ?? TaskExecutorMaxStep; - if (effectiveStepCount > maxSteps) { - throw new Error("The task failed to complete, max step count reached."); - } + return { effectiveStepCount, maxSteps }; } }