feat(core): Fully migrate packages/core to AgentLoopContext. (#22115)

packages/a2a-server/src/agent/task-event-driven.test.ts

@@ -26,7 +26,7 @@ describe('Task Event-Driven Scheduler', () => {
     mockConfig = createMockConfig({
       isEventDrivenSchedulerEnabled: () => true,
     }) as Config;
-    messageBus = mockConfig.getMessageBus();
+    messageBus = mockConfig.messageBus;
     mockEventBus = {
       publish: vi.fn(),
       on: vi.fn(),
@@ -360,7 +360,7 @@ describe('Task Event-Driven Scheduler', () => {
       isEventDrivenSchedulerEnabled: () => true,
       getApprovalMode: () => ApprovalMode.YOLO,
     }) as Config;
-    const yoloMessageBus = yoloConfig.getMessageBus();
+    const yoloMessageBus = yoloConfig.messageBus;
 
     // @ts-expect-error - Calling private constructor
     const task = new Task('task-id', 'context-id', yoloConfig, mockEventBus);

packages/a2a-server/src/agent/task.ts

@@ -5,6 +5,7 @@
  */
 
 import {
+  type AgentLoopContext,
   Scheduler,
   type GeminiClient,
   GeminiEventType,
@@ -114,7 +115,8 @@ export class Task {
 
     this.scheduler = this.setupEventDrivenScheduler();
 
-    this.geminiClient = this.config.getGeminiClient();
+    const loopContext: AgentLoopContext = this.config;
+    this.geminiClient = loopContext.geminiClient;
     this.pendingToolConfirmationDetails = new Map();
     this.taskState = 'submitted';
     this.eventBus = eventBus;
@@ -143,7 +145,8 @@ export class Task {
   // process. This is not scoped to the individual task but reflects the global connection
   // state managed within the @gemini-cli/core module.
   async getMetadata(): Promise<TaskMetadata> {
-    const toolRegistry = this.config.getToolRegistry();
+    const loopContext: AgentLoopContext = this.config;
+    const toolRegistry = loopContext.toolRegistry;
     const mcpServers = this.config.getMcpClientManager()?.getMcpServers() || {};
     const serverStatuses = getAllMCPServerStatuses();
     const servers = Object.keys(mcpServers).map((serverName) => ({
@@ -376,7 +379,8 @@ export class Task {
   private messageBusListener?: (message: ToolCallsUpdateMessage) => void;
 
   private setupEventDrivenScheduler(): Scheduler {
-    const messageBus = this.config.getMessageBus();
+    const loopContext: AgentLoopContext = this.config;
+    const messageBus = loopContext.messageBus;
     const scheduler = new Scheduler({
       schedulerId: this.id,
       context: this.config,
@@ -395,9 +399,11 @@ export class Task {
 
   dispose(): void {
     if (this.messageBusListener) {
-      this.config
-        .getMessageBus()
-        .unsubscribe(MessageBusType.TOOL_CALLS_UPDATE, this.messageBusListener);
+      const loopContext: AgentLoopContext = this.config;
+      loopContext.messageBus.unsubscribe(
+        MessageBusType.TOOL_CALLS_UPDATE,
+        this.messageBusListener,
+      );
       this.messageBusListener = undefined;
     }
 
@@ -948,7 +954,8 @@ export class Task {
 
         try {
           if (correlationId) {
-            await this.config.getMessageBus().publish({
+            const loopContext: AgentLoopContext = this.config;
+            await loopContext.messageBus.publish({
               type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
               correlationId,
               confirmed:

packages/a2a-server/src/commands/memory.test.ts

@@ -59,6 +59,9 @@ describe('a2a-server memory commands', () => {
     } as unknown as ToolRegistry;
 
     mockConfig = {
+      get toolRegistry() {
+        return mockToolRegistry;
+      },
       getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
     } as unknown as Config;
 
@@ -168,7 +171,6 @@ describe('a2a-server memory commands', () => {
       ]);
 
       expect(mockAddMemory).toHaveBeenCalledWith(fact);
-      expect(mockConfig.getToolRegistry).toHaveBeenCalled();
       expect(mockToolRegistry.getTool).toHaveBeenCalledWith('save_memory');
       expect(mockSaveMemoryTool.buildAndExecute).toHaveBeenCalledWith(
         { fact },

packages/a2a-server/src/commands/memory.ts

@@ -15,6 +15,7 @@ import type {
   CommandContext,
   CommandExecutionResponse,
 } from './types.js';
+import type { AgentLoopContext } from '@google/gemini-cli-core';
 
 const DEFAULT_SANITIZATION_CONFIG = {
   allowedEnvironmentVariables: [],
@@ -95,7 +96,8 @@ export class AddMemoryCommand implements Command {
       return { name: this.name, data: result.content };
     }
 
-    const toolRegistry = context.config.getToolRegistry();
+    const loopContext: AgentLoopContext = context.config;
+    const toolRegistry = loopContext.toolRegistry;
     const tool = toolRegistry.getTool(result.toolName);
     if (tool) {
       const abortController = new AbortController();

packages/a2a-server/src/utils/testing_utils.ts

@@ -16,6 +16,7 @@ import {
   DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD,
   GeminiClient,
   HookSystem,
+  type MessageBus,
   PolicyDecision,
   tmpdir,
   type Config,
@@ -31,9 +32,27 @@ export function createMockConfig(
   const tmpDir = tmpdir();
   // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
   const mockConfig = {
-    get toolRegistry(): ToolRegistry {
+    get config() {
       // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
-      return (this as unknown as Config).getToolRegistry();
+      return this as unknown as Config;
+    },
+    get toolRegistry() {
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+      const config = this as unknown as Config;
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+      return config.getToolRegistry?.() as unknown as ToolRegistry;
+    },
+    get messageBus() {
+      return (
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+        (this as unknown as Config).getMessageBus?.() as unknown as MessageBus
+      );
+    },
+    get geminiClient() {
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+      const config = this as unknown as Config;
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+      return config.getGeminiClient?.() as unknown as GeminiClient;
     },
     getToolRegistry: vi.fn().mockReturnValue({
       getTool: vi.fn(),
@@ -81,9 +100,6 @@ export function createMockConfig(
     ...overrides,
   } as unknown as Config;
 
-  // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
-  (mockConfig as unknown as { config: Config; promptId: string }).config =
-    mockConfig;
   // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
   (mockConfig as unknown as { config: Config; promptId: string }).promptId =
     'test-prompt-id';