From fdac4dce03d329f60f621f79ee3191db27430b5d Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:18:47 +0000 Subject: [PATCH] fix: invalidate cached memory context on memory changes and context reset The per-session memory context cache (index snapshot + hot memories) was cleared only at compaction boundaries, so deleted or renamed memory files kept appearing in the advertised index for the rest of the segment. Now any MemoryService change event invalidates all live sessions' caches, and resetContext invalidates the affected session. Rebuild stays lazy on the next stream. Tradeoff: a memory write mid-segment costs one prompt cache miss on the next request; correct index beats cache stability. --- src/node/orpc/router.ts | 19 +- .../agentSession.memoryContext.test.ts | 111 +++ .../agentSession.startupAutoRetry.test.ts | 67 +- src/node/services/agentSession.testHarness.ts | 2 + src/node/services/agentSession.ts | 109 ++- src/node/services/aiService.test.ts | 591 ++++++++++++++- src/node/services/aiService.ts | 297 +++++++- src/node/services/coreServices.ts | 7 +- src/node/services/memoryService.ts | 45 +- src/node/services/streamManager.test.ts | 249 ++++++ src/node/services/streamManager.ts | 184 ++++- src/node/services/tools/memory.ts | 16 +- src/node/services/workspaceService.test.ts | 712 ++++++++++++++++++ src/node/services/workspaceService.ts | 132 +++- 14 files changed, 2457 insertions(+), 84 deletions(-) diff --git a/src/node/orpc/router.ts b/src/node/orpc/router.ts index 37e1d1c63e..a8d9d927d7 100644 --- a/src/node/orpc/router.ts +++ b/src/node/orpc/router.ts @@ -3879,13 +3879,18 @@ export const router = (authToken?: string) => { error: "Project memory is unavailable: no project is associated with this session", }; } - await context.memoryMetaService.setPinned( - memoryLogicalKey(scope, relPath, { - projectPath: resolved.projectPath, - workspaceId: input.workspaceId ?? "", - }), - input.pinned - ); + // Route through MemoryService so the pin change emits a change + // event and invalidates session-cached memory contexts. + try { + await context.memoryService.setPinned( + resolved.scopeCtx, + input.path, + input.pinned, + "user" + ); + } catch (error) { + return { success: false as const, error: getErrorMessage(error) }; + } return { success: true as const, data: undefined }; }), consolidationStatus: t diff --git a/src/node/services/agentSession.memoryContext.test.ts b/src/node/services/agentSession.memoryContext.test.ts index 3e33d93573..4c26c997c1 100644 --- a/src/node/services/agentSession.memoryContext.test.ts +++ b/src/node/services/agentSession.memoryContext.test.ts @@ -214,6 +214,117 @@ describe("AgentSession memory context", () => { } }); + test("recomputes the context after invalidateMemoryContext (memory change / context reset)", async () => { + using sessionDir = new DisposableTempDir("agent-session-memory-context-invalidate"); + const { historyService, cleanup } = await createTestHistoryService(); + historyCleanup = cleanup; + + let indexEntries = [{ path: "/memories/global/deleted.md", description: "stale" }]; + const buildMemorySessionContext = mock(() => + Promise.resolve({ indexEntries, hotMemoriesBlock: null }) + ); + const session = createSession({ + historyService, + sessionDir: sessionDir.path, + buildMemorySessionContext, + }); + const priv = session as unknown as PrivateSessionAccess; + + try { + expect((await priv.resolveMemoryContext("test-model"))?.indexEntries).toHaveLength(1); + + indexEntries = []; + session.invalidateMemoryContext(); + + expect((await priv.resolveMemoryContext("test-model"))?.indexEntries).toHaveLength(0); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(2); + } finally { + session.dispose(); + } + }); + + test("rebuilds mid-request when invalidation arrives during an in-flight build", async () => { + using sessionDir = new DisposableTempDir("agent-session-memory-context-race"); + const { historyService, cleanup } = await createTestHistoryService(); + historyCleanup = cleanup; + + let indexEntries = [{ path: "/memories/global/deleted.md", description: "stale" }]; + const pendingReleases: Array<() => void> = []; + const buildMemorySessionContext = mock(async () => { + const entries = indexEntries; + await new Promise((resolve) => { + pendingReleases.push(resolve); + }); + return { indexEntries: entries, hotMemoriesBlock: null }; + }); + const session = createSession({ + historyService, + sessionDir: sessionDir.path, + buildMemorySessionContext, + }); + const priv = session as unknown as PrivateSessionAccess; + + try { + const firstResolve = priv.resolveMemoryContext("test-model"); + // The memory change lands while the first build is still awaited. + indexEntries = []; + session.invalidateMemoryContext(); + pendingReleases.shift()?.(); + // The stale build result triggers a rebuild; release it once it starts. + while (pendingReleases.length === 0) { + await new Promise((resolve) => setTimeout(resolve, 0)); + } + pendingReleases.shift()?.(); + + // The current request already sees the post-invalidation state. + expect((await firstResolve)?.indexEntries).toHaveLength(0); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(2); + + // The clean rebuild was cached: no further build on the next resolve. + const secondResolve = priv.resolveMemoryContext("test-model"); + expect((await secondResolve)?.indexEntries).toHaveLength(0); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(2); + } finally { + session.dispose(); + } + }); + + test("persistent invalidation churn stops rebuilding at the attempt bound and stays uncached", async () => { + using sessionDir = new DisposableTempDir("agent-session-memory-context-churn"); + const { historyService, cleanup } = await createTestHistoryService(); + historyCleanup = cleanup; + + let version = 0; + // Every build observes a concurrent invalidation before it completes. + const buildMemorySessionContext = mock(() => { + version++; + session.invalidateMemoryContext(); + return Promise.resolve({ + indexEntries: [{ path: `/memories/global/v${version}.md`, description: `v${version}` }], + hotMemoriesBlock: null, + }); + }); + const session = createSession({ + historyService, + sessionDir: sessionDir.path, + buildMemorySessionContext, + }); + const priv = session as unknown as PrivateSessionAccess; + + try { + // Three attempts (the bound), then the freshest snapshot is served. + const first = await priv.resolveMemoryContext("test-model"); + expect(first?.indexEntries[0]?.path).toBe("/memories/global/v3.md"); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(3); + + // Nothing was cached, so the next request rebuilds. + await priv.resolveMemoryContext("test-model"); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(6); + } finally { + session.dispose(); + } + }); + test("recomputes the context after a compaction boundary is consumed", async () => { using sessionDir = new DisposableTempDir("agent-session-memory-context-compaction"); const { historyService, cleanup } = await createTestHistoryService(); diff --git a/src/node/services/agentSession.startupAutoRetry.test.ts b/src/node/services/agentSession.startupAutoRetry.test.ts index fb7e547db4..99059d772f 100644 --- a/src/node/services/agentSession.startupAutoRetry.test.ts +++ b/src/node/services/agentSession.startupAutoRetry.test.ts @@ -15,6 +15,7 @@ import type { WorkspaceMetadata } from "@/common/types/workspace"; import { Ok } from "@/common/types/result"; import { GOAL_CONTINUATION_KIND } from "@/constants/goals"; import { WORKSPACE_DEFAULTS } from "@/constants/workspaceDefaults"; +import { CONTEXT_BOUNDARY_KINDS } from "@/common/constants/contextBoundary"; interface AutoRetryResumeRequest { options: SendMessageOptions; @@ -33,7 +34,10 @@ interface SessionBundle { cleanup: () => Promise; } -async function createSessionBundle(workspaceId: string): Promise { +async function createSessionBundle( + workspaceId: string, + getContextResetState?: () => { resetting: boolean; boundaryEpoch: number } +): Promise { const workspaceMetadata: WorkspaceMetadata = { id: workspaceId, name: workspaceId, @@ -56,6 +60,7 @@ async function createSessionBundle(workspaceId: string): Promise initStateManagerOverrides: { replayInit: mock(() => Promise.resolve()), }, + getContextResetState, captureEvents: true, }); } @@ -124,6 +129,66 @@ describe("AgentSession startup auto-retry recovery", () => { session.dispose(); }); + test("defers the retry when a context reset boundary lands after the history reads", async () => { + const workspaceId = "startup-retry-reset-overlap"; + let boundaryEpoch = 0; + const { session, historyService, events, cleanup } = await createSessionBundle( + workspaceId, + () => ({ resetting: false, boundaryEpoch }) + ); + cleanups.push(cleanup); + + const appendResult = await historyService.appendToHistory( + workspaceId, + createMuxMessage("user-1", "user", "interrupted turn", { timestamp: Date.now() }) + ); + expect(appendResult.success).toBe(true); + + // Simulate resetContext committing its boundary between this check's + // history reads and its retry decision: the reads return the pre-reset + // tail, then the boundary lands and the epoch advances. Scheduling the + // retry from that stale tail would resume the context the user reset. + const realGetLastMessages = historyService.getLastMessages.bind(historyService); + let overlapped = false; + spyOn(historyService, "getLastMessages").mockImplementation(async (readWorkspaceId, count) => { + const result = await realGetLastMessages(readWorkspaceId, count); + if (!overlapped) { + overlapped = true; + const boundaryAppend = await historyService.appendToHistory( + readWorkspaceId, + createMuxMessage("reset-boundary", "assistant", "", { + timestamp: Date.now(), + contextBoundaryKind: CONTEXT_BOUNDARY_KINDS.RESET, + }) + ); + expect(boundaryAppend.success).toBe(true); + boundaryEpoch += 1; + } + return result; + }); + + session.ensureStartupAutoRetryCheck(); + + // The deferred check re-runs when idle and settles on the post-boundary + // history; wait for the pipeline to finish before asserting. + const priv = session as unknown as { + startupAutoRetryCheckScheduled: boolean; + startupAutoRetryCheckPromise: Promise | null; + }; + const deadline = Date.now() + 5000; + while (!(priv.startupAutoRetryCheckScheduled && priv.startupAutoRetryCheckPromise === null)) { + if (Date.now() > deadline) { + throw new Error("Startup auto-retry check did not settle"); + } + await new Promise((resolve) => setTimeout(resolve, 5)); + } + + expect(events.find((event) => event.type === "auto-retry-scheduled")).toBeUndefined(); + expect(session.hasPendingAutoRetry()).toBe(false); + + session.dispose(); + }); + test("startup auto-retry reuses workspace-turn metadata from the retry user message", async () => { const workspaceId = "startup-retry-workspace-turn-metadata"; const { session, historyService, aiService, cleanup } = await createSessionBundle(workspaceId); diff --git a/src/node/services/agentSession.testHarness.ts b/src/node/services/agentSession.testHarness.ts index d332353c5d..d991a57a8c 100644 --- a/src/node/services/agentSession.testHarness.ts +++ b/src/node/services/agentSession.testHarness.ts @@ -66,6 +66,7 @@ export interface AgentSessionHarnessOptions { backgroundProcessManager?: BackgroundProcessManager; backgroundProcessManagerOverrides?: Partial; onCompactionComplete?: (metadata: CompactionCompletionMetadata) => void; + getContextResetState?: () => { resetting: boolean; boundaryEpoch: number }; captureEvents?: boolean; } @@ -108,6 +109,7 @@ export async function createAgentSessionHarness( initStateManager, backgroundProcessManager, onCompactionComplete: options.onCompactionComplete, + getContextResetState: options.getContextResetState, }); const events: WorkspaceChatMessage[] = []; diff --git a/src/node/services/agentSession.ts b/src/node/services/agentSession.ts index 72bfb6af0f..12dee85d33 100644 --- a/src/node/services/agentSession.ts +++ b/src/node/services/agentSession.ts @@ -335,6 +335,13 @@ interface AgentSessionOptions { onIdleCompactionOutcome?: (success: boolean) => void; /** Called when post-compaction context state may have changed (plan/file edits) */ onPostCompactionStateChange?: () => void; + /** + * Snapshot of WorkspaceService.resetContext activity for this workspace. + * boundaryEpoch counts appended reset boundaries, so a boundary landing + * between the startup auto-retry check's history reads and its retry + * decision is detectable even after the reset finished. + */ + getContextResetState?: () => { resetting: boolean; boundaryEpoch: number }; } enum TurnPhase { @@ -361,6 +368,7 @@ export class AgentSession { private readonly workspaceGoalService?: WorkspaceGoalService; private readonly keepBackgroundProcesses: boolean; private readonly onPostCompactionStateChange?: () => void; + private readonly getContextResetState: () => { resetting: boolean; boundaryEpoch: number }; private readonly emitter = new EventEmitter(); private readonly aiListeners: Array<{ event: string; handler: (...args: unknown[]) => void }> = []; @@ -399,6 +407,7 @@ export class AgentSession { private startupRecoveryPromise: Promise | null = null; private startupAutoRetryCheckScheduled = false; private startupAutoRetryCheckPromise: Promise | null = null; + private startupAutoRetryCheckInFlight = false; private startupAutoRetryHistoryReadFailureCount = 0; private startupAutoRetryDeferredRetryDelayMs = 0; private autoRetryEnabledPreference: boolean | null = null; @@ -444,10 +453,15 @@ export class AgentSession { * the memory tool description plus an optional hot-memories block, keyed by * model because the hot set is token-budgeted with the active model's * tokenizer. Index-only entries can be upgraded once final tool policy keeps - * the memory tool; compaction clears the map so repeated turns keep - * prompt-cache-stable bytes without preserving stale files forever. + * the memory tool. Cleared at compaction boundaries, on context reset, and + * on any memory change so requests never advertise stale memory files. */ private readonly memoryContextByModelString = new Map(); + /** + * Bumped on every invalidation so builds that were already in flight when a + * memory change arrived cannot repopulate the cache with a stale snapshot. + */ + private memoryContextGeneration = 0; /** * Cache the last-known experiment state so we don't spam metadata refresh * when post-compaction context is disabled. @@ -550,6 +564,7 @@ export class AgentSession { onCompactionComplete, onIdleCompactionOutcome, onPostCompactionStateChange, + getContextResetState, } = options; assert(typeof workspaceId === "string", "workspaceId must be a string"); @@ -565,6 +580,8 @@ export class AgentSession { this.workspaceGoalService = workspaceGoalService; this.keepBackgroundProcesses = keepBackgroundProcesses ?? false; this.onPostCompactionStateChange = onPostCompactionStateChange; + this.getContextResetState = + getContextResetState ?? (() => ({ resetting: false, boundaryEpoch: 0 })); this.compactionHandler = new CompactionHandler({ workspaceId: this.workspaceId, @@ -1609,6 +1626,19 @@ export class AgentSession { } private async scheduleStartupAutoRetryIfNeeded(): Promise { + // Mid-check the session is neither busy nor holding a pending retry, so + // resetContext's guards need this in-flight marker to know a retry + // decision may be imminent. + this.startupAutoRetryCheckInFlight = true; + try { + return await this.evaluateStartupAutoRetry(); + } finally { + this.startupAutoRetryCheckInFlight = false; + } + } + + private async evaluateStartupAutoRetry(): Promise { + const contextResetStateAtStart = this.getContextResetState(); if (this.disposed || this.isBusy() || this.isAiStreaming()) { // Busy/streaming deferrals are state-driven; do not carry history-error backoff. this.startupAutoRetryDeferredRetryDelayMs = 0; @@ -1698,6 +1728,18 @@ export class AgentSession { this.startupAutoRetryDeferredRetryDelayMs = 0; return "deferred"; } + // A context reset overlapping this check invalidates the tail the reads + // above captured: its boundary can land between those reads and the + // retry start, and the retry would resume pre-reset context. Defer; the + // re-run reads history again and sees the boundary. + const contextResetState = this.getContextResetState(); + if ( + contextResetState.resetting || + contextResetState.boundaryEpoch !== contextResetStateAtStart.boundaryEpoch + ) { + this.startupAutoRetryDeferredRetryDelayMs = 0; + return "deferred"; + } await this.handleStreamFailureForAutoRetry({ type: "unknown", message: "startup_interrupted_stream", @@ -1889,6 +1931,15 @@ export class AgentSession { ); } + /** + * True while startup recovery is mid-decision: the session is not busy and + * holds no scheduled retry yet, but a retry can start as soon as pending + * reads settle. resetContext's guards treat this like a pending turn. + */ + isStartupRecoveryInFlight(): boolean { + return this.startupRecoveryPromise !== null || this.startupAutoRetryCheckInFlight; + } + scheduleStartupRecovery(): void { if (this.disposed || this.startupRecoveryScheduled || this.startupRecoveryPromise) { return; @@ -5784,6 +5835,11 @@ export class AgentSession { this.fileChangeTracker.clear(); } + invalidateMemoryContext(): void { + this.memoryContextGeneration++; + this.memoryContextByModelString.clear(); + } + /** * Resolve the memory session context (index snapshot + optional hot block) * for the current session segment. @@ -5791,10 +5847,11 @@ export class AgentSession { * Computed lazily on the first stream for each model. The first pass is * index-only so final tool policy can strip memory without paying hot-set * tokenization cost; if memory survives policy, the cache is upgraded with - * the token-budgeted hot block. Compaction clears the cache. Invoked by - * AIService.streamMessage after runtime.ensureReady(): caching before the - * runtime is started (stopped Docker/remote workspace) would pin an - * empty/partial context for the whole segment. + * the token-budgeted hot block. Compaction, context reset, and memory + * changes clear the cache. Invoked by AIService.streamMessage after + * runtime.ensureReady(): caching before the runtime is started (stopped + * Docker/remote workspace) would pin an empty/partial context for the + * whole segment. */ private async resolveMemoryContext( modelString: string, @@ -5807,18 +5864,32 @@ export class AgentSession { return cached.context ?? undefined; } - // Guard for test mocks that may not implement buildMemorySessionContext. - const context = - typeof this.aiService.buildMemorySessionContext === "function" - ? await this.aiService.buildMemorySessionContext(this.workspaceId, modelString, { - includeHotMemories, - }) - : null; - this.memoryContextByModelString.set(modelString, { - context, - includesHotMemories: includeHotMemories, - }); - return context ?? undefined; + // A memory change during a build makes that snapshot stale for the + // current request, not just later ones: rebuild until a build completes + // without a concurrent invalidation. Bounded so pathological invalidation + // churn cannot livelock the stream; the final snapshot is then served + // uncached so the next request rebuilds. + const maxBuildAttempts = 3; + for (let attempt = 1; ; attempt++) { + const generationAtBuildStart = this.memoryContextGeneration; + // Guard for test mocks that may not implement buildMemorySessionContext. + const context = + typeof this.aiService.buildMemorySessionContext === "function" + ? await this.aiService.buildMemorySessionContext(this.workspaceId, modelString, { + includeHotMemories, + }) + : null; + if (generationAtBuildStart === this.memoryContextGeneration) { + this.memoryContextByModelString.set(modelString, { + context, + includesHotMemories: includeHotMemories, + }); + return context ?? undefined; + } + if (attempt >= maxBuildAttempts) { + return context ?? undefined; + } + } } /** @@ -5841,7 +5912,7 @@ export class AgentSession { // Compaction boundary: invalidate the session-cached memory context so // the next stream recomputes the index and hot set from current // files/pins/usage stats. - this.memoryContextByModelString.clear(); + this.invalidateMemoryContext(); // Clear file state cache since history context is gone this.fileChangeTracker.clear(); diff --git a/src/node/services/aiService.test.ts b/src/node/services/aiService.test.ts index c7b1b3c8bf..6fa87f75e8 100644 --- a/src/node/services/aiService.test.ts +++ b/src/node/services/aiService.test.ts @@ -1154,6 +1154,7 @@ describe("AIService.streamMessage compaction boundary slicing", () => { interface StreamMessageHarness { config: Config; service: AIService; + historyService: HistoryService; planPayloadMessageIds: string[][]; preparedPayloadMessageIds: string[][]; preparedToolNamesForSentinel: string[][]; @@ -1277,6 +1278,7 @@ describe("AIService.streamMessage compaction boundary slicing", () => { return { config, service, + historyService, planPayloadMessageIds, preparedPayloadMessageIds, preparedToolNamesForSentinel, @@ -1626,6 +1628,576 @@ describe("AIService.streamMessage compaction boundary slicing", () => { ); }); + it("re-resolves the memory context before tool assembly so a mid-preparation invalidation is not advertised stale", async () => { + using muxHome = new DisposableTempDir("ai-service-memory-reresolve"); + const projectPath = path.join(muxHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + + const workspaceId = "workspace-memory-reresolve"; + const metadata = createLocalWorkspaceMetadata(workspaceId, projectPath); + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- stub for memory availability gating + const stubTool: Tool = {} as never; + const harness = createHarness(muxHome.path, metadata, { + allTools: { memory: stubTool }, + }); + harness.service.setMemoryService( + new MemoryService(harness.config, new MemoryMetaService(muxHome.path)) + ); + + // The session cache serves the stale snapshot until an invalidation + // lands mid-preparation; every resolve after that returns fresh state. + const staleEntries = [{ path: "/memories/global/deleted.md", description: "stale" }]; + const freshEntries = [{ path: "/memories/global/kept.md", description: "fresh" }]; + let resolveCount = 0; + const resolveMemoryContext = mock(() => { + resolveCount++; + return Promise.resolve({ + indexEntries: resolveCount === 1 ? staleEntries : freshEntries, + hotMemoriesBlock: null, + }); + }); + + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("latest-user", "user", "fix the issue")], + workspaceId, + modelString: KNOWN_MODELS.SONNET.id, + thinkingLevel: "off", + experiments: { memory: true }, + resolveMemoryContext, + }); + expect(result.success).toBe(true); + + // The memory tool description must be assembled from the re-resolved + // snapshot, not the pre-preparation capture. + const toolsConfig = harness.getToolsForModelSpy.mock.calls[0]?.[1]; + expect(toolsConfig?.memoryIndexEntries).toEqual(freshEntries); + }); + + it("refreshes the memory tool description when memory changes during tool assembly", async () => { + using muxHome = new DisposableTempDir("ai-service-memory-assembly-window"); + const projectPath = path.join(muxHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + + const workspaceId = "workspace-memory-assembly-window"; + const metadata = createLocalWorkspaceMetadata(workspaceId, projectPath); + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- stub for memory availability gating + const stubTool: Tool = { description: "memory base" } as never; + const harness = createHarness(muxHome.path, metadata, { + allTools: { memory: stubTool }, + }); + harness.service.setMemoryService( + new MemoryService(harness.config, new MemoryMetaService(muxHome.path)) + ); + + // The rename/delete lands while getToolsForModel is awaited: resolves + // before assembly see v1, resolves after see v2. + let memoryVersion = 1; + const resolveMemoryContext = mock(() => + Promise.resolve({ + indexEntries: [ + { path: `/memories/global/v${memoryVersion}.md`, description: `v${memoryVersion}` }, + ], + hotMemoriesBlock: null, + }) + ); + harness.getToolsForModelSpy.mockImplementation(() => { + memoryVersion = 2; + return Promise.resolve({ memory: stubTool }); + }); + + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("latest-user", "user", "fix the issue")], + workspaceId, + modelString: KNOWN_MODELS.SONNET.id, + thinkingLevel: "off", + experiments: { memory: true }, + resolveMemoryContext, + }); + expect(result.success).toBe(true); + + // The tool record handed to the stream must advertise the post-assembly + // snapshot, not the entries captured before getToolsForModel awaited. + const START_STREAM_TOOLS_INDEX = 9; + const streamTools = harness.startStreamCalls[0]?.[START_STREAM_TOOLS_INDEX] as + | Record + | undefined; + expect(streamTools?.memory?.description).toContain("/memories/global/v2.md"); + expect(streamTools?.memory?.description).not.toContain("/memories/global/v1.md"); + }); + + it("refreshes the memory tool description when memory changes during the hot-set upgrade", async () => { + using muxHome = new DisposableTempDir("ai-service-memory-hot-upgrade-window"); + const projectPath = path.join(muxHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + + const workspaceId = "workspace-memory-hot-upgrade-window"; + const metadata = createLocalWorkspaceMetadata(workspaceId, projectPath); + const experimentsService = new ExperimentsService({ + telemetryService: new TelemetryService(muxHome.path), + muxHome: muxHome.path, + }); + spyOn(experimentsService, "isExperimentEnabled").mockImplementation( + (experimentId) => + experimentId === EXPERIMENT_IDS.MEMORY || experimentId === EXPERIMENT_IDS.MEMORY_HOT_SET + ); + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- stub for memory availability gating + const stubTool: Tool = { description: "memory base" } as never; + const harness = createHarness(muxHome.path, metadata, { + allTools: { memory: stubTool }, + experimentsService, + }); + harness.service.setMemoryService( + new MemoryService(harness.config, new MemoryMetaService(muxHome.path)) + ); + + // Identity-stable index-only snapshot (like the session cache); a rename + // lands while the hot-set upgrade build is awaited, so only the upgrade + // resolve returns the post-invalidation entries. + const staleContext = { + indexEntries: [{ path: "/memories/global/v1.md", description: "v1" }], + hotMemoriesBlock: null, + }; + const resolveMemoryContext = mock( + (_modelString: string, options?: { includeHotMemories?: boolean }) => { + if (options?.includeHotMemories === false) { + return Promise.resolve(staleContext); + } + return Promise.resolve({ + indexEntries: [{ path: "/memories/global/v2.md", description: "v2" }], + hotMemoriesBlock: "fresh", + }); + } + ); + + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("latest-user", "user", "fix the issue")], + workspaceId, + modelString: KNOWN_MODELS.SONNET.id, + thinkingLevel: "off", + experiments: { memory: true }, + resolveMemoryContext, + }); + expect(result.success).toBe(true); + + // The provider request must not pair the fresh system block with a tool + // index still advertising the pre-upgrade snapshot. + const START_STREAM_TOOLS_INDEX = 9; + const streamTools = harness.startStreamCalls[0]?.[START_STREAM_TOOLS_INDEX] as + | Record + | undefined; + expect(streamTools?.memory?.description).toContain("/memories/global/v2.md"); + expect(streamTools?.memory?.description).not.toContain("/memories/global/v1.md"); + }); + + it("revalidates the memory snapshot at the request boundary when memory changes during message preparation", async () => { + using muxHome = new DisposableTempDir("ai-service-memory-request-boundary"); + const projectPath = path.join(muxHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + + const workspaceId = "workspace-memory-request-boundary"; + const metadata = createLocalWorkspaceMetadata(workspaceId, projectPath); + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- stub for memory availability gating + const stubTool: Tool = { description: "memory base" } as never; + const harness = createHarness(muxHome.path, metadata, { + allTools: { memory: stubTool }, + }); + harness.service.setMemoryService( + new MemoryService(harness.config, new MemoryMetaService(muxHome.path)) + ); + + // Identity-stable per version, like the session cache: resolves return + // the same reference until an invalidation bumps the version. + let memoryVersion = 1; + let cachedVersion = 0; + let cachedContext: + | { indexEntries: Array<{ path: string; description: string }>; hotMemoriesBlock: null } + | undefined; + const resolveMemoryContext = mock(() => { + if (cachedContext === undefined || cachedVersion !== memoryVersion) { + cachedVersion = memoryVersion; + cachedContext = { + indexEntries: [ + { path: `/memories/global/v${memoryVersion}.md`, description: `v${memoryVersion}` }, + ], + hotMemoriesBlock: null, + }; + } + return Promise.resolve(cachedContext); + }); + // The rename/delete lands while prepareMessagesForProvider is awaited, + // after every earlier re-resolve window has already passed. + spyOn(messagePipeline, "prepareMessagesForProvider").mockImplementation((pipelineArgs) => { + memoryVersion = 2; + return Promise.resolve( + pipelineArgs.messagesWithSentinel as unknown as Awaited< + ReturnType + > + ); + }); + + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("latest-user", "user", "fix the issue")], + workspaceId, + modelString: KNOWN_MODELS.SONNET.id, + thinkingLevel: "off", + experiments: { memory: true }, + resolveMemoryContext, + }); + expect(result.success).toBe(true); + + // The request that actually streams must advertise the post-invalidation + // snapshot, not the capture from before message preparation. + const START_STREAM_TOOLS_INDEX = 9; + const streamTools = harness.startStreamCalls[0]?.[START_STREAM_TOOLS_INDEX] as + | Record + | undefined; + expect(streamTools?.memory?.description).toContain("/memories/global/v2.md"); + expect(streamTools?.memory?.description).not.toContain("/memories/global/v1.md"); + }); + + it("deletes the assistant placeholder when the boundary revalidation rebuild fails", async () => { + using muxHome = new DisposableTempDir("ai-service-memory-boundary-failure"); + const projectPath = path.join(muxHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + + const workspaceId = "workspace-memory-boundary-failure"; + const metadata = createLocalWorkspaceMetadata(workspaceId, projectPath); + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- stub for memory availability gating + const stubTool: Tool = { description: "memory base" } as never; + const harness = createHarness(muxHome.path, metadata, { + allTools: { memory: stubTool }, + }); + harness.service.setMemoryService( + new MemoryService(harness.config, new MemoryMetaService(muxHome.path)) + ); + // The shared stub no-ops history writes; this test asserts on persisted + // history, so restore the real disk-backed implementations. + spyOn(harness.historyService, "appendToHistory").mockImplementation((wsId, message) => + HistoryService.prototype.appendToHistory.call(harness.historyService, wsId, message) + ); + spyOn(harness.historyService, "deleteMessage").mockImplementation((wsId, messageId) => + HistoryService.prototype.deleteMessage.call(harness.historyService, wsId, messageId) + ); + + let memoryVersion = 1; + let cachedVersion = 0; + let cachedContext: + | { indexEntries: Array<{ path: string; description: string }>; hotMemoriesBlock: null } + | undefined; + const resolveMemoryContext = mock(() => { + if (cachedContext === undefined || cachedVersion !== memoryVersion) { + cachedVersion = memoryVersion; + cachedContext = { + indexEntries: [ + { path: `/memories/global/v${memoryVersion}.md`, description: `v${memoryVersion}` }, + ], + hotMemoriesBlock: null, + }; + } + return Promise.resolve(cachedContext); + }); + // The memory change lands during message preparation, so the boundary + // revalidation loop performs a rebuild after the placeholder append; + // that rebuild then fails. + let failContextBuilds = false; + spyOn(messagePipeline, "prepareMessagesForProvider").mockImplementation((pipelineArgs) => { + memoryVersion = 2; + failContextBuilds = true; + return Promise.resolve( + pipelineArgs.messagesWithSentinel as unknown as Awaited< + ReturnType + > + ); + }); + spyOn(streamContextBuilder, "buildStreamSystemContext").mockImplementation(() => { + if (failContextBuilds) { + return Promise.reject(new Error("boundary rebuild failed")); + } + return Promise.resolve({ + agentSystemPromptSections: ["test-agent-prompt"], + systemMessage: "test-system-message", + systemMessageTokens: 1, + agentDefinitions: undefined, + availableSkills: undefined, + ancestorPlanFilePaths: [], + }); + }); + + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("latest-user", "user", "fix the issue")], + workspaceId, + modelString: KNOWN_MODELS.SONNET.id, + thinkingLevel: "off", + experiments: { memory: true }, + resolveMemoryContext, + }); + expect(result.success).toBe(false); + + // The persisted placeholder must not survive the failure: an empty + // non-partial assistant would be treated as a trailing assistant turn on + // the next load and get a synthetic [CONTINUE]. + const history = await harness.historyService.getHistoryFromLatestBoundary(workspaceId); + expect(history.success).toBe(true); + const assistants = (history.success ? history.data : []).filter( + (message) => message.role === "assistant" + ); + expect(assistants).toHaveLength(0); + }); + + it("provides a refresh hook that pulls memory changes during stream startup awaits", async () => { + using muxHome = new DisposableTempDir("ai-service-memory-start-refresh"); + const projectPath = path.join(muxHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + + const workspaceId = "workspace-memory-start-refresh"; + const metadata = createLocalWorkspaceMetadata(workspaceId, projectPath); + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- stub for memory availability gating + const stubTool: Tool = { description: "memory base" } as never; + const harness = createHarness(muxHome.path, metadata, { + allTools: { memory: stubTool }, + }); + harness.service.setMemoryService( + new MemoryService(harness.config, new MemoryMetaService(muxHome.path)) + ); + // The default harness stub returns a constant system message; embed the + // hot-memories block so refreshes are observable in the built prompt. + spyOn(streamContextBuilder, "buildStreamSystemContext").mockImplementation((contextArgs) => + Promise.resolve({ + agentSystemPromptSections: ["test-agent-prompt"], + systemMessage: `test-system-message ${contextArgs.hotMemoriesBlock ?? ""}`, + systemMessageTokens: 1, + agentDefinitions: undefined, + availableSkills: undefined, + ancestorPlanFilePaths: [], + }) + ); + // Custom "Tool: memory" instructions must survive every description + // refresh, not just the initial getToolsForModel augmentation. + spyOn(systemMessageModule, "readToolInstructions").mockResolvedValue({ + memory: "CUSTOM-MEMORY-RULES", + }); + + let memoryVersion = 1; + let cachedVersion = 0; + let cachedContext: + | { indexEntries: Array<{ path: string; description: string }>; hotMemoriesBlock: string } + | undefined; + const resolveMemoryContext = mock(() => { + if (cachedContext === undefined || cachedVersion !== memoryVersion) { + cachedVersion = memoryVersion; + cachedContext = { + indexEntries: [ + { path: `/memories/global/v${memoryVersion}.md`, description: `v${memoryVersion}` }, + ], + hotMemoriesBlock: `hot-memories-v${memoryVersion}`, + }; + } + return Promise.resolve(cachedContext); + }); + + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("latest-user", "user", "fix the issue")], + workspaceId, + modelString: KNOWN_MODELS.SONNET.id, + thinkingLevel: "off", + experiments: { memory: true }, + resolveMemoryContext, + }); + expect(result.success).toBe(true); + + // startStream awaits setup (mutex, safety, temp dir) before registering + // the request; the hook passed to it must surface a memory change that + // lands during those awaits. + const START_STREAM_REFRESH_INDEX = 28; + const refresh = harness.startStreamCalls[0]?.[START_STREAM_REFRESH_INDEX] as + | (() => Promise<{ system: string; tools?: Record } | null>) + | undefined; + expect(typeof refresh).toBe("function"); + if (!refresh) { + throw new Error("Expected refresh hook on startStream"); + } + + // Unchanged memory: no refreshed payload, the captured snapshot stands. + expect(await refresh()).toBeNull(); + + memoryVersion = 2; + const refreshed = await refresh(); + expect(refreshed).not.toBeNull(); + expect(refreshed?.tools?.memory?.description).toContain("/memories/global/v2.md"); + expect(refreshed?.tools?.memory?.description).not.toContain("/memories/global/v1.md"); + // The refresh rebuilds the description from the base definition; the + // configured "Tool: memory" instructions must be re-appended. + expect(refreshed?.tools?.memory?.description).toContain("CUSTOM-MEMORY-RULES"); + + // The debug snapshot was captured before the hook ran; it must follow the + // refresh or the modal would show a request that was never sent. + const debugSnapshot = harness.service.debugGetLastLlmRequest(workspaceId); + expect(debugSnapshot.success).toBe(true); + const snapshotSystem = debugSnapshot.success ? debugSnapshot.data?.systemMessage : undefined; + expect(snapshotSystem).toContain("hot-memories-v2"); + expect(snapshotSystem).not.toContain("hot-memories-v1"); + }); + + it("resolves a fresh memory context for the refusal fallback when hot-set preload is disabled", async () => { + using muxHome = new DisposableTempDir("ai-service-memory-fallback-no-hot-set"); + const projectPath = path.join(muxHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + + const workspaceId = "workspace-memory-fallback-no-hot-set"; + const sourceModel = KNOWN_MODELS.SONNET.id; + const fallbackModel = KNOWN_MODELS.GPT.id; + await writeMainConfig(muxHome.path, { + modelFallbacks: { + [sourceModel]: { models: [fallbackModel] }, + }, + }); + + const metadata = createLocalWorkspaceMetadata(workspaceId, projectPath); + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- stub for memory availability gating + const stubTool: Tool = { description: "memory base" } as never; + const harness = createHarness(muxHome.path, metadata, { + allTools: { memory: stubTool }, + useRequestedModelString: true, + }); + harness.service.setMemoryService( + new MemoryService(harness.config, new MemoryMetaService(muxHome.path)) + ); + + let memoryVersion = 1; + let cachedVersion = 0; + let cachedContext: + | { indexEntries: Array<{ path: string; description: string }>; hotMemoriesBlock: null } + | undefined; + const resolveMemoryContext = mock(() => { + if (cachedContext === undefined || cachedVersion !== memoryVersion) { + cachedVersion = memoryVersion; + cachedContext = { + indexEntries: [ + { path: `/memories/global/v${memoryVersion}.md`, description: `v${memoryVersion}` }, + ], + hotMemoriesBlock: null, + }; + } + return Promise.resolve(cachedContext); + }); + + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("latest-user", "user", "fix the issue")], + workspaceId, + modelString: sourceModel, + thinkingLevel: "off", + experiments: { memory: true }, + resolveMemoryContext, + }); + expect(result.success).toBe(true); + + const modelFallback = harness.startStreamCalls[0]?.[START_STREAM_MODEL_FALLBACK_INDEX] as + | ModelFallbackOptions + | undefined; + expect(modelFallback).toBeDefined(); + if (!modelFallback) { + throw new Error("Expected modelFallback options on startStream"); + } + + // A rename lands between the original assembly and the refusal fallback: + // the fallback prepare must advertise the current index even though hot + // preloading never upgrades the context. + memoryVersion = 2; + const prepared = await modelFallback.prepare(fallbackModel); + expect(prepared.success).toBe(true); + if (!prepared.success) { + throw new Error("Expected fallback prepare to succeed"); + } + const memoryDescription = prepared.data.tools?.memory?.description; + expect(memoryDescription).toContain("/memories/global/v2.md"); + expect(memoryDescription).not.toContain("/memories/global/v1.md"); + }); + + it("revalidates fallback memory at the request boundary when memory changes during fallback preparation", async () => { + using muxHome = new DisposableTempDir("ai-service-memory-fallback-boundary"); + const projectPath = path.join(muxHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + + const workspaceId = "workspace-memory-fallback-boundary"; + const sourceModel = KNOWN_MODELS.SONNET.id; + const fallbackModel = KNOWN_MODELS.GPT.id; + await writeMainConfig(muxHome.path, { + modelFallbacks: { + [sourceModel]: { models: [fallbackModel] }, + }, + }); + + const metadata = createLocalWorkspaceMetadata(workspaceId, projectPath); + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- stub for memory availability gating + const stubTool: Tool = { description: "memory base" } as never; + const harness = createHarness(muxHome.path, metadata, { + allTools: { memory: stubTool }, + useRequestedModelString: true, + }); + harness.service.setMemoryService( + new MemoryService(harness.config, new MemoryMetaService(muxHome.path)) + ); + + let memoryVersion = 1; + let cachedVersion = 0; + let cachedContext: + | { indexEntries: Array<{ path: string; description: string }>; hotMemoriesBlock: null } + | undefined; + const resolveMemoryContext = mock(() => { + if (cachedContext === undefined || cachedVersion !== memoryVersion) { + cachedVersion = memoryVersion; + cachedContext = { + indexEntries: [ + { path: `/memories/global/v${memoryVersion}.md`, description: `v${memoryVersion}` }, + ], + hotMemoriesBlock: null, + }; + } + return Promise.resolve(cachedContext); + }); + // The rename lands while the FALLBACK preparation awaits message prep, + // after prepare()'s own initial memory resolve already returned v1. + spyOn(messagePipeline, "prepareMessagesForProvider").mockImplementation((pipelineArgs) => { + if (pipelineArgs.modelString === fallbackModel) { + memoryVersion = 2; + } + return Promise.resolve( + pipelineArgs.messagesWithSentinel as unknown as Awaited< + ReturnType + > + ); + }); + + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("latest-user", "user", "fix the issue")], + workspaceId, + modelString: sourceModel, + thinkingLevel: "off", + experiments: { memory: true }, + resolveMemoryContext, + }); + expect(result.success).toBe(true); + + const modelFallback = harness.startStreamCalls[0]?.[START_STREAM_MODEL_FALLBACK_INDEX] as + | ModelFallbackOptions + | undefined; + expect(modelFallback).toBeDefined(); + if (!modelFallback) { + throw new Error("Expected modelFallback options on startStream"); + } + + const prepared = await modelFallback.prepare(fallbackModel); + expect(prepared.success).toBe(true); + if (!prepared.success) { + throw new Error("Expected fallback prepare to succeed"); + } + // The fallback request must advertise the post-invalidation snapshot, + // not the resolve captured before its awaited preparation stages. + const memoryDescription = prepared.data.tools?.memory?.description; + expect(memoryDescription).toContain("/memories/global/v2.md"); + expect(memoryDescription).not.toContain("/memories/global/v1.md"); + }); + // GPT-5.6 Chat Completions explicit-caching seam: fallback provider options // and route metadata must be rebuilt from the fallback model/route so cache // fields cannot leak across routes in either direction. @@ -1961,6 +2533,9 @@ describe("AIService.streamMessage compaction boundary slicing", () => { new MemoryService(harness.config, new MemoryMetaService(muxHome.path)) ); const memoryCalls: Array<{ includeHotMemories: boolean }> = []; + // Identity-stable like the session cache: revalidation resolves must + // observe an unchanged snapshot, not fabricate a memory change. + const cachedContext = { indexEntries: [], hotMemoriesBlock: null }; const result = await harness.service.streamMessage({ messages: [createMuxMessage("latest-user", "user", "hello")], @@ -1970,13 +2545,15 @@ describe("AIService.streamMessage compaction boundary slicing", () => { experiments: { memory: true }, resolveMemoryContext: (_modelString, options) => { memoryCalls.push({ includeHotMemories: options?.includeHotMemories !== false }); - return Promise.resolve({ indexEntries: [], hotMemoriesBlock: null }); + return Promise.resolve(cachedContext); }, }); expect(result.success).toBe(true); expect(harness.streamSystemContextMemoryToolFlags).toEqual([true, false]); - expect(memoryCalls).toEqual([{ includeHotMemories: false }]); + // The stripped memory tool must never trigger a hot-set upgrade. + expect(memoryCalls.length).toBeGreaterThan(0); + expect(memoryCalls.some((call) => call.includeHotMemories)).toBe(false); }); it("does not upgrade memory context when the hot-set sub-experiment is disabled", async () => { @@ -1995,6 +2572,9 @@ describe("AIService.streamMessage compaction boundary slicing", () => { new MemoryService(harness.config, new MemoryMetaService(muxHome.path)) ); const memoryCalls: Array<{ includeHotMemories: boolean }> = []; + // Identity-stable like the session cache: revalidation resolves must + // observe an unchanged snapshot, not fabricate a memory change. + const cachedContext = { indexEntries: [], hotMemoriesBlock: null }; const result = await harness.service.streamMessage({ messages: [createMuxMessage("latest-user", "user", "hello")], @@ -2004,13 +2584,16 @@ describe("AIService.streamMessage compaction boundary slicing", () => { experiments: { memory: true }, resolveMemoryContext: (_modelString, options) => { memoryCalls.push({ includeHotMemories: options?.includeHotMemories !== false }); - return Promise.resolve({ indexEntries: [], hotMemoriesBlock: null }); + return Promise.resolve(cachedContext); }, }); expect(result.success).toBe(true); + // No rebuild: the snapshot never changed, so the pre-policy prompt stands. expect(harness.streamSystemContextMemoryToolFlags).toEqual([true]); - expect(memoryCalls).toEqual([{ includeHotMemories: false }]); + // Hot-set disabled: no resolve may request the hot block. + expect(memoryCalls.length).toBeGreaterThan(0); + expect(memoryCalls.some((call) => call.includeHotMemories)).toBe(false); }); it("anchors the memory index by project identity and gates hot preloading on the memory-hot-set experiment", async () => { diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts index 4a23481baa..1d7492e833 100644 --- a/src/node/services/aiService.ts +++ b/src/node/services/aiService.ts @@ -22,7 +22,12 @@ import type { GoalRecordV1 } from "@/common/types/goal"; import type { ModelMessage, MuxMessage, MuxMessageMetadata } from "@/common/types/message"; import { createMuxMessage } from "@/common/types/message"; import type { Config } from "@/node/config"; -import { StreamManager, type ModelFallbackOptions, type StreamTextOnChunk } from "./streamManager"; +import { + StreamManager, + type ModelFallbackOptions, + type RefreshMemoryDependentContext, + type StreamTextOnChunk, +} from "./streamManager"; import { runLanguageModelCleanup } from "./languageModelCleanup"; import type { InitStateManager } from "./initStateManager"; import type { SendMessageError } from "@/common/types/errors"; @@ -92,7 +97,10 @@ import { type MemorySessionContext, } from "@/node/services/memoryService"; import { formatHotMemoriesBlock } from "@/node/services/memoryHotSet"; -import { resolveMemoryAccessPolicy } from "@/node/services/tools/memory"; +import { + buildMemoryToolDescription, + resolveMemoryAccessPolicy, +} from "@/node/services/tools/memory"; import { isExecLikeEditingCapableInResolvedChain } from "@/common/utils/agentTools"; import { buildProviderOptions, @@ -214,6 +222,14 @@ export function prepareProviderRequestMessages( }; } +/** + * Replace a tool's advertised description without widening the Tool union + * (a plain spread of the union does not narrow back to Tool). + */ +function withToolDescription(toolValue: T, description: string): T { + return { ...toolValue, description }; +} + function replaceOrAppendMessageById(messages: MuxMessage[], replacement: MuxMessage): MuxMessage[] { const index = messages.findIndex((message) => message.id === replacement.id); if (index === -1) { @@ -587,12 +603,8 @@ export class AIService extends EventEmitter { * frequently used memory files; memory-hot-set sub-experiment). Returns * null when the memory experiment is off. * - * Callers (AgentSession) cache the result per model and recompute it only - * on the first use of a model in a session segment, or at compaction - * boundaries, so repeated turns keep prompt-cache-stable bytes. Memories - * written mid-segment surface in the next segment's index for cached models - * (the writing agent already has its own tool calls in context, and `view` - * lists live state). + * AgentSession invalidates cached results at compaction, context reset, and + * memory changes. */ async buildMemorySessionContext( workspaceId: string, @@ -1466,21 +1478,25 @@ export class AIService extends EventEmitter { // claude-skills-compat is host-evaluated (like memory-hot-set): sub-agents share the // host ExperimentsService, so it is not inherited through SendMessageOptions.experiments. const claudeSkillsCompatExperimentEnabled = this.isClaudeSkillsCompatEnabled(); - // Once final tool policy keeps the memory tool, upgrade the index-only - // memory context (resolved pre-policy with includeHotMemories: false) to - // the token-budgeted hot block for the model that will actually stream. - // Returns the unchanged pre-policy `memoryContext` reference when hot - // preloading is off or the memory tool was stripped, so callers can use - // identity comparison to decide whether the system prompt must be rebuilt. + // Once final tool policy keeps the memory tool, re-resolve the memory + // context for the model that will actually stream; the hot-set + // experiment gates only the token-budgeted hot block, not the resolve + // itself, so a memory change since an earlier capture is picked up even + // with hot preloading off (the session cache makes the no-change case + // an identity-stable cache hit). Returns the unchanged `fallbackContext` + // reference when the memory tool was stripped, so callers can use + // identity comparison to decide whether the system prompt must be + // rebuilt. const upgradeMemoryContextForModel = async ( memoryToolAvailableForModel: boolean, - modelStringForContext: string + modelStringForContext: string, + fallbackContext: MemorySessionContext | undefined ): Promise => - memoryToolAvailableForModel && - memoryHotSetExperimentEnabled && - resolveMemoryContext !== undefined - ? await resolveMemoryContext(modelStringForContext, { includeHotMemories: true }) - : memoryContext; + memoryToolAvailableForModel && resolveMemoryContext !== undefined + ? await resolveMemoryContext(modelStringForContext, { + includeHotMemories: memoryHotSetExperimentEnabled, + }) + : fallbackContext; emitStartupBreadcrumb("loading_workspace_context"); const resolveAgentForStreamStartedAt = Date.now(); const agentResult = await resolveAgentForStream({ @@ -2027,6 +2043,15 @@ export class AIService extends EventEmitter { // stay scoped to this specific assistant turn. The placeholder is appended to history below // (after the abort check). const assistantMessageId = createAssistantMessageId(); + // Memory changes during the async preparation above (agent/skill + // discovery, MCP startup) invalidate the session cache but cannot + // update the earlier `memoryContext` capture. Re-resolve immediately + // before tool assembly: a cache hit returns the identical reference, + // an invalidation forces a rebuild, and the identity change makes the + // final system-prompt comparison below rebuild too. + const memoryContextForTools = resolveMemoryContext + ? await resolveMemoryContext(modelString, { includeHotMemories: false }) + : undefined; const allowLegacyInvalidWorkflowAgentOutputSchema = await this.shouldAllowLegacyInvalidWorkflowAgentOutputSchema(metadata); // Hoisted so the refusal-fallback prepare() can rebuild the toolset for a @@ -2227,7 +2252,7 @@ export class AIService extends EventEmitter { availableSkills, // Session-segment memory index advertised in the memory tool // description (same disclosure mechanic as skills). - memoryIndexEntries: memoryContext?.indexEntries, + memoryIndexEntries: memoryContextForTools?.indexEntries, // Trust gating: only run hooks/scripts when the full shared workspace runtime is trusted. trusted: sharedExecutionTrusted, }; @@ -2293,15 +2318,59 @@ export class AIService extends EventEmitter { } } + // getToolsForModel appended any configured "Tool: memory" instructions + // to the built description, so every later description refresh (memory + // index changes, hot-set upgrades, fallback rebuilds) must re-append + // them or the refreshed tool silently drops the agent's custom rules. + const refreshedMemoryToolDescription = ( + indexEntries: Parameters[0] + ): string => { + const rebuilt = buildMemoryToolDescription(indexEntries); + const memoryInstructions = toolInstructions.memory; + return memoryInstructions ? `${rebuilt}\n\n${memoryInstructions}` : rebuilt; + }; + const advisorToolAvailable = tools.advisor !== undefined; const memoryToolAvailable = tools.memory !== undefined; + // Tool assembly awaited above (getToolsForModel, policy/experiments) + // is another invalidation window: re-resolve once more and refresh the + // already-created memory tool's description if the snapshot moved, so + // the request cannot advertise paths deleted during assembly. + const memoryContextAfterAssembly = resolveMemoryContext + ? await resolveMemoryContext(modelString, { includeHotMemories: false }) + : undefined; + if (tools.memory !== undefined && memoryContextAfterAssembly !== memoryContextForTools) { + tools = { + ...tools, + memory: withToolDescription( + tools.memory, + refreshedMemoryToolDescription(memoryContextAfterAssembly?.indexEntries) + ), + }; + } const finalMemoryContext = await upgradeMemoryContextForModel( memoryToolAvailable, - modelString + modelString, + memoryContextAfterAssembly ); + // The hot-set upgrade is one more awaited window: a memory change + // during it lands in finalMemoryContext (and the system prompt rebuild + // below), so the already-created memory tool must advertise the same + // snapshot instead of the pre-upgrade capture. + if (tools.memory !== undefined && finalMemoryContext !== memoryContextAfterAssembly) { + tools = { + ...tools, + memory: withToolDescription( + tools.memory, + refreshedMemoryToolDescription(finalMemoryContext?.indexEntries) + ), + }; + } const finalStreamSystemContext = advisorToolAvailable === advisorToolEligible && memoryToolAvailable === memoryToolEligible && + // Identity against the context the pre-policy prompt was built from: + // a mid-preparation invalidation produced a fresh reference above. finalMemoryContext === memoryContext ? prePolicyStreamSystemContext : await (async () => { @@ -2451,6 +2520,63 @@ export class AIService extends EventEmitter { return Ok(undefined); } + // Memory can also change while the later awaited stages run (final + // system-prompt rebuild, tokenizer loading, message preparation, history + // append). Everything from here to startStream is synchronous, so + // re-resolve once more and repeat the memory-dependent assembly while + // the snapshot keeps moving; bounded like resolveMemoryContext's rebuild + // so invalidation churn cannot livelock the stream. + let memoryContextAtBoundary = finalMemoryContext; + const revalidateMemoryDependentState = async (): Promise => { + let changed = false; + for (let revalidation = 0; revalidation < 3; revalidation++) { + const revalidatedMemoryContext = await upgradeMemoryContextForModel( + memoryToolAvailable, + modelString, + memoryContextAtBoundary + ); + if (revalidatedMemoryContext === memoryContextAtBoundary) { + break; + } + changed = true; + memoryContextAtBoundary = revalidatedMemoryContext; + if (tools.memory !== undefined) { + tools = { + ...tools, + memory: withToolDescription( + tools.memory, + refreshedMemoryToolDescription(revalidatedMemoryContext?.indexEntries) + ), + }; + } + const revalidatedStreamSystemContext = await buildStreamSystemContextForToolset( + { advisorToolAvailable, memoryToolAvailable }, + modelString, + revalidatedMemoryContext + ); + systemMessage = revalidatedStreamSystemContext.systemMessage; + systemMessageTokens = revalidatedStreamSystemContext.systemMessageTokens; + if (mcpWarningPrefix != null) { + systemMessage = `${mcpWarningPrefix}${systemMessage}`; + const tokenizer = await getTokenizerForModel( + modelString, + resolveModelForMetadata(modelString, this.providerService.getConfig()) + ); + systemMessageTokens = await tokenizer.countTokens(systemMessage); + } + } + return changed; + }; + try { + await revalidateMemoryDependentState(); + } catch (error) { + // The placeholder was persisted above; failing without deleting it + // would leave an empty non-partial assistant in history that the next + // load treats as a trailing assistant turn (synthetic [CONTINUE]). + await deleteAbortedPlaceholder(assistantMessageId); + throw error; + } + // Build provider options based on thinking level and request-sliced message history. const truncationMode = openaiTruncationModeOverride; // Use the same boundary-sliced payload history that we send to the provider. @@ -2655,6 +2781,38 @@ export class AIService extends EventEmitter { } const toolsForStream = tools; + // startStream awaits its own setup (workspace mutex, stream safety, + // temp-dir creation) before registering the request; a memory change + // during those awaits would be invisible in the system/tools captured + // below, so this hook re-runs the bounded revalidation immediately + // before registration. Best-effort: the request-boundary snapshot is + // complete and valid, so a failed rebuild keeps it rather than failing + // the stream (and leaving the persisted placeholder behind). + const refreshMemoryDependentContext: RefreshMemoryDependentContext = async () => { + try { + const changed = await revalidateMemoryDependentState(); + if (!changed) { + return null; + } + // The debug snapshot above captured the pre-refresh system message; + // update it so the modal shows the request that was actually sent. + const capturedSnapshot = this.lastLlmRequestByWorkspace.get(workspaceId); + if (capturedSnapshot?.messageId === assistantMessageId) { + this.lastLlmRequestByWorkspace.set(workspaceId, { + ...capturedSnapshot, + systemMessage, + }); + } + return { system: systemMessage, tools, systemMessageTokens }; + } catch (error) { + workspaceLog.warn( + "Memory revalidation at stream registration failed; keeping the request-boundary snapshot", + { error: getErrorMessage(error) } + ); + return null; + } + }; + const canQueueDevToolsRunMetadata = this.devToolsService?.enabled === true && typeof modelResult.data.model !== "string" && @@ -2790,8 +2948,21 @@ export class AIService extends EventEmitter { const nextMemoryToolAvailable = nextTools.memory !== undefined; const nextMemoryContext = await upgradeMemoryContextForModel( nextMemoryToolAvailable, - next.canonicalModelString + next.canonicalModelString, + memoryContextAfterAssembly ); + // The rebuilt toolset used the config's captured index + // entries; refresh the description from the snapshot just + // resolved for the fallback model. + if (nextTools.memory !== undefined) { + nextTools = { + ...nextTools, + memory: withToolDescription( + nextTools.memory, + refreshedMemoryToolDescription(nextMemoryContext?.indexEntries) + ), + }; + } // Rebuild the system prompt for the fallback model (tool // instructions and "Model:" sections are model-keyed), keeping @@ -2842,6 +3013,82 @@ export class AIService extends EventEmitter { workspaceId, }); + // Fallback request boundary: the awaited stages above + // (system rebuild, tokenizer, message preparation) are + // invalidation windows too, and everything after this loop + // is synchronous until prepare() returns. Mirror the + // primary path's bounded revalidation. + let fallbackMemoryContext = nextMemoryContext; + const revalidateNextMemoryDependentState = async (): Promise => { + let changed = false; + for (let revalidation = 0; revalidation < 3; revalidation++) { + const revalidatedMemoryContext = await upgradeMemoryContextForModel( + nextMemoryToolAvailable, + next.canonicalModelString, + fallbackMemoryContext + ); + if (revalidatedMemoryContext === fallbackMemoryContext) { + break; + } + changed = true; + fallbackMemoryContext = revalidatedMemoryContext; + if (nextTools.memory !== undefined) { + nextTools = { + ...nextTools, + memory: withToolDescription( + nextTools.memory, + refreshedMemoryToolDescription(revalidatedMemoryContext?.indexEntries) + ), + }; + } + const revalidatedSystemContext = await buildStreamSystemContextForToolset( + { + advisorToolAvailable: nextTools.advisor !== undefined, + memoryToolAvailable: nextMemoryToolAvailable, + }, + next.canonicalModelString, + revalidatedMemoryContext + ); + nextSystem = revalidatedSystemContext.systemMessage; + nextSystemTokens = revalidatedSystemContext.systemMessageTokens; + if (mcpWarningPrefix != null) { + nextSystem = `${mcpWarningPrefix}${nextSystem}`; + const revalidationTokenizer = await getTokenizerForModel( + next.canonicalModelString, + resolveModelForMetadata( + next.canonicalModelString, + this.providerService.getConfig() + ) + ); + nextSystemTokens = await revalidationTokenizer.countTokens(nextSystem); + } + } + return changed; + }; + await revalidateNextMemoryDependentState(); + // Per-step hook for the fallback stream, mirroring the + // primary path's refreshMemoryDependentContext. + const refreshNextMemoryDependentContext: RefreshMemoryDependentContext = + async () => { + try { + const changed = await revalidateNextMemoryDependentState(); + if (!changed) { + return null; + } + return { + system: nextSystem, + tools: nextTools, + systemMessageTokens: nextSystemTokens, + }; + } catch (error) { + workspaceLog.warn( + "Fallback memory revalidation failed; keeping the current snapshot", + { error: getErrorMessage(error) } + ); + return null; + } + }; + const nextProviderOptions = buildProviderOptions( next.canonicalModelString, nextThinkingLevel, @@ -2967,6 +3214,7 @@ export class AIService extends EventEmitter { thinkingLevel: nextThinkingLevel, rebuildProviderOptionsForThinkingLevel: rebuildNextProviderOptionsForThinkingLevel, + refreshMemoryDependentContext: refreshNextMemoryDependentContext, initialMetadataPatch: { routedThroughGateway: next.routedThroughGateway, ...(next.routeProvider != null ? { routeProvider: next.routeProvider } : {}), @@ -3036,7 +3284,8 @@ export class AIService extends EventEmitter { modelFallback, toolSearchRuntime?.state, activeTurnThinkingOverride, - rebuildProviderOptionsForThinkingLevel + rebuildProviderOptionsForThinkingLevel, + refreshMemoryDependentContext ); recordStartupPhaseTiming("startStreamMs", startStreamStartedAt); diff --git a/src/node/services/coreServices.ts b/src/node/services/coreServices.ts index 2fc3918342..871971e5b7 100644 --- a/src/node/services/coreServices.ts +++ b/src/node/services/coreServices.ts @@ -27,7 +27,7 @@ import type { WorkspaceMcpOverridesService } from "@/node/services/workspaceMcpO import type { PolicyService } from "@/node/services/policyService"; import type { TelemetryService } from "@/node/services/telemetryService"; import type { ExperimentsService } from "@/node/services/experimentsService"; -import { MemoryService } from "@/node/services/memoryService"; +import { MemoryService, type MemoryChangeEvent } from "@/node/services/memoryService"; import { MemoryConsolidationService } from "@/node/services/memoryConsolidationService"; import { MemoryMetaService } from "@/node/services/memoryMeta"; import type { SessionTimingService } from "@/node/services/sessionTimingService"; @@ -157,6 +157,11 @@ export function createCoreServices(opts: CoreServicesOptions): CoreServices { opts.opResolver ); aiService.setWorkspaceHeartbeatService(workspaceService); + // Memory changes invalidate live sessions' cached contexts; the event scopes + // the fan-out (global -> all, workspace/project -> matching sessions only). + memoryService.on("change", (event: MemoryChangeEvent) => + workspaceService.invalidateMemoryContexts(event) + ); // Tool-started workflows share the same sidebar activity cache as ORPC-started workflows, // so terminal updates must prune active run counts regardless of launch path. aiService.setWorkflowRunStatusChangedHandler((event) => diff --git a/src/node/services/memoryService.ts b/src/node/services/memoryService.ts index b7e159bf05..1b701acb55 100644 --- a/src/node/services/memoryService.ts +++ b/src/node/services/memoryService.ts @@ -39,7 +39,7 @@ import { getErrorMessage } from "@/common/utils/errors"; import { isMultiProject } from "@/common/utils/multiProject"; import { MULTI_PROJECT_CONFIG_KEY } from "@/common/constants/multiProject"; import type { WorkspaceMetadata } from "@/common/types/workspace"; -import type { Config } from "@/node/config"; +import type { Config, Workspace as WorkspaceConfigEntry } from "@/node/config"; import type { Runtime } from "@/node/runtime/Runtime"; import { MutexMap } from "@/node/utils/concurrency/mutexMap"; import { memoryLogicalKey, type MemoryMetaService } from "@/node/services/memoryMeta"; @@ -293,6 +293,26 @@ export function resolveMemoryProjectIdentity(metadata: WorkspaceMetadata): strin return isMultiProject(metadata) ? "" : metadata.projectPath; } +/** + * Same identity as resolveMemoryProjectIdentity, computed synchronously from + * a config workspace entry: mutation paths must invalidate matching caches + * before yielding, so they cannot go through async metadata resolution. + * Mirrors Config.getAllWorkspaceMetadata's projectPath resolution (scratch + * identity is the scratch directory, not the "_scratch" bucket key). + */ +export function resolveMemoryProjectIdentityFromConfigEntry( + bucketProjectPath: string, + workspace: Pick +): string { + if ((workspace.projects?.length ?? 0) > 1) { + return ""; + } + if (workspace.kind === "scratch") { + return workspace.path; + } + return workspace.projects?.[0]?.projectPath ?? bucketProjectPath; +} + class LocalMemoryStore implements MemoryStore { constructor(readonly physicalRoot: string) {} @@ -643,6 +663,29 @@ export class MemoryService extends EventEmitter { // Commands // ------------------------------------------------------------------------- + /** + * Pins live in the meta sidecar, not the file stores, but they feed hot-set + * selection, so they must flow through MemoryService to emit a change event + * that invalidates session-cached memory contexts. + */ + async setPinned( + ctx: MemoryScopeContext, + virtualPath: string, + pinned: boolean, + actor: MemoryActor + ): Promise { + const parsed = parseMemoryPath(virtualPath); + const scope = this.requireFilePath(parsed, virtualPath); + const key = this.logicalKeyFor(ctx, scope, parsed.relPath); + if (key === null) { + throw new MemoryCommandError( + `Cannot pin ${virtualPath}: this scope has no stable identity in the current context` + ); + } + await this.metaService.setPinned(key, pinned); + this.emitChange(ctx, scope, parsed.relPath, actor); + } + async view( ctx: MemoryScopeContext, virtualPath: string, diff --git a/src/node/services/streamManager.test.ts b/src/node/services/streamManager.test.ts index ca7a41decb..fa0463f1b6 100644 --- a/src/node/services/streamManager.test.ts +++ b/src/node/services/streamManager.test.ts @@ -1111,6 +1111,194 @@ describe("StreamManager - sequential tool execution", () => { }); }); +describe("StreamManager - per-step memory refresh", () => { + interface MemoryRefreshResult { + system: string; + tools: Record | undefined; + systemMessageTokens: number; + } + type PrepareStepFn = (options: { + messages: ModelMessage[]; + }) => Promise< + { messages?: ModelMessage[]; instructions?: string | { content: string } } | undefined + >; + + function buildRequestWithHook( + system: string, + modelString: string, + tools: Record | undefined, + hook: () => Promise + ): { + request: { + system?: + | string + | { role: "system"; content: string; providerOptions?: Record }; + tools?: Record; + }; + options: Parameters[0]; + prepareStep: PrepareStepFn; + } { + const streamManager = new StreamManager(historyService); + const buildRequestConfig = Reflect.get(streamManager, "buildStreamRequestConfig") as ( + ...args: unknown[] + ) => { + system?: + | string + | { role: "system"; content: string; providerOptions?: Record }; + tools?: Record; + }; + const createStreamResultMethod = Reflect.get(streamManager, "createStreamResult") as ( + request: unknown, + abortController: AbortController + ) => unknown; + const model = createAnthropic({ apiKey: "test" })("claude-sonnet-4-5"); + const streamTextSpy = spyOn(aiSdk, "streamText").mockReturnValue({ + fullStream: (async function* asyncGenerator() { + yield* [] as unknown[]; + await Promise.resolve(); + })(), + usage: Promise.resolve(undefined), + providerMetadata: Promise.resolve(undefined), + totalUsage: Promise.resolve(undefined), + steps: Promise.resolve([]), + } as unknown as ReturnType); + const request = buildRequestConfig.call( + streamManager, + model, + modelString, + [{ role: "user", content: "hello" }], + system, + undefined, // routeProvider + tools, + undefined, // providerOptions + undefined, // maxOutputTokens + undefined, // callSettingsOverrides + undefined, // toolPolicy + undefined, // hasQueuedMessages + undefined, // headers + undefined, // anthropicCacheTtlOverride + undefined, // onChunk + undefined, // onStepMessages + undefined, // toolSearchState + undefined, // onToolExecutionStart + undefined, // thinkingOverrideState + undefined, // rebuildProviderOptionsForThinkingLevel + hook + ); + createStreamResultMethod.call(streamManager, request, new AbortController()); + // The spy may carry calls from earlier suites in the same process; only + // the call this helper just triggered matters. + expect(streamTextSpy.mock.calls.length).toBeGreaterThan(0); + const options = streamTextSpy.mock.calls[streamTextSpy.mock.calls.length - 1][0]; + const prepareStep = options.prepareStep as unknown as PrepareStepFn; + expect(typeof prepareStep).toBe("function"); + return { request, options, prepareStep }; + } + + function makeMemoryTool(description: string): Tool { + return tool({ + description, + inputSchema: z.object({}), + execute: () => Promise.resolve({ ok: true }), + }); + } + + afterEach(() => { + mock.restore(); + }); + + test("replaces the cached system message and memory tool description between steps", async () => { + let hookResult: MemoryRefreshResult | null = null; + const { options, prepareStep } = buildRequestWithHook( + "SYS-V1", + KNOWN_MODELS.SONNET.id, + { memory: makeMemoryTool("MEM-V1") }, + () => Promise.resolve(hookResult) + ); + // Anthropic path: system prompt travels as a cached system message. + const initialMessages = options.messages!; + expect(initialMessages[0]?.role).toBe("system"); + expect(initialMessages[0]?.content).toBe("SYS-V1"); + expect(options.system).toBeUndefined(); + + expect(await prepareStep({ messages: initialMessages })).toBeUndefined(); + + hookResult = { + system: "SYS-V2", + tools: { memory: makeMemoryTool("MEM-V2") }, + systemMessageTokens: 3, + }; + const result = await prepareStep({ messages: initialMessages }); + expect(result?.messages?.[0]?.content).toBe("SYS-V2"); + expect(result?.messages?.[0]?.providerOptions).toEqual(initialMessages[0]?.providerOptions); + expect(result?.instructions).toBeUndefined(); + // In-place update: streamText keeps reading the same tools reference. + const streamTextTools = options.tools as Record; + expect(streamTextTools.memory.description).toBe("MEM-V2"); + + // Already applied: the next step gets no redundant override. + expect(await prepareStep({ messages: result!.messages! })).toBeUndefined(); + }); + + test("overrides the system argument through per-step instructions", async () => { + let hookResult: MemoryRefreshResult | null = null; + const { request, prepareStep } = buildRequestWithHook( + "SYS-V1", + KNOWN_MODELS.GPT.id, + undefined, + () => Promise.resolve(hookResult) + ); + expect(request.system).toBe("SYS-V1"); + + const stepMessages: ModelMessage[] = [{ role: "user", content: "hello" }]; + expect(await prepareStep({ messages: stepMessages })).toBeUndefined(); + + hookResult = { system: "SYS-V2", tools: undefined, systemMessageTokens: 3 }; + const result = await prepareStep({ messages: stepMessages }); + expect(result?.instructions).toBe("SYS-V2"); + expect(request.system).toBe("SYS-V2"); + + expect(await prepareStep({ messages: stepMessages })).toBeUndefined(); + }); + + test("keeps structured system providerOptions and reports refreshed token counts", async () => { + let hookResult: MemoryRefreshResult | null = null; + const { request, prepareStep } = buildRequestWithHook( + "SYS-V1", + KNOWN_MODELS.GPT.id, + undefined, + () => Promise.resolve(hookResult) + ); + // Simulate the OpenAI explicit prompt-cache breakpoint shape: a + // structured system message whose providerOptions must survive refreshes. + const cacheOptions = { openai: { promptCacheBreakpoint: true } }; + request.system = { + role: "system", + content: "SYS-V1", + providerOptions: cacheOptions, + }; + const refreshState = ( + request as unknown as { + memoryRefreshState: { onApplied?: (r: { systemMessageTokens: number }) => void }; + } + ).memoryRefreshState; + const appliedTokenCounts: number[] = []; + refreshState.onApplied = ({ systemMessageTokens }) => { + appliedTokenCounts.push(systemMessageTokens); + }; + + const stepMessages: ModelMessage[] = [{ role: "user", content: "hello" }]; + hookResult = { system: "SYS-V2", tools: undefined, systemMessageTokens: 7 }; + const result = await prepareStep({ messages: stepMessages }); + const structured = result?.instructions as + | { role: string; content: string; providerOptions?: unknown } + | undefined; + expect(structured?.content).toBe("SYS-V2"); + expect(structured?.providerOptions).toEqual(cacheOptions); + expect(appliedTokenCounts).toEqual([7]); + }); +}); + describe("StreamManager - call settings overrides", () => { interface StreamRequestConfigForTests { model: unknown; @@ -1485,6 +1673,67 @@ describe("StreamManager - language model cleanup", () => { expect(getCleanupCalls()).toBe(1); }); + test("applies the refreshed memory-dependent context before registering the request", async () => { + const streamManager = new StreamManager(historyService); + const { model } = createCleanupModel("refresh-context-model"); + let capturedSystem: unknown; + const replaceCreateStreamResult = Reflect.set( + streamManager, + "createStreamResult", + (request: { system?: unknown }) => { + capturedSystem = request.system; + throw new Error("stop before processing"); + } + ); + expect(replaceCreateStreamResult).toBe(true); + + // The hook runs after startStream's awaited setup; its refreshed system + // must be what the registered request carries, not the call argument. + let refreshCalls = 0; + const result = await streamManager.startStream( + "refresh-context-workspace", + [{ role: "user", content: "hello" }], + model, + "openai:gpt-4.1-mini", + 1, + "stale-system", + runtime, + "refresh-context-message", + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + () => { + refreshCalls += 1; + return Promise.resolve({ + system: "refreshed-system", + tools: undefined, + systemMessageTokens: 42, + }); + } + ); + + expect(result.success).toBe(false); + expect(refreshCalls).toBe(1); + expect(capturedSystem).toBe("refreshed-system"); + }); + test("runs model cleanup when stream creation throws before processing", async () => { const streamManager = new StreamManager(historyService); const { model, getCleanupCalls } = createCleanupModel("cleanup-create-throw-model"); diff --git a/src/node/services/streamManager.ts b/src/node/services/streamManager.ts index 395fba973c..0cb31ee1aa 100644 --- a/src/node/services/streamManager.ts +++ b/src/node/services/streamManager.ts @@ -208,6 +208,13 @@ interface StreamRequestConfig { * options for the stream's model. `null` ⇒ not applicable / no-op. */ rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel; + /** + * Per-step memory revalidation: prepareStep consults the hook before every + * model step so a memory change while earlier steps ran (e.g. during tool + * execution) reaches the next provider request instead of waiting for the + * next turn. + */ + memoryRefreshState?: MemoryRefreshStepState; } /** @@ -231,6 +238,11 @@ export interface PreparedModelFallback { * that OpenAI rejects). */ tools: Record | undefined; + /** + * Per-step memory revalidation hook bound to the fallback model's rebuilt + * system/toolset, mirroring the primary request's hook. + */ + refreshMemoryDependentContext?: RefreshMemoryDependentContext; providerOptions?: Record; headers?: Record; callSettingsOverrides?: ResolvedCallSettingsOverrides; @@ -246,6 +258,38 @@ export interface PreparedModelFallback { rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel; } +/** + * Built by AIService to pull a fresh memory-dependent snapshot (system prompt, + * tool descriptions, token count) immediately before the stream request is + * registered: startStream's own awaited setup (workspace mutex, stream safety, + * temp-dir creation) opens a window in which a memory change would otherwise + * be invisible because system/tools were captured as call arguments. Returns + * null when nothing changed (or on a failed rebuild, where the request-boundary + * snapshot, which is complete and valid, is kept instead of failing the stream). + */ +export type RefreshMemoryDependentContext = () => Promise<{ + system: string; + tools: Record | undefined; + systemMessageTokens: number; +} | null>; + +/** + * Per-step memory revalidation state. `appliedSystemText` tracks the system + * prompt text currently in effect for this request so the cached system + * message embedded in `messages` (Anthropic path) can be replaced by exact + * match only. + */ +interface MemoryRefreshStepState { + hook: RefreshMemoryDependentContext; + appliedSystemText: string; + /** + * Wired once the stream registration exists; propagates the refreshed + * system token count into the live stream metadata so stream-end token + * attribution matches the system prompt the later steps actually sent. + */ + onApplied?: (refreshed: { systemMessageTokens: number }) => void; +} + export interface ModelFallbackPrepareOptions { /** * In-memory partial assistant turn to include when continuing after a @@ -1605,7 +1649,8 @@ export class StreamManager extends EventEmitter { toolSearchState?: ToolSearchStreamState, onToolExecutionStart?: (toolCallId: string) => void, thinkingOverrideState?: ActiveTurnThinkingOverride, - rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel + rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel, + refreshMemoryDependentContext?: RefreshMemoryDependentContext ): StreamRequestConfig { // Mid-turn thinking overrides mutate providerOptions IN PLACE (the SDK's // per-step deep-merge reads the object passed at streamText() time, so @@ -1687,6 +1732,10 @@ export class StreamManager extends EventEmitter { toolSearchState, thinkingOverrideState, rebuildProviderOptionsForThinkingLevel, + memoryRefreshState: + refreshMemoryDependentContext != null + ? { hook: refreshMemoryDependentContext, appliedSystemText: system } + : undefined, }; } @@ -1793,6 +1842,75 @@ export class StreamManager extends EventEmitter { return rebuilt.providerOptions; } + /** + * Re-run AIService's memory revalidation before a model step: steps are + * separate provider requests, so a memory change while earlier steps ran + * (e.g. during a long tool execution) would otherwise keep advertising + * renamed or deleted memory paths until the next turn. The SDK has no + * per-step tools override, so the memory tool description is updated IN + * PLACE (prepareTools re-reads descriptions from the same reference on + * every step, like the providerOptions deep-merge). The system prompt is + * swapped through the step return value: an `instructions` override when + * streamText received a system argument (the SDK persists a returned + * instructions override across subsequent steps, so later no-change steps + * keep it), or an exact-match replacement of the cached system message for + * providers that carry it inside `messages`. + * The hook resolves null on both no-change and failed rebuilds, so a cache + * rebuild failure keeps the current snapshot instead of failing the loop. + */ + private async applyPendingMemoryRefresh( + request: StreamRequestConfig, + stepMessages: ModelMessage[] + ): Promise< + { instructions?: string | SystemModelMessage; messages?: ModelMessage[] } | undefined + > { + const state = request.memoryRefreshState; + if (state == null) { + return undefined; + } + const refreshed = await state.hook(); + if (refreshed == null) { + return undefined; + } + const refreshedDescription = refreshed.tools?.memory?.description; + const requestMemoryTool = request.tools?.memory; + if (typeof refreshedDescription === "string" && requestMemoryTool != null) { + requestMemoryTool.description = refreshedDescription; + } + if (refreshed.system === state.appliedSystemText) { + return undefined; + } + if (request.system != null) { + // Structured system messages (OpenAI explicit prompt-cache breakpoint) + // keep their providerOptions: only the content is swapped, and the + // structured object is returned so the per-step override does not + // downgrade the request to a plain string. + request.system = + typeof request.system === "string" + ? refreshed.system + : { ...request.system, content: refreshed.system }; + state.appliedSystemText = refreshed.system; + state.onApplied?.({ systemMessageTokens: refreshed.systemMessageTokens }); + return { instructions: request.system }; + } + let replaced = false; + const nextMessages = stepMessages.map((message) => { + if (!replaced && message.role === "system" && message.content === state.appliedSystemText) { + replaced = true; + return { ...message, content: refreshed.system }; + } + return message; + }); + if (!replaced) { + // No exact-match cached system message; leave appliedSystemText alone + // so a later step can still find and replace the original. + return undefined; + } + state.appliedSystemText = refreshed.system; + state.onApplied?.({ systemMessageTokens: refreshed.systemMessageTokens }); + return { messages: nextMessages }; + } + private createStreamResult( request: StreamRequestConfig, abortController: AbortController, @@ -1819,7 +1937,10 @@ export class StreamManager extends EventEmitter { const withoutWorkflowRunRecords = stripWorkflowRunRecordsFromModelMessages(stepMessages); const rewritten = await extractToolMediaAsUserMessagesFromModelMessages(withoutWorkflowRunRecords); - const effectiveMessages = rewritten === stepMessages ? stepMessages : rewritten; + // Memory revalidation before this step's provider request is built: + // may replace the cached system message inside the step messages. + const memoryRefresh = await this.applyPendingMemoryRefresh(request, rewritten); + const effectiveMessages = memoryRefresh?.messages ?? rewritten; if (stepTracker) { stepTracker.latestMessages = effectiveMessages; } @@ -1834,14 +1955,18 @@ export class StreamManager extends EventEmitter { // this step's provider request is built. const thinkingOverride = this.applyPendingThinkingOverride(request); if ( - rewritten === stepMessages && + effectiveMessages === stepMessages && + memoryRefresh === undefined && activeTools === undefined && thinkingOverride === undefined ) { return undefined; } return { - ...(rewritten === stepMessages ? {} : { messages: rewritten }), + ...(effectiveMessages === stepMessages ? {} : { messages: effectiveMessages }), + ...(memoryRefresh?.instructions !== undefined + ? { instructions: memoryRefresh.instructions } + : {}), ...(activeTools !== undefined ? { activeTools } : {}), // Defense in depth: the in-place request mutation is authoritative // (per-step deep-merge cannot delete keys); returning the rebuilt @@ -1893,7 +2018,8 @@ export class StreamManager extends EventEmitter { modelFallback?: ModelFallbackOptions, toolSearchState?: ToolSearchStreamState, thinkingOverrideState?: ActiveTurnThinkingOverride, - rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel + rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel, + refreshMemoryDependentContext?: RefreshMemoryDependentContext ): WorkspaceStreamInfo { // abortController is created and linked to the caller-provided abortSignal in startStream(). @@ -1918,7 +2044,8 @@ export class StreamManager extends EventEmitter { toolSearchState, (toolCallId) => this.handleToolExecutionStart(workspaceId, messageId, toolCallId), thinkingOverrideState, - rebuildProviderOptionsForThinkingLevel + rebuildProviderOptionsForThinkingLevel, + refreshMemoryDependentContext ); // Start streaming - this can throw immediately if API key is missing @@ -1994,6 +2121,15 @@ export class StreamManager extends EventEmitter { } } + // Per-step memory refresh: route the refreshed system token count into + // this stream's metadata so stream-end attribution matches the system + // prompt the later steps actually sent. + if (request.memoryRefreshState) { + request.memoryRefreshState.onApplied = ({ systemMessageTokens }) => { + streamInfo.initialMetadata = { ...streamInfo.initialMetadata, systemMessageTokens }; + }; + } + // Atomically register the stream this.workspaceStreams.set(workspaceId, streamInfo); @@ -2630,8 +2766,17 @@ export class StreamManager extends EventEmitter { // hop) with a closure bound to the FALLBACK model. Attached before // createStreamResult below in case the SDK eagerly prepares step 1. streamInfo.request.thinkingOverrideState, - prepared.data.rebuildProviderOptionsForThinkingLevel + prepared.data.rebuildProviderOptionsForThinkingLevel, + prepared.data.refreshMemoryDependentContext ); + // Rebind the metadata sink for the swapped request, mirroring + // createStreamAtomically. Attached before createStreamResult below in + // case the SDK eagerly prepares step 1. + if (nextRequest.memoryRefreshState) { + nextRequest.memoryRefreshState.onApplied = ({ systemMessageTokens }) => { + streamInfo.initialMetadata = { ...streamInfo.initialMetadata, systemMessageTokens }; + }; + } // createStreamResult may eagerly prepare the first fallback step and update // latestMessages. Clear stale source-step messages before starting it so a // later disk-reset await cannot wipe freshly prepared fallback messages. @@ -4082,7 +4227,8 @@ export class StreamManager extends EventEmitter { modelFallback?: ModelFallbackOptions, toolSearchState?: ToolSearchStreamState, thinkingOverrideState?: ActiveTurnThinkingOverride, - rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel + rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel, + refreshMemoryDependentContext?: RefreshMemoryDependentContext ): Promise> { const typedWorkspaceId = workspaceId as WorkspaceId; @@ -4139,6 +4285,25 @@ export class StreamManager extends EventEmitter { return Ok(streamToken); } + // Memory can change while the awaited setup above runs; pull a fresh + // memory-dependent snapshot now, while everything from here to + // registration is synchronous, so the provider request cannot + // advertise a renamed or deleted memory path. + if (refreshMemoryDependentContext) { + const refreshed = await refreshMemoryDependentContext(); + if (streamAbortController.signal.aborted) { + return Ok(streamToken); + } + if (refreshed != null) { + system = refreshed.system; + tools = refreshed.tools; + initialMetadata = { + ...initialMetadata, + systemMessageTokens: refreshed.systemMessageTokens, + }; + } + } + // Step 4: Atomic stream creation and registration const streamInfo = this.createStreamAtomically( typedWorkspaceId, @@ -4168,7 +4333,8 @@ export class StreamManager extends EventEmitter { modelFallback, toolSearchState, thinkingOverrideState, - rebuildProviderOptionsForThinkingLevel + rebuildProviderOptionsForThinkingLevel, + refreshMemoryDependentContext ); // Guard against a narrow race: diff --git a/src/node/services/tools/memory.ts b/src/node/services/tools/memory.ts index cbf44f6921..6a35a3d540 100644 --- a/src/node/services/tools/memory.ts +++ b/src/node/services/tools/memory.ts @@ -49,14 +49,22 @@ export function resolveMemoryAccessPolicy(options: { * Build the dynamic memory tool description: the base description plus the * session-segment memory index (same disclosure mechanic as skills — index * advertised next to the tool schema, contents fetched on demand). Falls back - * to the base description when no snapshot was resolved. + * to the base description when no snapshot was resolved. Exported so the + * stream setup can refresh an already-created tool's description when a + * memory change lands during async tool assembly. */ -function buildMemoryDescription(config: ToolConfiguration): string { +export function buildMemoryToolDescription( + memoryIndexEntries: ToolConfiguration["memoryIndexEntries"] +): string { const baseDescription = TOOL_DEFINITIONS.memory.description; - if (config.memoryIndexEntries == null) { + if (memoryIndexEntries == null) { return baseDescription; } - return `${baseDescription}\n\n${formatMemoryIndexForToolDescription(config.memoryIndexEntries)}`; + return `${baseDescription}\n\n${formatMemoryIndexForToolDescription(memoryIndexEntries)}`; +} + +function buildMemoryDescription(config: ToolConfiguration): string { + return buildMemoryToolDescription(config.memoryIndexEntries); } /** diff --git a/src/node/services/workspaceService.test.ts b/src/node/services/workspaceService.test.ts index 44fa53583c..8c73257b01 100644 --- a/src/node/services/workspaceService.test.ts +++ b/src/node/services/workspaceService.test.ts @@ -17,6 +17,8 @@ import type { ProjectsConfig } from "@/common/types/project"; import type { Config } from "@/node/config"; import type { HistoryService } from "./historyService"; import { createTestHistoryService } from "./testHistoryService"; +import { MemoryService } from "./memoryService"; +import { MemoryMetaService } from "./memoryMeta"; import type { SessionTimingService } from "./sessionTimingService"; import { SessionUsageService } from "./sessionUsageService"; import type { AIService } from "./aiService"; @@ -3181,6 +3183,7 @@ describe("WorkspaceService truncateHistory goal acknowledgment", () => { } as unknown as AIService); const initStateManager = { on: mock(() => undefined), + off: mock(() => undefined), getInitState: mock(() => null), } as unknown as InitStateManager; const workspaceService = new WorkspaceService( @@ -3368,6 +3371,715 @@ describe("WorkspaceService truncateHistory goal acknowledgment", () => { } }); + function createAIServiceForSessionDispose(): AIService { + return { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + } as unknown as AIService; + } + + test("context reset produces a fresh memory context snapshot", async () => { + let indexEntries = [{ path: "/memories/global/phantom.md", description: "stale" }]; + const buildMemorySessionContext = mock(() => + Promise.resolve({ indexEntries, hotMemoriesBlock: null }) + ); + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + buildMemorySessionContext, + } as unknown as AIService; + const { config, historyService, workspaceService, cleanup } = await createServices(aiService); + const workspaceId = "context-reset-memory-invalidation"; + try { + await config.addWorkspace("/tmp/context-reset-memory-project", { + id: workspaceId, + name: workspaceId, + projectName: "context-reset-memory-project", + projectPath: "/tmp/context-reset-memory-project", + runtimeConfig: { type: "local" }, + }); + expect( + ( + await historyService.appendToHistory( + workspaceId, + createMuxMessage("pre-reset-user", "user", "before reset", {}) + ) + ).success + ).toBe(true); + + const session = workspaceService.getOrCreateSession(workspaceId); + const priv = session as unknown as { + resolveMemoryContext( + modelString: string + ): Promise<{ indexEntries: Array<{ path: string }> } | undefined>; + }; + expect((await priv.resolveMemoryContext("test-model"))?.indexEntries).toHaveLength(1); + + indexEntries = []; + expect(await workspaceService.resetContext(workspaceId)).toEqual({ + success: true, + data: "reset", + }); + + expect((await priv.resolveMemoryContext("test-model"))?.indexEntries).toHaveLength(0); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(2); + workspaceService.disposeSession(workspaceId); + } finally { + await cleanup(); + } + }); + + test("context reset invalidates a live startup-recovery session's memory context", async () => { + let indexEntries = [{ path: "/memories/global/phantom.md", description: "stale" }]; + const buildMemorySessionContext = mock(() => + Promise.resolve({ indexEntries, hotMemoriesBlock: null }) + ); + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + buildMemorySessionContext, + } as unknown as AIService; + const { config, historyService, workspaceService, cleanup } = await createServices(aiService); + const workspaceId = "context-reset-recovery-session"; + try { + await config.addWorkspace("/tmp/context-reset-recovery-project", { + id: workspaceId, + name: workspaceId, + projectName: "context-reset-recovery-project", + projectPath: "/tmp/context-reset-recovery-project", + runtimeConfig: { type: "local" }, + }); + expect( + ( + await historyService.appendToHistory( + workspaceId, + createMuxMessage("pre-reset-user", "user", "before reset", {}) + ) + ).success + ).toBe(true); + + const session = workspaceService.getOrCreateSession(workspaceId); + const maps = workspaceService as unknown as { + sessions: Map; + transientStartupRecoverySessions: Map; + }; + // Model a live startup-recovery session, which is tracked outside `sessions`. + maps.sessions.delete(workspaceId); + maps.transientStartupRecoverySessions.set(workspaceId, session); + const priv = session as unknown as { + resolveMemoryContext( + modelString: string + ): Promise<{ indexEntries: Array<{ path: string }> } | undefined>; + }; + expect((await priv.resolveMemoryContext("test-model"))?.indexEntries).toHaveLength(1); + + indexEntries = []; + expect(await workspaceService.resetContext(workspaceId)).toEqual({ + success: true, + data: "reset", + }); + + // The recovery session's cached snapshot was invalidated by the reset. + expect((await priv.resolveMemoryContext("test-model"))?.indexEntries).toHaveLength(0); + maps.transientStartupRecoverySessions.delete(workspaceId); + maps.sessions.set(workspaceId, session); + workspaceService.disposeSession(workspaceId); + } finally { + await cleanup(); + } + }); + + test("reset is rejected when a startup-recovery retry turns busy during the history read", async () => { + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + } as unknown as AIService; + const { config, historyService, workspaceService, cleanup } = await createServices(aiService); + const workspaceId = "context-reset-busy-recovery"; + try { + await config.addWorkspace("/tmp/context-reset-busy-recovery-project", { + id: workspaceId, + name: workspaceId, + projectName: "context-reset-busy-recovery-project", + projectPath: "/tmp/context-reset-busy-recovery-project", + runtimeConfig: { type: "local" }, + }); + expect( + ( + await historyService.appendToHistory( + workspaceId, + createMuxMessage("pre-reset-user", "user", "before reset", {}) + ) + ).success + ).toBe(true); + + const session = workspaceService.getOrCreateSession(workspaceId); + const maps = workspaceService as unknown as { + sessions: Map; + transientStartupRecoverySessions: Map; + }; + // Model a live startup-recovery session, which is tracked outside `sessions`. + maps.sessions.delete(workspaceId); + maps.transientStartupRecoverySessions.set(workspaceId, session); + + // The retry starts (session turns busy) while the reset awaits the + // history read: its request already contains pre-reset history, so the + // boundary must not be appended into that active turn. + let recoveryBusy = false; + spyOn(session, "isBusy").mockImplementation(() => recoveryBusy); + const realGetHistory = historyService.getHistoryFromLatestBoundary.bind(historyService); + spyOn(historyService, "getHistoryFromLatestBoundary").mockImplementation( + async (readWorkspaceId) => { + const result = await realGetHistory(readWorkspaceId); + recoveryBusy = true; + return result; + } + ); + + const result = await workspaceService.resetContext(workspaceId); + expect(result.success).toBe(false); + expect(result.success ? "" : result.error).toContain("turn is active"); + + // No reset boundary landed in history. + const history = await realGetHistory(workspaceId); + expect(history.success).toBe(true); + const boundaries = (history.success ? history.data : []).filter( + (message) => message.metadata?.contextBoundaryKind != null + ); + expect(boundaries).toHaveLength(0); + + maps.transientStartupRecoverySessions.delete(workspaceId); + maps.sessions.set(workspaceId, session); + workspaceService.disposeSession(workspaceId); + } finally { + await cleanup(); + } + }); + + test("reset is rejected while a startup-recovery retry is scheduled but not yet running", async () => { + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + } as unknown as AIService; + const { config, historyService, workspaceService, cleanup } = await createServices(aiService); + const workspaceId = "context-reset-scheduled-recovery"; + try { + await config.addWorkspace("/tmp/context-reset-scheduled-recovery-project", { + id: workspaceId, + name: workspaceId, + projectName: "context-reset-scheduled-recovery-project", + projectPath: "/tmp/context-reset-scheduled-recovery-project", + runtimeConfig: { type: "local" }, + }); + expect( + ( + await historyService.appendToHistory( + workspaceId, + createMuxMessage("pre-reset-user", "user", "before reset", {}) + ) + ).success + ).toBe(true); + + const session = workspaceService.getOrCreateSession(workspaceId); + const maps = workspaceService as unknown as { + sessions: Map; + transientStartupRecoverySessions: Map; + }; + maps.sessions.delete(workspaceId); + maps.transientStartupRecoverySessions.set(workspaceId, session); + + // The retry timer is armed but has not fired: the session is not busy, + // yet the retry could start mid-reset (e.g. during the boundary append) + // and build its request from pre-reset history. + spyOn(session, "hasPendingAutoRetry").mockReturnValue(true); + + const result = await workspaceService.resetContext(workspaceId); + expect(result.success).toBe(false); + expect(result.success ? "" : result.error).toContain("queued user input is pending"); + + const history = await historyService.getHistoryFromLatestBoundary(workspaceId); + expect(history.success).toBe(true); + const boundaries = (history.success ? history.data : []).filter( + (message) => message.metadata?.contextBoundaryKind != null + ); + expect(boundaries).toHaveLength(0); + + maps.transientStartupRecoverySessions.delete(workspaceId); + maps.sessions.set(workspaceId, session); + workspaceService.disposeSession(workspaceId); + } finally { + await cleanup(); + } + }); + + test("reset is rejected while a startup-recovery check is still deciding", async () => { + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + } as unknown as AIService; + const { config, historyService, workspaceService, cleanup } = await createServices(aiService); + const workspaceId = "context-reset-inflight-recovery"; + try { + await config.addWorkspace("/tmp/context-reset-inflight-recovery-project", { + id: workspaceId, + name: workspaceId, + projectName: "context-reset-inflight-recovery-project", + projectPath: "/tmp/context-reset-inflight-recovery-project", + runtimeConfig: { type: "local" }, + }); + expect( + ( + await historyService.appendToHistory( + workspaceId, + createMuxMessage("pre-reset-user", "user", "before reset", {}) + ) + ).success + ).toBe(true); + + const session = workspaceService.getOrCreateSession(workspaceId); + const maps = workspaceService as unknown as { + sessions: Map; + transientStartupRecoverySessions: Map; + }; + maps.sessions.delete(workspaceId); + maps.transientStartupRecoverySessions.set(workspaceId, session); + + // Hold the check inside its history reads: the session is not busy and + // no retry is scheduled yet, so only the in-flight state can tell the + // reset that a retry decision from pre-reset history may be imminent. + let releaseRead: () => void = () => undefined; + const readGate = new Promise((resolve) => { + releaseRead = resolve; + }); + spyOn(historyService, "getLastMessages").mockImplementation(async () => { + await readGate; + // A settled assistant tail concludes the released check without + // scheduling a retry. + return { + success: true as const, + data: [createMuxMessage("done", "assistant", "done", {})], + }; + }); + + session.ensureStartupAutoRetryCheck(); + expect(session.isStartupRecoveryInFlight()).toBe(true); + + const result = await workspaceService.resetContext(workspaceId); + expect(result.success).toBe(false); + expect(result.success ? "" : result.error).toContain("startup recovery"); + + releaseRead(); + await (session as unknown as { startupAutoRetryCheckPromise: Promise | null }) + .startupAutoRetryCheckPromise; + expect(session.isStartupRecoveryInFlight()).toBe(false); + + const history = await historyService.getHistoryFromLatestBoundary(workspaceId); + expect(history.success).toBe(true); + const boundaries = (history.success ? history.data : []).filter( + (message) => message.metadata?.contextBoundaryKind != null + ); + expect(boundaries).toHaveLength(0); + + maps.transientStartupRecoverySessions.delete(workspaceId); + maps.sessions.set(workspaceId, session); + workspaceService.disposeSession(workspaceId); + } finally { + await cleanup(); + } + }); + + test("a retry racing the reset's acknowledgment sees fresh memory state", async () => { + let indexEntries = [{ path: "/memories/global/phantom.md", description: "stale" }]; + const buildMemorySessionContext = mock(() => + Promise.resolve({ indexEntries, hotMemoriesBlock: null }) + ); + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + buildMemorySessionContext, + } as unknown as AIService; + const { config, historyService, workspaceService, cleanup } = await createServices(aiService); + const workspaceId = "context-reset-retry-race"; + try { + await config.addWorkspace("/tmp/context-reset-retry-race-project", { + id: workspaceId, + name: workspaceId, + projectName: "context-reset-retry-race-project", + projectPath: "/tmp/context-reset-retry-race-project", + runtimeConfig: { type: "local" }, + }); + expect( + ( + await historyService.appendToHistory( + workspaceId, + createMuxMessage("pre-reset-user", "user", "before reset", {}) + ) + ).success + ).toBe(true); + + const session = workspaceService.getOrCreateSession(workspaceId); + const maps = workspaceService as unknown as { + sessions: Map; + transientStartupRecoverySessions: Map; + }; + // Model a live startup-recovery session, invisible to the busy guards. + maps.sessions.delete(workspaceId); + maps.transientStartupRecoverySessions.set(workspaceId, session); + const priv = session as unknown as { + resolveMemoryContext( + modelString: string + ): Promise<{ indexEntries: Array<{ path: string }> } | undefined>; + }; + expect((await priv.resolveMemoryContext("test-model"))?.indexEntries).toHaveLength(1); + indexEntries = []; + + // Model a startup-recovery retry firing while the reset is awaiting + // goal acknowledgment: the boundary is already appended, so the retry + // must not build from the pre-reset memory snapshot. + let entriesSeenDuringAck: number | undefined; + workspaceService.setWorkspaceGoalService({ + requireUserAcknowledgment: mock(async () => { + entriesSeenDuringAck = (await priv.resolveMemoryContext("test-model"))?.indexEntries + .length; + }), + } as unknown as WorkspaceGoalService); + + expect(await workspaceService.resetContext(workspaceId)).toEqual({ + success: true, + data: "reset", + }); + expect(entriesSeenDuringAck).toBe(0); + + maps.transientStartupRecoverySessions.delete(workspaceId); + maps.sessions.set(workspaceId, session); + workspaceService.disposeSession(workspaceId); + } finally { + await cleanup(); + } + }); + + test("memory delete drops phantom index entries from the next-built context", async () => { + let memoryService: MemoryService | undefined; + let memoryCtx: Parameters[0] | undefined; + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + buildMemorySessionContext: mock(async () => ({ + indexEntries: await memoryService!.listIndexEntries(memoryCtx!), + hotMemoriesBlock: null, + })), + } as unknown as AIService; + const { config, workspaceService, cleanup } = await createServices(aiService); + const workspaceId = "memory-change-invalidation"; + try { + memoryService = new MemoryService(config, new MemoryMetaService(config.rootDir)); + memoryCtx = { runtime: null, checkoutCwd: "", workspaceId, projectPath: "" }; + // Mirrors the coreServices wiring. + memoryService.on("change", () => workspaceService.invalidateMemoryContexts()); + + expect( + (await memoryService.create(memoryCtx, "/memories/global/phantom.md", "stale", "agent")) + .success + ).toBe(true); + + const session = workspaceService.getOrCreateSession(workspaceId); + const priv = session as unknown as { + resolveMemoryContext( + modelString: string + ): Promise<{ indexEntries: Array<{ path: string }> } | undefined>; + }; + expect( + (await priv.resolveMemoryContext("test-model"))?.indexEntries.map((entry) => entry.path) + ).toEqual(["/memories/global/phantom.md"]); + + expect( + (await memoryService.deletePath(memoryCtx, "/memories/global/phantom.md", "agent")).success + ).toBe(true); + + expect((await priv.resolveMemoryContext("test-model"))?.indexEntries).toEqual([]); + workspaceService.disposeSession(workspaceId); + } finally { + await cleanup(); + } + }); + + test("memory rename retargets index entries in the next-built context", async () => { + let memoryService: MemoryService | undefined; + let memoryCtx: Parameters[0] | undefined; + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + buildMemorySessionContext: mock(async () => ({ + indexEntries: await memoryService!.listIndexEntries(memoryCtx!), + hotMemoriesBlock: null, + })), + } as unknown as AIService; + const { config, workspaceService, cleanup } = await createServices(aiService); + const workspaceId = "memory-rename-invalidation"; + try { + memoryService = new MemoryService(config, new MemoryMetaService(config.rootDir)); + memoryCtx = { runtime: null, checkoutCwd: "", workspaceId, projectPath: "" }; + // Mirrors the coreServices wiring. + memoryService.on("change", () => workspaceService.invalidateMemoryContexts()); + + expect( + (await memoryService.create(memoryCtx, "/memories/global/old-name.md", "keep", "agent")) + .success + ).toBe(true); + + const session = workspaceService.getOrCreateSession(workspaceId); + const priv = session as unknown as { + resolveMemoryContext( + modelString: string + ): Promise<{ indexEntries: Array<{ path: string }> } | undefined>; + }; + expect( + (await priv.resolveMemoryContext("test-model"))?.indexEntries.map((entry) => entry.path) + ).toEqual(["/memories/global/old-name.md"]); + + expect( + ( + await memoryService.rename( + memoryCtx, + "/memories/global/old-name.md", + "/memories/global/new-name.md", + "agent" + ) + ).success + ).toBe(true); + + expect( + (await priv.resolveMemoryContext("test-model"))?.indexEntries.map((entry) => entry.path) + ).toEqual(["/memories/global/new-name.md"]); + workspaceService.disposeSession(workspaceId); + } finally { + await cleanup(); + } + }); + + test("memory pin changes invalidate cached contexts", async () => { + const buildMemorySessionContext = mock(() => + Promise.resolve({ indexEntries: [], hotMemoriesBlock: null }) + ); + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + buildMemorySessionContext, + } as unknown as AIService; + const { config, workspaceService, cleanup } = await createServices(aiService); + const workspaceId = "memory-pin-invalidation"; + try { + const memoryService = new MemoryService(config, new MemoryMetaService(config.rootDir)); + const memoryCtx = { runtime: null, checkoutCwd: "", workspaceId, projectPath: "" }; + // Mirrors the coreServices wiring. + memoryService.on("change", () => workspaceService.invalidateMemoryContexts()); + + expect( + (await memoryService.create(memoryCtx, "/memories/global/pinnable.md", "body", "agent")) + .success + ).toBe(true); + + const session = workspaceService.getOrCreateSession(workspaceId); + const priv = session as unknown as { + resolveMemoryContext(modelString: string): Promise; + }; + await priv.resolveMemoryContext("test-model"); + await priv.resolveMemoryContext("test-model"); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(1); + + // Pinning feeds hot-set selection, so it must drop the cached context. + await memoryService.setPinned(memoryCtx, "/memories/global/pinnable.md", true, "user"); + + await priv.resolveMemoryContext("test-model"); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(2); + workspaceService.disposeSession(workspaceId); + } finally { + await cleanup(); + } + }); + + test("workspace-scoped memory changes do not invalidate unrelated sessions", async () => { + const buildMemorySessionContext = mock(() => + Promise.resolve({ indexEntries: [], hotMemoriesBlock: null }) + ); + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + buildMemorySessionContext, + } as unknown as AIService; + const { workspaceService, cleanup } = await createServices(aiService); + try { + const targetSession = workspaceService.getOrCreateSession("ws-target"); + const otherSession = workspaceService.getOrCreateSession("ws-other"); + const priv = (session: unknown) => + session as { resolveMemoryContext(modelString: string): Promise }; + await priv(targetSession).resolveMemoryContext("test-model"); + await priv(otherSession).resolveMemoryContext("test-model"); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(2); + + workspaceService.invalidateMemoryContexts({ + scope: "workspace", + workspaceId: "ws-target", + projectPath: "", + }); + + // Only the target session rebuilds; the other stays cached. + await priv(targetSession).resolveMemoryContext("test-model"); + await priv(otherSession).resolveMemoryContext("test-model"); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(3); + + // Global changes still fan out to everyone. + workspaceService.invalidateMemoryContexts({ + scope: "global", + workspaceId: "", + projectPath: "", + }); + await priv(targetSession).resolveMemoryContext("test-model"); + await priv(otherSession).resolveMemoryContext("test-model"); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(5); + + workspaceService.disposeSession("ws-target"); + workspaceService.disposeSession("ws-other"); + } finally { + await cleanup(); + } + }); + + test("project-scoped memory changes match the memory identity, including scratch", async () => { + const buildMemorySessionContext = mock(() => + Promise.resolve({ indexEntries: [], hotMemoriesBlock: null }) + ); + const aiService = { + on: mock(() => undefined), + off: mock(() => undefined), + isStreaming: mock(() => false), + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + buildMemorySessionContext, + } as unknown as AIService; + const { config, workspaceService, cleanup } = await createServices(aiService); + try { + const projectPath = "/tmp/memory-scope-project"; + await config.addWorkspace(projectPath, { + id: "ws-project-scope", + name: "ws-project-scope", + projectName: "memory-scope-project", + projectPath, + runtimeConfig: { type: "local" }, + }); + const scratchDir = "/tmp/mux/scratch/ws-scratch-scope"; + await config.addWorkspace(SCRATCH_PROJECT_CONFIG_KEY, { + kind: "scratch", + id: "ws-scratch-scope", + name: "ws-scratch-scope", + projectName: "Scratch", + projectPath: scratchDir, + runtimeConfig: { type: "local" }, + namedWorkspacePath: scratchDir, + }); + + const projectSession = workspaceService.getOrCreateSession("ws-project-scope"); + const scratchSession = workspaceService.getOrCreateSession("ws-scratch-scope"); + const priv = (session: unknown) => + session as { resolveMemoryContext(modelString: string): Promise }; + await priv(projectSession).resolveMemoryContext("test-model"); + await priv(scratchSession).resolveMemoryContext("test-model"); + expect(buildMemorySessionContext).toHaveBeenCalledTimes(2); + + // Scratch chats emit project events keyed by the scratch directory + // (their memory project identity), not the "_scratch" config bucket key. + workspaceService.invalidateMemoryContexts({ + scope: "project", + workspaceId: "", + projectPath: scratchDir, + }); + await priv(projectSession).resolveMemoryContext("test-model"); + await priv(scratchSession).resolveMemoryContext("test-model"); + // Only the scratch session rebuilt. + expect(buildMemorySessionContext).toHaveBeenCalledTimes(3); + + workspaceService.invalidateMemoryContexts({ + scope: "project", + workspaceId: "", + projectPath, + }); + await priv(projectSession).resolveMemoryContext("test-model"); + await priv(scratchSession).resolveMemoryContext("test-model"); + // Only the project session rebuilt. + expect(buildMemorySessionContext).toHaveBeenCalledTimes(4); + + workspaceService.disposeSession("ws-project-scope"); + workspaceService.disposeSession("ws-scratch-scope"); + } finally { + await cleanup(); + } + }); + + test("invalidateMemoryContexts reaches live startup-recovery sessions", async () => { + const { workspaceService, cleanup } = await createServices(createAIServiceForSessionDispose()); + const workspaceId = "startup-recovery-memory-invalidation"; + try { + const session = workspaceService.getOrCreateSession(workspaceId); + const maps = workspaceService as unknown as { + sessions: Map; + transientStartupRecoverySessions: Map; + }; + // Model a live startup-recovery session, which is tracked outside `sessions`. + maps.sessions.delete(workspaceId); + maps.transientStartupRecoverySessions.set(workspaceId, session); + const invalidateSpy = spyOn(session, "invalidateMemoryContext"); + + workspaceService.invalidateMemoryContexts(); + + expect(invalidateSpy).toHaveBeenCalledTimes(1); + maps.transientStartupRecoverySessions.delete(workspaceId); + maps.sessions.set(workspaceId, session); + workspaceService.disposeSession(workspaceId); + } finally { + await cleanup(); + } + }); + + test("invalidateMemoryContexts invalidates every live session", async () => { + const { workspaceService, cleanup } = await createServices(createAIServiceForSessionDispose()); + try { + const sessionA = workspaceService.getOrCreateSession("memory-invalidate-a"); + const sessionB = workspaceService.getOrCreateSession("memory-invalidate-b"); + const spyA = spyOn(sessionA, "invalidateMemoryContext"); + const spyB = spyOn(sessionB, "invalidateMemoryContext"); + + workspaceService.invalidateMemoryContexts(); + + expect(spyA).toHaveBeenCalledTimes(1); + expect(spyB).toHaveBeenCalledTimes(1); + workspaceService.disposeSession("memory-invalidate-a"); + workspaceService.disposeSession("memory-invalidate-b"); + } finally { + await cleanup(); + } + }); + test("context reset is a no-op when repeated without provider-eligible messages", async () => { const { config, historyService, workspaceService, cleanup } = await createServices(); const workspaceId = "context-reset-noop"; diff --git a/src/node/services/workspaceService.ts b/src/node/services/workspaceService.ts index b15d74f2b3..5cdd21a04f 100644 --- a/src/node/services/workspaceService.ts +++ b/src/node/services/workspaceService.ts @@ -257,6 +257,8 @@ import { import type { WorkspaceLifecycleHooks } from "@/node/services/workspaceLifecycleHooks"; import type { TaskService } from "@/node/services/taskService"; import { findWorkspaceEntry } from "@/node/services/taskUtils"; +import type { MemoryChangeEvent } from "@/node/services/memoryService"; +import { resolveMemoryProjectIdentityFromConfigEntry } from "@/node/services/memoryService"; import type { WorktreeArchiveSnapshotService } from "@/node/services/worktreeArchiveSnapshotService"; import { DisposableTempDir } from "@/node/services/tempDir"; @@ -1856,6 +1858,10 @@ export class WorkspaceService extends EventEmitter { // Blocks new sends while a context reset is committing its durable boundary and cleanup. private readonly resettingContextWorkspaces = new Set(); + // Counts appended reset boundaries per workspace so startup auto-retry + // checks can detect a boundary that landed after their history reads. + private readonly contextResetBoundaryEpochs = new Map(); + // Tracks in-flight fork auto-title generations so only the first accepted continue // message can claim the workspace title. private readonly autoTitlingWorkspaces = new Set(); @@ -3142,6 +3148,10 @@ export class WorkspaceService extends EventEmitter { initStateManager: this.initStateManager, workspaceGoalService: this.workspaceGoalService, backgroundProcessManager: this.backgroundProcessManager, + getContextResetState: () => ({ + resetting: this.resettingContextWorkspaces.has(workspaceId), + boundaryEpoch: this.contextResetBoundaryEpochs.get(workspaceId) ?? 0, + }), onCompactionComplete: (metadata) => { this.schedulePostCompactionMetadataRefresh(workspaceId); // Compaction marks a long session with accumulated learnings: harvest @@ -9218,7 +9228,11 @@ export class WorkspaceService extends EventEmitter { } hasPendingQueuedOrPreparingTurn(workspaceId: string): boolean { - const session = this.sessions.get(workspaceId.trim()); + const trimmed = workspaceId.trim(); + // Startup-recovery sessions live outside `sessions` but can hold a + // scheduled auto-retry that will start a turn on its own. + const session = + this.sessions.get(trimmed) ?? this.transientStartupRecoverySessions.get(trimmed); if (!session) { return false; } @@ -9358,6 +9372,53 @@ export class WorkspaceService extends EventEmitter { return Ok(undefined); } + /** + * Invalidates session-cached memory contexts. When a change event is given, + * only sessions that can observe the changed file are invalidated: global + * memory fans out to every session, workspace memory to that workspace, and + * project memory to sessions in the same project (project identity, not + * checkout path, per MemoryChangeEvent). + */ + invalidateMemoryContexts( + event?: Pick + ): void { + // Synchronous on purpose: this runs inside MemoryService's change emit, + // so mutation routes cannot return (and the next stream cannot start) + // with a stale cache still in place. + let cfg: ReturnType | undefined; + const affects = (workspaceId: string): boolean => { + if (!event || event.scope === "global") { + return true; + } + if (event.scope === "workspace") { + return workspaceId === event.workspaceId; + } + cfg ??= this.config.loadConfigOrDefault(); + const entry = findWorkspaceEntry(cfg, workspaceId); + if (entry == null) { + return false; + } + // Project events carry the memory project identity (scratch workspaces + // use their scratch directory, not the config bucket key), so matching + // must resolve the same identity the event was built from. + return ( + resolveMemoryProjectIdentityFromConfigEntry(entry.projectPath, entry.workspace) === + event.projectPath + ); + }; + + for (const [workspaceId, session] of this.sessions) { + if (affects(workspaceId)) { + session.invalidateMemoryContext(); + } + } + for (const [workspaceId, session] of this.transientStartupRecoverySessions) { + if (affects(workspaceId)) { + session.invalidateMemoryContext(); + } + } + } + async resetContext(workspaceId: string): Promise> { if (this.resettingContextWorkspaces.has(workspaceId)) { return Err("Context reset is already in progress for this workspace."); @@ -9365,17 +9426,34 @@ export class WorkspaceService extends EventEmitter { this.resettingContextWorkspaces.add(workspaceId); try { - const session = this.sessions.get(workspaceId); - if (session?.isBusy() || this.aiService.isStreaming(workspaceId)) { - return Err( - "Cannot reset context while a turn is active. Press Esc to stop the stream first." - ); - } - - if (this.hasPendingQueuedOrPreparingTurn(workspaceId)) { - return Err( - "Cannot reset context while queued user input is pending. Send or clear the queued message first." - ); + // Startup-recovery retries run on a session tracked outside `sessions`, + // so the busy guard must resolve both maps, freshly on each check. + const isTurnActive = (): boolean => + this.sessions.get(workspaceId)?.isBusy() === true || + this.transientStartupRecoverySessions.get(workspaceId)?.isBusy() === true || + this.aiService.isStreaming(workspaceId); + const turnBlockError = (): string | undefined => { + if (isTurnActive()) { + return "Cannot reset context while a turn is active. Press Esc to stop the stream first."; + } + if (this.hasPendingQueuedOrPreparingTurn(workspaceId)) { + return "Cannot reset context while queued user input is pending. Send or clear the queued message first."; + } + // A startup-recovery check that is mid-decision is neither busy nor + // holding a scheduled retry yet, but it can start a retry from + // pre-reset history as soon as its reads settle. + if ( + this.sessions.get(workspaceId)?.isStartupRecoveryInFlight() === true || + this.transientStartupRecoverySessions.get(workspaceId)?.isStartupRecoveryInFlight() === + true + ) { + return "Cannot reset context while startup recovery is checking for an interrupted turn. Try again in a moment."; + } + return undefined; + }; + const entryBlockError = turnBlockError(); + if (entryBlockError != null) { + return Err(entryBlockError); } const historyResult = await this.historyService.getHistoryFromLatestBoundary(workspaceId); @@ -9390,6 +9468,16 @@ export class WorkspaceService extends EventEmitter { return Ok("noop"); } + // The history read above is an await window: a startup-recovery retry + // starting during it already built its request from pre-reset history, + // and appending the boundary would land it inside that active turn. A + // retry that is merely scheduled must block too: its timer can fire + // during the boundary append below, which has no later recheck. + const preAppendBlockError = turnBlockError(); + if (preAppendBlockError != null) { + return Err(preAppendBlockError); + } + const boundaryMessage = createMuxMessage( createContextResetBoundaryMessageId(), "assistant", @@ -9404,14 +9492,30 @@ export class WorkspaceService extends EventEmitter { if (!appendResult.success) { return Err(`Failed to append context reset boundary: ${appendResult.error}`); } + // Bump before anything else observes the append: a startup auto-retry + // check whose history reads predate this boundary compares epochs at + // its retry decision and defers instead of resuming pre-reset context. + this.contextResetBoundaryEpochs.set( + workspaceId, + (this.contextResetBoundaryEpochs.get(workspaceId) ?? 0) + 1 + ); const typedBoundaryMessage = { ...boundaryMessage, type: "message" as const }; - if (session) { - session.emitChatEvent(typedBoundaryMessage); + const liveSession = this.sessions.get(workspaceId); + if (liveSession) { + liveSession.emitChatEvent(typedBoundaryMessage); } else { this.emit("chat", { workspaceId, message: typedBoundaryMessage }); } + // Invalidate synchronously before any further await: a startup-recovery + // retry (tracked outside `sessions`, rejected by the busy rechecks above + // only while it is running) could otherwise start during the + // acknowledgment await and build its provider request from the + // pre-reset memory snapshot. + const resetSession = liveSession ?? this.transientStartupRecoverySessions.get(workspaceId); + resetSession?.invalidateMemoryContext(); + try { await this.workspaceGoalService?.requireUserAcknowledgment(workspaceId); } catch (error) {