Reload gemini memory on extension load/unload + memory refresh refactor (#12651)

This commit is contained in:
Jacob MacDonald
2025-11-07 09:07:25 -08:00
committed by GitHub
parent ef4030331a
commit 47603ef8e1
15 changed files with 276 additions and 259 deletions
+46 -11
View File
@@ -6,6 +6,7 @@
import type { EventEmitter } from 'node:events';
import type { Config, GeminiCLIExtension } from '../config/config.js';
import { refreshServerHierarchicalMemory } from './memoryDiscovery.js';
export abstract class ExtensionLoader {
// Assigned in `start`.
@@ -18,6 +19,9 @@ export abstract class ExtensionLoader {
protected stoppingCount: number = 0;
protected stopCompletedCount: number = 0;
// Whether or not we are currently executing `start`
private isStarting: boolean = false;
constructor(private readonly eventEmitter?: EventEmitter<ExtensionEvents>) {}
/**
@@ -32,16 +36,21 @@ export abstract class ExtensionLoader {
* McpClientManager, PromptRegistry, and GeminiChat set up.
*/
async start(config: Config): Promise<void> {
if (!this.config) {
this.config = config;
} else {
throw new Error('Already started, you may only call `start` once.');
this.isStarting = true;
try {
if (!this.config) {
this.config = config;
} else {
throw new Error('Already started, you may only call `start` once.');
}
await Promise.all(
this.getExtensions()
.filter((e) => e.isActive)
.map(this.startExtension.bind(this)),
);
} finally {
this.isStarting = false;
}
await Promise.all(
this.getExtensions()
.filter((e) => e.isActive)
.map(this.startExtension.bind(this)),
);
}
/**
@@ -64,12 +73,15 @@ export abstract class ExtensionLoader {
});
try {
await this.config.getMcpClientManager()!.startExtension(extension);
// Note: Context files are loaded only once all extensions are done
// loading/unloading to reduce churn, see the `maybeRefreshMemories` call
// below.
// TODO: Update custom command updating away from the event based system
// and call directly into a custom command manager here. See the
// useSlashCommandProcessor hook which responds to events fired here today.
// TODO: Move all enablement of extension features here, including at least:
// - context file loading
// - excluded tool configuration
} finally {
this.startCompletedCount++;
@@ -81,6 +93,25 @@ export abstract class ExtensionLoader {
this.startingCount = 0;
this.startCompletedCount = 0;
}
await this.maybeRefreshMemories();
}
}
private async maybeRefreshMemories(): Promise<void> {
if (!this.config) {
throw new Error(
'Cannot refresh gemini memories prior to calling `start`.',
);
}
if (
!this.isStarting && // Don't refresh memories on the first call to `start`.
this.startingCount === this.startCompletedCount &&
this.stoppingCount === this.stopCompletedCount
) {
// Wait until all extensions are done starting and stopping before we
// reload memory, this is somewhat expensive and also busts the context
// cache, we want to only do it once.
await refreshServerHierarchicalMemory(this.config);
}
}
@@ -119,12 +150,15 @@ export abstract class ExtensionLoader {
try {
await this.config.getMcpClientManager()!.stopExtension(extension);
// Note: Context files are loaded only once all extensions are done
// loading/unloading to reduce churn, see the `maybeRefreshMemories` call
// below.
// TODO: Update custom command updating away from the event based system
// and call directly into a custom command manager here. See the
// useSlashCommandProcessor hook which responds to events fired here today.
// TODO: Remove all extension features here, including at least:
// - context files
// - excluded tools
} finally {
this.stopCompletedCount++;
@@ -136,6 +170,7 @@ export abstract class ExtensionLoader {
this.stoppingCount = 0;
this.stopCompletedCount = 0;
}
await this.maybeRefreshMemories();
}
}