Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion packages/a2a-server/src/agent/task.ts
Original file line number Diff line number Diff line change
Expand Up @@ -542,7 +542,7 @@ export class Task {
const messageBus = this.config.getMessageBus();
const scheduler = new Scheduler({
schedulerId: this.id,
config: this.config,
context: this.config,
messageBus,
getPreferredEditor: () => DEFAULT_GUI_EDITOR,
});
Expand Down
2 changes: 1 addition & 1 deletion packages/cli/src/nonInteractiveCli.ts
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,7 @@ export async function runNonInteractive({

const geminiClient = config.getGeminiClient();
const scheduler = new Scheduler({
config,
context: config,
messageBus: config.getMessageBus(),
getPreferredEditor: () => undefined,
schedulerId: ROOT_SCHEDULER_ID,
Expand Down
2 changes: 1 addition & 1 deletion packages/cli/src/ui/hooks/useToolScheduler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ export function useToolScheduler(
const scheduler = useMemo(
() =>
new Scheduler({
config,
context: config,
messageBus,
getPreferredEditor: () => getPreferredEditorRef.current(),
schedulerId: ROOT_SCHEDULER_ID,
Expand Down
15 changes: 12 additions & 3 deletions packages/core/src/agents/agent-scheduler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ vi.mock('../scheduler/scheduler.js', () => ({

describe('agent-scheduler', () => {
let mockToolRegistry: Mocked<ToolRegistry>;
let mockConfig: Mocked<Config>;
let mockMessageBus: Mocked<MessageBus>;

beforeEach(() => {
Expand All @@ -29,6 +30,14 @@ describe('agent-scheduler', () => {
getTool: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
} as unknown as Mocked<ToolRegistry>;
mockConfig = {
getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
toolRegistry: mockToolRegistry,
} as unknown as Mocked<Config>;
(mockConfig as unknown as { messageBus: MessageBus }).messageBus =
mockMessageBus;
(mockConfig as unknown as { toolRegistry: ToolRegistry }).toolRegistry =
mockToolRegistry;
});

it('should create a scheduler with agent-specific config', async () => {
Expand Down Expand Up @@ -69,7 +78,8 @@ describe('agent-scheduler', () => {
}),
);

const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].config;
// Verify that the scheduler's context has the overridden tool registry
const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].context;
expect(schedulerConfig.toolRegistry).toBe(mockToolRegistry);
});

Expand Down Expand Up @@ -106,9 +116,8 @@ describe('agent-scheduler', () => {
},
);

const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].config;
const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].context;
expect(schedulerConfig.toolRegistry).toBe(agentRegistry);
expect(schedulerConfig.toolRegistry).not.toBe(mainRegistry);
expect(schedulerConfig.getToolRegistry()).toBe(agentRegistry);
});
});
2 changes: 1 addition & 1 deletion packages/core/src/agents/agent-scheduler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ export async function scheduleAgentTools(
});

const scheduler = new Scheduler({
config: agentConfig,
context: agentConfig,
messageBus: toolRegistry.getMessageBus(),
getPreferredEditor: getPreferredEditor ?? (() => undefined),
schedulerId,
Expand Down
9 changes: 7 additions & 2 deletions packages/core/src/agents/browser/browserAgentInvocation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import { randomUUID } from 'node:crypto';
import type { Config } from '../../config/config.js';
import { type AgentLoopContext } from '../../config/agent-loop-context.js';
import { LocalAgentExecutor } from '../local-executor.js';
import { safeJsonToMarkdown } from '../../utils/markdownUtils.js';
import {
Expand Down Expand Up @@ -179,7 +180,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation<
ToolResult
> {
constructor(
private readonly config: Config,
private readonly context: AgentLoopContext,
params: AgentInputs,
messageBus: MessageBus,
_toolName?: string,
Expand All @@ -194,6 +195,10 @@ export class BrowserAgentInvocation extends BaseToolInvocation<
);
}

private get config(): Config {
return this.context.config;
}

/**
* Returns a concise, human-readable description of the invocation.
*/
Expand Down Expand Up @@ -409,7 +414,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation<
// Create and run executor with the configured definition
const executor = await LocalAgentExecutor.create(
definition,
this.config,
this.context,
onActivity,
);

Expand Down
16 changes: 13 additions & 3 deletions packages/core/src/agents/local-executor.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -307,6 +307,11 @@ describe('LocalAgentExecutor', () => {
vi.useFakeTimers();

mockConfig = makeFakeConfig();
// .config is already set correctly by the getter on the instance.
Object.defineProperty(mockConfig, 'promptId', {
get: () => 'test-prompt-id',
configurable: true,
});
parentToolRegistry = new ToolRegistry(
mockConfig,
mockConfig.getMessageBus(),
Expand All @@ -319,7 +324,9 @@ describe('LocalAgentExecutor', () => {
);
parentToolRegistry.registerTool(MOCK_TOOL_NOT_ALLOWED);

vi.spyOn(mockConfig, 'getToolRegistry').mockReturnValue(parentToolRegistry);
vi.spyOn(mockConfig, 'toolRegistry', 'get').mockReturnValue(
parentToolRegistry,
);
vi.spyOn(mockConfig, 'getAgentRegistry').mockReturnValue({
getAllAgentNames: () => [],
} as unknown as AgentRegistry);
Expand Down Expand Up @@ -382,7 +389,10 @@ describe('LocalAgentExecutor', () => {

it('should use parentPromptId from context to create agentId', async () => {
const parentId = 'parent-id';
mockedPromptIdContext.getStore.mockReturnValue(parentId);
Object.defineProperty(mockConfig, 'promptId', {
get: () => parentId,
configurable: true,
});

const definition = createTestDefinition();
const executor = await LocalAgentExecutor.create(
Expand Down Expand Up @@ -2052,7 +2062,7 @@ describe('LocalAgentExecutor', () => {
vi.spyOn(configWithHints, 'getAgentRegistry').mockReturnValue({
getAllAgentNames: () => [],
} as unknown as AgentRegistry);
vi.spyOn(configWithHints, 'getToolRegistry').mockReturnValue(
vi.spyOn(configWithHints, 'toolRegistry', 'get').mockReturnValue(
parentToolRegistry,
);
});
Expand Down
63 changes: 33 additions & 30 deletions packages/core/src/agents/local-executor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
*/

import type { Config } from '../config/config.js';
import { type AgentLoopContext } from '../config/agent-loop-context.js';
import { reportError } from '../utils/errorReporting.js';
import { GeminiChat, StreamEventType } from '../core/geminiChat.js';
import {
Expand Down Expand Up @@ -92,29 +93,33 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {

private readonly agentId: string;
private readonly toolRegistry: ToolRegistry;
private readonly runtimeContext: Config;
private readonly context: AgentLoopContext;
private readonly onActivity?: ActivityCallback;
private readonly compressionService: ChatCompressionService;
private readonly parentCallId?: string;
private hasFailedCompressionAttempt = false;

private get config(): Config {
return this.context.config;
}

/**
* Creates and validates a new `AgentExecutor` instance.
*
* This method ensures that all tools specified in the agent's definition are
* safe for non-interactive use before creating the executor.
*
* @param definition The definition object for the agent.
* @param runtimeContext The global runtime configuration.
* @param context The execution context.
* @param onActivity An optional callback to receive activity events.
* @returns A promise that resolves to a new `LocalAgentExecutor` instance.
*/
static async create<TOutput extends z.ZodTypeAny>(
definition: LocalAgentDefinition<TOutput>,
runtimeContext: Config,
context: AgentLoopContext,
onActivity?: ActivityCallback,
): Promise<LocalAgentExecutor<TOutput>> {
const parentMessageBus = runtimeContext.getMessageBus();
const parentMessageBus = context.messageBus;

// Create an override object to inject the subagent name into tool confirmation requests
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
Expand All @@ -133,12 +138,12 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {

// Create an isolated tool registry for this agent instance.
const agentToolRegistry = new ToolRegistry(
runtimeContext,
context.config,
subagentMessageBus,
);
const parentToolRegistry = runtimeContext.getToolRegistry();
const parentToolRegistry = context.toolRegistry;
const allAgentNames = new Set(
runtimeContext.getAgentRegistry().getAllAgentNames(),
context.config.getAgentRegistry().getAllAgentNames(),
);

const registerToolByName = (toolName: string) => {
Expand Down Expand Up @@ -190,15 +195,15 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
agentToolRegistry.sortTools();

// Get the parent prompt ID from context
const parentPromptId = promptIdContext.getStore();
const parentPromptId = context.promptId;

// Get the parent tool call ID from context
const toolContext = getToolCallContext();
const parentCallId = toolContext?.callId;

return new LocalAgentExecutor(
definition,
runtimeContext,
context,
agentToolRegistry,
parentPromptId,
parentCallId,
Expand All @@ -214,14 +219,14 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
*/
private constructor(
definition: LocalAgentDefinition<TOutput>,
runtimeContext: Config,
context: AgentLoopContext,
toolRegistry: ToolRegistry,
parentPromptId: string | undefined,
parentCallId: string | undefined,
onActivity?: ActivityCallback,
) {
this.definition = definition;
this.runtimeContext = runtimeContext;
this.context = context;
this.toolRegistry = toolRegistry;
this.onActivity = onActivity;
this.compressionService = new ChatCompressionService();
Expand Down Expand Up @@ -418,7 +423,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
} finally {
clearTimeout(graceTimeoutId);
logRecoveryAttempt(
this.runtimeContext,
this.config,
new RecoveryAttemptEvent(
this.agentId,
this.definition.name,
Expand Down Expand Up @@ -466,7 +471,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
const combinedSignal = AbortSignal.any([signal, deadlineTimer.signal]);

logAgentStart(
this.runtimeContext,
this.config,
new AgentStartEvent(this.agentId, this.definition.name),
);

Expand All @@ -477,7 +482,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
const augmentedInputs = {
...inputs,
cliVersion: await getVersion(),
activeModel: this.runtimeContext.getActiveModel(),
activeModel: this.config.getActiveModel(),
today: new Date().toLocaleDateString(),
};

Expand All @@ -494,13 +499,12 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
// Capture the index of the last hint before starting to avoid re-injecting old hints.
// NOTE: Hints added AFTER this point will be broadcast to all currently running
// local agents via the listener below.
const startIndex =
this.runtimeContext.userHintService.getLatestHintIndex();
this.runtimeContext.userHintService.onUserHint(hintListener);
const startIndex = this.config.userHintService.getLatestHintIndex();
this.config.userHintService.onUserHint(hintListener);

try {
const initialHints =
this.runtimeContext.userHintService.getUserHintsAfter(startIndex);
this.config.userHintService.getUserHintsAfter(startIndex);
const formattedInitialHints = formatUserHintsForModel(initialHints);

let currentMessage: Content = formattedInitialHints
Expand Down Expand Up @@ -561,7 +565,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
}
}
} finally {
this.runtimeContext.userHintService.offUserHint(hintListener);
this.config.userHintService.offUserHint(hintListener);
}

// === UNIFIED RECOVERY BLOCK ===
Expand Down Expand Up @@ -674,7 +678,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
} finally {
deadlineTimer.abort();
logAgentFinish(
this.runtimeContext,
this.config,
new AgentFinishEvent(
this.agentId,
this.definition.name,
Expand All @@ -697,7 +701,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
prompt_id,
false,
model,
this.runtimeContext,
this.config,
this.hasFailedCompressionAttempt,
);

Expand Down Expand Up @@ -735,11 +739,10 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
const modelConfigAlias = getModelConfigAlias(this.definition);

// Resolve the model config early to get the concrete model string (which may be `auto`).
const resolvedConfig =
this.runtimeContext.modelConfigService.getResolvedConfig({
model: modelConfigAlias,
overrideScope: this.definition.name,
});
const resolvedConfig = this.config.modelConfigService.getResolvedConfig({
model: modelConfigAlias,
overrideScope: this.definition.name,
});
const requestedModel = resolvedConfig.model;

let modelToUse: string;
Expand All @@ -756,7 +759,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
signal,
requestedModel,
};
const router = this.runtimeContext.getModelRouterService();
const router = this.config.getModelRouterService();
const decision = await router.route(routingContext);
modelToUse = decision.model;
} catch (error) {
Expand Down Expand Up @@ -844,7 +847,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {

try {
return new GeminiChat(
this.runtimeContext,
this.config,
systemInstruction,
[{ functionDeclarations: tools }],
startHistory,
Expand Down Expand Up @@ -1092,7 +1095,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
// Execute standard tool calls using the new scheduler
if (toolRequests.length > 0) {
const completedCalls = await scheduleAgentTools(
this.runtimeContext,
this.config,
toolRequests,
{
schedulerId: this.agentId,
Expand Down Expand Up @@ -1240,7 +1243,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
let finalPrompt = templateString(promptConfig.systemPrompt, inputs);

// Append environment context (CWD and folder structure).
const dirContext = await getDirectoryContextString(this.runtimeContext);
const dirContext = await getDirectoryContextString(this.config);
finalPrompt += `\n\n# Environment Context\n${dirContext}`;

// Append standard rules for non-interactive execution.
Expand Down
5 changes: 5 additions & 0 deletions packages/core/src/agents/local-invocation.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,11 @@ describe('LocalSubagentInvocation', () => {
beforeEach(() => {
vi.clearAllMocks();
mockConfig = makeFakeConfig();
// .config is already set correctly by the getter on the instance.
Object.defineProperty(mockConfig, 'promptId', {
get: () => 'test-prompt-id',
configurable: true,
});
mockMessageBus = createMockMessageBus();

mockExecutorInstance = {
Expand Down
Loading
Loading