From 16768c08f2fa1f0fe2dca32021f81d9f6c8e3316 Mon Sep 17 00:00:00 2001
From: Michael Bleigh <mbleigh@mbleigh.com>
Date: Tue, 7 Apr 2026 16:45:22 -0700
Subject: [PATCH] refactor(core): make LegacyAgentSession dependencies optional
 (#24287)

Co-authored-by: Adam Weidman <adamfweidman@gmail.com>
Co-authored-by: Adam Weidman <adamfweidman@google.com>
---
 .../src/agent/legacy-agent-session.test.ts    | 23 +++++++-----
 .../core/src/agent/legacy-agent-session.ts    | 36 ++++++++++++++-----
 2 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/packages/core/src/agent/legacy-agent-session.test.ts b/packages/core/src/agent/legacy-agent-session.test.ts
index 38bea34910..1de5d90e20 100644
--- a/packages/core/src/agent/legacy-agent-session.test.ts
+++ b/packages/core/src/agent/legacy-agent-session.test.ts
@@ -17,6 +17,9 @@ import type {
   ToolCallRequestInfo,
 } from '../scheduler/types.js';
 import { CoreToolCallStatus } from '../scheduler/types.js';
+import type { GeminiClient } from '../core/client.js';
+import type { Scheduler } from '../scheduler/scheduler.js';
+import type { Config } from '../config/config.js';
 
 // ---------------------------------------------------------------------------
 // Mock helpers
@@ -24,7 +27,7 @@ import { CoreToolCallStatus } from '../scheduler/types.js';
 
 function createMockDeps(
   overrides?: Partial<LegacyAgentSessionDeps>,
-): LegacyAgentSessionDeps {
+): Required<LegacyAgentSessionDeps> {
   const mockClient = {
     sendMessageStream: vi.fn(),
     getChat: vi.fn().mockReturnValue({
@@ -40,18 +43,22 @@ function createMockDeps(
   const mockConfig = {
     getMaxSessionTurns: vi.fn().mockReturnValue(-1),
     getModel: vi.fn().mockReturnValue('gemini-2.5-pro'),
+    getGeminiClient: vi.fn().mockReturnValue(mockClient),
+    getMessageBus: vi.fn().mockImplementation(() => ({
+      subscribe: vi.fn(),
+      unsubscribe: vi.fn(),
+    })),
   };
 
   return {
-    client: mockClient as unknown as LegacyAgentSessionDeps['client'],
-
-    scheduler: mockScheduler as unknown as LegacyAgentSessionDeps['scheduler'],
-
-    config: mockConfig as unknown as LegacyAgentSessionDeps['config'],
+    client: mockClient as unknown as GeminiClient,
+    scheduler: mockScheduler as unknown as Scheduler,
+    config: mockConfig as unknown as Config,
     promptId: 'test-prompt',
     streamId: 'test-stream',
+    getPreferredEditor: vi.fn().mockReturnValue(undefined),
     ...overrides,
-  };
+  } as Required<LegacyAgentSessionDeps>;
 }
 
 async function* makeStream(
@@ -129,7 +136,7 @@ async function collectEvents(
 // ---------------------------------------------------------------------------
 
 describe('LegacyAgentSession', () => {
-  let deps: LegacyAgentSessionDeps;
+  let deps: Required<LegacyAgentSessionDeps>;
 
   beforeEach(() => {
     deps = createMockDeps();
diff --git a/packages/core/src/agent/legacy-agent-session.ts b/packages/core/src/agent/legacy-agent-session.ts
index 667c85f5ed..757dbdb952 100644
--- a/packages/core/src/agent/legacy-agent-session.ts
+++ b/packages/core/src/agent/legacy-agent-session.ts
@@ -14,10 +14,11 @@ import type { Part } from '@google/genai';
 import type { GeminiClient } from '../core/client.js';
 import type { Config } from '../config/config.js';
 import type { ToolCallRequestInfo } from '../scheduler/types.js';
-import type { Scheduler } from '../scheduler/scheduler.js';
+import { Scheduler } from '../scheduler/scheduler.js';
 import { recordToolCallInteractions } from '../code_assist/telemetry.js';
 import { ToolErrorType, isFatalToolError } from '../tools/tool-error.js';
 import { debugLogger } from '../utils/debugLogger.js';
+import type { EditorType } from '../utils/editor.js';
 import {
   buildToolResponseData,
   contentPartsToGeminiParts,
@@ -45,14 +46,17 @@ function isAbortLikeError(err: unknown): boolean {
 }
 
 export interface LegacyAgentSessionDeps {
-  client: GeminiClient;
-  scheduler: Scheduler;
   config: Config;
-  promptId: string;
+  client?: GeminiClient;
+  scheduler?: Scheduler;
+  promptId?: string;
   streamId?: string;
+  getPreferredEditor?: () => EditorType | undefined;
 }
 
-class LegacyAgentProtocol implements AgentProtocol {
+const schedulerMap = new WeakMap<Config, Scheduler>();
+
+export class LegacyAgentProtocol implements AgentProtocol {
   private _events: AgentEvent[] = [];
   private _subscribers = new Set<(event: AgentEvent) => void>();
   private _translationState: TranslationState;
@@ -69,10 +73,26 @@ class LegacyAgentProtocol implements AgentProtocol {
   constructor(deps: LegacyAgentSessionDeps) {
     this._translationState = createTranslationState(deps.streamId);
     this._nextStreamIdOverride = deps.streamId;
-    this._client = deps.client;
-    this._scheduler = deps.scheduler;
     this._config = deps.config;
-    this._promptId = deps.promptId;
+    this._client = deps.client ?? deps.config.getGeminiClient();
+    this._promptId = deps.promptId ?? deps.config.promptId ?? '';
+
+    if (deps.scheduler) {
+      this._scheduler = deps.scheduler;
+    } else {
+      let scheduler = schedulerMap.get(deps.config);
+      if (!scheduler) {
+        const sessionId = deps.config.getSessionId();
+        const schedulerId = `legacy-agent-scheduler-${sessionId}`;
+        scheduler = new Scheduler({
+          context: deps.config,
+          schedulerId,
+          getPreferredEditor: deps.getPreferredEditor ?? (() => undefined),
+        });
+        schedulerMap.set(deps.config, scheduler);
+      }
+      this._scheduler = scheduler;
+    }
   }
 
   get events(): readonly AgentEvent[] {