From 8fcb18996a07ab5fbcbc17762c20d3f9546fb0e8 Mon Sep 17 00:00:00 2001 From: Adam Weidman Date: Sun, 15 Mar 2026 21:38:11 -0400 Subject: [PATCH] refactor(core): unify InjectionService API to single onInjection interface Remove legacy onUserHint/offUserHint/addUserHint methods. All callers now use addInjection(text, source) and onInjection/offInjection with source-based filtering where needed. --- packages/cli/src/test-utils/AppRig.tsx | 2 +- packages/cli/src/ui/AppContainer.tsx | 14 ++-- .../core/src/agents/local-executor.test.ts | 15 ++++- .../core/src/agents/subagent-tool.test.ts | 14 ++-- .../core/src/config/injectionService.test.ts | 65 ++++++------------- packages/core/src/config/injectionService.ts | 40 ++---------- 6 files changed, 52 insertions(+), 98 deletions(-) diff --git a/packages/cli/src/test-utils/AppRig.tsx b/packages/cli/src/test-utils/AppRig.tsx index 921d3ceaf8..cf22bec1ae 100644 --- a/packages/cli/src/test-utils/AppRig.tsx +++ b/packages/cli/src/test-utils/AppRig.tsx @@ -611,7 +611,7 @@ export class AppRig { async addUserHint(hint: string) { if (!this.config) throw new Error('AppRig not initialized'); await act(async () => { - this.config!.injectionService.addUserHint(hint); + this.config!.injectionService.addInjection(hint, 'user_steering'); }); } diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 259913f161..610106810c 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -85,6 +85,7 @@ import { buildUserSteeringHintPrompt, logBillingEvent, ApiKeyUpdatedEvent, + type InjectionSource, } from '@google/gemini-cli-core'; import { validateAuthMethod } from '../config/auth.js'; import process from 'node:process'; @@ -1089,13 +1090,16 @@ Logging in with Google... Restarting Gemini CLI to continue. }, []); useEffect(() => { - const hintListener = (hint: string) => { - pendingHintsRef.current.push(hint); + const hintListener = (text: string, source: InjectionSource) => { + if (source !== 'user_steering') { + return; + } + pendingHintsRef.current.push(text); setPendingHintCount((prev) => prev + 1); }; - config.injectionService.onUserHint(hintListener); + config.injectionService.onInjection(hintListener); return () => { - config.injectionService.offUserHint(hintListener); + config.injectionService.offInjection(hintListener); }; }, [config]); @@ -1259,7 +1263,7 @@ Logging in with Google... Restarting Gemini CLI to continue. if (!trimmed) { return; } - config.injectionService.addUserHint(trimmed); + config.injectionService.addInjection(trimmed, 'user_steering'); // Render hints with a distinct style. historyManager.addItem({ type: 'hint', diff --git a/packages/core/src/agents/local-executor.test.ts b/packages/core/src/agents/local-executor.test.ts index f64e79590c..79a69d8850 100644 --- a/packages/core/src/agents/local-executor.test.ts +++ b/packages/core/src/agents/local-executor.test.ts @@ -2105,7 +2105,10 @@ describe('LocalAgentExecutor', () => { // Give the loop a chance to start and register the listener await vi.advanceTimersByTimeAsync(1); - configWithHints.injectionService.addUserHint('Initial Hint'); + configWithHints.injectionService.addInjection( + 'Initial Hint', + 'user_steering', + ); // Resolve the tool call to complete Turn 1 resolveToolCall!([ @@ -2151,7 +2154,10 @@ describe('LocalAgentExecutor', () => { it('should NOT inject legacy hints added before executor was created', async () => { const definition = createTestDefinition(); - configWithHints.injectionService.addUserHint('Legacy Hint'); + configWithHints.injectionService.addInjection( + 'Legacy Hint', + 'user_steering', + ); const executor = await LocalAgentExecutor.create( definition, @@ -2218,7 +2224,10 @@ describe('LocalAgentExecutor', () => { await vi.advanceTimersByTimeAsync(1); // Add the hint while the tool call is pending - configWithHints.injectionService.addUserHint('Corrective Hint'); + configWithHints.injectionService.addInjection( + 'Corrective Hint', + 'user_steering', + ); // Now resolve the tool call to complete Turn 1 resolveToolCall!([ diff --git a/packages/core/src/agents/subagent-tool.test.ts b/packages/core/src/agents/subagent-tool.test.ts index 761431c2ba..438df59cd3 100644 --- a/packages/core/src/agents/subagent-tool.test.ts +++ b/packages/core/src/agents/subagent-tool.test.ts @@ -214,7 +214,7 @@ describe('SubAgentInvocation', () => { describe('withUserHints', () => { it('should NOT modify query for local agents', async () => { mockConfig = makeFakeConfig({ modelSteering: true }); - mockConfig.injectionService.addUserHint('Test Hint'); + mockConfig.injectionService.addInjection('Test Hint', 'user_steering'); const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus); const params = { query: 'original query' }; @@ -229,7 +229,7 @@ describe('SubAgentInvocation', () => { it('should NOT modify query for remote agents if model steering is disabled', async () => { mockConfig = makeFakeConfig({ modelSteering: false }); - mockConfig.injectionService.addUserHint('Test Hint'); + mockConfig.injectionService.addInjection('Test Hint', 'user_steering'); const tool = new SubagentTool( testRemoteDefinition, @@ -276,8 +276,8 @@ describe('SubAgentInvocation', () => { // @ts-expect-error - accessing private method for testing const invocation = tool.createInvocation(params, mockMessageBus); - mockConfig.injectionService.addUserHint('Hint 1'); - mockConfig.injectionService.addUserHint('Hint 2'); + mockConfig.injectionService.addInjection('Hint 1', 'user_steering'); + mockConfig.injectionService.addInjection('Hint 2', 'user_steering'); // @ts-expect-error - accessing private method for testing const hintedParams = invocation.withUserHints(params); @@ -289,7 +289,7 @@ describe('SubAgentInvocation', () => { it('should NOT include legacy hints added before the invocation was created', async () => { mockConfig = makeFakeConfig({ modelSteering: true }); - mockConfig.injectionService.addUserHint('Legacy Hint'); + mockConfig.injectionService.addInjection('Legacy Hint', 'user_steering'); const tool = new SubagentTool( testRemoteDefinition, @@ -308,7 +308,7 @@ describe('SubAgentInvocation', () => { expect(hintedParams.query).toBe('original query'); // Add a new hint after creation - mockConfig.injectionService.addUserHint('New Hint'); + mockConfig.injectionService.addInjection('New Hint', 'user_steering'); // @ts-expect-error - accessing private method for testing hintedParams = invocation.withUserHints(params); @@ -318,7 +318,7 @@ describe('SubAgentInvocation', () => { it('should NOT modify query if query is missing or not a string', async () => { mockConfig = makeFakeConfig({ modelSteering: true }); - mockConfig.injectionService.addUserHint('Hint'); + mockConfig.injectionService.addInjection('Hint', 'user_steering'); const tool = new SubagentTool( testRemoteDefinition, diff --git a/packages/core/src/config/injectionService.test.ts b/packages/core/src/config/injectionService.test.ts index 4803b49e7d..482eca6403 100644 --- a/packages/core/src/config/injectionService.test.ts +++ b/packages/core/src/config/injectionService.test.ts @@ -10,7 +10,7 @@ import { InjectionService } from './injectionService.js'; describe('InjectionService', () => { it('is disabled by default and ignores user_steering injections', () => { const service = new InjectionService(() => false); - service.addUserHint('this hint should be ignored'); + service.addInjection('this hint should be ignored', 'user_steering'); expect(service.getUserHints()).toEqual([]); expect(service.getLatestHintIndex()).toBe(-1); }); @@ -18,9 +18,9 @@ describe('InjectionService', () => { it('stores trimmed injections and exposes them via indexing when enabled', () => { const service = new InjectionService(() => true); - service.addUserHint(' first hint '); - service.addUserHint('second hint'); - service.addUserHint(' '); + service.addInjection(' first hint ', 'user_steering'); + service.addInjection('second hint', 'user_steering'); + service.addInjection(' ', 'user_steering'); expect(service.getUserHints()).toEqual(['first hint', 'second hint']); expect(service.getLatestHintIndex()).toBe(1); @@ -36,38 +36,38 @@ describe('InjectionService', () => { const service = new InjectionService(() => true); expect(service.getLastUserHintAt()).toBeNull(); - service.addUserHint('hint'); + service.addInjection('hint', 'user_steering'); const timestamp = service.getLastUserHintAt(); expect(timestamp).not.toBeNull(); expect(typeof timestamp).toBe('number'); }); - it('notifies user hint listeners when a user_steering injection is added', () => { + it('notifies listeners when an injection is added', () => { const service = new InjectionService(() => true); const listener = vi.fn(); - service.onUserHint(listener); + service.onInjection(listener); - service.addUserHint('new hint'); + service.addInjection('new hint', 'user_steering'); - expect(listener).toHaveBeenCalledWith('new hint'); + expect(listener).toHaveBeenCalledWith('new hint', 'user_steering'); }); - it('does NOT notify user hint listeners after they are unregistered', () => { + it('does NOT notify listeners after they are unregistered', () => { const service = new InjectionService(() => true); const listener = vi.fn(); - service.onUserHint(listener); - service.offUserHint(listener); + service.onInjection(listener); + service.offInjection(listener); - service.addUserHint('ignored hint'); + service.addInjection('ignored hint', 'user_steering'); expect(listener).not.toHaveBeenCalled(); }); it('should clear all injections', () => { const service = new InjectionService(() => true); - service.addUserHint('hint 1'); - service.addUserHint('hint 2'); + service.addInjection('hint 1', 'user_steering'); + service.addInjection('hint 2', 'user_steering'); expect(service.getUserHints()).toHaveLength(2); service.clear(); @@ -75,18 +75,18 @@ describe('InjectionService', () => { expect(service.getLatestHintIndex()).toBe(-1); }); - describe('typed injection API', () => { - it('notifies typed listeners with source for user_steering', () => { + describe('source-specific behavior', () => { + it('notifies listeners with source for user_steering', () => { const service = new InjectionService(() => true); const listener = vi.fn(); service.onInjection(listener); - service.addUserHint('steering hint'); + service.addInjection('steering hint', 'user_steering'); expect(listener).toHaveBeenCalledWith('steering hint', 'user_steering'); }); - it('notifies typed listeners with source for background_completion', () => { + it('notifies listeners with source for background_completion', () => { const service = new InjectionService(() => true); const listener = vi.fn(); service.onInjection(listener); @@ -99,22 +99,6 @@ describe('InjectionService', () => { ); }); - it('does NOT notify user hint listeners for background_completion', () => { - const service = new InjectionService(() => true); - const userListener = vi.fn(); - const typedListener = vi.fn(); - service.onUserHint(userListener); - service.onInjection(typedListener); - - service.addInjection('bg output', 'background_completion'); - - expect(typedListener).toHaveBeenCalledWith( - 'bg output', - 'background_completion', - ); - expect(userListener).not.toHaveBeenCalled(); - }); - it('accepts background_completion even when model steering is disabled', () => { const service = new InjectionService(() => false); const listener = vi.fn(); @@ -139,16 +123,5 @@ describe('InjectionService', () => { expect(listener).not.toHaveBeenCalled(); expect(service.getUserHints()).toEqual([]); }); - - it('unregisters typed listeners correctly', () => { - const service = new InjectionService(() => true); - const listener = vi.fn(); - service.onInjection(listener); - service.offInjection(listener); - - service.addInjection('bg output', 'background_completion'); - - expect(listener).not.toHaveBeenCalled(); - }); }); }); diff --git a/packages/core/src/config/injectionService.ts b/packages/core/src/config/injectionService.ts index 7991a95bea..7018f34e1c 100644 --- a/packages/core/src/config/injectionService.ts +++ b/packages/core/src/config/injectionService.ts @@ -20,9 +20,8 @@ export type InjectionListener = (text: string, source: InjectionSource) => void; * Service for managing injections into the model conversation. * * Multiple sources (user steering, background execution completions, etc.) - * can feed into this service. Consumers register typed listeners via - * {@link onInjection} to receive injections with source information, or use the - * legacy {@link onUserHint} API for backward compatibility. + * can feed into this service. Consumers register listeners via + * {@link onInjection} to receive injections with source information. */ export class InjectionService { private readonly injections: Array<{ @@ -31,7 +30,6 @@ export class InjectionService { timestamp: number; }> = []; private readonly injectionListeners: Set = new Set(); - private readonly userHintListeners: Set<(hint: string) => void> = new Set(); constructor(private readonly isEnabled: () => boolean) {} @@ -51,55 +49,25 @@ export class InjectionService { } this.injections.push({ text: trimmed, source, timestamp: Date.now() }); - // Fire typed listeners (new API) for (const listener of this.injectionListeners) { listener(trimmed, source); } - - // Fire legacy listeners (user_steering only) - if (source === 'user_steering') { - for (const listener of this.userHintListeners) { - listener(trimmed); - } - } } /** - * Adds a new steering hint from the user. - * Convenience wrapper around {@link addInjection} with `user_steering` source. - */ - addUserHint(hint: string): void { - this.addInjection(hint, 'user_steering'); - } - - /** - * Registers a typed listener for injections from any source. + * Registers a listener for injections from any source. */ onInjection(listener: InjectionListener): void { this.injectionListeners.add(listener); } /** - * Unregisters a typed injection listener. + * Unregisters an injection listener. */ offInjection(listener: InjectionListener): void { this.injectionListeners.delete(listener); } - /** - * Registers a listener for user steering hints only. - */ - onUserHint(listener: (hint: string) => void): void { - this.userHintListeners.add(listener); - } - - /** - * Unregisters a user steering hint listener. - */ - offUserHint(listener: (hint: string) => void): void { - this.userHintListeners.delete(listener); - } - /** * Returns all collected injection texts (all sources). */