diff --git a/src/common/schemas/project.ts b/src/common/schemas/project.ts index bd1c33dc78e..9d575be479b 100644 --- a/src/common/schemas/project.ts +++ b/src/common/schemas/project.ts @@ -105,6 +105,10 @@ export const WorkspaceConfigSchema = z.object({ description: "If set, this workspace is a child workspace spawned from the parent workspaceId (enables nesting in UI and backend orchestration).", }), + memoryOwnerWorkspaceId: z.string().optional().meta({ + description: + "Memory owner pinned when an intermediate ancestor was removed while this descendant stayed alive: the parentWorkspaceId chain no longer reaches the task-tree root, so this keeps /memories/workspace bound to the root's store (memoryWorkspaceOwner.ts). Set only by workspace removal.", + }), agentType: z.string().optional().meta({ description: 'If set, selects an agent preset for this workspace (e.g., "explore" or "exec").', }), diff --git a/src/node/services/memoryConsolidationService.test.ts b/src/node/services/memoryConsolidationService.test.ts index f2d31f3bbc3..00d20304066 100644 --- a/src/node/services/memoryConsolidationService.test.ts +++ b/src/node/services/memoryConsolidationService.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from "bun:test"; +import { Effect } from "effect"; import * as fsPromises from "node:fs/promises"; import * as path from "node:path"; @@ -7,7 +8,10 @@ import type { LanguageModelV3CallOptions, LanguageModelV3StreamPart } from "@ai- import type { CompactionCompletionMetadata } from "@/common/types/compaction"; import { createMuxMessage } from "@/common/types/message"; -import type { MemoryConsolidationStatusChangeEventPayload } from "@/common/orpc/schemas/memory"; +import type { + MemoryConsolidationStatusChangeEventPayload, + MemoryHarvestRecordPayload, +} from "@/common/orpc/schemas/memory"; import { MULTI_PROJECT_CONFIG_KEY } from "@/common/constants/multiProject"; import { EXPERIMENT_IDS } from "@/common/constants/experiments"; import { @@ -17,6 +21,7 @@ import { import { Ok } from "@/common/types/result"; import { Config } from "@/node/config"; import { + HARVEST_MAX_ATTEMPTS, MemoryConsolidationService, resolveDreamAgentBody, resolveDreamModelString, @@ -1503,6 +1508,102 @@ describe("MemoryConsolidationService", () => { expect(fixture.modelCalls).toHaveLength(3); }); + it("releases the teardown gate when a removal aborts before its point of no return", async () => { + using fixture = await createFixture(); + // The removal drain marks the workspace; every trigger is refused... + await fixture.service.cancelInFlightConsolidation("ws-dream"); + const refused = await fixture.service.maybeRun("ws-dream", "manual"); + expect(refused.success).toBe(false); + if (!refused.success) expect(refused.error).toContain("being removed"); + expect(fixture.modelCalls).toHaveLength(0); + // ...until the aborted removal (no tombstone, workspace intact) lifts it. + fixture.service.releaseRemovalCancellation("ws-dream"); + expect((await fixture.service.maybeRun("ws-dream", "manual")).success).toBe(true); + expect(fixture.modelCalls).toHaveLength(1); + }); + + it("finalizes a removed workspace's retryable harvest records so they are never retried", async () => { + using fixture = await createFixture({ modelFactory: harvestCandidateModel }); + await fixture.addWorkspace("ws-sub", { parentWorkspaceId: "ws-dream" }); + const metadata = await seedCompactionEpoch(fixture, "ws-sub"); + await fsPromises.writeFile( + path.join(fixture.xumHome, "memory-consolidation.json"), + JSON.stringify({ + workspaces: {}, + harvestsByWorkspace: { + "ws-sub": { + [metadata.summaryMessageId]: { + status: "failed", + startedAt: Date.now() - 10_000, + completedAt: Date.now() - 9_000, + attemptCount: 1, + boundaryKey: metadata.summaryMessageId, + compactionEpoch: metadata.compactionEpoch, + acceptedCandidates: 0, + skippedCandidates: 0, + error: "crashed mid-harvest", + completionMetadata: metadata, + }, + }, + }, + }) + ); + await fixture.service.finalizeHarvestsForRemoval("ws-sub"); + const record = (await fixture.service.getStatus("ws-sub")).latestHarvestRecord; + expect(record?.status).toBe("failed"); + expect(record?.attemptCount).toBe(HARVEST_MAX_ATTEMPTS); + // The owner's run no longer sees a retryable child bucket. + expect((await fixture.service.maybeRun("ws-dream", "manual")).success).toBe(true); + expect((await fixture.service.getStatus("ws-sub")).latestHarvestRecord?.attemptCount).toBe( + HARVEST_MAX_ATTEMPTS + ); + }); + + it("keeps a removal-finalized harvest record terminal against residual retryable writes", async () => { + using fixture = await createFixture({ modelFactory: harvestCandidateModel }); + await fixture.addWorkspace("ws-sub", { parentWorkspaceId: "ws-dream" }); + const metadata = await seedCompactionEpoch(fixture, "ws-sub"); + const boundaryKey = metadata.summaryMessageId; + const base = { + startedAt: Date.now() - 10_000, + attemptCount: 1, + boundaryKey, + compactionEpoch: metadata.compactionEpoch, + acceptedCandidates: 0, + skippedCandidates: 0, + completionMetadata: metadata, + }; + // Residual runs of the bounded cancellation drain record through the same + // path as the live harvest; reach it directly to interleave with finalization. + const save = (record: MemoryHarvestRecordPayload) => + Effect.runPromise( + ( + fixture.service as unknown as { + saveHarvestRecordEffect: ( + workspaceId: string, + boundaryKey: string, + record: MemoryHarvestRecordPayload, + projectPath: string + ) => Effect.Effect; + } + ).saveHarvestRecordEffect("ws-sub", boundaryKey, record, "") + ); + await save({ ...base, status: "pending" }); + await fixture.service.finalizeHarvestsForRemoval("ws-sub"); + const latest = async () => (await fixture.service.getStatus("ws-sub")).latestHarvestRecord; + expect((await latest())?.attemptCount).toBe(HARVEST_MAX_ATTEMPTS); + + // A residual retryable failure landing after finalization must not reopen the bucket... + await save({ ...base, status: "failed", completedAt: Date.now(), error: "residual failure" }); + expect((await latest())?.attemptCount).toBe(HARVEST_MAX_ATTEMPTS); + expect((await latest())?.error).toContain("workspace removed"); + // ...while a residual completion (its writes really landed) is kept as the truth, + // and finalization never demotes a completed record. + await save({ ...base, status: "completed", completedAt: Date.now(), acceptedCandidates: 1 }); + await fixture.service.finalizeHarvestsForRemoval("ws-sub"); + expect((await latest())?.status).toBe("completed"); + }); + it("launch sweep skips archived workspaces and caps runs per launch", async () => { using fixture = await createFixture(); const dayAgo = Date.now() - 25 * 60 * 60 * 1000; diff --git a/src/node/services/memoryConsolidationService.ts b/src/node/services/memoryConsolidationService.ts index 33fe3676308..f7c5a7a7fbf 100644 --- a/src/node/services/memoryConsolidationService.ts +++ b/src/node/services/memoryConsolidationService.ts @@ -291,7 +291,26 @@ function pruneHarvestRecords(records: Record): void } } -const HARVEST_MAX_ATTEMPTS = 3; +export const HARVEST_MAX_ATTEMPTS = 3; + +/** Completed, or failed with retries exhausted: nothing may retry it. */ +function isTerminalHarvestRecord(record: MemoryHarvestRecord): boolean { + return ( + record.status === "completed" || + (record.status === "failed" && record.attemptCount >= HARVEST_MAX_ATTEMPTS) + ); +} + +/** Terminal marker for a bucket whose transcript is being deleted (see finalizeHarvestsForRemoval). */ +function finalizeHarvestRecordForRemoval(record: MemoryHarvestRecord): MemoryHarvestRecord { + return { + ...record, + status: "failed", + completedAt: record.completedAt ?? Date.now(), + attemptCount: HARVEST_MAX_ATTEMPTS, + error: "workspace removed before the harvest could be retried; transcript no longer available", + }; +} export class MemoryConsolidationService extends EventEmitter { private readonly sidecarPath: string; @@ -324,10 +343,9 @@ export class MemoryConsolidationService extends EventEmitter { * post-harvest sweep, and a cancelled run still starts retryable-harvest * recovery, each with a fresh un-aborted signal. Entry points refuse and * new controllers start pre-aborted while a workspace is in this set. - * Entries are never cleared: removal is terminal, and if a force=false - * removal fails after the drain, losing background consolidation for the - * surviving workspace (until restart) matches the documented drained- - * producers tradeoff in WorkspaceService.removeWorkspace. Cross-PROCESS + * Entries are cleared only when removal aborts before its point of no + * return (releaseRemovalCancellation); once the tombstone is published, + * removal is terminal. Cross-PROCESS * teardown is covered by the durable removal tombstone instead (see * workspaceRemoval.ts), checked at memory mutation commit points. */ @@ -485,16 +503,27 @@ export class MemoryConsolidationService extends EventEmitter { const self = this; return Effect.uninterruptible( Effect.gen(function* () { - yield* Effect.promise(() => + const saved = yield* Effect.promise(() => self.locks.withLock(self.sidecarPath, async () => { const file = await self.load(); file.harvestsByWorkspace[workspaceId] ??= {}; + const existing = file.harvestsByWorkspace[workspaceId][boundaryKey]; + // A terminal record is never reopened: removal finalization + // (finalizeHarvestsForRemoval) races the bounded cancellation + // drain's residual harvest runs on this file, and a residual + // pending/retryable-failure write landing afterwards would turn + // a bucket whose transcript is gone back into a retry candidate. + // Only a genuine completion may replace it (the writes happened). + if (existing !== undefined && isTerminalHarvestRecord(existing)) { + if (record.status !== "completed") return false; + } file.harvestsByWorkspace[workspaceId][boundaryKey] = record; pruneHarvestRecords(file.harvestsByWorkspace[workspaceId]); await writeFileAtomic(self.sidecarPath, JSON.stringify(file, null, 2)); + return true; }) ); - self.emitStatusChange(workspaceId, projectPath); + if (saved) self.emitStatusChange(workspaceId, projectPath); }) ); } @@ -626,6 +655,51 @@ export class MemoryConsolidationService extends EventEmitter { return Effect.runPromise(this.cancelInFlightConsolidationEffect(workspaceId)); } + /** + * Removal aborted BEFORE its point of no return (no tombstone published, the + * workspace stays registered and intact — e.g. a non-forced removal whose + * checkout deletion was refused): lift the teardown gate again, or the + * surviving workspace would refuse every Dream run and post-compaction + * harvest until restart. The drained in-flight runs are gone regardless + * (retryable harvests recover on the next trigger). + */ + releaseRemovalCancellation(workspaceId: string): void { + this.removalCancelled.delete(workspaceId); + } + + /** + * Removal teardown for harvest state: the workspace's transcript is about + * to be deleted, so its failed/stale-pending harvest records can never be + * retried (recovery needs the compaction epoch's messages) — and once the + * config entry is gone they could not even be associated with the memory + * owner. Mark them terminal now so nothing lingers as "retryable". + */ + async finalizeHarvestsForRemoval(workspaceId: string): Promise { + // One read-check-write under the sidecar lock: residual harvest runs + // (cancelInFlightConsolidation's drain is bounded) may still be recording + // outcomes, and a completion landing between an unlocked read and this + // write must not be overwritten with a failure. + const finalized = await this.locks.withLock(this.sidecarPath, async () => { + const file = await this.load(); + const records = file.harvestsByWorkspace[workspaceId]; + if (records === undefined) return false; + let changed = false; + for (const [boundaryKey, record] of Object.entries(records)) { + if (isTerminalHarvestRecord(record)) continue; + records[boundaryKey] = finalizeHarvestRecordForRemoval(record); + changed = true; + } + if (changed) await writeFileAtomic(this.sidecarPath, JSON.stringify(file, null, 2)); + return changed; + }); + if (!finalized) return; + const workspace = this.config.findWorkspace(workspaceId); + this.emitStatusChange( + workspaceId, + workspace == null ? "" : resolveConsolidationProjectPath(workspace) + ); + } + /** * Teardown pipeline: uninterruptible end-to-end so the r61 mark, the abort * loop, and the residual-run handoff can never be separated, with the @@ -846,11 +920,18 @@ export class MemoryConsolidationService extends EventEmitter { } const projectPath = resolveConsolidationProjectPath(workspace); + // A child's redirected run sweeps under the owner's identity; the + // child's own removal tombstone still refuses every read and commit of + // the run (MemoryScopeContext.guardedWorkspaceId) — a remover in another + // backend cannot abort this controller. const ctx: MemoryScopeContext = { runtime: null, checkoutCwd: "", workspaceId, projectPath, + ...(options.actingWorkspaceId !== undefined && options.actingWorkspaceId !== workspaceId + ? { guardedWorkspaceId: options.actingWorkspaceId } + : {}), }; const result = yield* Effect.promise(async () => diff --git a/src/node/services/memoryService.test.ts b/src/node/services/memoryService.test.ts index f8d7bc5f8ae..5c9b3b85f15 100644 --- a/src/node/services/memoryService.test.ts +++ b/src/node/services/memoryService.test.ts @@ -6,6 +6,7 @@ import { createHash } from "node:crypto"; import * as fsPromises from "node:fs/promises"; import * as path from "node:path"; import { Config } from "@/node/config"; +import { getErrorMessage } from "@/common/utils/errors"; import { LocalRuntime } from "@/node/runtime/LocalRuntime"; import { extractMemoryDescription, @@ -1244,6 +1245,249 @@ describe("MemoryService", () => { unreadable.mockRestore(); expect(fixture.service.resolveWorkspaceMemoryOwnerId("ws-child")).toBe("ws-owner"); }); + it("keeps a grandchild on the root store via the pinned owner after its parent is removed", async () => { + using fixture = await createFixture("ws-grandchild"); + await registerTaskTree(fixture); + // Removal of the intermediate "ws-child" pins memoryOwnerWorkspaceId on + // its children before deregistering it (WorkspaceService.remove). + await fixture.config.editConfig((cfg) => { + const project = cfg.projects.get(FIXTURE_PROJECT_PATH)!; + for (const ws of project.workspaces) { + if (ws.parentWorkspaceId === "ws-child") ws.memoryOwnerWorkspaceId = "ws-owner"; + } + project.workspaces = project.workspaces.filter((ws) => ws.id !== "ws-child"); + return cfg; + }); + expect(fixture.service.resolveWorkspaceMemoryOwnerId("ws-grandchild")).toBe("ws-owner"); + const created = await fixture.service.create( + fixture.ctx, + "/memories/workspace/still-shared.md", + "root store", + "agent" + ); + expect(created.success).toBe(true); + expect( + await pathExists( + path.join(fixture.config.sessionsDir, "ws-owner", "memory", "still-shared.md") + ) + ).toBe(true); + // A pinned owner that is itself gone falls back to self. + await fixture.config.editConfig((cfg) => { + const project = cfg.projects.get(FIXTURE_PROJECT_PATH)!; + project.workspaces = project.workspaces.filter((ws) => ws.id !== "ws-owner"); + return cfg; + }); + expect(fixture.service.resolveWorkspaceMemoryOwnerId("ws-grandchild")).toBe("ws-grandchild"); + }); + + it("re-resolves the owner per command and refuses reads once the owner is tombstoned", async () => { + using fixture = await createFixture("ws-child"); + await registerTaskTree(fixture); + await fixture.service.create(fixture.ctx, "/memories/workspace/n.md", "shared", "agent"); + // One context serves a whole stream (createMemoryTool): a cached owner + // must not outlive the command that resolved it. + expect(fixture.service.ownerWorkspaceIdFor(fixture.ctx)).toBe("ws-owner"); + const resolve = spyOn(fixture.service, "resolveWorkspaceMemoryOwnerId"); + expect((await fixture.service.view(fixture.ctx, "/memories/workspace/n.md")).success).toBe( + true + ); + expect(resolve).toHaveBeenCalled(); + + // Another backend removed the owner: its durable tombstone (no local + // event) must stop the child's reads of the shared notebook. + const tombstonePath = workspaceRemovalTombstonePath(fixture.xumHome, "ws-owner"); + await fsPromises.mkdir(path.dirname(tombstonePath), { recursive: true }); + await fsPromises.writeFile(tombstonePath, JSON.stringify({ workspaceId: "ws-owner" })); + const refused = await fixture.service.view(fixture.ctx, "/memories/workspace/n.md"); + expect(refused.success).toBe(false); + if (!refused.success) expect(refused.error).toContain("was removed"); + const root = await fixture.service.view(fixture.ctx, "/memories"); + expect(root.success).toBe(true); + if (root.success) expect(root.output).toContain("unavailable"); + // The prompt-context path is guarded too: the index no longer lists + // the store. + expect( + (await fixture.service.listIndexEntries(fixture.ctx)).some( + (entry) => entry.scope === "workspace" + ) + ).toBe(false); + }); + + it("guards a context acting on a removed child's behalf like the child itself", async () => { + using fixture = await createFixture("ws-child"); + await registerTaskTree(fixture); + // A child's consolidation run sweeps under the OWNER's identity; the + // child's removal by another backend (tombstone, no local signal) must + // still refuse that run's reads and commits in every scope. + const ownerCtx = { ...fixture.ctx, workspaceId: "ws-owner", guardedWorkspaceId: "ws-child" }; + await fixture.service.create(ownerCtx, "/memories/workspace/n.md", "shared", "agent"); + await fixture.service.create(ownerCtx, "/memories/global/g.md", "global", "agent"); + const tombstonePath = workspaceRemovalTombstonePath(fixture.xumHome, "ws-child"); + await fsPromises.mkdir(path.dirname(tombstonePath), { recursive: true }); + await fsPromises.writeFile(tombstonePath, JSON.stringify({ workspaceId: "ws-child" })); + for (const attempt of [ + () => fixture.service.view(ownerCtx, "/memories/workspace/n.md"), + () => + fixture.service.strReplace(ownerCtx, "/memories/workspace/n.md", "shared", "x", "agent"), + () => fixture.service.strReplace(ownerCtx, "/memories/global/g.md", "global", "x", "agent"), + () => fixture.service.create(ownerCtx, "/memories/project/p.md", "p", "agent"), + ]) { + const result = await attempt(); + expect(result.success).toBe(false); + if (!result.success) expect(result.error).toContain("ws-child was removed"); + } + expect( + await fsPromises.readFile( + path.join(fixture.config.sessionsDir, "ws-owner", "memory", "n.md"), + "utf-8" + ) + ).toBe("shared"); + // The owner's own contexts are unaffected. + const plainOwner = { ...fixture.ctx, workspaceId: "ws-owner" }; + expect((await fixture.service.view(plainOwner, "/memories/workspace/n.md")).success).toBe( + true + ); + }); + + it("withholds a read whose workspace was tombstoned after the pre-read check", async () => { + using fixture = await createFixture("ws-child"); + await registerTaskTree(fixture); + await fixture.service.create(fixture.ctx, "/memories/workspace/n.md", "shared", "agent"); + await fixture.service.create(fixture.ctx, "/memories/global/g.md", "global", "agent"); + const tombstonePath = workspaceRemovalTombstonePath(fixture.xumHome, "ws-child"); + // Another backend's removal lands between the check that opened the + // store and the read itself: every path that exposes the store's bytes + // or listing re-checks before returning them. + const service = fixture.service as unknown as { + openWorkspaceStore: (...args: unknown[]) => Promise; + }; + const original = service.openWorkspaceStore.bind(fixture.service); + const tombstoneAfterOpen = () => + spyOn(service, "openWorkspaceStore").mockImplementationOnce(async (...args) => { + await original(...args); + await fsPromises.mkdir(path.dirname(tombstonePath), { recursive: true }); + await fsPromises.writeFile(tombstonePath, JSON.stringify({ workspaceId: "ws-child" })); + }); + const untombstone = () => fsPromises.rm(tombstonePath, { force: true }); + + tombstoneAfterOpen(); + const file = await fixture.service.view(fixture.ctx, "/memories/workspace/n.md"); + expect(file.success).toBe(false); + if (!file.success) expect(file.error).toContain("was removed"); + await untombstone(); + + // The usage record waits for the owner-store lock, which a removal + // holds while it publishes the tombstone: landing there, after the + // bytes were read, must still withhold them. + const usageService = fixture.service as unknown as { + recordUsage: (...args: unknown[]) => Promise; + }; + const originalUsage = usageService.recordUsage.bind(fixture.service); + const usage = spyOn(usageService, "recordUsage").mockImplementationOnce(async (...args) => { + await originalUsage(...args); + await fsPromises.mkdir(path.dirname(tombstonePath), { recursive: true }); + await fsPromises.writeFile(tombstonePath, JSON.stringify({ workspaceId: "ws-child" })); + }); + const lateFile = await fixture.service.view(fixture.ctx, "/memories/workspace/n.md"); + expect(usage).toHaveBeenCalledTimes(1); + expect(lateFile.success).toBe(false); + if (!lateFile.success) expect(lateFile.error).toContain("was removed"); + usage.mockRestore(); + await untombstone(); + + tombstoneAfterOpen(); + const dir = await fixture.service.view(fixture.ctx, "/memories/workspace"); + expect(dir.success).toBe(false); + if (!dir.success) expect(dir.error).toContain("was removed"); + await untombstone(); + + tombstoneAfterOpen(); + const root = await fixture.service.view(fixture.ctx, "/memories"); + expect(root.success).toBe(true); + if (root.success) { + expect(root.output).toContain("unavailable"); + expect(root.output).not.toContain("n.md"); + } + await untombstone(); + + tombstoneAfterOpen(); + const entries = await fixture.service.listIndexEntries(fixture.ctx); + expect(entries.map((entry) => entry.scope)).toEqual(["global"]); + await untombstone(); + + tombstoneAfterOpen(); + const ui = await fixture.service.readFileWithSha(fixture.ctx, "/memories/workspace/n.md"); + expect(ui.success).toBe(false); + await untombstone(); + + // Hot-set reads happen after the index enumeration passed: the + // tombstone landing before the file read drops the item. + const hotBefore = await fixture.service.listHotMemories(fixture.ctx, { + countTokens: (text) => Promise.resolve(text.length), + }); + expect(hotBefore.some((item) => item.path === "/memories/workspace/n.md")).toBe(true); + // Interleaving: the tombstone lands after listIndexEntries built the + // candidate list and before the hot-set file reads. + const originalList = fixture.service.listIndexEntries.bind(fixture.service); + const listIndex = spyOn(fixture.service, "listIndexEntries").mockImplementationOnce( + async (ctx) => { + const result = await originalList(ctx); + await fsPromises.mkdir(path.dirname(tombstonePath), { recursive: true }); + await fsPromises.writeFile(tombstonePath, JSON.stringify({ workspaceId: "ws-child" })); + return result; + } + ); + try { + const hot = await fixture.service.listHotMemories(fixture.ctx, { + countTokens: (text) => Promise.resolve(text.length), + }); + expect(hot.some((item) => item.path === "/memories/workspace/n.md")).toBe(false); + expect(hot.some((item) => item.path === "/memories/global/g.md")).toBe(true); + } finally { + listIndex.mockRestore(); + await untombstone(); + } + // Selection keeps awaiting token counts after the file reads: a + // tombstone landing there still withholds the workspace items. + let counted = 0; + const hotAfterCount = await fixture.service.listHotMemories(fixture.ctx, { + countTokens: async (text) => { + if (counted++ === 0) { + await fsPromises.mkdir(path.dirname(tombstonePath), { recursive: true }); + await fsPromises.writeFile(tombstonePath, JSON.stringify({ workspaceId: "ws-child" })); + } + return text.length; + }, + }); + expect(counted).toBeGreaterThan(0); + expect(hotAfterCount.some((item) => item.path === "/memories/workspace/n.md")).toBe(false); + expect(hotAfterCount.some((item) => item.path === "/memories/global/g.md")).toBe(true); + await untombstone(); + }); + + it("refuses a pin toggle once the owner it was bound to is tombstoned", async () => { + using fixture = await createFixture("ws-child"); + await registerTaskTree(fixture); + await fixture.service.create(fixture.ctx, "/memories/workspace/n.md", "shared", "agent"); + const events: unknown[] = []; + fixture.service.on("change", (event) => events.push(event)); + // Owner removed by another backend between the tab's owner resolution + // and the pin's lock acquisition: the pin must not be committed under + // the dead owner's logical key while the route reports success. + const tombstonePath = workspaceRemovalTombstonePath(fixture.xumHome, "ws-owner"); + await fsPromises.mkdir(path.dirname(tombstonePath), { recursive: true }); + await fsPromises.writeFile(tombstonePath, JSON.stringify({ workspaceId: "ws-owner" })); + const refused = await fixture.service + .setPinned(fixture.ctx, "/memories/workspace/n.md", true) + .then( + () => null, + (error: unknown) => error + ); + expect(refused).toBeInstanceOf(Error); + expect(getErrorMessage(refused)).toContain("was removed"); + expect((await fixture.metaService.getPinnedKeys()).size).toBe(0); + expect(events).toEqual([]); + }); }); describe("memory index entries", () => { diff --git a/src/node/services/memoryService.ts b/src/node/services/memoryService.ts index b23c08f4d34..bd6ee6d380f 100644 --- a/src/node/services/memoryService.ts +++ b/src/node/services/memoryService.ts @@ -88,6 +88,15 @@ export interface MemoryScopeContext { * and sidecar logical keys; empty when no project identity is available. */ projectPath: string; + /** + * A further workspace on whose behalf this context acts, guarded like the + * acting one: a sub-agent's consolidation run sweeps the OWNER's notebook + * under the owner's identity (`workspaceId`), and the child's removal — + * possibly by another backend, which cannot abort this run — must refuse + * every read and commit of that run at the tombstone check, not only its + * start (r77). + */ + guardedWorkspaceId?: string; } export type MemoryActor = "agent" | "user"; @@ -739,7 +748,7 @@ export class MemoryService extends EventEmitter { * index+hot-set build and reused for every entry within it, so the stamp * stat behind resolveWorkspaceMemoryOwnerId runs once per operation instead * of once per candidate file. Staleness is bounded to that one operation; - * writes are still gated by the commit check (assertMutationCommittable). + * writes are still gated by the store-bound tombstone check. */ private readonly ownerByContext = new WeakMap(); @@ -947,11 +956,85 @@ export class MemoryService extends EventEmitter { relPath: string ): Promise { const store = this.getStore(ctx, scope); + if (scope === "workspace") await this.openWorkspaceStore(ctx, store); await store.assertRootSafe(); await store.assertContained(relPath); return store; } + /** + * Every workspace-scope entry point (commands, root listing, index build) + * goes through here: refuse revoked access. + */ + private async openWorkspaceStore(ctx: MemoryScopeContext, store: MemoryStore): Promise { + await this.assertWorkspaceStoreReadable(ctx, store); + } + + /** + * The workspace whose session dir physically holds `store` (the memory + * owner a workspace-scope store was resolved to), or null for stores that + * are not session-bound (global/project), which live elsewhere. + */ + private storeOwnerWorkspaceId(store: MemoryStore): string | null { + const rel = path.relative(this.config.sessionsDir, store.physicalRoot); + if (rel === "" || rel.startsWith("..") || path.isAbsolute(rel)) return null; + return rel.split(path.sep)[0]; + } + + /** Acting workspace plus the store's owner: both must be alive to touch the store. */ + private guardedWorkspaceIds(ctx: MemoryScopeContext, store: MemoryStore): string[] { + const owner = this.storeOwnerWorkspaceId(store); + return [ + ...new Set([ + ctx.workspaceId, + ...(owner === null ? [] : [owner]), + ...(ctx.guardedWorkspaceId === undefined || ctx.guardedWorkspaceId === "" + ? [] + : [ctx.guardedWorkspaceId]), + ]), + ]; + } + + /** + * Reads have no commit guard, so a removed child's stream in ANOTHER backend + * (which the remover cannot cancel) could keep viewing its former owner's + * notebook — including notes written after the removal — through the + * shared store. Refuse workspace-scope reads once the acting workspace or + * the store's owner is tombstoned (the tombstone is durable and + * cross-process; see workspaceRemoval.ts). + */ + private async assertWorkspaceStoreReadable( + ctx: MemoryScopeContext, + store: MemoryStore + ): Promise { + if (ctx.workspaceId === "") return; + for (const workspaceId of this.guardedWorkspaceIds(ctx, store)) { + if (await isWorkspaceRemovalTombstoned(this.config.rootDir, workspaceId)) { + throw new MemoryCommandError( + `Workspace ${workspaceId} was removed; the workspace memory store is no longer available` + ); + } + } + } + + /** + * Post-read gate for workspace-scope reads. openWorkspaceStore checks the + * tombstones BEFORE the read; another backend can publish the acting + * workspace's (or the owner's) removal tombstone while the read is in + * flight, and the bytes would then be exposed on behalf of a workspace that + * no longer exists. Re-checked after every read whose result leaves the + * service (view, listings, index/hot-set builds, UI reads), before the + * result is returned. Other scopes are never shared and have no tombstone. + */ + private async assertWorkspaceReadExposable( + ctx: MemoryScopeContext, + scope: MemoryScope, + store: MemoryStore + ): Promise { + if (scope !== "workspace") return; + await this.assertWorkspaceStoreReadable(ctx, store); + } + private requireFilePath(parsed: ParsedMemoryPath, virtualPath: string): MemoryScope { if (parsed.scope === null || parsed.relPath === "") { throw new MemoryCommandError( @@ -1021,15 +1104,24 @@ export class MemoryService extends EventEmitter { * removal here at commit time and refuses instead of recreating the * deleted session directory via its write or journal append. * - * The owner the command's store was bound to is then compared with a fresh - * resolution: the per-context cache (ownerWorkspaceIdFor) may hold a - * self-fallback taken while config.json was missing or malformed, and if - * the file recovers before this command commits, the write would land in - * the child's private store although the tree is shared again. Refused as - * a recoverable error; the retried command resolves the owner anew. + * Both the acting workspace and the workspace that physically owns the + * RESOLVED store are checked: a removed sub-agent must not keep writing + * into its parent's notebook, and a removed owner must not have its session + * directory recreated by a lingering child's write. The owner is derived + * from the store the command already bound to — not re-resolved — so an + * ownership change between resolution and lock acquisition cannot make the + * check pass for the new owner while the write lands in the old one. + * + * The bound owner is then compared with a fresh resolution: the + * per-context cache (ownerWorkspaceIdFor) may hold a self-fallback taken + * while config.json was missing or malformed, and if the file recovers + * before this command commits, the write would land in the child's private + * store although the tree is shared again. Refused as a recoverable error; + * the retried command resolves the owner anew. */ private async assertMutationCommittable( ctx: MemoryScopeContext, + store: MemoryStore, signal: AbortSignal | undefined, virtualPath: string ): Promise { @@ -1039,8 +1131,8 @@ export class MemoryService extends EventEmitter { ); } if (ctx.workspaceId === "") return; - const boundOwner = this.ownerByContext.get(ctx); - if (boundOwner !== undefined) { + const boundOwner = this.storeOwnerWorkspaceId(store); + if (boundOwner !== null) { const currentOwner = this.resolveWorkspaceMemoryOwnerId(ctx.workspaceId); if (currentOwner !== boundOwner) { throw new MemoryCommandError( @@ -1048,16 +1140,7 @@ export class MemoryService extends EventEmitter { ); } } - // The acting workspace AND, for the workspace scope, the store owner: a - // child's mutation waits on the owner's store lock while the owner is - // removed (tombstone published, session dir deleted), then resumes and - // would recreate the owner's directory. Global/project stores are not - // the owner's, so a removed owner does not refuse those. - const storeOwner = - parseMemoryPath(virtualPath).scope === "workspace" - ? this.ownerWorkspaceIdFor(ctx) - : undefined; - for (const workspaceId of new Set([ctx.workspaceId, storeOwner ?? ctx.workspaceId])) { + for (const workspaceId of this.guardedWorkspaceIds(ctx, store)) { if (await isWorkspaceRemovalTombstoned(this.config.rootDir, workspaceId)) { throw new MemoryCommandError( `Workspace ${workspaceId} was removed; refusing to commit the mutation of ${virtualPath}` @@ -1216,7 +1299,7 @@ export class MemoryService extends EventEmitter { // the owner the context resolved BEFORE the lock. If ownership moved // meanwhile, the pin would land under a dead logical key and the // route would still report success. Refuse instead. - await this.assertMutationCommittable(ctx, undefined, virtualPath); + await this.assertMutationCommittable(ctx, store, undefined, virtualPath); await this.metaService.setPinned(key, pinned); }); } else { @@ -1264,9 +1347,11 @@ export class MemoryService extends EventEmitter { sections.push(`- ${scope}/`); try { const store = this.getStore(ctx, scope); + if (scope === "workspace") await this.openWorkspaceStore(ctx, store); // Read-only: never create roots just to list (missing ⇒ empty). await store.assertRootSafe(); const files = await store.listFiles(); + await this.assertWorkspaceReadExposable(ctx, scope, store); sections.push(...renderTree(files, MEMORY_VIEW_MAX_DEPTH - 1, " ")); } catch (error) { // Self-healing: an unavailable scope must not break the whole view. @@ -1283,6 +1368,7 @@ export class MemoryService extends EventEmitter { // write — but the scope itself always exists in the protocol. if (kind === "dir" || (kind === null && parsed.relPath === "")) { const files = await store.listFiles(); + await this.assertWorkspaceReadExposable(ctx, parsed.scope, store); const prefix = parsed.relPath === "" ? "" : `${parsed.relPath}/`; const scopedFiles = files .filter((file) => file.startsWith(prefix)) @@ -1300,6 +1386,9 @@ export class MemoryService extends EventEmitter { const content = await this.readBoundedTextFile(store, parsed.relPath, virtualPath); const output = renderFileView(content, options); await this.recordUsage(ctx, parsed.scope, parsed.relPath, { write: false }); + // AFTER recordUsage — the last await before the content leaves: a + // tombstone published meanwhile must still withhold the bytes. + await this.assertWorkspaceReadExposable(ctx, parsed.scope, store); return { success: true, output }; }); } @@ -1322,7 +1411,7 @@ export class MemoryService extends EventEmitter { // only INSIDE the lock and after the removal check (r62), so the // mkdir serializes with removal's locked deletion and cannot // recreate a removed session directory. - await this.assertMutationCommittable(ctx, abortSignal, virtualPath); + await this.assertMutationCommittable(ctx, store, abortSignal, virtualPath); await store.ensureRoot(); const existing = await store.kind(parsed.relPath); if (existing !== null) { @@ -1336,7 +1425,7 @@ export class MemoryService extends EventEmitter { `The ${scope} memory scope is full (${MEMORY_MAX_FILES_PER_SCOPE} files); delete unused files first` ); } - await this.assertMutationCommittable(ctx, abortSignal, virtualPath); + await this.assertMutationCommittable(ctx, store, abortSignal, virtualPath); await store.writeFile(parsed.relPath, fileText); // Row is written before the create is acknowledged (mutation → row → ack). await this.journalRefinement( @@ -1377,7 +1466,7 @@ export class MemoryService extends EventEmitter { const content = await this.readTextFileForEdit(store, parsed.relPath, virtualPath); const updated = computeStrReplaceUpdate(content, oldStr, newStr, virtualPath); assertWithinFileSizeCap(updated); - await this.assertMutationCommittable(ctx, abortSignal, virtualPath); + await this.assertMutationCommittable(ctx, store, abortSignal, virtualPath); await store.writeFile(parsed.relPath, updated); // Row is written before the edit is acknowledged (mutation → row → ack). await this.journalRefinement( @@ -1430,7 +1519,7 @@ export class MemoryService extends EventEmitter { const content = await this.readTextFileForEdit(store, parsed.relPath, virtualPath); const { updated, insertedLineCount } = computeInsertUpdate(content, insertLine, insertText); assertWithinFileSizeCap(updated); - await this.assertMutationCommittable(ctx, abortSignal, virtualPath); + await this.assertMutationCommittable(ctx, store, abortSignal, virtualPath); await store.writeFile(parsed.relPath, updated); // Row is written before the edit is acknowledged (mutation → row → ack). await this.journalRefinement( @@ -1490,7 +1579,7 @@ export class MemoryService extends EventEmitter { } const store = await this.resolveStore(ctx, scope, parsed.relPath); return withTargetMutationLock(this.config.rootDir, this.storeLockKey(store), async () => { - await this.assertMutationCommittable(ctx, abortSignal, virtualPath); + await this.assertMutationCommittable(ctx, store, abortSignal, virtualPath); await store.ensureRoot(); const kind = await store.kind(parsed.relPath); if (kind === "dir") { @@ -1525,7 +1614,7 @@ export class MemoryService extends EventEmitter { mutation.insertText ).updated; assertWithinFileSizeCap(updated, maxFileBytes); - await this.assertMutationCommittable(ctx, abortSignal, virtualPath); + await this.assertMutationCommittable(ctx, store, abortSignal, virtualPath); await store.writeFile(parsed.relPath, updated); const physicalPath = store.physicalPath(parsed.relPath); // Row is written before the write is acknowledged (mutation → row → ack). @@ -1702,7 +1791,7 @@ export class MemoryService extends EventEmitter { // Prior contents must be captured before removal; the row itself is // written after the mutation succeeds and before it is acknowledged. const inverse = await this.captureDeleteInverse(store, parsed.relPath, kind); - await this.assertMutationCommittable(ctx, abortSignal, virtualPath); + await this.assertMutationCommittable(ctx, store, abortSignal, virtualPath); await store.remove(parsed.relPath); if (inverse !== null) { await this.journalRefinement( @@ -1768,7 +1857,7 @@ export class MemoryService extends EventEmitter { if (newKind !== null) { throw new MemoryCommandError(`Destination ${newVirtualPath} already exists`); } - await this.assertMutationCommittable(ctx, abortSignal, oldVirtualPath); + await this.assertMutationCommittable(ctx, store, abortSignal, oldVirtualPath); await store.rename(oldParsed.relPath, newParsed.relPath); // Row is written before the rename is acknowledged (mutation → row → ack). await this.journalRefinement( @@ -1851,6 +1940,7 @@ export class MemoryService extends EventEmitter { const scope = this.requireFilePath(parsed, virtualPath); const store = await this.resolveStore(ctx, scope, parsed.relPath); const content = await this.readTextFileForEdit(store, parsed.relPath, virtualPath); + await this.assertWorkspaceReadExposable(ctx, scope, store); // Deliberately NOT recorded as a use: this is a human browsing the // Memory tab/settings, and usage stats must reflect agent reads only so // UI browsing never inflates hot-set ranking. (UI saves still count — @@ -1892,7 +1982,7 @@ export class MemoryService extends EventEmitter { async () => { // UI save can create new files: materialize the scope root on // first use — in-lock, after the removal check (r62; see create). - await this.assertMutationCommittable(ctx, abortSignal, virtualPath); + await this.assertMutationCommittable(ctx, store, abortSignal, virtualPath); await store.ensureRoot(); const kind = await store.kind(parsed.relPath); if (kind === "dir") { @@ -1919,7 +2009,7 @@ export class MemoryService extends EventEmitter { ); } } - await this.assertMutationCommittable(ctx, abortSignal, virtualPath); + await this.assertMutationCommittable(ctx, store, abortSignal, virtualPath); await store.writeFile(parsed.relPath, content); await this.recordUsage(ctx, scope, parsed.relPath, { write: true }); this.emitChange(ctx, scope, parsed.relPath, actor); @@ -1947,8 +2037,17 @@ export class MemoryService extends EventEmitter { async listIndexEntries(ctx: MemoryScopeContext): Promise { const entries: MemoryIndexEntry[] = []; for (const scope of MEMORY_SCOPES) { + // Per-scope buffer: the scope's entries join the result only once the + // post-read gate below passed, so a tombstone published mid-enumeration + // drops the whole scope rather than a prefix of it. + const scopeEntries: MemoryIndexEntry[] = []; try { const store = this.getStore(ctx, scope); + // Prompt context is a read of the (possibly shared) store: a removed + // child's stream in another backend must not keep indexing / hot-set + // reading its former owner's notes. Refused here (skipped below) like + // any other scope failure. + if (scope === "workspace") await this.openWorkspaceStore(ctx, store); // Read-only enumeration (stream startup, Memory tab) must not create // scope roots unnecessarily. Missing roots list as empty. await store.assertRootSafe(); @@ -1996,8 +2095,10 @@ export class MemoryService extends EventEmitter { } catch { // Unreadable file: list it without a description. } - entries.push({ path: toVirtualPath(scope, relPath), scope, relPath, description }); + scopeEntries.push({ path: toVirtualPath(scope, relPath), scope, relPath, description }); } + await this.assertWorkspaceReadExposable(ctx, scope, store); + entries.push(...scopeEntries); } catch (error) { log.debug("[MemoryService] skipping scope in memory index", { scope, error }); } @@ -2031,24 +2132,41 @@ export class MemoryService extends EventEmitter { lastAccessedAt: stats?.lastAccessedAt ?? null, }; }); - return selectHotMemories({ + const selected = await selectHotMemories({ candidates, countTokens: options.countTokens, tokenBudgetActive: options.tokenBudgetActive, onlyContextNotes: options.onlyContextNotes, - readFile: (virtualPath) => { + readFile: async (virtualPath) => { const parsed = parseMemoryPath(virtualPath); const scope = this.requireFilePath(parsed, virtualPath); // Paths come from listIndexEntries (already enumerated under the scope // roots), so no extra containment walk is needed for these reads. // Bounded prefix: selection truncates to MEMORY_HOT_SET_MAX_ITEM_BYTES // anyway; +1 byte preserves its over-budget (truncation marker) check. - return this.getStore(ctx, scope).readFilePrefix( + const store = this.getStore(ctx, scope); + const content = await store.readFilePrefix( parsed.relPath, MEMORY_HOT_SET_MAX_ITEM_BYTES + 1 ); + await this.assertWorkspaceReadExposable(ctx, scope, store); + return content; }, }); + // Selection keeps awaiting (token counting, repeatedly) after the last + // per-file gate: a tombstone published meanwhile must still withhold the + // buffered owner notes. Final check once selection is done; the workspace + // items are dropped (the scope reads as unavailable, like in the index). + const isWorkspaceItem = (item: MemoryHotSetItem): boolean => + parseMemoryPath(item.path).scope === "workspace"; + if (selected.some(isWorkspaceItem)) { + try { + await this.assertWorkspaceStoreReadable(ctx, this.getStore(ctx, "workspace")); + } catch { + return selected.filter((item) => !isWorkspaceItem(item)); + } + } + return selected; } } diff --git a/src/node/services/memoryWorkspaceOwner.test.ts b/src/node/services/memoryWorkspaceOwner.test.ts index f11101fbeca..a11f54c54e2 100644 --- a/src/node/services/memoryWorkspaceOwner.test.ts +++ b/src/node/services/memoryWorkspaceOwner.test.ts @@ -1,13 +1,16 @@ import { describe, expect, it } from "bun:test"; import type { Config } from "@/node/config"; import { + pinDescendantWorkspaceMemoryOwners, resolveWorkspaceMemoryOwnerId, workspaceMemoryOwnerResolver, } from "./memoryWorkspaceOwner"; type ProjectsConfig = ReturnType; -function topology(workspaces: Array<{ id: string; parentWorkspaceId?: string }>): ProjectsConfig { +function topology( + workspaces: Array<{ id: string; parentWorkspaceId?: string; memoryOwnerWorkspaceId?: string }> +): ProjectsConfig { return { projects: new Map([ ["/tmp/project", { workspaces: workspaces.map((ws) => ({ path: `/tmp/${ws.id}`, ...ws })) }], @@ -68,3 +71,87 @@ describe("resolveWorkspaceMemoryOwnerId", () => { expect(workspaceMemoryOwnerResolver(topology([{ id: "ws-owner" }]))).not.toBe(resolve); }); }); + +describe("pinDescendantWorkspaceMemoryOwners", () => { + it("pins each surviving child to the owner it resolves to now", () => { + const cfg = topology([ + { id: "ws-owner" }, + { id: "ws-other" }, + { id: "ws-mid", parentWorkspaceId: "ws-owner" }, + // No pin: the walk through ws-mid reaches ws-owner. + { id: "ws-plain", parentWorkspaceId: "ws-mid" }, + // Stale pin (its owner is gone): the resolver walks past it today, but + // once ws-mid is removed that walk would dangle — replaced. + { id: "ws-stale", parentWorkspaceId: "ws-mid", memoryOwnerWorkspaceId: "ws-gone" }, + // Pin to another live notebook while the parent is still registered: + // a state this code never writes (pins are recorded as an ancestor is + // removed). The live chain wins — the child has been using ws-owner's + // notebook — and the removal re-pins it to that (r84), rather than + // letting corrupt raw config redirect it across task trees. + { id: "ws-pinned", parentWorkspaceId: "ws-mid", memoryOwnerWorkspaceId: "ws-other" }, + // Not a child of the removed node: untouched. + { id: "ws-sibling", parentWorkspaceId: "ws-owner" }, + ]); + const before = Object.fromEntries( + ["ws-plain", "ws-stale", "ws-pinned"].map((id) => [ + id, + resolveWorkspaceMemoryOwnerId(cfg, id), + ]) + ); + expect(before).toEqual({ + "ws-plain": "ws-owner", + "ws-stale": "ws-owner", + "ws-pinned": "ws-owner", + }); + + const pinned = pinDescendantWorkspaceMemoryOwners(cfg, "ws-mid"); + expect(Object.fromEntries(pinned)).toEqual(before); + const entries = [...cfg.projects.values()][0].workspaces; + const pinOf = (id: string) => entries.find((ws) => ws.id === id)!.memoryOwnerWorkspaceId; + expect(pinOf("ws-plain")).toBe("ws-owner"); + expect(pinOf("ws-stale")).toBe("ws-owner"); + expect(pinOf("ws-pinned")).toBe("ws-owner"); + expect(pinOf("ws-sibling")).toBeUndefined(); + + // With ws-mid gone, every pinned child still resolves as before. + const after = topology( + entries + .filter((ws) => ws.id !== "ws-mid") + .map((ws) => ({ + id: ws.id!, + ...(ws.parentWorkspaceId === undefined + ? {} + : { parentWorkspaceId: ws.parentWorkspaceId }), + ...(ws.memoryOwnerWorkspaceId === undefined + ? {} + : { memoryOwnerWorkspaceId: ws.memoryOwnerWorkspaceId }), + })) + ); + for (const [id, owner] of Object.entries(before)) { + expect(resolveWorkspaceMemoryOwnerId(after, id)).toBe(owner); + } + }); + + it("honors a pin only once the recorded parent is gone", () => { + const live = topology([ + { id: "ws-owner" }, + { id: "ws-other" }, + { id: "ws-child", parentWorkspaceId: "ws-owner", memoryOwnerWorkspaceId: "ws-other" }, + { id: "ws-grand", parentWorkspaceId: "ws-child" }, + ]); + // Parent registered: the chain decides, for the child and everything below it. + expect(resolveWorkspaceMemoryOwnerId(live, "ws-child")).toBe("ws-owner"); + expect(resolveWorkspaceMemoryOwnerId(live, "ws-grand")).toBe("ws-owner"); + // Parent gone: the (live) pin decides; a pin whose owner is gone too + // leaves the child on its own store. + const dangling = topology([ + { id: "ws-other" }, + { id: "ws-child", parentWorkspaceId: "ws-owner", memoryOwnerWorkspaceId: "ws-other" }, + { id: "ws-grand", parentWorkspaceId: "ws-child" }, + { id: "ws-orphan", parentWorkspaceId: "ws-owner", memoryOwnerWorkspaceId: "ws-gone" }, + ]); + expect(resolveWorkspaceMemoryOwnerId(dangling, "ws-child")).toBe("ws-other"); + expect(resolveWorkspaceMemoryOwnerId(dangling, "ws-grand")).toBe("ws-other"); + expect(resolveWorkspaceMemoryOwnerId(dangling, "ws-orphan")).toBe("ws-orphan"); + }); +}); diff --git a/src/node/services/memoryWorkspaceOwner.ts b/src/node/services/memoryWorkspaceOwner.ts index 0092cca81cb..7cd5a75ef0f 100644 --- a/src/node/services/memoryWorkspaceOwner.ts +++ b/src/node/services/memoryWorkspaceOwner.ts @@ -19,7 +19,8 @@ type ProjectsConfig = ReturnType; * parentWorkspaceId, so those fallbacks keep their private store usable. * * Pure over one config snapshot (indexed once per snapshot, see - * workspaceMemoryOwnerResolver); MemoryService memoizes it. + * workspaceMemoryOwnerResolver); MemoryService memoizes it, removal calls it + * directly. */ export function resolveWorkspaceMemoryOwnerId(cfg: ProjectsConfig, workspaceId: string): string { return workspaceMemoryOwnerResolver(cfg)(workspaceId); @@ -68,7 +69,21 @@ export function workspaceMemoryOwnerResolver(cfg: ProjectsConfig): (workspaceId: } return workspaceId; } + // A pinned owner is recorded when an intermediate ancestor is removed + // (pinDescendantWorkspaceMemoryOwners), so it only speaks for a chain + // that DANGLES: while the recorded parent is still registered the walk + // follows it, and a pin that disagrees with a live parent (raw config, + // never produced by this code) heals on the next removal instead of + // redirecting the child into an unrelated tree's notebook. With the + // parent gone, a live pin decides; a pin whose owner is gone too leaves + // the child on its own store. const parentWorkspaceId = entry.parentWorkspaceId; + const parentLive = + parentWorkspaceId !== undefined && parentWorkspaceId !== "" && byId.has(parentWorkspaceId); + if (!parentLive) { + const pinned = entry.memoryOwnerWorkspaceId; + if (pinned !== undefined && pinned !== "" && byId.has(pinned)) return pinned; + } if (parentWorkspaceId === undefined || parentWorkspaceId === "") return current; current = parentWorkspaceId; } @@ -80,3 +95,35 @@ export function workspaceMemoryOwnerResolver(cfg: ProjectsConfig): (workspaceId: resolversBySnapshot.set(cfg, resolver); return resolver; } + +/** + * Removal of `removedWorkspaceId`: pin each surviving direct child to the + * owner it resolves to NOW, so the notebook it uses stays the same once the + * chain through the removed node dangles. The pin is whatever the walk + * resolves to while the node is still registered — an existing pin is + * overwritten by it (a live parent takes precedence over a pin in the + * resolver, so that IS the notebook the child has been using), and a stale + * one (its owner gone) is replaced likewise. Mutates the entries in place; + * returns the pins written, for the caller's verified read-back. + */ +export function pinDescendantWorkspaceMemoryOwners( + cfg: ProjectsConfig, + removedWorkspaceId: string +): Map { + assert(removedWorkspaceId.length > 0, "pinDescendantWorkspaceMemoryOwners requires an id"); + const resolve = workspaceMemoryOwnerResolver(cfg); + const pinned = new Map(); + for (const project of cfg.projects.values()) { + for (const workspace of project.workspaces) { + if (workspace.parentWorkspaceId !== removedWorkspaceId || workspace.id === undefined) { + continue; + } + // Resolved before this loop mutates anything: every child's chain runs + // through the removed node, never through a sibling being pinned. + const owner = resolve(workspace.id); + workspace.memoryOwnerWorkspaceId = owner; + pinned.set(workspace.id, owner); + } + } + return pinned; +} diff --git a/src/node/services/workspaceService.test.ts b/src/node/services/workspaceService.test.ts index c6567ed2857..67270f9fb35 100644 --- a/src/node/services/workspaceService.test.ts +++ b/src/node/services/workspaceService.test.ts @@ -58,6 +58,7 @@ import type { WorkspaceMetadata, } from "@/common/types/workspace"; import { makeAgentTaskIntegrationFake } from "./taskWorkspaceSeam.testUtils"; +import { resolveWorkspaceMemoryOwnerId } from "./memoryWorkspaceOwner"; import type { BackgroundProcessManager } from "./backgroundProcessManager"; import type { TerminalService } from "@/node/services/terminalService"; import type { DesktopSessionManager } from "@/node/services/desktop/DesktopSessionManager"; @@ -14898,6 +14899,146 @@ describe("WorkspaceService remove shared-workspace guard", () => { }); }); +describe("WorkspaceService remove shared memory owner pinning", () => { + const projectPath = "/tmp/proj-memory-pin"; + const runtimeConfig = { type: "worktree" as const, srcBaseDir: "/tmp/src" }; + interface Entry { + id: string; + name: string; + path: string; + runtimeConfig: typeof runtimeConfig; + parentWorkspaceId?: string; + memoryOwnerWorkspaceId?: string; + } + + /** owner → mid → grand: removing `mid` must keep `grand` on the owner's notebook. */ + function buildTopology(): { projects: Map } { + return { + projects: new Map([ + [ + projectPath, + { + trusted: true, + workspaces: [ + { id: "ws-owner", name: "owner", path: `${projectPath}/owner`, runtimeConfig }, + { + id: "ws-mid", + name: "mid", + path: `${projectPath}/mid`, + runtimeConfig, + parentWorkspaceId: "ws-owner", + }, + { + id: "ws-grand", + name: "grand", + path: `${projectPath}/grand`, + runtimeConfig, + parentWorkspaceId: "ws-mid", + }, + ], + }, + ], + ]), + }; + } + + function buildConfig(options: { persistPins: boolean }): { + config: Partial; + topology: ReturnType; + } { + const topology = buildTopology(); + const config = { + rootDir: path.join(tmpdir(), "mux-memory-pin", `root-${crypto.randomUUID()}`), + srcDir: "/tmp/src", + sessionsDir: path.join(tmpdir(), "mux-memory-pin", `sessions-${crypto.randomUUID()}`), + removeWorkspace: mock(() => Promise.resolve()), + findWorkspace: mock(() => ({ workspacePath: `${projectPath}/mid`, projectPath })), + loadConfigOrDefault: mock(() => topology), + // Config swallows write failures: a pin that does not land must be + // caught by the removal's verified read-back, so the no-persist variant + // applies the edit to a throwaway copy. + editConfig: mock((edit: (cfg: ReturnType) => unknown) => { + edit(options.persistPins ? topology : buildTopology()); + return Promise.resolve(); + }), + } as unknown as Partial; + return { config, topology }; + } + + function buildAiService(): AIService { + class FakeAIService extends EventEmitter { + isStreaming = mock(() => false); + stopStream = mock(() => Promise.resolve({ success: true as const, data: undefined })); + getWorkspaceMetadata = mock(() => + Promise.resolve({ + success: true as const, + data: { id: "ws-mid", name: "mid", projectPath, runtimeConfig }, + }) + ); + } + return new FakeAIService() as unknown as AIService; + } + + test("pins surviving descendants to the root owner before tearing the middle node down", async () => { + const deleteWorkspace = mock(() => + Promise.resolve({ success: true as const, deletedPath: `${projectPath}/mid` }) + ); + const createRuntimeSpy = spyOn(runtimeFactory, "createRuntime").mockReturnValue({ + deleteWorkspace, + } as unknown as ReturnType); + const { config, topology } = buildConfig({ persistPins: true }); + try { + const workspaceService = createWorkspaceServiceForTest({ + config, + aiService: buildAiService(), + }); + const result = await workspaceService.remove("ws-mid"); + expect(result.success).toBe(true); + expect(deleteWorkspace).toHaveBeenCalledTimes(1); + const grand = topology.projects + .get(projectPath)! + .workspaces.find((ws) => ws.id === "ws-grand"); + expect(grand?.memoryOwnerWorkspaceId).toBe("ws-owner"); + // Once ws-mid is gone the pin keeps ws-grand on the root's notebook. + topology.projects.get(projectPath)!.workspaces = topology.projects + .get(projectPath)! + .workspaces.filter((ws) => ws.id !== "ws-mid"); + expect(resolveWorkspaceMemoryOwnerId(topology as never, "ws-grand")).toBe("ws-owner"); + } finally { + createRuntimeSpy.mockRestore(); + } + }); + + test("aborts a non-forced removal (workspace intact) when the descendant pin does not persist", async () => { + const deleteWorkspace = mock(() => + Promise.resolve({ success: true as const, deletedPath: `${projectPath}/mid` }) + ); + const createRuntimeSpy = spyOn(runtimeFactory, "createRuntime").mockReturnValue({ + deleteWorkspace, + } as unknown as ReturnType); + const { config } = buildConfig({ persistPins: false }); + try { + const workspaceService = createWorkspaceServiceForTest({ + config, + aiService: buildAiService(), + }); + const refused = await workspaceService.remove("ws-mid"); + expect(refused.success).toBe(false); + if (!refused.success) expect(refused.error).toContain("retry the removal"); + // Nothing destructive ran: no checkout deletion, no deregistration. + expect(deleteWorkspace).not.toHaveBeenCalled(); + expect(config.removeWorkspace).not.toHaveBeenCalled(); + + // Forced removal accepts the loss and proceeds. + const forced = await workspaceService.remove("ws-mid", true); + expect(forced.success).toBe(true); + expect(deleteWorkspace).toHaveBeenCalledTimes(1); + } finally { + createRuntimeSpy.mockRestore(); + } + }); +}); + describe("WorkspaceService remove desktop session cleanup", () => { const workspaceId = "ws-remove-desktop"; @@ -14991,6 +15132,47 @@ describe("WorkspaceService remove desktop session cleanup", () => { expect(reopened).toEqual([workspaceId]); }); + test("remove() lifts the consolidation teardown gate only when it aborts before committing", async () => { + const calls: string[] = []; + workspaceService.setMemoryConsolidationService({ + triggerInBackground: () => undefined, + triggerHarvestThenSweepInBackground: () => undefined, + cancelInFlightConsolidation: () => { + calls.push("cancel"); + return Promise.resolve(); + }, + releaseRemovalCancellation: () => { + calls.push("release"); + }, + finalizeHarvestsForRemoval: () => { + calls.push("finalize"); + return Promise.resolve(); + }, + }); + // Aborted before the point of no return (live descendant tasks): the + // workspace stays intact, so any teardown gate is lifted again and no + // harvest state is finalized. + let descendants = true; + workspaceService.setAgentTaskIntegration( + makeAgentTaskIntegrationFake({ hasDescendantAgentTasks: () => descendants }) + ); + const aborted = await workspaceService.remove(workspaceId); + expect(aborted.success).toBe(false); + expect(calls).toEqual(["release"]); + // Committed removal: cancelled (drained), harvest records finalized once + // the session directory is gone, and never released. + descendants = false; + calls.length = 0; + const sessionDir = path.join(tempRoot, "sessions", workspaceId); + await fsPromises.mkdir(sessionDir, { recursive: true }); + const removed = await workspaceService.remove(workspaceId); + expect(removed.success).toBe(true); + expect(existsSync(sessionDir)).toBe(false); + expect(calls.filter((call) => call === "cancel").length).toBeGreaterThan(0); + expect(calls).toContain("finalize"); + expect(calls).not.toContain("release"); + }); + test("remove() flushes the timeline before deleting the session directory", async () => { const sessionDir = path.join(tempRoot, "sessions", workspaceId); await fsPromises.mkdir(sessionDir, { recursive: true }); diff --git a/src/node/services/workspaceService.ts b/src/node/services/workspaceService.ts index 4baa0e055bc..fa478f984c4 100644 --- a/src/node/services/workspaceService.ts +++ b/src/node/services/workspaceService.ts @@ -130,7 +130,10 @@ import { deriveSideChannelModelCandidates, startAbandonedBranchSummaryInBackground, } from "@/node/services/branchSummary"; -import { resolveWorkspaceMemoryOwnerId } from "@/node/services/memoryWorkspaceOwner"; +import { + pinDescendantWorkspaceMemoryOwners, + resolveWorkspaceMemoryOwnerId, +} from "@/node/services/memoryWorkspaceOwner"; import { healRemovalTombstonesForRegisteredWorkspaces, removeSessionDirUnderMemoryLocks, @@ -2743,6 +2746,8 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { triggerInBackground(workspaceId: string, trigger: "compaction" | "archive"): void; triggerHarvestThenSweepInBackground(metadata: CompactionCompletionMetadata): void; cancelInFlightConsolidation(workspaceId: string): Promise; + releaseRemovalCancellation(workspaceId: string): void; + finalizeHarvestsForRemoval(workspaceId: string): Promise; }; private worktreeArchiveSnapshotService?: WorktreeArchiveSnapshotLifecycleService; private agentTaskIntegration?: AgentTaskIntegration; @@ -3086,6 +3091,8 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { triggerInBackground(workspaceId: string, trigger: "compaction" | "archive"): void; triggerHarvestThenSweepInBackground(metadata: CompactionCompletionMetadata): void; cancelInFlightConsolidation(workspaceId: string): Promise; + releaseRemovalCancellation(workspaceId: string): void; + finalizeHarvestsForRemoval(workspaceId: string): Promise; }): void { this.memoryConsolidationService = service; } @@ -5807,6 +5814,11 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { this.removingWorkspaces.add(workspaceId); let timelineClosed = false; let removedFromConfig = false; + // Set once removal passes its point of no return (session teardown and + // tombstone follow unconditionally); an abort before that leaves the + // workspace registered and intact, so the finally lifts the consolidation + // teardown gate the drains below installed. + let removalCommitted = false; // If this workspace is mid-init, cancel the fire-and-forget init work (postCreateSetup, // sync/checkout, .xum/init hook, etc.) so removal doesn't leave orphaned background work. @@ -5883,6 +5895,9 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { await this.sessionTimingService?.waitForIdle(workspaceId); let parentWorkspaceId: string | null = null; + // Memory owner resolved while the workspace was still fully registered + // (metadata path); reused for the destructive step below. + let verifiedSharedMemoryOwnerId: string | null = null; let childTaskModelString: string | undefined; let childTaskThinkingLevel: ThinkingLevel | undefined; @@ -5952,6 +5967,45 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { await clearPendingBranchSummary(workspaceId); await this.refinePassCanceller?.cancelInFlightRefinePass(workspaceId); + // Shared workspace memory (sub-agents write into their task-tree + // owner's store). BEFORE any destructive step — so a failure leaves a + // fully intact, retryable workspace: pin the owner on surviving + // descendants (their parent chain is about to lose this node), + // verified by reading the config back because Config swallows write + // failures. + const sharedMemoryOwnerId = resolveWorkspaceMemoryOwnerId( + this.config.loadConfigOrDefault(), + workspaceId + ); + verifiedSharedMemoryOwnerId = sharedMemoryOwnerId; + if (sharedMemoryOwnerId !== workspaceId) { + try { + let pinnedOwners = new Map(); + await this.config.editConfig((cfg) => { + pinnedOwners = pinDescendantWorkspaceMemoryOwners(cfg, workspaceId); + return cfg; + }); + const persisted = this.config.loadConfigOrDefault(); + for (const [id, owner] of pinnedOwners) { + const entry = findWorkspaceEntry(persisted, id); + if (entry?.workspace.memoryOwnerWorkspaceId !== owner) { + throw new Error(`memory owner pin for descendant ${id} did not persist`); + } + } + } catch (error) { + if (!force) { + return Err( + `Failed to hand this sub-agent's shared workspace memory over to its owner (${getErrorMessage(error)}); the workspace was left intact — retry the removal` + ); + } + log.warn("Forced removal: shared-memory handover to the owner failed", { + workspaceId, + sharedMemoryOwnerId, + error: getErrorMessage(error), + }); + } + } + if (isMultiProject(metadata)) { const projects = getProjects(metadata); const deleteErrors: string[] = []; @@ -6247,6 +6301,7 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { // // Intentionally deferred until we're committed to removal: if runtime deletion fails with // force=false we return early and keep init state intact so init-end can refresh metadata. + removalCommitted = true; this.initStateManager.clearInMemoryState(workspaceId); // Dispose the session before deleting its directory: disposal aborts the active stream, and @@ -6335,10 +6390,9 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { // hold that store's lock too, so a child mutation admitted under the // owner key cannot commit after this tombstone. The workspace is // still registered here, so its parent chain resolves. - const memoryOwnerId = resolveWorkspaceMemoryOwnerId( - this.config.loadConfigOrDefault(), - workspaceId - ); + const memoryOwnerId = + verifiedSharedMemoryOwnerId ?? + resolveWorkspaceMemoryOwnerId(this.config.loadConfigOrDefault(), workspaceId); await removeSessionDirUnderMemoryLocks({ rootDir: this.config.rootDir, sessionDir, @@ -6349,12 +6403,21 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { ? undefined : path.join(this.config.sessionsDir, memoryOwnerId), }); + // Only once the session (and with it the transcript) is gone are the + // retryable harvest records truly unrecoverable; an aborted removal + // above must leave them retryable. + await this.memoryConsolidationService?.finalizeHarvestsForRemoval(workspaceId); } catch (error) { // r63: without a durable tombstone the retained orphan stays // writable by foreign backends forever — abort the removal (the // workspace stays registered and retryable) instead of proceeding // to deregistration below. if (error instanceof TombstoneNotDurableError) { + // No durable tombstone was published: the workspace stays + // registered with its session directory intact, so the + // consolidation teardown gate is lifted again in the finally like + // any pre-commit abort. + removalCommitted = false; throw error; } log.error(`Failed to remove session directory for ${workspaceId}:`, error); @@ -6480,6 +6543,9 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { const message = getErrorMessage(error); return Err(`Failed to remove workspace: ${message}`); } finally { + if (!removalCommitted) { + this.memoryConsolidationService?.releaseRemovalCancellation(workspaceId); + } if (releaseOverridesLock !== undefined) { try { await releaseOverridesLock();