fix(agent-core-v2): cache workspace alias resolution across calls

packages/agent-core-v2/src/app/workspace/fileWorkspacePersistence.ts

@@ -1,6 +1,8 @@
 import { LifecycleScope } from '#/app/scopes';
 
+import { Disposable } from '#/_base/di/lifecycle';
 import { ScopeActivation, registerScopedService } from '#/_base/di/scope';
+import { Emitter, type Event } from '#/_base/event';
 import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore';
 
 import type { Workspace } from './workspace';
@@ -15,10 +17,20 @@ const WORKSPACE_CATALOG_VERSION = 1;
 const WORKSPACE_CATALOG_SCOPE = '';
 const WORKSPACE_CATALOG_KEY = 'workspaces.json';
 
-export class FileWorkspacePersistence implements IWorkspacePersistence {
+export class FileWorkspacePersistence extends Disposable implements IWorkspacePersistence {
   declare readonly _serviceBrand: undefined;
 
-  constructor(@IAtomicDocumentStore private readonly docs: IAtomicDocumentStore) {}
+  private readonly changeEmitter = this._register(new Emitter<void>());
+  readonly onDidChange: Event<void> = this.changeEmitter.event;
+
+  constructor(@IAtomicDocumentStore private readonly docs: IAtomicDocumentStore) {
+    super();
+    this._register(
+      this.docs.watch(WORKSPACE_CATALOG_SCOPE, WORKSPACE_CATALOG_KEY)(() => {
+        this.changeEmitter.fire();
+      }),
+    );
+  }
 
   async load(): Promise<WorkspaceCatalog | undefined> {
     const file = await this.docs.get<PersistedWorkspaceFile>(
@@ -70,6 +82,7 @@ export class FileWorkspacePersistence implements IWorkspacePersistence {
       deleted_workspace_ids: [...catalog.deletedIds],
     };
     await this.docs.set(WORKSPACE_CATALOG_SCOPE, WORKSPACE_CATALOG_KEY, file);
+    this.changeEmitter.fire();
   }
 }
 

packages/agent-core-v2/src/app/workspace/workspacePersistence.ts

@@ -1,4 +1,5 @@
 import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation';
+import type { Event } from '#/_base/event';
 
 import type { Workspace } from './workspace';
 
@@ -23,6 +24,8 @@ export interface WorkspaceCatalog {
 export interface IWorkspacePersistence {
   readonly _serviceBrand: undefined;
 
+  readonly onDidChange: Event<void>;
+
   load(): Promise<WorkspaceCatalog | undefined>;
   save(catalog: WorkspaceCatalog): Promise<void>;
 }

packages/agent-core-v2/src/app/workspaceAliases/workspaceAliasesService.ts

@@ -1,34 +1,171 @@
 import { LifecycleScope } from '#/app/scopes';
 
+import { Disposable } from '#/_base/di/lifecycle';
 import { ScopeActivation, registerScopedService } from '#/_base/di/scope';
-import { IWorkspaceService } from '#/app/workspace/workspace';
+import { encodeWorkDirKey, workspaceRootKey } from '#/_base/utils/workdir-slug';
+import { IWorkspaceService, type Workspace } from '#/app/workspace/workspace';
 import {
-  collectAliasIds,
   readSessionIndexEntries,
+  SESSION_INDEX_KEY,
+  SESSION_INDEX_SCOPE,
 } from '#/app/workspace/workspaceAlias';
 import { IWorkspacePersistence } from '#/app/workspace/workspacePersistence';
+import { IAppendLogStore } from '#/persistence/interface/appendLogStore';
 import { IFileSystemStorageService } from '#/persistence/interface/storage';
 
 import { IWorkspaceAliases } from './workspaceAliases';
 
-export class WorkspaceAliasesService implements IWorkspaceAliases {
+interface CatalogSnapshot {
+  readonly byId: ReadonlyMap<string, Workspace>;
+  readonly idsByRootKey: ReadonlyMap<string, readonly string[]>;
+}
+
+interface SessionIndexSnapshot {
+  readonly idsByRootKey: ReadonlyMap<string, readonly string[]>;
+}
+
+function rootKeyIndex<T>(
+  items: readonly T[],
+  rootOf: (item: T) => string,
+  idOf: (item: T) => string,
+): Map<string, readonly string[]> {
+  const map = new Map<string, string[]>();
+  for (const item of items) {
+    const key = workspaceRootKey(rootOf(item));
+    const id = idOf(item);
+    const bucket = map.get(key);
+    if (bucket === undefined) {
+      map.set(key, [id]);
+    } else if (!bucket.includes(id)) {
+      bucket.push(id);
+    }
+  }
+  return map;
+}
+
+export class WorkspaceAliasesService extends Disposable implements IWorkspaceAliases {
   declare readonly _serviceBrand: undefined;
 
+  private catalogCache: CatalogSnapshot | undefined;
+  private sessionIndexCache: { snapshot: SessionIndexSnapshot; size: number | undefined } | undefined;
+  private catalogPromise:
+    | Promise<{ snapshot: CatalogSnapshot; generation: number }>
+    | undefined;
+  private sessionIndexPromise:
+    | Promise<{ snapshot: SessionIndexSnapshot; generation: number }>
+    | undefined;
+  private invalidationGeneration = 0;
+  private catalogMergePrimed = false;
+
   constructor(
     @IWorkspaceService private readonly workspaces: IWorkspaceService,
     @IWorkspacePersistence private readonly store: IWorkspacePersistence,
     @IFileSystemStorageService private readonly storage: IFileSystemStorageService,
-  ) {}
+    @IAppendLogStore private readonly appendLogs: IAppendLogStore,
+  ) {
+    super();
+    this._register(
+      this.store.onDidChange(() => {
+        this.invalidationGeneration += 1;
+        this.catalogCache = undefined;
+      }),
+    );
+    this._register(
+      this.appendLogs.onDidWrite((write) => {
+        if (write.scope === SESSION_INDEX_SCOPE && write.key === SESSION_INDEX_KEY) {
+          this.invalidationGeneration += 1;
+          this.sessionIndexCache = undefined;
+        }
+      }),
+    );
+  }
 
   async resolveAliasIds(id: string): Promise<readonly string[]> {
-    const entry = await this.workspaces.get(id);
-    if (entry === undefined) return [id];
-    const catalog = (await this.store.load()) ?? { workspaces: [], deletedIds: [] };
-    return collectAliasIds(
-      catalog.workspaces,
-      await readSessionIndexEntries(this.storage),
-      entry.root,
-    );
+    for (;;) {
+      const generation = this.invalidationGeneration;
+      const [catalog, index] = await Promise.all([this.catalog(), this.sessionIndex()]);
+      if (generation !== this.invalidationGeneration) continue;
+      const entry = catalog.byId.get(id);
+      if (entry === undefined) return [id];
+      const rootKey = workspaceRootKey(entry.root);
+      const fromCatalog = catalog.idsByRootKey.get(rootKey);
+      const fromIndex = index.idsByRootKey.get(rootKey);
+      if (fromCatalog === undefined) return fromIndex ?? [id];
+      if (fromIndex === undefined) return fromCatalog;
+      const merged = [...fromCatalog];
+      for (const alias of fromIndex) {
+        if (!merged.includes(alias)) merged.push(alias);
+      }
+      return merged;
+    }
+  }
+
+  private async catalog(): Promise<CatalogSnapshot> {
+    if (this.catalogCache !== undefined) return this.catalogCache;
+    this.catalogPromise ??= this.loadCatalog();
+    const { snapshot, generation } = await this.catalogPromise;
+    if (generation !== this.invalidationGeneration) return this.catalog();
+    return snapshot;
+  }
+
+  private async loadCatalog(): Promise<{ snapshot: CatalogSnapshot; generation: number }> {
+    try {
+      if (!this.catalogMergePrimed) {
+        await this.workspaces.list();
+        this.catalogMergePrimed = true;
+      }
+      const generation = this.invalidationGeneration;
+      const workspaces = (await this.store.load())?.workspaces ?? [];
+      const snapshot: CatalogSnapshot = {
+        byId: new Map(workspaces.map((ws) => [ws.id, ws] as const)),
+        idsByRootKey: rootKeyIndex(
+          workspaces,
+          (ws) => ws.root,
+          (ws) => ws.id,
+        ),
+      };
+      if (generation === this.invalidationGeneration) {
+        this.catalogCache = snapshot;
+      }
+      return { snapshot, generation };
+    } finally {
+      this.catalogPromise = undefined;
+    }
+  }
+
+  private async sessionIndex(): Promise<SessionIndexSnapshot> {
+    const cache = this.sessionIndexCache;
+    if (
+      cache !== undefined &&
+      (await this.storage.size(SESSION_INDEX_SCOPE, SESSION_INDEX_KEY)) === cache.size
+    ) {
+      return cache.snapshot;
+    }
+    this.sessionIndexPromise ??= this.loadSessionIndex();
+    const { snapshot, generation } = await this.sessionIndexPromise;
+    if (generation !== this.invalidationGeneration) return this.sessionIndex();
+    return snapshot;
+  }
+
+  private async loadSessionIndex(): Promise<{ snapshot: SessionIndexSnapshot; generation: number }> {
+    try {
+      const generation = this.invalidationGeneration;
+      const entries = await readSessionIndexEntries(this.storage);
+      const snapshot: SessionIndexSnapshot = {
+        idsByRootKey: rootKeyIndex(entries, (entry) => entry.workDir, (entry) =>
+          encodeWorkDirKey(entry.workDir),
+        ),
+      };
+      if (generation === this.invalidationGeneration) {
+        this.sessionIndexCache = {
+          snapshot,
+          size: await this.storage.size(SESSION_INDEX_SCOPE, SESSION_INDEX_KEY),
+        };
+      }
+      return { snapshot, generation };
+    } finally {
+      this.sessionIndexPromise = undefined;
+    }
   }
 }
 

packages/agent-core-v2/src/persistence/backends/node-fs/appendLogStore.ts

@@ -1,13 +1,15 @@
-import { toDisposable, type IDisposable } from '#/_base/di/lifecycle';
+import { Disposable, toDisposable, type IDisposable } from '#/_base/di/lifecycle';
 import { LifecycleScope } from '#/app/scopes';
 import { ScopeActivation, registerScopedService } from '#/_base/di/scope';
+import { Emitter, type Event } from '#/_base/event';
 
 import { IFileSystemStorageService } from '#/persistence/interface/storage';
 import {
   AppendLogCorruptedError,
   IAppendLogStore,
   type AppendLogOptions,
   type AppendLogReadOptions,
+  type AppendLogWrite,
 } from '#/persistence/interface/appendLogStore';
 
 const textEncoder = new TextEncoder();
@@ -33,12 +35,16 @@ interface LogState {
   onError?: (error: unknown) => void;
 }
 
-export class AppendLogStore implements IAppendLogStore {
+export class AppendLogStore extends Disposable implements IAppendLogStore {
   declare readonly _serviceBrand: undefined;
 
   private readonly logs = new Map<string, LogState>();
+  private readonly writeEmitter = this._register(new Emitter<AppendLogWrite>());
+  readonly onDidWrite: Event<AppendLogWrite> = this.writeEmitter.event;
 
-  constructor(@IFileSystemStorageService private readonly storage: IFileSystemStorageService) {}
+  constructor(@IFileSystemStorageService private readonly storage: IFileSystemStorageService) {
+    super();
+  }
 
   append<R>(scope: string, key: string, record: R, options?: AppendLogOptions): void {
     const state = this.state(scope, key);
@@ -117,12 +123,13 @@ export class AppendLogStore implements IAppendLogStore {
       try {
         await this.storage.write(scope, key, encoded, { atomic: true });
         state.storageFailure = undefined;
+        return true;
       } catch (error) {
         state.storageFailure = { error };
         throw error;
       }
     });
-    await this.ownFlush(scope, key, state, rewrite);
+    await this.ownFlush(scope, key, state, rewrite, { value: false });
   }
 
   async flush(): Promise<void> {
@@ -190,7 +197,8 @@ export class AppendLogStore implements IAppendLogStore {
   private flushState(scope: string, key: string, state: LogState): Promise<void> {
     if (state.flushPromise !== undefined) return state.flushPromise;
     if (state.storageFailure !== undefined) return Promise.reject(state.storageFailure.error);
-    return this.ownFlush(scope, key, state, this.drain(scope, key, state));
+    const wroteBox = { value: false };
+    return this.ownFlush(scope, key, state, this.drain(scope, key, state, wroteBox), wroteBox);
   }
 
   private release(scope: string, key: string, state: LogState): void {
@@ -216,10 +224,11 @@ export class AppendLogStore implements IAppendLogStore {
     scope: string,
     key: string,
     state: LogState,
-    operation: Promise<void>,
+    operation: Promise<boolean>,
+    wroteBox: { value: boolean },
   ): Promise<void> {
     let owned!: Promise<void>;
-    owned = this.finishOwnedFlush(scope, key, state, operation, () => owned);
+    owned = this.finishOwnedFlush(scope, key, state, operation, wroteBox, () => owned);
     state.flushPromise = owned;
     return owned;
   }
@@ -228,47 +237,55 @@ export class AppendLogStore implements IAppendLogStore {
     scope: string,
     key: string,
     state: LogState,
-    operation: Promise<void>,
+    operation: Promise<boolean>,
+    wroteBox: { value: boolean },
     owner: () => Promise<void>,
   ): Promise<void> {
     let failure: { readonly error: unknown } | undefined;
     try {
-      await operation;
-    } catch (error) {
-      failure = { error };
-    }
-    const owned = owner();
-    if (state.flushPromise === owned) {
-      try {
-        if (failure === undefined) {
-          while (state.flushPromise === owned && state.pending.length > 0) {
-            await this.drain(scope, key, state);
-          }
-        }
-      } finally {
-        if (state.flushPromise === owned) {
-          state.flushPromise = undefined;
+      if (await operation) wroteBox.value = true;
+      const owned = owner();
+      if (state.flushPromise === owned) {
+        while (state.flushPromise === owned && state.pending.length > 0) {
+          await this.drain(scope, key, state, wroteBox);
         }
       }
+    } catch (error) {
+      failure ??= { error };
+    } finally {
+      const owned = owner();
+      if (state.flushPromise === owned) {
+        state.flushPromise = undefined;
+      }
     }
+    if (wroteBox.value) this.writeEmitter.fire({ scope, key });
     if (failure !== undefined) throw failure.error;
   }
 
-  private async drain(scope: string, key: string, state: LogState): Promise<void> {
+  private async drain(
+    scope: string,
+    key: string,
+    state: LogState,
+    wroteBox?: { value: boolean },
+  ): Promise<boolean> {
     const cutoverEpoch = state.cutoverEpoch;
     await state.ready;
-    if (state.cutoverEpoch !== cutoverEpoch) return;
+    if (state.cutoverEpoch !== cutoverEpoch) return false;
+    let wrote = false;
     while (state.pending.length > 0) {
       const batch = state.pending.slice();
       try {
         await this.storage.append(scope, key, encodeBatch(batch), { durable: true });
+        wrote = true;
+        if (wroteBox !== undefined) wroteBox.value = true;
       } catch (error) {
         const failure = (state.storageFailure ??= { error });
         throw failure.error;
       }
-      if (state.cutoverEpoch !== cutoverEpoch) return;
+      if (state.cutoverEpoch !== cutoverEpoch) return wrote;
       state.pending.splice(0, batch.length);
     }
+    return wrote;
   }
 }