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) {