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/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/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/__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..a6a934c500 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( @@ -50,6 +54,7 @@ describe("task-memory adaptor", () => { expect(await stateStore.read(task.id)).toMatchObject({ parentTaskId: "parent", useCase: "task-memory", + maxSteps: 3, }); }); @@ -78,6 +83,7 @@ describe("task-memory adaptor", () => { }, parentTaskId: "parent", parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, }); await expect( @@ -135,6 +141,7 @@ describe("task-memory adaptor", () => { }, parentTaskId: "parent", parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, }); await adaptor.update({ @@ -155,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", @@ -180,6 +192,7 @@ describe("task-memory adaptor", () => { }, parentTaskId: "parent", parentCwd: "/repo", + getCompactThreshold: () => TestCompactThreshold, }); await adaptor.update({ @@ -198,18 +211,226 @@ 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 () => { + 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 expect(adaptor.takeCompactionBoundaryMessageId()).resolves.toBe( + "assistant-1", + ); + expect(taskMemoryState).toMatchObject({ + extractionAttemptsSinceCompact: 0, + extractedSinceCompact: false, + lastExtractionMessageId: undefined, + }); + 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); }); }); @@ -238,7 +459,7 @@ describe("auto-memory adaptor", () => { await expect( adaptor.update({ - messages: makeParentMessages(), + messages: makeAutoMemoryParentMessages(), status: "completed", }), ).resolves.toBe(true); @@ -252,6 +473,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/**)", @@ -300,7 +522,7 @@ describe("auto-memory adaptor", () => { await expect( adaptor.update({ - messages: makeParentMessages(), + messages: makeAutoMemoryParentMessages(), status: "completed", }), ).resolves.toBe(true); @@ -341,30 +563,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); @@ -456,7 +669,7 @@ describe("auto-memory adaptor", () => { }); await adaptor.update({ - messages: makeParentMessages(), + messages: makeAutoMemoryParentMessages(), status: "completed", }); const [extractionTask] = store.backgroundTasks(); @@ -478,6 +691,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"); @@ -490,8 +707,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", @@ -547,6 +883,7 @@ function createTestBackgroundTask({ parentTaskId: agent.parentTaskId, tools: agent.tools, useCase: agent.label, + maxSteps: agent.maxSteps, baselineStepCount: agent.baselineStepCount, }); store.commit( @@ -701,6 +1038,48 @@ 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) { return { system: tokens, diff --git a/packages/livekit/src/background-task/memory/auto-memory.ts b/packages/livekit/src/background-task/memory/auto-memory.ts index b6075fd231..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, @@ -31,6 +34,9 @@ const MemoryReadToolNames = [ "searchFiles", ] as const; const MemoryAgentWriteToolNames = ["writeToFile", "applyDiff"] as const; +const AutoMemoryMaxSteps = 5; +const AutoMemoryDreamMaxSteps = 20; +const MinNewUserTurnsPerExtraction = 3; const MaxSessionTranscriptChars = 24_000; const MaxPartChars = 4_000; @@ -76,6 +82,7 @@ async function startAutoMemoryExtraction({ previousMessageCount, }), tools: buildMemoryTools(context), + maxSteps: AutoMemoryMaxSteps, }); const handle = await startForkAgent(agent); @@ -148,6 +155,7 @@ async function startAutoMemoryDream({ sessions, }), tools: buildMemoryTools(run.context), + maxSteps: AutoMemoryDreamMaxSteps, }); const handle = await startForkAgent(agent); @@ -174,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: { @@ -192,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, }, @@ -282,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 @@ -473,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), @@ -528,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/background-task/memory/task-memory.ts b/packages/livekit/src/background-task/memory/task-memory.ts index 61106b567e..e1e8e84df5 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, }; @@ -43,6 +40,7 @@ const TaskMemoryAllowedTools: readonly ToolSpecInput[] = [ "readFile", `writeToFile(${TaskMemoryFileUri})`, ]; +const TaskMemoryMaxSteps = 3; const TaskMemoryStoreFilePath = new URL(TaskMemoryFileUri).pathname; @@ -53,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 ); } @@ -83,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, }; } @@ -124,6 +125,7 @@ async function startTaskMemoryExtraction({ parentCwd, directive: prompts.taskMemory.buildExtractionDirective(existingMemory), tools: TaskMemoryAllowedTools, + maxSteps: TaskMemoryMaxSteps, }); const handle = await startForkAgent(agent); @@ -163,14 +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: 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, @@ -187,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"; } } @@ -213,19 +226,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: { @@ -236,10 +236,13 @@ type TaskMemoryAdaptorOptions = { parentTaskId: string; parentCwd: string | undefined | (() => string | undefined); isSubTask?: boolean; + getCompactThreshold?: () => number | undefined; }; export class TaskMemoryAdaptor { private readonly stateStore: MemoryStateStore; + private state: TaskMemoryState | undefined; + private transitionQueue = Promise.resolve(); constructor(private readonly options: TaskMemoryAdaptorOptions) { this.stateStore = @@ -248,34 +251,62 @@ export class TaskMemoryAdaptor { } getState() { - return this.stateStore.get() ?? { ...DefaultTaskMemoryState }; + return this.state ?? this.stateStore.get() ?? { ...DefaultTaskMemoryState }; } - resetTokenBaseline() { - return this.stateStore.set({ - ...this.getState(), - lastExtractionTokens: 0, + 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); + 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( @@ -284,11 +315,11 @@ 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, - parentMessages: data.messages, + parentMessages: messages, parentCwd, parentTaskTitle: task.title ?? undefined, existingMemory: memoryFile?.content ?? undefined, @@ -302,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; @@ -314,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/background-task/task-executor/__tests__/task-executor.test.ts b/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts index 0dd53729fd..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 = { @@ -123,6 +129,136 @@ describe("TaskExecutor", () => { await executor.dispose(); }); + it("uses the background task's configured max steps", async () => { + const store = new FakeLiveKitStore([ + makeTask({ id: "task", status: "pending-tool" }), + ]); + store.setMessages("task", [ + makeAssistantMessage([ + { type: "step-start" }, + { type: "step-start" }, + makeToolPart("readFile", "read", { path: "a.ts" }), + ]), + ]); + 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: { + kind: "InternalError", + message: "The task failed to complete, max step count reached.", + }, + }); + 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 4aba34fe98..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,15 +618,30 @@ 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, stepCount - (this.taskState.baselineStepCount ?? 0), ); + const maxSteps = this.taskState.maxSteps ?? TaskExecutorMaxStep; - if (effectiveStepCount > TaskExecutorMaxStep) { - throw new Error("The task failed to complete, max step count reached."); - } + return { effectiveStepCount, maxSteps }; } } 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/__tests__/live-chat-kit-memory.test.ts b/packages/livekit/src/chat/__tests__/live-chat-kit-memory.test.ts index 31aece827b..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,8 +393,9 @@ describe("LiveChatKit memory lifecycle", () => { }, }); - chatKit.chat.messages = [userMessage(), assistantMessage()]; - setLatestRequestSnapshot(chatKit, 20_000, 0); + chatKit.chat.messages = threeUserTurns(); + // Above 80% of the 67k auto-compact threshold. + setLatestRequestSnapshot(chatKit, 60_000, 0); chatKit.chat.finish(assistantMessage()); await chatKit.drainBackgroundTasksAndSettleMemory(); @@ -444,7 +445,7 @@ describe("LiveChatKit memory lifecycle", () => { }, }); - chatKit.chat.messages = [userMessage(), assistantMessage()]; + chatKit.chat.messages = threeUserTurns(); 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(); @@ -867,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, 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 12635a8623..b65da62e7d 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"; @@ -177,6 +178,7 @@ async function createBackgroundTaskFromForkAgent({ parentTaskId: agent.parentTaskId, tools: agent.tools, useCase: agent.label, + maxSteps: agent.maxSteps, baselineStepCount: agent.baselineStepCount, }); @@ -213,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({ @@ -496,6 +500,7 @@ export class LiveChatKit< parentTaskId: taskId, parentCwd: defaultMemoryParentCwd, isSubTask, + getCompactThreshold: () => this.getAutoCompactThreshold(), }) : undefined; this.autoMemoryAdaptor = @@ -514,8 +519,6 @@ export class LiveChatKit< }) : undefined; - const readEffectiveTaskMemoryState = () => - this.taskMemoryAdaptor?.getState(); this.transport = new FlexibleChatTransport({ store, blobStore, @@ -598,10 +601,12 @@ 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() }); + const taskMemoryBoundaryMessageId = + await this.taskMemoryAdaptor?.takeCompactionBoundaryMessageId(); if (isAutoCompact) { logger.info( `Auto-compact triggered (totalTokens=${ @@ -618,8 +623,7 @@ export class LiveChatKit< recentFiles: await readRecentFilesForCompact( getRecentFilesForCompact, ), - taskMemoryBoundaryMessageId: - readEffectiveTaskMemoryState()?.lastExtractionMessageId, + taskMemoryBoundaryMessageId, abortSignal, inline: true, store: this.store, @@ -673,10 +677,12 @@ 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() }); + const taskMemoryBoundaryMessageId = + await this.taskMemoryAdaptor?.takeCompactionBoundaryMessageId(); const summary = await compactTask({ blobStore: this.blobStore, taskId: this.taskId, @@ -686,8 +692,7 @@ export class LiveChatKit< recentFiles: await readRecentFilesForCompact( getRecentFilesForCompact, ), - taskMemoryBoundaryMessageId: - readEffectiveTaskMemoryState()?.lastExtractionMessageId, + taskMemoryBoundaryMessageId, store: this.store, }); @@ -1069,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; @@ -1116,10 +1133,6 @@ export class LiveChatKit< success: boolean, onCompactFinish: ((success: boolean) => MaybePromise) | undefined, ) { - if (success) { - await this.taskMemoryAdaptor?.resetTokenBaseline(); - } - 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) { 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, };