diff --git a/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts b/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts index 64448c8a6a2..f133ac93a69 100644 --- a/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts +++ b/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts @@ -100,6 +100,59 @@ describe('StreamJsonOutputAdapter', () => { }); }); + it('should emit active goal stream events', () => { + adapter.processEvent({ + type: GeminiEventType.ActiveGoal, + value: { + condition: 'finish the refactor', + iterations: 2, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing verification', + }, + }); + + adapter.processEvent({ + type: GeminiEventType.ActiveGoal, + value: null, + }); + + const activeGoalEvents = stdoutWriteSpy.mock.calls + .map((call: unknown[]) => JSON.parse(call[0] as string)) + .filter( + (message: { type?: string; event?: { type?: string } }) => + message.type === 'stream_event' && + message.event?.type === 'active_goal', + ); + + expect(activeGoalEvents).toEqual([ + expect.objectContaining({ + session_id: 'test-session-id', + parent_tool_use_id: null, + event: { + type: 'active_goal', + active_goal: { + condition: 'finish the refactor', + iterations: 2, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing verification', + }, + }, + }), + expect.objectContaining({ + session_id: 'test-session-id', + parent_tool_use_id: null, + event: { + type: 'active_goal', + active_goal: null, + }, + }), + ]); + }); + it('should emit message_start event on first content', () => { adapter.processEvent({ type: GeminiEventType.Content, diff --git a/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.ts b/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.ts index 58095221ac1..1eac3e26018 100644 --- a/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.ts +++ b/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.ts @@ -7,9 +7,11 @@ import { randomUUID } from 'node:crypto'; import type { Config, + ServerGeminiStreamEvent, ToolCallRequestInfo, McpToolProgressData, } from '@qwen-code/qwen-code-core'; +import { GeminiEventType } from '@qwen-code/qwen-code-core'; import type { CLIAssistantMessage, CLIMessage, @@ -122,6 +124,21 @@ export class StreamJsonOutputAdapter this.emitMessage(message); } + override processEvent(event: ServerGeminiStreamEvent): void { + if (event.type === GeminiEventType.ActiveGoal) { + this.emitStreamEventIfEnabled( + { + type: 'active_goal', + active_goal: event.value, + }, + null, + ); + return; + } + + super.processEvent(event); + } + /** * Overrides base class hook to emit stream event when text block is created. */ diff --git a/packages/cli/src/nonInteractive/types.ts b/packages/cli/src/nonInteractive/types.ts index 84efda11ea8..774bea2ba81 100644 --- a/packages/cli/src/nonInteractive/types.ts +++ b/packages/cli/src/nonInteractive/types.ts @@ -1,5 +1,6 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ import type { + ActiveGoal, SubagentConfig, McpToolProgressData, } from '@qwen-code/qwen-code-core'; @@ -246,13 +247,19 @@ export interface ToolProgressStreamEvent { content: McpToolProgressData; } +export interface ActiveGoalStreamEvent { + type: 'active_goal'; + active_goal: ActiveGoal | null; +} + export type StreamEvent = | MessageStartStreamEvent | ContentBlockStartEvent | ContentBlockDeltaEvent | ContentBlockStopEvent | MessageStopStreamEvent - | ToolProgressStreamEvent; + | ToolProgressStreamEvent + | ActiveGoalStreamEvent; export interface CLIPartialAssistantMessage { type: 'stream_event'; diff --git a/packages/cli/src/nonInteractiveCliCommands.test.ts b/packages/cli/src/nonInteractiveCliCommands.test.ts index 0813d566bb1..3f448bbb1f3 100644 --- a/packages/cli/src/nonInteractiveCliCommands.test.ts +++ b/packages/cli/src/nonInteractiveCliCommands.test.ts @@ -4,15 +4,19 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { getAvailableCommands, handleSlashCommand, } from './nonInteractiveCliCommands.js'; -import type { Config } from '@qwen-code/qwen-code-core'; +import { + __resetActiveGoalStoreForTests, + type Config, +} from '@qwen-code/qwen-code-core'; import type { LoadedSettings } from './config/settings.js'; import { CommandKind, type ExecutionMode } from './ui/commands/types.js'; import { filterCommandsForMode } from './services/commandUtils.js'; +import { goalCommand } from './ui/commands/goalCommand.js'; // Mock the CommandService const mockGetCommands = vi.hoisted(() => vi.fn()); @@ -32,6 +36,7 @@ describe('handleSlashCommand', () => { beforeEach(() => { vi.clearAllMocks(); + __resetActiveGoalStoreForTests(); // getCommandsForMode applies real mode filtering on top of getCommands() mockGetCommandsForMode.mockImplementation((mode: ExecutionMode) => filterCommandsForMode(mockGetCommands(), mode), @@ -55,6 +60,12 @@ describe('handleSlashCommand', () => { getFolderTrustFeature: vi.fn().mockReturnValue(false), getFolderTrust: vi.fn().mockReturnValue(false), getProjectRoot: vi.fn().mockReturnValue('/test/project'), + isTrustedFolder: vi.fn().mockReturnValue(true), + getDisableAllHooks: vi.fn().mockReturnValue(false), + getHookSystem: vi.fn().mockReturnValue({ + addFunctionHook: vi.fn().mockReturnValue('goal-hook-id'), + removeFunctionHook: vi.fn().mockReturnValue(true), + }), setModelInvocableCommandsProvider: vi.fn(), setModelInvocableCommandsExecutor: vi.fn(), getDisabledSlashCommands: vi.fn().mockReturnValue([]), @@ -71,6 +82,10 @@ describe('handleSlashCommand', () => { abortController = new AbortController(); }); + afterEach(() => { + __resetActiveGoalStoreForTests(); + }); + it('should return no_command for non-slash input', async () => { const result = await handleSlashCommand( 'regular text', @@ -199,6 +214,94 @@ describe('handleSlashCommand', () => { } }); + it('should execute /goal in non-interactive mode as a submit_prompt command', async () => { + mockGetCommands.mockReturnValue([goalCommand]); + + const result = await handleSlashCommand( + '/goal write a hello world script', + abortController, + mockConfig, + mockSettings, + ); + + expect(result.type).toBe('submit_prompt'); + if (result.type === 'submit_prompt') { + expect(result.content).toEqual([ + expect.objectContaining({ + text: expect.stringContaining('write a hello world script'), + }), + ]); + } + }); + + it('should report no active goal for empty non-interactive /goal', async () => { + mockGetCommands.mockReturnValue([goalCommand]); + + const result = await handleSlashCommand( + '/goal', + abortController, + mockConfig, + mockSettings, + ); + + expect(result).toMatchObject({ + type: 'message', + messageType: 'info', + content: 'No goal set. Usage: `/goal ` (or `/goal clear`).', + }); + }); + + it('should report active goal status after setting a non-interactive /goal', async () => { + mockGetCommands.mockReturnValue([goalCommand]); + + await handleSlashCommand( + '/goal write a hello world script', + abortController, + mockConfig, + mockSettings, + ); + const result = await handleSlashCommand( + '/goal', + abortController, + mockConfig, + mockSettings, + ); + + expect(result).toMatchObject({ + type: 'message', + messageType: 'info', + }); + if (result.type === 'message') { + expect(result.content).toContain( + 'Goal active: write a hello world script', + ); + expect(result.content).toContain('not yet evaluated'); + } + }); + + it('should report cleared goal for non-interactive /goal clear', async () => { + mockGetCommands.mockReturnValue([goalCommand]); + + await handleSlashCommand( + '/goal write a hello world script', + abortController, + mockConfig, + mockSettings, + ); + const result = await handleSlashCommand( + '/goal clear', + abortController, + mockConfig, + mockSettings, + ); + + expect(result).toMatchObject({ + type: 'message', + messageType: 'info', + content: 'Goal cleared: write a hello world script', + }); + }); + it('should execute FILE commands in any mode without explicit supportedModes', async () => { const mockFileCommand = { name: 'custom', diff --git a/packages/cli/src/ui/commands/goalCommand.test.ts b/packages/cli/src/ui/commands/goalCommand.test.ts index 42aa2f2910f..286bf8f0a01 100644 --- a/packages/cli/src/ui/commands/goalCommand.test.ts +++ b/packages/cli/src/ui/commands/goalCommand.test.ts @@ -31,8 +31,11 @@ describe('goalCommand', () => { beforeEach(() => __resetActiveGoalStoreForTests()); afterEach(() => __resetActiveGoalStoreForTests()); - it('is currently limited to interactive mode', () => { - expect(goalCommand.supportedModes).toEqual(['interactive']); + it('is available in interactive and non-interactive modes', () => { + expect(goalCommand.supportedModes).toEqual([ + 'interactive', + 'non_interactive', + ]); }); it('rejects when config is missing', async () => { diff --git a/packages/cli/src/ui/commands/goalCommand.ts b/packages/cli/src/ui/commands/goalCommand.ts index 626bd7a2c92..dbf37465ca4 100644 --- a/packages/cli/src/ui/commands/goalCommand.ts +++ b/packages/cli/src/ui/commands/goalCommand.ts @@ -100,7 +100,7 @@ export const goalCommand: SlashCommand = { }, argumentHint: '[ | clear]', kind: CommandKind.BUILT_IN, - supportedModes: ['interactive'] as const, + supportedModes: ['interactive', 'non_interactive'] as const, action: async ( context: CommandContext, args: string, @@ -158,6 +158,9 @@ export const goalCommand: SlashCommand = { durationMs: Date.now() - cleared.setAt, }; context.ui.addItem(clearedItem, Date.now()); + if (context.executionMode === 'non_interactive') { + return infoMessage(`Goal cleared: ${cleared.condition}`); + } return; } diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 8517598e315..749a8510bd4 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -78,6 +78,9 @@ const mockParseAndFormatApiError = vi.hoisted(() => ); const mockLogApiCancel = vi.hoisted(() => vi.fn()); const mockGetActiveGoal = vi.hoisted(() => vi.fn()); +const mockActiveGoalEquals = vi.hoisted(() => vi.fn()); +const mockSetActiveGoal = vi.hoisted(() => vi.fn()); +const mockClearActiveGoal = vi.hoisted(() => vi.fn()); vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { const actualCoreModule = (await importOriginal()) as any; @@ -90,6 +93,9 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { parseAndFormatApiError: mockParseAndFormatApiError, logApiCancel: mockLogApiCancel, getActiveGoal: mockGetActiveGoal, + activeGoalEquals: mockActiveGoalEquals, + setActiveGoal: mockSetActiveGoal, + clearActiveGoal: mockClearActiveGoal, }; }); @@ -153,6 +159,7 @@ describe('useGeminiStream', () => { beforeEach(() => { vi.clearAllMocks(); // Clear mocks before each test mockGetActiveGoal.mockReturnValue(undefined); + mockActiveGoalEquals.mockReturnValue(false); vi.mocked(findLastSafeSplitPoint).mockImplementation( (s: string) => s.length, ); @@ -4638,6 +4645,81 @@ describe('useGeminiStream', () => { }); describe('StopHookLoop Event', () => { + it('syncs active_goal events into the active goal store', async () => { + const activeGoal = { + condition: 'finish the refactor', + iterations: 1, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing verification', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.ActiveGoal, + value: activeGoal, + }; + yield { + type: ServerGeminiEventType.ActiveGoal, + value: null, + }; + })(), + ); + mockGetActiveGoal + .mockReturnValueOnce(undefined) + .mockReturnValueOnce(activeGoal); + mockActiveGoalEquals.mockReturnValue(false); + + const { result } = renderTestHook(); + + await act(async () => { + await result.current.submitQuery('continue goal'); + }); + + expect(mockSetActiveGoal).toHaveBeenCalledWith( + 'test-session-id', + activeGoal, + ); + expect(mockClearActiveGoal).toHaveBeenCalledWith('test-session-id'); + }); + + it('skips redundant active_goal store updates', async () => { + const activeGoal = { + condition: 'finish the refactor', + iterations: 1, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing verification', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.ActiveGoal, + value: activeGoal, + }; + yield { + type: ServerGeminiEventType.ActiveGoal, + value: null, + }; + })(), + ); + mockGetActiveGoal + .mockReturnValueOnce(activeGoal) + .mockReturnValueOnce(undefined); + mockActiveGoalEquals.mockReturnValue(true); + + const { result } = renderTestHook(); + + await act(async () => { + await result.current.submitQuery('continue goal'); + }); + + expect(mockSetActiveGoal).not.toHaveBeenCalled(); + expect(mockClearActiveGoal).not.toHaveBeenCalled(); + }); + it('should handle StopHookLoop event and add stop hook loop history item', async () => { mockSendMessageStream.mockReturnValue( (async function* () { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 0bd888d44e6..13f3ab4264b 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -26,6 +26,7 @@ import type { ToolCallRequestInfo, GeminiErrorEventValue, StopFailureErrorType, + ActiveGoal, } from '@qwen-code/qwen-code-core'; import { GeminiEventType as ServerGeminiEventType, @@ -52,6 +53,9 @@ import { getUnsupportedImageFormatWarning, generateToolUseSummary, getActiveGoal, + activeGoalEquals, + setActiveGoal, + clearActiveGoal, } from '@qwen-code/qwen-code-core'; import { type Part, type PartListUnion, FinishReason } from '@google/genai'; import type { @@ -1339,6 +1343,25 @@ export const useGeminiStream = ( [addItem, config, pendingHistoryItemRef, setPendingHistoryItem], ); + const handleActiveGoalEvent = useCallback( + (activeGoal: ActiveGoal | null) => { + const sessionId = config.getSessionId(); + const currentActiveGoal = getActiveGoal(sessionId); + if (activeGoal) { + if (activeGoalEquals(currentActiveGoal, activeGoal)) { + return; + } + setActiveGoal(sessionId, activeGoal); + return; + } + if (!currentActiveGoal) { + return; + } + clearActiveGoal(sessionId); + }, + [config], + ); + const processGeminiStreamEvents = useCallback( async ( stream: AsyncIterable, @@ -1573,6 +1596,9 @@ export const useGeminiStream = ( flushBufferedStreamEvents(); handleStopHookLoopEvent(event.value, userMessageTimestamp); break; + case ServerGeminiEventType.ActiveGoal: + handleActiveGoalEvent(event.value); + break; default: { // enforces exhaustive switch-case const unreachable: never = event; @@ -1609,6 +1635,7 @@ export const useGeminiStream = ( setPendingHistoryItem, handleUserPromptSubmitBlockedEvent, handleStopHookLoopEvent, + handleActiveGoalEvent, addItem, dualOutput, ], diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 7e968b1326e..a45c51eda1a 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -52,6 +52,11 @@ import { promptIdContext } from '../utils/promptIdContext.js'; import { setSimulate429 } from '../utils/testUtils.js'; import { ideContextStore } from '../ide/ideContext.js'; import { uiTelemetryService } from '../telemetry/uiTelemetry.js'; +import { + __resetActiveGoalStoreForTests, + clearActiveGoal, + setActiveGoal, +} from '../goals/activeGoalStore.js'; // Mock fs module to prevent actual file system operations during tests const mockFileSystem = new Map(); @@ -514,6 +519,7 @@ describe('Gemini Client (client.ts)', () => { afterEach(() => { vi.restoreAllMocks(); + __resetActiveGoalStoreForTests(); }); describe('initialize', () => { @@ -4501,6 +4507,245 @@ Other open files: client['chat'] = mockChat as GeminiChat; }); + it('emits active_goal when a goal is active for the turn', async () => { + setActiveGoal('test-session-id', { + condition: 'finish the refactor', + iterations: 2, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing verification', + }); + + const events = await fromAsync( + client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-active-goal', + ), + ); + + expect(events[0]).toEqual({ + type: GeminiEventType.ActiveGoal, + value: { + condition: 'finish the refactor', + iterations: 2, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing verification', + }, + }); + }); + + it('emits active_goal null when the Stop hook clears the goal', async () => { + setActiveGoal('test-session-id', { + condition: 'finish the refactor', + iterations: 2, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing verification', + }); + const mockMessageBus = { + request: vi.fn().mockImplementation(async () => { + clearActiveGoal('test-session-id'); + return {}; + }), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'Stop', + ); + client['chat'] = { + addHistory: vi.fn(), + getHistory: vi.fn().mockReturnValue([ + { + role: 'model', + parts: [{ text: 'done' }], + }, + ]), + } as unknown as GeminiChat; + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'done' }; + })(), + ); + + const events = await fromAsync( + client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-cleared-active-goal', + ), + ); + + expect(events).toContainEqual({ + type: GeminiEventType.ActiveGoal, + value: null, + }); + }); + + it('emits active_goal null when the Stop hook clears the goal before aborting', async () => { + setActiveGoal('test-session-id', { + condition: 'finish the refactor', + iterations: 2, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing verification', + }); + const abortController = new AbortController(); + const mockMessageBus = { + request: vi.fn().mockImplementation(async () => { + clearActiveGoal('test-session-id'); + abortController.abort(); + return {}; + }), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'Stop', + ); + client['chat'] = { + addHistory: vi.fn(), + getHistory: vi.fn().mockReturnValue([ + { + role: 'model', + parts: [{ text: 'done' }], + }, + ]), + } as unknown as GeminiChat; + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'done' }; + })(), + ); + + const events = await fromAsync( + client.sendMessageStream( + [{ text: 'Hi' }], + abortController.signal, + 'prompt-cleared-active-goal-then-aborted', + ), + ); + + const activeGoalEvents = events.filter( + (event) => event.type === GeminiEventType.ActiveGoal, + ); + + expect(activeGoalEvents).toEqual([ + { + type: GeminiEventType.ActiveGoal, + value: expect.objectContaining({ + condition: 'finish the refactor', + }), + }, + { + type: GeminiEventType.ActiveGoal, + value: null, + }, + ]); + }); + + it('emits active_goal changes when aborting before Stop hook continuation', async () => { + setActiveGoal('test-session-id', { + condition: 'finish the refactor', + iterations: 2, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing verification', + }); + const abortController = new AbortController(); + const mockMessageBus = { + request: vi.fn().mockImplementation(async () => { + setActiveGoal('test-session-id', { + condition: 'finish the refactor', + iterations: 3, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing validation', + }); + return { + output: { + get decision() { + abortController.abort(); + return 'block'; + }, + reason: 'Keep working', + }, + stopHookCount: 1, + }; + }), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'Stop', + ); + client['chat'] = { + addHistory: vi.fn(), + getHistory: vi.fn().mockReturnValue([ + { + role: 'model', + parts: [{ text: 'done' }], + }, + ]), + } as unknown as GeminiChat; + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'done' }; + })(), + ); + + const events = await fromAsync( + client.sendMessageStream( + [{ text: 'Hi' }], + abortController.signal, + 'prompt-stop-hook-continuation-aborted', + ), + ); + const activeGoalEvents = events.filter( + (event) => event.type === GeminiEventType.ActiveGoal, + ); + + expect(activeGoalEvents).toEqual([ + { + type: GeminiEventType.ActiveGoal, + value: expect.objectContaining({ + condition: 'finish the refactor', + iterations: 2, + }), + }, + { + type: GeminiEventType.ActiveGoal, + value: expect.objectContaining({ + condition: 'finish the refactor', + iterations: 3, + lastReason: 'still missing validation', + }), + }, + ]); + expect(events).not.toContainEqual( + expect.objectContaining({ + type: GeminiEventType.StopHookLoop, + }), + ); + }); + it('should skip messageBus.request for UserPromptSubmit when hasHooksForEvent returns false', async () => { // Enable hooks and provide messageBus const mockMessageBus = { @@ -4606,6 +4851,74 @@ Other open files: }); }); + it('emits one active_goal null when the blocking cap aborts an active goal', async () => { + setActiveGoal('test-session-id', { + condition: 'finish the refactor', + iterations: 2, + setAt: 123, + tokensAtStart: 456, + hookId: 'goal-hook-id', + lastReason: 'still missing verification', + }); + const mockMessageBus = { + request: vi.fn().mockResolvedValue({ + output: { + decision: 'block', + reason: 'Keep working', + }, + stopHookCount: 1, + }), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'Stop', + ); + vi.mocked(mockConfig.getStopHookBlockingCap).mockReturnValue(1); + + client['chat'] = { + addHistory: vi.fn(), + getHistory: vi.fn().mockReturnValue([ + { + role: 'model', + parts: [{ text: 'not done' }], + }, + ]), + } as unknown as GeminiChat; + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'not done' }; + })(), + ); + + const events = await fromAsync( + client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-stop-cap-active-goal', + ), + ); + const activeGoalEvents = events.filter( + (event) => event.type === GeminiEventType.ActiveGoal, + ); + + expect(activeGoalEvents).toEqual([ + { + type: GeminiEventType.ActiveGoal, + value: expect.objectContaining({ + condition: 'finish the refactor', + }), + }, + { + type: GeminiEventType.ActiveGoal, + value: null, + }, + ]); + }); + it('should not skip hooks when hasHooksForEvent returns true', async () => { const mockMessageBus = { request: vi.fn().mockResolvedValue({ modifiedPrompt: undefined }), diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index b50e5aa411d..2ff9f1bc6f7 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -18,7 +18,11 @@ import { ApprovalMode, type Config } from '../config/config.js'; import { createDebugLogger } from '../utils/debugLogger.js'; import { recordStartupEvent } from '../utils/startupEventSink.js'; import { microcompactHistory } from '../services/microcompaction/microcompact.js'; -import { getActiveGoal } from '../goals/activeGoalStore.js'; +import { + activeGoalEquals, + getActiveGoal, + type ActiveGoal, +} from '../goals/activeGoalStore.js'; import { abortGoalForStopHookCap } from '../goals/goalHook.js'; import { formatStopHookBlockingCapWarning } from '../hooks/stopHookCap.js'; @@ -1454,6 +1458,29 @@ export class GeminiClient { requestToSend = [...systemReminders, ...requestToSend]; } + const activeGoalAtTurnStart = getActiveGoal(this.config.getSessionId()); + if (activeGoalAtTurnStart) { + yield { + type: GeminiEventType.ActiveGoal, + value: activeGoalAtTurnStart, + }; + } + let lastEmittedActiveGoal: ActiveGoal | undefined = activeGoalAtTurnStart; + // Tracks the last emitted goal value to suppress duplicate events. + // Mutates `lastEmittedActiveGoal` when an event is returned. + const maybeEmitActiveGoalChange = ( + nextActiveGoal: ActiveGoal | undefined, + ): ServerGeminiStreamEvent | undefined => { + if (activeGoalEquals(lastEmittedActiveGoal, nextActiveGoal)) { + return undefined; + } + lastEmittedActiveGoal = nextActiveGoal; + return { + type: GeminiEventType.ActiveGoal, + value: nextActiveGoal ?? null, + }; + }; + const resultStream = turn.run(model, requestToSend, signal); let didUpdateIdeContextState = false; for await (const event of resultStream) { @@ -1570,8 +1597,20 @@ export class GeminiClient { MessageBusType.HOOK_EXECUTION_RESPONSE, ); + // Stop hook callbacks can mutate active goal state during request(). + // Capture it before cancellation returns so clear events are not lost. + const activeGoalAfterStopHook = getActiveGoal( + this.config.getSessionId(), + ); + // Check if aborted after hook execution if (signal.aborted) { + const activeGoalEvent = maybeEmitActiveGoalChange( + activeGoalAfterStopHook, + ); + if (activeGoalEvent) { + yield activeGoalEvent; + } if (isTopLevelInteraction) endInteractionSpan('cancelled'); return turn; } @@ -1597,6 +1636,12 @@ export class GeminiClient { ) { // Check if aborted before continuing if (signal.aborted) { + const activeGoalEvent = maybeEmitActiveGoalChange( + activeGoalAfterStopHook, + ); + if (activeGoalEvent) { + yield activeGoalEvent; + } if (isTopLevelInteraction) endInteractionSpan('cancelled'); return turn; } @@ -1626,6 +1671,14 @@ export class GeminiClient { this.config.getSessionId(), warning, ); + const activeGoalAfterCap = getActiveGoal( + this.config.getSessionId(), + ); + const activeGoalEvent = + maybeEmitActiveGoalChange(activeGoalAfterCap); + if (activeGoalEvent) { + yield activeGoalEvent; + } yield { type: GeminiEventType.HookSystemMessage, value: warning, @@ -1635,6 +1688,13 @@ export class GeminiClient { return turn; } + const activeGoalEvent = maybeEmitActiveGoalChange( + activeGoalAfterStopHook, + ); + if (activeGoalEvent) { + yield activeGoalEvent; + } + yield { type: GeminiEventType.StopHookLoop, value: { @@ -1665,6 +1725,13 @@ export class GeminiClient { endInteractionSpan(signal.aborted ? 'cancelled' : 'ok'); return hookTurn; } + + const activeGoalEvent = maybeEmitActiveGoalChange( + activeGoalAfterStopHook, + ); + if (activeGoalEvent) { + yield activeGoalEvent; + } } if (!turn.pendingToolCalls.length && signal && !signal.aborted) { diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 34f78bcd063..8847120a843 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -34,6 +34,7 @@ import { type ThoughtSummary, } from '../utils/thoughtUtils.js'; import type { LoopType } from '../telemetry/types.js'; +import type { ActiveGoal } from '../goals/activeGoalStore.js'; // Define a structure for tools passed to the server export interface ServerTool { @@ -64,6 +65,7 @@ export enum GeminiEventType { HookSystemMessage = 'hook_system_message', UserPromptSubmitBlocked = 'user_prompt_submit_blocked', StopHookLoop = 'stop_hook_loop', + ActiveGoal = 'active_goal', } export type ServerGeminiRetryEvent = { @@ -233,8 +235,14 @@ export type ServerGeminiStopHookLoopEvent = { }; }; +export type ServerGeminiActiveGoalEvent = { + type: GeminiEventType.ActiveGoal; + value: ActiveGoal | null; +}; + // The original union type, now composed of the individual types export type ServerGeminiStreamEvent = + | ServerGeminiActiveGoalEvent | ServerGeminiChatCompressedEvent | ServerGeminiCitationEvent | ServerGeminiContentEvent diff --git a/packages/core/src/goals/activeGoalStore.test.ts b/packages/core/src/goals/activeGoalStore.test.ts index 435a8071be0..6cfe1a91ac1 100644 --- a/packages/core/src/goals/activeGoalStore.test.ts +++ b/packages/core/src/goals/activeGoalStore.test.ts @@ -7,6 +7,7 @@ import { beforeEach, describe, expect, it } from 'vitest'; import { __resetActiveGoalStoreForTests, + activeGoalEquals, clearActiveGoal, getActiveGoal, recordGoalIteration, @@ -60,4 +61,19 @@ describe('activeGoalStore', () => { it('recordGoalIteration is a no-op when no goal exists', () => { expect(recordGoalIteration('sess-missing', 'noop')).toBeUndefined(); }); + + it('compares active goal snapshots by value', () => { + expect(activeGoalEquals(undefined, undefined)).toBe(true); + expect(activeGoalEquals(makeGoal(), makeGoal())).toBe(true); + expect( + activeGoalEquals(makeGoal(), makeGoal({ lastReason: undefined })), + ).toBe(true); + expect( + activeGoalEquals( + makeGoal({ iterations: 1 }), + makeGoal({ iterations: 2 }), + ), + ).toBe(false); + expect(activeGoalEquals(makeGoal(), undefined)).toBe(false); + }); }); diff --git a/packages/core/src/goals/activeGoalStore.ts b/packages/core/src/goals/activeGoalStore.ts index 162a1221ab9..184015c8b2d 100644 --- a/packages/core/src/goals/activeGoalStore.ts +++ b/packages/core/src/goals/activeGoalStore.ts @@ -20,6 +20,26 @@ export interface ActiveGoal { const store = new Map(); +export function activeGoalEquals( + left: ActiveGoal | undefined, + right: ActiveGoal | undefined, +): boolean { + if (left === right) return true; + if (!left || !right) return false; + return stableActiveGoalKey(left) === stableActiveGoalKey(right); +} + +function stableActiveGoalKey(goal: ActiveGoal): string { + const comparable: Record = {}; + for (const key of Object.keys(goal).sort() as Array) { + const value = goal[key]; + if (value !== undefined) { + comparable[key] = value; + } + } + return JSON.stringify(comparable); +} + export function getActiveGoal(sessionId: string): ActiveGoal | undefined { return store.get(sessionId); } diff --git a/packages/core/src/goals/index.ts b/packages/core/src/goals/index.ts index eb382b58b95..471afdf6f36 100644 --- a/packages/core/src/goals/index.ts +++ b/packages/core/src/goals/index.ts @@ -11,6 +11,7 @@ export type { GoalTerminalObserver, } from './activeGoalStore.js'; export { + activeGoalEquals, getActiveGoal, setActiveGoal, clearActiveGoal, diff --git a/packages/core/src/hooks/promptHookRunner.test.ts b/packages/core/src/hooks/promptHookRunner.test.ts index c93b1347771..8fef034118e 100644 --- a/packages/core/src/hooks/promptHookRunner.test.ts +++ b/packages/core/src/hooks/promptHookRunner.test.ts @@ -472,6 +472,9 @@ describe('PromptHookRunner', () => { }); it('should track duration correctly', async () => { + vi.useFakeTimers(); + vi.setSystemTime(0); + const mockResponse = createMockResponse('{"ok": true}'); mockGenerateContent.mockImplementation( () => @@ -483,13 +486,20 @@ describe('PromptHookRunner', () => { const config = createMockConfig(); const input = createMockInput(); - const result = await promptRunner.execute( - config, - HookEventName.PreToolUse, - input, - ); + try { + const execution = promptRunner.execute( + config, + HookEventName.PreToolUse, + input, + ); + + await vi.advanceTimersByTimeAsync(50); + const result = await execution; - expect(result.duration).toBeGreaterThanOrEqual(50); + expect(result.duration).toBe(50); + } finally { + vi.useRealTimers(); + } }); it('should handle multiple $ARGUMENTS placeholders', async () => {