diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index e7d8078579..f4966c099f 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -77,6 +77,7 @@ import type { InjectionSource } from '../config/injectionService.js'; import { createScopedWorkspaceContext, runWithScopedWorkspaceContext, + runWithScopedActiveExtension, } from '../config/scoped-config.js'; import { CompleteTaskTool } from '../tools/complete-task.js'; import { @@ -523,21 +524,27 @@ export class LocalAgentExecutor { * @returns A promise that resolves to the agent's final output. */ async run(inputs: AgentInputs, signal: AbortSignal): Promise { - // If the agent definition declares additional workspace directories, - // wrap execution in a scoped workspace context. All calls to - // Config.getWorkspaceContext() within this scope will see the extended - // directories, without mutating the shared Config. - const dirs = this.definition.workspaceDirectories; - if (dirs && dirs.length > 0) { - const scopedCtx = createScopedWorkspaceContext( - this.context.config.getWorkspaceContext(), - dirs, - ); - return runWithScopedWorkspaceContext(scopedCtx, () => - this.runInternal(inputs, signal), - ); - } - return this.runInternal(inputs, signal); + // Isolate activeExtensionName for sub-agents to prevent leaking context switches + return runWithScopedActiveExtension( + this.context.config.activeExtensionName ?? null, + () => { + // If the agent definition declares additional workspace directories, + // wrap execution in a scoped workspace context. All calls to + // Config.getWorkspaceContext() within this scope will see the extended + // directories, without mutating the shared Config. + const dirs = this.definition.workspaceDirectories; + if (dirs && dirs.length > 0) { + const scopedCtx = createScopedWorkspaceContext( + this.context.config.getWorkspaceContext(), + dirs, + ); + return runWithScopedWorkspaceContext(scopedCtx, () => + this.runInternal(inputs, signal), + ); + } + return this.runInternal(inputs, signal); + }, + ); } private async runInternal( diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index f7d97e1c8f..c19de3b5d8 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -133,7 +133,10 @@ import type { GenerateContentParameters } from '@google/genai'; export type { MCPOAuthConfig, AnyToolInvocation, AnyDeclarativeTool }; import type { AnyToolInvocation, AnyDeclarativeTool } from '../tools/tools.js'; import { WorkspaceContext } from '../utils/workspaceContext.js'; -import { getWorkspaceContextOverride } from './scoped-config.js'; +import { + getWorkspaceContextOverride, + getActiveExtensionOverride, +} from './scoped-config.js'; import { Storage } from './storage.js'; import type { ShellExecutionConfig } from '../services/shellExecutionService.js'; import { FileExclusions } from '../utils/ignorePatterns.js'; @@ -740,13 +743,22 @@ export class Config implements McpContext, AgentLoopContext { private _activeExtensionName?: string; get activeExtensionName(): string | undefined { + const override = getActiveExtensionOverride(); + if (override !== undefined) { + return override.name === null ? undefined : override.name; + } return ( this._activeExtensionName || process.env['GEMINI_CLI_ACTIVE_EXTENSION'] ); } setActiveExtensionName(name: string | undefined): void { - this._activeExtensionName = name; + const override = getActiveExtensionOverride(); + if (override !== undefined) { + override.name = name ?? null; + } else { + this._activeExtensionName = name; + } } private _resourceRegistry!: ResourceRegistry; private agentRegistry!: AgentRegistry; diff --git a/packages/core/src/config/scoped-config.test.ts b/packages/core/src/config/scoped-config.test.ts index 59689e25a9..30e5d07239 100644 --- a/packages/core/src/config/scoped-config.test.ts +++ b/packages/core/src/config/scoped-config.test.ts @@ -12,6 +12,8 @@ import { createScopedWorkspaceContext, runWithScopedWorkspaceContext, getWorkspaceContextOverride, + runWithScopedActiveExtension, + getActiveExtensionOverride, } from './scoped-config.js'; import { Config } from './config.js'; @@ -204,3 +206,65 @@ describe('runWithScopedWorkspaceContext', () => { }); }); }); + +describe('runWithScopedActiveExtension', () => { + let config: Config; + let tempDir: string; + + beforeEach(() => { + tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'scoped-run-')); + config = new Config({ + targetDir: tempDir, + sessionId: 'test-session', + debugMode: false, + cwd: tempDir, + model: 'test-model', + }); + }); + + afterEach(() => { + fs.rmSync(tempDir, { recursive: true, force: true }); + }); + + it('should override Config.activeExtensionName within scope', () => { + config.setActiveExtensionName('global-ext'); + + runWithScopedActiveExtension('scoped-ext', () => { + expect(config.activeExtensionName).toBe('scoped-ext'); + }); + + expect(config.activeExtensionName).toBe('global-ext'); + }); + + it('should handle null to mask the global extension', () => { + config.setActiveExtensionName('global-ext'); + + runWithScopedActiveExtension(null, () => { + expect(config.activeExtensionName).toBeUndefined(); + }); + + expect(config.activeExtensionName).toBe('global-ext'); + }); + + it('should allow mutating the scoped extension using Config.setActiveExtensionName', () => { + config.setActiveExtensionName('global-ext'); + + runWithScopedActiveExtension('scoped-ext', () => { + config.setActiveExtensionName('mutated-scoped-ext'); + expect(config.activeExtensionName).toBe('mutated-scoped-ext'); + }); + + // The global state should remain untouched + expect(config.activeExtensionName).toBe('global-ext'); + }); + + it('should return undefined from getActiveExtensionOverride outside scope', () => { + expect(getActiveExtensionOverride()).toBeUndefined(); + }); + + it('should return the object from getActiveExtensionOverride inside scope', () => { + runWithScopedActiveExtension('scoped-ext', () => { + expect(getActiveExtensionOverride()).toEqual({ name: 'scoped-ext' }); + }); + }); +}); diff --git a/packages/core/src/config/scoped-config.ts b/packages/core/src/config/scoped-config.ts index 90cdea2da6..2434ac0ed0 100644 --- a/packages/core/src/config/scoped-config.ts +++ b/packages/core/src/config/scoped-config.ts @@ -19,6 +19,9 @@ import { WorkspaceContext } from '../utils/workspaceContext.js'; * This follows the same pattern as `toolCallContext` and `promptIdContext`. */ const workspaceContextOverride = new AsyncLocalStorage(); +const activeExtensionOverride = new AsyncLocalStorage<{ + name: string | null; +}>(); /** * Returns the current workspace context override, if any. @@ -28,6 +31,16 @@ export function getWorkspaceContextOverride(): WorkspaceContext | undefined { return workspaceContextOverride.getStore(); } +/** + * Returns the current active extension name override, if any. + * Called by `Config.activeExtensionName` getter/setter to check for isolated scoped execution. + */ +export function getActiveExtensionOverride(): + | { name: string | null } + | undefined { + return activeExtensionOverride.getStore(); +} + /** * Runs a function with a scoped workspace context override. * Any calls to `Config.getWorkspaceContext()` within `fn` will return @@ -44,6 +57,22 @@ export function runWithScopedWorkspaceContext( return workspaceContextOverride.run(scopedContext, fn); } +/** + * Runs a function with a scoped active extension context override. + * Any calls to `Config.activeExtensionName` within `fn` will return + * the scoped context instead of the inherited default. + * + * @param scopedExtension The active extension name to use within the scope. + * @param fn The function to run. + * @returns The result of the function. + */ +export function runWithScopedActiveExtension( + scopedExtension: string | null, + fn: () => T, +): T { + return activeExtensionOverride.run({ name: scopedExtension }, fn); +} + /** * Creates a {@link WorkspaceContext} that extends a parent's directories * with additional ones.