From 1a6f73043960893e9e78969579b654784260dc31 Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Tue, 17 Mar 2026 21:34:50 -0700 Subject: [PATCH] feat: add Forever Mode with time-based scheduled work and A2A listener Add --forever CLI flag that enables autonomous agent operation with scheduled work, context management, A2A protocol support, and session optimization. Core features: - Time-based WorkScheduler: manages a sorted list of scheduled prompts that fire at absolute or relative times, persisted across sessions - schedule_work tool: add/cancel scheduled prompts with 'at' (local time) or 'inMinutes' (relative) params; current time and schedule prepended to every model turn - A2A HTTP listener: JSON-RPC 2.0 server bridges external messages into the session (message/send, tasks/get, responses/poll) - PreCompress hook: hooks can return newHistory to replace built-in LLM compression - Idle hook: fires after configurable inactivity, can auto-submit prompts - Forever mode disables MemoryTool, EnterPlanModeTool, interactive shell UI: - ScheduledWorkDisplay component shows all pending items above the context summary bar - A2A port shown in StatusDisplay when active Session optimization for long-running sessions: - Record lastCompressionIndex on ConversationRecord; on resume, only load post-compression messages - Restore scheduled work items on session resume (past-due fire immediately) - Skip file I/O in updateMessagesFromHistory when no tool results to sync - Prune UI history to last 50 items after each context compression --- packages/cli/src/acp/acpClient.ts | 5 +- packages/cli/src/config/config.ts | 8 + packages/cli/src/config/settingsSchema.ts | 12 + packages/cli/src/external-listener.ts | 453 ++++++++++++++++ packages/cli/src/gemini.test.tsx | 1 + packages/cli/src/interactiveCli.tsx | 20 + packages/cli/src/nonInteractiveCli.ts | 1 + packages/cli/src/test-utils/mockConfig.ts | 8 + packages/cli/src/ui/App.test.tsx | 1 + packages/cli/src/ui/AppContainer.tsx | 87 +++ .../cli/src/ui/components/Composer.test.tsx | 5 + packages/cli/src/ui/components/Composer.tsx | 2 + .../ui/components/ScheduledWorkDisplay.tsx | 66 +++ .../src/ui/components/StatusDisplay.test.tsx | 13 + .../cli/src/ui/components/StatusDisplay.tsx | 6 + .../cli/src/ui/contexts/UIStateContext.tsx | 1 + .../src/ui/hooks/creditsFlowHandler.test.ts | 1 + packages/cli/src/ui/hooks/useGeminiStream.ts | 81 ++- .../src/ui/hooks/useHistoryManager.test.ts | 80 ++- .../cli/src/ui/hooks/useHistoryManager.ts | 29 +- .../src/ui/hooks/useIncludeDirsTrust.test.tsx | 1 + .../src/ui/hooks/useQuotaAndFallback.test.ts | 1 + .../cli/src/ui/hooks/useSessionBrowser.ts | 7 +- .../cli/src/ui/hooks/useSessionResume.test.ts | 56 ++ packages/cli/src/ui/hooks/useSessionResume.ts | 18 +- packages/cli/src/utils/events.ts | 4 + packages/cli/src/utils/sessionUtils.ts | 14 +- packages/core/src/config/config.ts | 39 +- packages/core/src/core/client.test.ts | 3 + packages/core/src/core/client.ts | 7 + packages/core/src/core/prompts.test.ts | 2 + packages/core/src/core/turn.ts | 3 + .../core/src/hooks/hookAggregator.test.ts | 28 + packages/core/src/hooks/hookAggregator.ts | 7 +- .../core/src/hooks/hookEventHandler.test.ts | 97 ++++ packages/core/src/hooks/hookEventHandler.ts | 20 + packages/core/src/hooks/hookRegistry.ts | 2 + packages/core/src/hooks/hookSystem.ts | 9 +- packages/core/src/hooks/types.test.ts | 3 + packages/core/src/hooks/types.ts | 32 +- packages/core/src/index.ts | 1 + .../core/src/prompts/promptProvider.test.ts | 1 + .../services/chatCompressionService.test.ts | 149 +++++- .../src/services/chatCompressionService.ts | 66 ++- .../core/src/services/chatRecordingService.ts | 93 +++- .../core/src/services/work-scheduler.test.ts | 498 ++++++++++++++++++ packages/core/src/services/work-scheduler.ts | 223 ++++++++ packages/core/src/tools/schedule-work.test.ts | 308 +++++++++++ packages/core/src/tools/schedule-work.ts | 210 ++++++++ packages/core/src/tools/tool-names.ts | 2 + packages/core/src/utils/sessionUtils.ts | 8 +- 51 files changed, 2737 insertions(+), 55 deletions(-) create mode 100644 packages/cli/src/external-listener.ts create mode 100644 packages/cli/src/ui/components/ScheduledWorkDisplay.tsx create mode 100644 packages/core/src/services/work-scheduler.test.ts create mode 100644 packages/core/src/services/work-scheduler.ts create mode 100644 packages/core/src/tools/schedule-work.test.ts create mode 100644 packages/core/src/tools/schedule-work.ts diff --git a/packages/cli/src/acp/acpClient.ts b/packages/cli/src/acp/acpClient.ts index 072d91c20a..357ef3b567 100644 --- a/packages/cli/src/acp/acpClient.ts +++ b/packages/cli/src/acp/acpClient.ts @@ -360,7 +360,10 @@ export class GeminiAgent { config.setFileSystemService(acpFileSystemService); } - const clientHistory = convertSessionToClientHistory(sessionData.messages); + const clientHistory = convertSessionToClientHistory( + sessionData.messages, + sessionData.lastCompressionIndex, + ); const geminiClient = config.getGeminiClient(); await geminiClient.initialize(); diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 80c1e19443..35016db4ea 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -97,6 +97,7 @@ export interface CliArgs { rawOutput: boolean | undefined; acceptRawOutputRisk: boolean | undefined; isCommand: boolean | undefined; + forever: boolean | undefined; } /** @@ -298,6 +299,12 @@ export async function parseArguments( .option('accept-raw-output-risk', { type: 'boolean', description: 'Suppress the security warning when using --raw-output.', + }) + .option('forever', { + type: 'boolean', + description: + 'Run as a long-running autonomous agent with auto-resume and schedule_work support.', + default: false, }), ) // Register MCP subcommands @@ -893,6 +900,7 @@ export async function loadCliConfig( }; }, enableConseca: settings.security?.enableConseca, + isForeverMode: !!argv.forever, }); } diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index 8a107c4d47..876bb5b684 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -2312,6 +2312,18 @@ const SETTINGS_SCHEMA = { ref: 'HookDefinitionArray', mergeStrategy: MergeStrategy.CONCAT, }, + Idle: { + type: 'array', + label: 'Idle Hooks', + category: 'Advanced', + requiresRestart: false, + default: [], + description: + 'Hooks that execute after a period of inactivity. Can trigger maintenance tasks like memory consolidation.', + showInDialog: false, + ref: 'HookDefinitionArray', + mergeStrategy: MergeStrategy.CONCAT, + }, }, additionalProperties: { type: 'array', diff --git a/packages/cli/src/external-listener.ts b/packages/cli/src/external-listener.ts new file mode 100644 index 0000000000..d9f5a5851b --- /dev/null +++ b/packages/cli/src/external-listener.ts @@ -0,0 +1,453 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import http from 'node:http'; +import { writeFileSync, mkdirSync, unlinkSync } from 'node:fs'; +import { join } from 'node:path'; +import os from 'node:os'; +import crypto from 'node:crypto'; +import { appEvents, AppEvent } from './utils/events.js'; + +// --- A2A Task management --- + +interface A2AResponseMessage { + kind: 'message'; + role: 'agent'; + parts: Array<{ kind: 'text'; text: string }>; + messageId: string; +} + +interface A2ATask { + id: string; + contextId: string; + status: { + state: 'submitted' | 'working' | 'completed' | 'failed'; + timestamp: string; + message?: A2AResponseMessage; + }; +} + +const tasks = new Map(); + +const TASK_CLEANUP_DELAY_MS = 10 * 60 * 1000; // 10 minutes +const DEFAULT_BLOCKING_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes + +interface ResponseWaiter { + taskId: string; + resolve: (text: string) => void; +} + +const responseWaiters: ResponseWaiter[] = []; + +// Queue for unsolicited responses (e.g. forever mode auto-resume output) +const unsolicitedResponses: string[] = []; + +/** + * Called by AppContainer when streaming transitions from non-Idle to Idle. + * If there's a pending A2A task, resolves it. Otherwise queues as unsolicited. + */ +export function notifyResponse(responseText: string): void { + if (!responseText) return; + + const waiter = responseWaiters.shift(); + if (!waiter) { + // No A2A task waiting — queue as unsolicited (forever mode, etc.) + unsolicitedResponses.push(responseText); + return; + } + + const task = tasks.get(waiter.taskId); + if (task) { + task.status = { + state: 'completed', + timestamp: new Date().toISOString(), + message: { + kind: 'message', + role: 'agent', + parts: [{ kind: 'text', text: responseText }], + messageId: crypto.randomUUID(), + }, + }; + scheduleTaskCleanup(task.id); + } + + waiter.resolve(responseText); +} + +/** + * Drain all unsolicited responses (from forever mode auto-resume, etc.). + */ +export function drainUnsolicitedResponses(): string[] { + return unsolicitedResponses.splice(0, unsolicitedResponses.length); +} + +/** + * Returns true if there are any in-flight tasks waiting for a response. + */ +export function hasPendingTasks(): boolean { + return responseWaiters.length > 0; +} + +/** + * Called when streaming starts (Idle -> non-Idle) to mark the oldest + * submitted task as "working". + */ +export function markTasksWorking(): void { + const waiter = responseWaiters[0]; + if (!waiter) return; + const task = tasks.get(waiter.taskId); + if (task && task.status.state === 'submitted') { + task.status = { + state: 'working', + timestamp: new Date().toISOString(), + }; + } +} + +function scheduleTaskCleanup(taskId: string): void { + setTimeout(() => { + tasks.delete(taskId); + }, TASK_CLEANUP_DELAY_MS); +} + +function createTask(): A2ATask { + const task: A2ATask = { + id: crypto.randomUUID(), + contextId: `session-${process.pid}`, + status: { + state: 'submitted', + timestamp: new Date().toISOString(), + }, + }; + tasks.set(task.id, task); + return task; +} + +function formatTaskResult(task: A2ATask): object { + return { + kind: 'task', + id: task.id, + contextId: task.contextId, + status: task.status, + }; +} + +// --- JSON-RPC helpers --- + +interface JsonRpcRequest { + jsonrpc?: string; + id?: string | number | null; + method?: string; + params?: Record; +} + +function jsonRpcSuccess(id: string | number | null, result: object): object { + return { jsonrpc: '2.0', id, result }; +} + +function jsonRpcError( + id: string | number | null, + code: number, + message: string, +): object { + return { jsonrpc: '2.0', id, error: { code, message } }; +} + +// --- HTTP utilities --- + +function getSessionsDir(): string { + return join(os.homedir(), '.gemini', 'sessions'); +} + +function getPortFilePath(): string { + return join(getSessionsDir(), `interactive-${process.pid}.port`); +} + +function buildAgentCard(port: number): object { + return { + name: 'Gemini CLI Interactive Session', + url: `http://localhost:${port}/`, + protocolVersion: '0.3.0', + provider: { organization: 'Google', url: 'https://google.com' }, + capabilities: { streaming: false, pushNotifications: false }, + defaultInputModes: ['text'], + defaultOutputModes: ['text'], + skills: [ + { + id: 'interactive_session', + name: 'Interactive Session', + description: 'Send messages to the live interactive Gemini CLI session', + }, + ], + }; +} + +interface A2AMessagePart { + kind?: string; + text?: string; +} + +function extractTextFromParts( + parts: A2AMessagePart[] | undefined, +): string | null { + if (!Array.isArray(parts)) { + return null; + } + const texts: string[] = []; + for (const part of parts) { + if (part.kind === 'text' && typeof part.text === 'string') { + texts.push(part.text); + } + } + return texts.length > 0 ? texts.join('\n') : null; +} + +function sendJson( + res: http.ServerResponse, + statusCode: number, + data: object, +): void { + const body = JSON.stringify(data); + res.writeHead(statusCode, { + 'Content-Type': 'application/json', + 'Content-Length': Buffer.byteLength(body), + }); + res.end(body); +} + +function readBody(req: http.IncomingMessage): Promise { + return new Promise((resolve, reject) => { + const chunks: Buffer[] = []; + let size = 0; + const maxSize = 1024 * 1024; // 1MB limit + req.on('data', (chunk: Buffer) => { + size += chunk.length; + if (size > maxSize) { + req.destroy(); + reject(new Error('Request body too large')); + return; + } + chunks.push(chunk); + }); + req.on('end', () => resolve(Buffer.concat(chunks).toString('utf-8'))); + req.on('error', reject); + }); +} + +// --- JSON-RPC request handlers --- + +function handleMessageSend( + rpcId: string | number | null, + params: Record, + res: http.ServerResponse, +): void { + const messageVal = params['message']; + const message = + messageVal && typeof messageVal === 'object' + ? (messageVal as { role?: string; parts?: A2AMessagePart[] }) + : undefined; + const text = extractTextFromParts(message?.parts); + if (!text) { + sendJson( + res, + 200, + jsonRpcError( + rpcId, + -32602, + 'Missing or empty text. Expected: params.message.parts with kind "text".', + ), + ); + return; + } + + const task = createTask(); + + // Inject message into the session + appEvents.emit(AppEvent.ExternalMessage, text); + + // Block until response (standard A2A message/send semantics) + const timer = setTimeout(() => { + const idx = responseWaiters.findIndex((w) => w.taskId === task.id); + if (idx !== -1) { + responseWaiters.splice(idx, 1); + } + task.status = { + state: 'failed', + timestamp: new Date().toISOString(), + }; + scheduleTaskCleanup(task.id); + sendJson(res, 200, jsonRpcError(rpcId, -32000, 'Request timed out')); + }, DEFAULT_BLOCKING_TIMEOUT_MS); + + responseWaiters.push({ + taskId: task.id, + resolve: () => { + clearTimeout(timer); + // Task is already updated in notifyResponse + const updatedTask = tasks.get(task.id); + sendJson( + res, + 200, + jsonRpcSuccess(rpcId, formatTaskResult(updatedTask ?? task)), + ); + }, + }); +} + +function handleResponsesPoll( + rpcId: string | number | null, + res: http.ServerResponse, +): void { + const responses = drainUnsolicitedResponses(); + sendJson(res, 200, jsonRpcSuccess(rpcId, { responses })); +} + +function handleTasksGet( + rpcId: string | number | null, + params: Record, + res: http.ServerResponse, +): void { + const taskId = params['id']; + if (typeof taskId !== 'string') { + sendJson( + res, + 200, + jsonRpcError(rpcId, -32602, 'Missing or invalid params.id'), + ); + return; + } + + const task = tasks.get(taskId); + if (!task) { + sendJson(res, 200, jsonRpcError(rpcId, -32001, 'Task not found')); + return; + } + + sendJson(res, 200, jsonRpcSuccess(rpcId, formatTaskResult(task))); +} + +// --- Server --- + +export interface ExternalListenerResult { + port: number; + cleanup: () => void; +} + +/** + * Start an embedded HTTP server that accepts A2A-format JSON-RPC messages + * and bridges them into the interactive session's message queue. + */ +export function startExternalListener(options?: { + port?: number; +}): Promise { + const port = options?.port ?? 0; + + return new Promise((resolve, reject) => { + const server = http.createServer( + (req: http.IncomingMessage, res: http.ServerResponse) => { + const url = new URL(req.url ?? '/', `http://localhost`); + + // GET /.well-known/agent-card.json + if ( + req.method === 'GET' && + url.pathname === '/.well-known/agent-card.json' + ) { + const address = server.address(); + const actualPort = + typeof address === 'object' && address ? address.port : port; + sendJson(res, 200, buildAgentCard(actualPort)); + return; + } + + // POST / — JSON-RPC 2.0 routing + if (req.method === 'POST' && url.pathname === '/') { + readBody(req) + .then((rawBody) => { + let parsed: JsonRpcRequest; + try { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + parsed = JSON.parse(rawBody) as JsonRpcRequest; + } catch { + sendJson( + res, + 200, + jsonRpcError(null, -32700, 'Parse error: invalid JSON'), + ); + return; + } + + const rpcId = parsed.id ?? null; + const method = parsed.method; + const params = parsed.params ?? {}; + + switch (method) { + case 'message/send': + handleMessageSend(rpcId, params, res); + break; + case 'tasks/get': + handleTasksGet(rpcId, params, res); + break; + case 'responses/poll': + handleResponsesPoll(rpcId, res); + break; + default: + sendJson( + res, + 200, + jsonRpcError( + rpcId, + -32601, + `Method not found: ${method ?? '(none)'}`, + ), + ); + } + }) + .catch(() => { + sendJson( + res, + 200, + jsonRpcError(null, -32603, 'Failed to read request body'), + ); + }); + return; + } + + // 404 for everything else + sendJson(res, 404, { error: 'Not found' }); + }, + ); + + server.listen(port, '127.0.0.1', () => { + const address = server.address(); + const actualPort = + typeof address === 'object' && address ? address.port : port; + + // Write port file + try { + const sessionsDir = getSessionsDir(); + mkdirSync(sessionsDir, { recursive: true }); + writeFileSync(getPortFilePath(), String(actualPort), 'utf-8'); + } catch { + // Non-fatal: port file is a convenience, not a requirement + } + + const cleanup = () => { + server.close(); + try { + unlinkSync(getPortFilePath()); + } catch { + // Ignore: file may already be deleted + } + }; + + resolve({ port: actualPort, cleanup }); + }); + + server.on('error', (err) => { + reject(err); + }); + }); +} diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index 31fec36db0..ccf58f6f42 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/packages/cli/src/gemini.test.tsx @@ -513,6 +513,7 @@ describe('gemini.tsx main function kitty protocol', () => { rawOutput: undefined, acceptRawOutputRisk: undefined, isCommand: undefined, + forever: undefined, }); await act(async () => { diff --git a/packages/cli/src/interactiveCli.tsx b/packages/cli/src/interactiveCli.tsx index a6337ef29c..8b4aba1251 100644 --- a/packages/cli/src/interactiveCli.tsx +++ b/packages/cli/src/interactiveCli.tsx @@ -10,6 +10,8 @@ import { basename } from 'node:path'; import { AppContainer } from './ui/AppContainer.js'; import { ConsolePatcher } from './ui/utils/ConsolePatcher.js'; import { registerCleanup, setupTtyCheck } from './utils/cleanup.js'; +import { startExternalListener } from './external-listener.js'; +import { appEvents, AppEvent } from './utils/events.js'; import { type StartupWarning, type Config, @@ -181,6 +183,24 @@ export async function startInteractiveUI( registerCleanup(() => instance.unmount()); + // Auto-start A2A HTTP listener in Forever Mode + if (config.getIsForeverMode()) { + try { + const listener = await startExternalListener({ port: 0 }); + registerCleanup(listener.cleanup); + appEvents.emit(AppEvent.A2AListenerStarted, listener.port); + coreEvents.emitFeedback( + 'info', + `A2A endpoint listening on port ${listener.port}`, + ); + } catch (err) { + coreEvents.emitFeedback( + 'warning', + `Failed to start A2A listener: ${err instanceof Error ? err.message : String(err)}`, + ); + } + } + registerCleanup(setupTtyCheck()); } diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 891e3d0ee9..43fb8a3021 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -222,6 +222,7 @@ export async function runNonInteractive({ await geminiClient.resumeChat( convertSessionToClientHistory( resumedSessionData.conversation.messages, + resumedSessionData.conversation.lastCompressionIndex, ), resumedSessionData, ); diff --git a/packages/cli/src/test-utils/mockConfig.ts b/packages/cli/src/test-utils/mockConfig.ts index d4f11212e3..631cc20c30 100644 --- a/packages/cli/src/test-utils/mockConfig.ts +++ b/packages/cli/src/test-utils/mockConfig.ts @@ -173,6 +173,14 @@ export const createMockConfig = (overrides: Partial = {}): Config => getHasAccessToPreviewModel: vi.fn().mockReturnValue(false), validatePathAccess: vi.fn().mockReturnValue(null), getUseAlternateBuffer: vi.fn().mockReturnValue(false), + getIsForeverMode: vi.fn().mockReturnValue(false), + getWorkScheduler: vi.fn().mockReturnValue({ + on: vi.fn(), + off: vi.fn(), + serialize: vi.fn().mockReturnValue([]), + restore: vi.fn(), + dispose: vi.fn(), + }), ...overrides, }) as unknown as Config; diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index 4e59ab854e..d6adb12836 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -82,6 +82,7 @@ describe('App', () => { updateItem: vi.fn(), clearItems: vi.fn(), loadHistory: vi.fn(), + pruneItems: vi.fn(), }, history: [], pendingHistoryItems: [], diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index b2402f9fe9..bbbdd4279d 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -127,6 +127,7 @@ import { useFolderTrust } from './hooks/useFolderTrust.js'; import { useIdeTrustListener } from './hooks/useIdeTrustListener.js'; import { type IdeIntegrationNudgeResult } from './IdeIntegrationNudge.js'; import { appEvents, AppEvent, TransientMessageType } from '../utils/events.js'; +import { notifyResponse, markTasksWorking } from '../external-listener.js'; import { type UpdateObject } from './utils/updateCheck.js'; import { setUpdateHandler } from '../utils/handleAutoUpdate.js'; import { registerCleanup, runExitCleanup } from '../utils/cleanup.js'; @@ -1142,6 +1143,7 @@ Logging in with Google... Restarting Gemini CLI to continue. terminalHeight, embeddedShellFocused, consumePendingHints, + historyManager.pruneItems, ); toggleBackgroundShellRef.current = toggleBackgroundShell; @@ -1213,6 +1215,89 @@ Logging in with Google... Restarting Gemini CLI to continue. isMcpReady, }); + // --- A2A listener integration --- + const [a2aListenerPort, setA2aListenerPort] = useState(null); + + useEffect(() => { + const handler = (port: number) => { + setA2aListenerPort(port); + }; + appEvents.on(AppEvent.A2AListenerStarted, handler); + return () => { + appEvents.off(AppEvent.A2AListenerStarted, handler); + }; + }, []); + + // Bridge external messages from A2A HTTP listener to message queue + useEffect(() => { + const handler = (text: string) => { + addMessage(text); + }; + appEvents.on(AppEvent.ExternalMessage, handler); + return () => { + appEvents.off(AppEvent.ExternalMessage, handler); + }; + }, [addMessage]); + + // Wire WorkScheduler: inject fired prompts and persist schedule changes + useEffect(() => { + const scheduler = config.getWorkScheduler(); + + const onFire = (prompt: string) => { + appEvents.emit(AppEvent.ExternalMessage, prompt); + }; + + const onChanged = () => { + // Persist pending items to the session file + const recordingService = config + .getGeminiClient() + ?.getChatRecordingService(); + if (recordingService) { + recordingService.recordScheduledWork(scheduler.serialize()); + } + }; + + scheduler.on('fire', onFire); + scheduler.on('changed', onChanged); + return () => { + scheduler.off('fire', onFire); + scheduler.off('changed', onChanged); + }; + }, [config]); + + // Track streaming state transitions for A2A response capture + const prevStreamingStateRef = useRef(streamingState); + + useEffect(() => { + const prev = prevStreamingStateRef.current; + prevStreamingStateRef.current = streamingState; + + // Mark tasks as "working" when streaming starts + if ( + prev === StreamingState.Idle && + streamingState !== StreamingState.Idle + ) { + markTasksWorking(); + } + + // Capture response when streaming ends (for A2A tasks or unsolicited output) + if ( + prev !== StreamingState.Idle && + streamingState === StreamingState.Idle + ) { + const history = historyManager.history; + const parts: string[] = []; + for (let i = history.length - 1; i >= 0; i--) { + const item = history[i]; + if (item.type !== 'gemini' && item.type !== 'gemini_content') break; + if (typeof item.text === 'string' && item.text) { + parts.unshift(item.text); + } + } + notifyResponse(parts.join('\n')); + } + }, [streamingState, historyManager.history]); + cancelHandlerRef.current = useCallback( (shouldRestorePrompt: boolean = true) => { const pendingHistoryItems = [ @@ -2300,6 +2385,7 @@ Logging in with Google... Restarting Gemini CLI to continue. adminSettingsChanged, newAgents, showIsExpandableHint, + a2aListenerPort, hintMode: config.isModelSteeringEnabled() && isToolExecuting([ @@ -2428,6 +2514,7 @@ Logging in with Google... Restarting Gemini CLI to continue. adminSettingsChanged, newAgents, showIsExpandableHint, + a2aListenerPort, ], ); diff --git a/packages/cli/src/ui/components/Composer.test.tsx b/packages/cli/src/ui/components/Composer.test.tsx index e0919947fb..d58700150d 100644 --- a/packages/cli/src/ui/components/Composer.test.tsx +++ b/packages/cli/src/ui/components/Composer.test.tsx @@ -61,6 +61,10 @@ vi.mock('./StatusDisplay.js', () => ({ StatusDisplay: () => StatusDisplay, })); +vi.mock('./ScheduledWorkDisplay.js', () => ({ + ScheduledWorkDisplay: () => null, +})); + vi.mock('./ToastDisplay.js', () => ({ ToastDisplay: () => ToastDisplay, shouldShowToast: (uiState: UIState) => @@ -202,6 +206,7 @@ const createMockUIState = (overrides: Partial = {}): UIState => activeHooks: [], isBackgroundShellVisible: false, embeddedShellFocused: false, + a2aListenerPort: null, quota: { userTier: undefined, stats: undefined, diff --git a/packages/cli/src/ui/components/Composer.tsx b/packages/cli/src/ui/components/Composer.tsx index 0864b8f02b..37fa66d51a 100644 --- a/packages/cli/src/ui/components/Composer.tsx +++ b/packages/cli/src/ui/components/Composer.tsx @@ -13,6 +13,7 @@ import { } from '@google/gemini-cli-core'; import { LoadingIndicator } from './LoadingIndicator.js'; import { StatusDisplay } from './StatusDisplay.js'; +import { ScheduledWorkDisplay } from './ScheduledWorkDisplay.js'; import { ToastDisplay, shouldShowToast } from './ToastDisplay.js'; import { ApprovalModeIndicator } from './ApprovalModeIndicator.js'; import { ShellModeIndicator } from './ShellModeIndicator.js'; @@ -336,6 +337,7 @@ export const Composer = ({ isFocused = true }: { isFocused?: boolean }) => { )} {showShortcutsHelp && } {showUiDetails && } + {showUiDetails && ( { + const config = useConfig(); + const [items, setItems] = useState([]); + + useEffect(() => { + const scheduler = config.getWorkScheduler(); + + const update = () => { + setItems(scheduler.getPendingItems()); + }; + + scheduler.on('changed', update); + update(); + + return () => { + scheduler.off('changed', update); + }; + }, [config]); + + if (items.length === 0) { + return null; + } + + return ( + + + ⏰ Scheduled work ({items.length}): + + {items.map((item) => { + const timeStr = item.fireAt.toLocaleTimeString([], { + hour: '2-digit', + minute: '2-digit', + }); + const diffMs = item.fireAt.getTime() - Date.now(); + const diffMins = Math.max(0, Math.ceil(diffMs / 60000)); + const truncatedPrompt = + item.prompt.length > 60 + ? item.prompt.slice(0, 57) + '...' + : item.prompt; + return ( + + {' '} + {timeStr} (in {diffMins}m) — {truncatedPrompt} + + ); + })} + + ); +}; diff --git a/packages/cli/src/ui/components/StatusDisplay.test.tsx b/packages/cli/src/ui/components/StatusDisplay.test.tsx index fcb66ea0b2..2a063df688 100644 --- a/packages/cli/src/ui/components/StatusDisplay.test.tsx +++ b/packages/cli/src/ui/components/StatusDisplay.test.tsx @@ -52,6 +52,7 @@ const createMockUIState = (overrides: UIStateOverrides = {}): UIState => geminiMdFileCount: 0, contextFileNames: [], backgroundShellCount: 0, + a2aListenerPort: null, buffer: { text: '' }, history: [{ id: 1, type: 'user', text: 'test' }], ...overrides, @@ -171,4 +172,16 @@ describe('StatusDisplay', () => { expect(lastFrame()).toContain('Shells: 3'); unmount(); }); + + it('renders A2A listener port when a2aListenerPort is set', async () => { + const uiState = createMockUIState({ + a2aListenerPort: 8080, + }); + const { lastFrame, unmount } = await renderStatusDisplay( + { hideContextSummary: false }, + uiState, + ); + expect(lastFrame()).toContain('A2A :8080'); + unmount(); + }); }); diff --git a/packages/cli/src/ui/components/StatusDisplay.tsx b/packages/cli/src/ui/components/StatusDisplay.tsx index 223340c039..7a0baa524a 100644 --- a/packages/cli/src/ui/components/StatusDisplay.tsx +++ b/packages/cli/src/ui/components/StatusDisplay.tsx @@ -28,6 +28,12 @@ export const StatusDisplay: React.FC = ({ return |⌐■_■|; } + if (uiState.a2aListenerPort !== null) { + return ( + ⚡ A2A :{uiState.a2aListenerPort} + ); + } + if ( uiState.activeHooks.length > 0 && settings.merged.hooksConfig.notifications diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index ea9025aa6b..9b9e138aae 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -223,6 +223,7 @@ export interface UIState { showIsExpandableHint: boolean; hintMode: boolean; hintBuffer: string; + a2aListenerPort: number | null; transientMessage: { text: string; type: TransientMessageType; diff --git a/packages/cli/src/ui/hooks/creditsFlowHandler.test.ts b/packages/cli/src/ui/hooks/creditsFlowHandler.test.ts index 37a6294010..60effa0c4a 100644 --- a/packages/cli/src/ui/hooks/creditsFlowHandler.test.ts +++ b/packages/cli/src/ui/hooks/creditsFlowHandler.test.ts @@ -57,6 +57,7 @@ describe('handleCreditsFlow', () => { updateItem: vi.fn(), clearItems: vi.fn(), loadHistory: vi.fn(), + pruneItems: vi.fn(), }; isDialogPending = { current: false }; mockSetOverageMenuRequest = vi.fn(); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 2034e14b87..e15fa087a4 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -220,6 +220,7 @@ export const useGeminiStream = ( terminalHeight: number, isShellFocused?: boolean, consumeUserHint?: () => string | null, + pruneItems?: () => void, ) => { const [initError, setInitError] = useState(null); const [retryStatus, setRetryStatus] = useState( @@ -256,6 +257,7 @@ export const useGeminiStream = ( const [_isFirstToolInGroup, isFirstToolInGroupRef, setIsFirstToolInGroup] = useStateAndRef(true); const processedMemoryToolsRef = useRef>(new Set()); + const { startNewPrompt, getPromptCount } = useSessionStats(); const storage = config.storage; const logger = useLogger(storage); @@ -1163,8 +1165,12 @@ export const useGeminiStream = ( } as HistoryItemInfo, userMessageTimestamp, ); + + // Prune old UI history items to prevent unbounded memory growth + // in long-running sessions. + pruneItems?.(); }, - [addItem, pendingHistoryItemRef, setPendingHistoryItem, config], + [addItem, pendingHistoryItemRef, setPendingHistoryItem, config, pruneItems], ); const handleMaxSessionTurnsEvent = useCallback( @@ -1506,9 +1512,20 @@ export const useGeminiStream = ( lastQueryRef.current = queryToSend; lastPromptIdRef.current = prompt_id!; + // Prepend current time (and schedule if items exist) so the + // model can reason about time and scheduling in forever mode. + let queryWithContext = queryToSend; + if ( + config.getIsForeverMode() && + typeof queryWithContext === 'string' + ) { + const scheduler = config.getWorkScheduler(); + queryWithContext = `[${scheduler.formatScheduleSummary()}]\n\n${queryWithContext}`; + } + try { const stream = geminiClient.sendMessageStream( - queryToSend, + queryWithContext, abortSignal, prompt_id!, undefined, @@ -1946,6 +1963,66 @@ export const useGeminiStream = ( storage, ]); + // Idle hook timer: fires after idleTimeout seconds of no activity. + // The timeout is read from the Idle hook definitions themselves. + // If hooks exist but don't declare a timeout, fall back to 300 seconds. + const DEFAULT_IDLE_TIMEOUT = 300; + const idleTimerRef = useRef | null>(null); + useEffect(() => { + // Clear any existing timer + if (idleTimerRef.current) { + clearTimeout(idleTimerRef.current); + idleTimerRef.current = null; + } + + if (streamingState !== StreamingState.Idle || !config.getEnableHooks()) { + return; + } + + // Derive timeout from registered Idle hook definitions. + const hookSystem = config.getHookSystem(); + const idleHooks = hookSystem + ?.getAllHooks() + .filter((h) => h.eventName === 'Idle' && h.enabled); + + if (!idleHooks || idleHooks.length === 0) { + return; + } + + // Use the max idleTimeout declared by any Idle hook, or the default. + const declaredTimeouts = idleHooks + .map((h) => h.idleTimeout) + .filter((t): t is number => typeof t === 'number' && t > 0); + const idleTimeoutSeconds = + declaredTimeouts.length > 0 + ? Math.max(...declaredTimeouts) + : DEFAULT_IDLE_TIMEOUT; + + const startTime = Date.now(); + idleTimerRef.current = setTimeout(async () => { + if (!hookSystem) return; + + const elapsed = Math.round((Date.now() - startTime) / 1000); + try { + const result = await hookSystem.fireIdleEvent(elapsed); + const prompt = result?.finalOutput?.hookSpecificOutput?.['prompt']; + if (typeof prompt === 'string' && prompt.trim()) { + // Auto-submit the prompt returned by the hook + void submitQuery(prompt); + } + } catch { + // Idle hook failures are non-fatal + } + }, idleTimeoutSeconds * 1000); + + return () => { + if (idleTimerRef.current) { + clearTimeout(idleTimerRef.current); + idleTimerRef.current = null; + } + }; + }, [streamingState, config, submitQuery]); + const lastOutputTime = Math.max( lastToolOutputTime, lastShellOutputTime, diff --git a/packages/cli/src/ui/hooks/useHistoryManager.test.ts b/packages/cli/src/ui/hooks/useHistoryManager.test.ts index 696f9d60c0..37e134f308 100644 --- a/packages/cli/src/ui/hooks/useHistoryManager.test.ts +++ b/packages/cli/src/ui/hooks/useHistoryManager.test.ts @@ -7,7 +7,7 @@ import { describe, it, expect } from 'vitest'; import { act } from 'react'; import { renderHook } from '../../test-utils/render.js'; -import { useHistory } from './useHistoryManager.js'; +import { useHistory, PRUNE_KEEP_COUNT } from './useHistoryManager.js'; import type { HistoryItem } from '../types.js'; describe('useHistoryManager', () => { @@ -255,4 +255,82 @@ describe('useHistoryManager', () => { expect(result.current.history[0].type).toBe('info'); }); }); + + describe('pruneItems', () => { + it('should prune history to PRUNE_KEEP_COUNT + 1 (marker) when over limit', () => { + const { result } = renderHook(() => useHistory()); + const itemCount = PRUNE_KEEP_COUNT + 20; + + act(() => { + for (let i = 0; i < itemCount; i++) { + result.current.addItem({ + type: 'user', + text: `Message ${i}`, + }); + } + }); + + expect(result.current.history).toHaveLength(itemCount); + + act(() => { + result.current.pruneItems(); + }); + + // PRUNE_KEEP_COUNT items + 1 prune marker + expect(result.current.history).toHaveLength(PRUNE_KEEP_COUNT + 1); + // First item should be the prune marker + expect(result.current.history[0].type).toBe('info'); + expect(result.current.history[0].text).toContain('pruned'); + // Last item should be the most recent message + expect( + result.current.history[result.current.history.length - 1].text, + ).toBe(`Message ${itemCount - 1}`); + }); + + it('should be a no-op when history is under the threshold', () => { + const { result } = renderHook(() => useHistory()); + const itemCount = 10; + + act(() => { + for (let i = 0; i < itemCount; i++) { + result.current.addItem({ + type: 'user', + text: `Message ${i}`, + }); + } + }); + + const historyBefore = result.current.history; + + act(() => { + result.current.pruneItems(); + }); + + // Should be unchanged (same reference) + expect(result.current.history).toBe(historyBefore); + expect(result.current.history).toHaveLength(itemCount); + }); + + it('should be a no-op when history is exactly at the threshold', () => { + const { result } = renderHook(() => useHistory()); + + act(() => { + for (let i = 0; i < PRUNE_KEEP_COUNT; i++) { + result.current.addItem({ + type: 'user', + text: `Message ${i}`, + }); + } + }); + + const historyBefore = result.current.history; + + act(() => { + result.current.pruneItems(); + }); + + expect(result.current.history).toBe(historyBefore); + expect(result.current.history).toHaveLength(PRUNE_KEEP_COUNT); + }); + }); }); diff --git a/packages/cli/src/ui/hooks/useHistoryManager.ts b/packages/cli/src/ui/hooks/useHistoryManager.ts index 93f7f01f28..6931fb535c 100644 --- a/packages/cli/src/ui/hooks/useHistoryManager.ts +++ b/packages/cli/src/ui/hooks/useHistoryManager.ts @@ -8,6 +8,12 @@ import { useState, useRef, useCallback, useMemo } from 'react'; import type { HistoryItem } from '../types.js'; import type { ChatRecordingService } from '@google/gemini-cli-core/src/services/chatRecordingService.js'; +/** + * Number of history items to keep when pruning after context compression. + * Exported for testing purposes. + */ +export const PRUNE_KEEP_COUNT = 50; + // Type for the updater function passed to updateHistoryItem type HistoryItemUpdater = ( prevItem: HistoryItem, @@ -26,6 +32,7 @@ export interface UseHistoryManagerReturn { ) => void; clearItems: () => void; loadHistory: (newHistory: HistoryItem[]) => void; + pruneItems: () => void; } /** @@ -156,6 +163,25 @@ export function useHistory({ messageIdCounterRef.current = 0; }, []); + // Prunes old history items, keeping only the most recent PRUNE_KEEP_COUNT. + // Intended to be called after context compression to free memory in + // long-running sessions. + const pruneItems = useCallback(() => { + setHistory((prevHistory) => { + if (prevHistory.length <= PRUNE_KEEP_COUNT) { + return prevHistory; + } + const kept = prevHistory.slice(-PRUNE_KEEP_COUNT); + + const marker = { + id: getNextMessageId(Date.now()), + type: 'info', + text: `ℹ️ Earlier history was pruned after context compression.`, + } as HistoryItem; + return [marker, ...kept]; + }); + }, [getNextMessageId]); + return useMemo( () => ({ history, @@ -163,7 +189,8 @@ export function useHistory({ updateItem, clearItems, loadHistory, + pruneItems, }), - [history, addItem, updateItem, clearItems, loadHistory], + [history, addItem, updateItem, clearItems, loadHistory, pruneItems], ); } diff --git a/packages/cli/src/ui/hooks/useIncludeDirsTrust.test.tsx b/packages/cli/src/ui/hooks/useIncludeDirsTrust.test.tsx index 3f9c656048..5980084fbb 100644 --- a/packages/cli/src/ui/hooks/useIncludeDirsTrust.test.tsx +++ b/packages/cli/src/ui/hooks/useIncludeDirsTrust.test.tsx @@ -91,6 +91,7 @@ describe('useIncludeDirsTrust', () => { updateItem: vi.fn(), clearItems: vi.fn(), loadHistory: vi.fn(), + pruneItems: vi.fn(), }; mockSetCustomDialog = vi.fn(); }); diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts index ea4234bd10..269134adac 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts @@ -85,6 +85,7 @@ describe('useQuotaAndFallback', () => { updateItem: vi.fn(), clearItems: vi.fn(), loadHistory: vi.fn(), + pruneItems: vi.fn(), }; mockSetModelSwitchedFromQuotaError = vi.fn(); mockOnShowAuthSelection = vi.fn(); diff --git a/packages/cli/src/ui/hooks/useSessionBrowser.ts b/packages/cli/src/ui/hooks/useSessionBrowser.ts index 9a34f68e0b..ee97003f46 100644 --- a/packages/cli/src/ui/hooks/useSessionBrowser.ts +++ b/packages/cli/src/ui/hooks/useSessionBrowser.ts @@ -78,12 +78,17 @@ export const useSessionBrowser = ( // We've loaded it; tell the UI about it. setIsSessionBrowserOpen(false); + const compressionIndex = conversation.lastCompressionIndex; const historyData = convertSessionToHistoryFormats( conversation.messages, + compressionIndex, ); await onLoadHistory( historyData.uiHistory, - convertSessionToClientHistory(conversation.messages), + convertSessionToClientHistory( + conversation.messages, + compressionIndex, + ), resumedSessionData, ); } catch (error) { diff --git a/packages/cli/src/ui/hooks/useSessionResume.test.ts b/packages/cli/src/ui/hooks/useSessionResume.test.ts index 9350cc167a..c378de13a0 100644 --- a/packages/cli/src/ui/hooks/useSessionResume.test.ts +++ b/packages/cli/src/ui/hooks/useSessionResume.test.ts @@ -33,6 +33,7 @@ describe('useSessionResume', () => { updateItem: vi.fn(), clearItems: vi.fn(), loadHistory: vi.fn(), + pruneItems: vi.fn(), }); let mockHistoryManager: UseHistoryManagerReturn; @@ -529,5 +530,60 @@ describe('useSessionResume', () => { // But UI history should have both expect(mockHistoryManager.addItem).toHaveBeenCalledTimes(2); }); + + it('should restore scheduled work from resumed session data', async () => { + const mockRestore = vi.fn(); + const configWithScheduler = { + ...mockConfig, + getWorkScheduler: vi.fn().mockReturnValue({ + restore: mockRestore, + }), + }; + + const scheduledWork = [ + { + id: 'test-1', + prompt: 'check status', + fireAt: new Date(Date.now() + 60_000).toISOString(), + createdAt: new Date().toISOString(), + }, + ]; + + const conversation: ConversationRecord = { + sessionId: 'auto-resume-scheduled', + projectHash: 'project-123', + startTime: '2025-01-01T00:00:00Z', + lastUpdated: '2025-01-01T01:00:00Z', + messages: [ + { + id: 'msg-1', + timestamp: '2025-01-01T00:01:00Z', + content: 'Hello', + type: 'user', + }, + ] as MessageRecord[], + scheduledWork, + }; + + await act(async () => { + renderHook(() => + useSessionResume({ + ...getDefaultProps(), + config: configWithScheduler as unknown as Config, + resumedSessionData: { + conversation, + filePath: '/path/to/session.json', + }, + }), + ); + }); + + await waitFor(() => { + expect(mockHistoryManager.clearItems).toHaveBeenCalled(); + }); + + expect(configWithScheduler.getWorkScheduler).toHaveBeenCalled(); + expect(mockRestore).toHaveBeenCalledWith(scheduledWork); + }); }); }); diff --git a/packages/cli/src/ui/hooks/useSessionResume.ts b/packages/cli/src/ui/hooks/useSessionResume.ts index 055686773b..c927c64f54 100644 --- a/packages/cli/src/ui/hooks/useSessionResume.ts +++ b/packages/cli/src/ui/hooks/useSessionResume.ts @@ -83,6 +83,16 @@ export function useSessionResume({ workspaceContext.addDirectories(resumedData.conversation.directories); } + // Restore scheduled work items from the resumed session. + // Past-due items fire immediately; future items get timers re-armed. + if ( + resumedData.conversation.scheduledWork && + resumedData.conversation.scheduledWork.length > 0 + ) { + const scheduler = config.getWorkScheduler(); + scheduler.restore(resumedData.conversation.scheduledWork); + } + // Give the history to the Gemini client. await config.getGeminiClient()?.resumeChat(clientHistory, resumedData); } catch (error) { @@ -109,12 +119,18 @@ export function useSessionResume({ !hasLoadedResumedSession.current ) { hasLoadedResumedSession.current = true; + const compressionIndex = + resumedSessionData.conversation.lastCompressionIndex; const historyData = convertSessionToHistoryFormats( resumedSessionData.conversation.messages, + compressionIndex, ); void loadHistoryForResume( historyData.uiHistory, - convertSessionToClientHistory(resumedSessionData.conversation.messages), + convertSessionToClientHistory( + resumedSessionData.conversation.messages, + compressionIndex, + ), resumedSessionData, ); } diff --git a/packages/cli/src/utils/events.ts b/packages/cli/src/utils/events.ts index 8291528ac1..42340ee5c1 100644 --- a/packages/cli/src/utils/events.ts +++ b/packages/cli/src/utils/events.ts @@ -23,6 +23,8 @@ export enum AppEvent { PasteTimeout = 'paste-timeout', TerminalBackground = 'terminal-background', TransientMessage = 'transient-message', + ExternalMessage = 'external-message', + A2AListenerStarted = 'a2a-listener-started', } export interface AppEvents { @@ -32,6 +34,8 @@ export interface AppEvents { [AppEvent.PasteTimeout]: never[]; [AppEvent.TerminalBackground]: [string]; [AppEvent.TransientMessage]: [TransientMessagePayload]; + [AppEvent.ExternalMessage]: [string]; + [AppEvent.A2AListenerStarted]: [number]; } export const appEvents = new EventEmitter(); diff --git a/packages/cli/src/utils/sessionUtils.ts b/packages/cli/src/utils/sessionUtils.ts index ca6685f47d..0312d973b6 100644 --- a/packages/cli/src/utils/sessionUtils.ts +++ b/packages/cli/src/utils/sessionUtils.ts @@ -539,12 +539,24 @@ export class SessionSelector { */ export function convertSessionToHistoryFormats( messages: ConversationRecord['messages'], + startIndex?: number, ): { uiHistory: HistoryItemWithoutId[]; } { const uiHistory: HistoryItemWithoutId[] = []; - for (const msg of messages) { + const hasCompressedHistory = + startIndex != null && startIndex > 0 && startIndex < messages.length; + const slice = hasCompressedHistory ? messages.slice(startIndex) : messages; + + if (hasCompressedHistory) { + uiHistory.push({ + type: MessageType.INFO, + text: `ℹ️ Earlier history (${startIndex} messages) was compressed. Showing post-compression messages only.`, + }); + } + + for (const msg of slice) { // Add thoughts if present if (msg.type === 'gemini' && msg.thoughts && msg.thoughts.length > 0) { for (const thought of msg.thoughts) { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index aa3e9aa5b6..c2c8e42037 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -33,6 +33,8 @@ import { WebFetchTool } from '../tools/web-fetch.js'; import { MemoryTool, setGeminiMdFilename } from '../tools/memoryTool.js'; import { WebSearchTool } from '../tools/web-search.js'; import { AskUserTool } from '../tools/ask-user.js'; +import { ScheduleWorkTool } from '../tools/schedule-work.js'; +import { WorkScheduler } from '../services/work-scheduler.js'; import { ExitPlanModeTool } from '../tools/exit-plan-mode.js'; import { EnterPlanModeTool } from '../tools/enter-plan-mode.js'; import { GeminiClient } from '../core/client.js'; @@ -640,6 +642,7 @@ export interface ConfigParameters { mcpEnabled?: boolean; extensionsEnabled?: boolean; agents?: AgentSettings; + isForeverMode?: boolean; onReload?: () => Promise<{ disabledSkills?: string[]; adminSkillsEnabled?: boolean; @@ -847,6 +850,8 @@ export class Config implements McpContext, AgentLoopContext { private readonly enableAgents: boolean; private agents: AgentSettings; + private readonly isForeverMode: boolean; + private readonly workScheduler: WorkScheduler; private readonly enableEventDrivenScheduler: boolean; private readonly skillsSupport: boolean; private disabledSkills: string[]; @@ -959,6 +964,8 @@ export class Config implements McpContext, AgentLoopContext { this._activeModel = params.model; this.enableAgents = params.enableAgents ?? true; this.agents = params.agents ?? {}; + this.isForeverMode = params.isForeverMode ?? false; + this.workScheduler = new WorkScheduler(); this.disableLLMCorrection = params.disableLLMCorrection ?? true; this.planEnabled = params.plan ?? true; this.trackerEnabled = params.tracker ?? false; @@ -2712,6 +2719,14 @@ export class Config implements McpContext, AgentLoopContext { return remoteThreshold; } + getIsForeverMode(): boolean { + return this.isForeverMode; + } + + getWorkScheduler(): WorkScheduler { + return this.workScheduler; + } + async getUserCaching(): Promise { await this.ensureExperimentsLoaded(); @@ -2863,6 +2878,7 @@ export class Config implements McpContext, AgentLoopContext { } isInteractiveShellEnabled(): boolean { + if (this.isForeverMode) return false; return ( this.interactive && this.ptyInfo !== 'child_process' && @@ -3184,15 +3200,24 @@ export class Config implements McpContext, AgentLoopContext { maybeRegister(ShellTool, () => registry.registerTool(new ShellTool(this, this.messageBus)), ); - maybeRegister(MemoryTool, () => - registry.registerTool(new MemoryTool(this.messageBus)), - ); + if (!this.isForeverMode) { + maybeRegister(MemoryTool, () => + registry.registerTool(new MemoryTool(this.messageBus)), + ); + } maybeRegister(WebSearchTool, () => registry.registerTool(new WebSearchTool(this, this.messageBus)), ); maybeRegister(AskUserTool, () => registry.registerTool(new AskUserTool(this.messageBus)), ); + if (this.isForeverMode) { + maybeRegister(ScheduleWorkTool, () => + registry.registerTool( + new ScheduleWorkTool(this.messageBus, this.workScheduler), + ), + ); + } if (this.getUseWriteTodos()) { maybeRegister(WriteTodosTool, () => registry.registerTool(new WriteTodosTool(this.messageBus)), @@ -3202,9 +3227,11 @@ export class Config implements McpContext, AgentLoopContext { maybeRegister(ExitPlanModeTool, () => registry.registerTool(new ExitPlanModeTool(this, this.messageBus)), ); - maybeRegister(EnterPlanModeTool, () => - registry.registerTool(new EnterPlanModeTool(this, this.messageBus)), - ); + if (!this.isForeverMode) { + maybeRegister(EnterPlanModeTool, () => + registry.registerTool(new EnterPlanModeTool(this, this.messageBus)), + ); + } } if (this.isTrackerEnabled()) { diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 77c4a5a498..3b42ec4c4d 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -143,6 +143,7 @@ const mockHookSystem = { fireBeforeAgentEvent: vi.fn().mockResolvedValue(undefined), fireAfterAgentEvent: vi.fn().mockResolvedValue(undefined), firePreCompressEvent: vi.fn().mockResolvedValue(undefined), + fireIdleEvent: vi.fn().mockResolvedValue(undefined), }; /** @@ -452,6 +453,7 @@ describe('Gemini Client (client.ts)', () => { getChatRecordingService: vi.fn().mockReturnValue({ getConversation: vi.fn().mockReturnValue(null), getConversationFilePath: vi.fn().mockReturnValue(null), + recordCompressionPoint: vi.fn(), }), }; client['chat'] = mockOriginalChat as GeminiChat; @@ -686,6 +688,7 @@ describe('Gemini Client (client.ts)', () => { const mockRecordingService = { getConversation: vi.fn().mockReturnValue(mockConversation), getConversationFilePath: vi.fn().mockReturnValue(mockFilePath), + recordCompressionPoint: vi.fn(), }; vi.mocked(mockOriginalChat.getChatRecordingService!).mockReturnValue( mockRecordingService as unknown as ChatRecordingService, diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 01577452f4..e4080c4167 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -879,6 +879,12 @@ export class GeminiClient { this.hookStateMap.delete(this.lastPromptId); this.lastPromptId = prompt_id; this.currentSequenceModel = null; + + // In Forever Mode, refresh the system instruction so memory changes + // (e.g. GEMINI.md updates from hooks/agents) are picked up immediately. + if (this.config.getIsForeverMode()) { + this.updateSystemInstruction(); + } } if (hooksEnabled && messageBus) { @@ -1178,6 +1184,7 @@ export class GeminiClient { // capture current session data before resetting const currentRecordingService = this.getChat().getChatRecordingService(); + currentRecordingService.recordCompressionPoint(); const conversation = currentRecordingService.getConversation(); const filePath = currentRecordingService.getConversationFilePath(); diff --git a/packages/core/src/core/prompts.test.ts b/packages/core/src/core/prompts.test.ts index 82a7943de4..a8908c9b6c 100644 --- a/packages/core/src/core/prompts.test.ts +++ b/packages/core/src/core/prompts.test.ts @@ -116,6 +116,7 @@ describe('Core System Prompt (prompts.ts)', () => { getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT), getApprovedPlanPath: vi.fn().mockReturnValue(undefined), isTrackerEnabled: vi.fn().mockReturnValue(false), + getIsForeverMode: vi.fn().mockReturnValue(false), get config() { return this; }, @@ -436,6 +437,7 @@ describe('Core System Prompt (prompts.ts)', () => { }), getApprovedPlanPath: vi.fn().mockReturnValue(undefined), isTrackerEnabled: vi.fn().mockReturnValue(false), + getIsForeverMode: vi.fn().mockReturnValue(false), get config() { return this; }, diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 9c0e536c48..6b26699dc2 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -182,6 +182,9 @@ export enum CompressionStatus { /** The compression was skipped due to previous failure, but content was truncated to budget */ CONTENT_TRUNCATED, + + /** The compression was replaced by a PreCompress hook */ + HOOK_REPLACED, } export interface ChatCompressionInfo { diff --git a/packages/core/src/hooks/hookAggregator.test.ts b/packages/core/src/hooks/hookAggregator.test.ts index ee9ade9a87..4903d0351d 100644 --- a/packages/core/src/hooks/hookAggregator.test.ts +++ b/packages/core/src/hooks/hookAggregator.test.ts @@ -472,5 +472,33 @@ describe('HookAggregator', () => { aggregated.finalOutput?.hookSpecificOutput?.['additionalContext'], ).toBe('Context from hook 1\nContext from hook 2'); }); + + it('should propagate refreshContext with any-true-wins logic', () => { + const results: HookExecutionResult[] = [ + createHookExecutionResult({ refreshContext: false }), + createHookExecutionResult({ refreshContext: true }), + ]; + + const aggregated = aggregator.aggregateResults( + results, + HookEventName.AfterAgent, + ); + + expect(aggregated.finalOutput?.refreshContext).toBe(true); + }); + + it('should not set refreshContext when no hook requests it', () => { + const results: HookExecutionResult[] = [ + createHookExecutionResult({}), + createHookExecutionResult({}), + ]; + + const aggregated = aggregator.aggregateResults( + results, + HookEventName.AfterAgent, + ); + + expect(aggregated.finalOutput?.refreshContext).toBeUndefined(); + }); }); }); diff --git a/packages/core/src/hooks/hookAggregator.ts b/packages/core/src/hooks/hookAggregator.ts index 73e814702e..7d503129e1 100644 --- a/packages/core/src/hooks/hookAggregator.ts +++ b/packages/core/src/hooks/hookAggregator.ts @@ -158,6 +158,11 @@ export class HookAggregator { merged.suppressOutput = true; } + // Handle refreshContext (any true wins) + if (output.refreshContext) { + merged.refreshContext = true; + } + // Handle clearContext (any true wins) - for AfterAgent hooks if (output.hookSpecificOutput?.['clearContext'] === true) { merged.hookSpecificOutput = { @@ -355,7 +360,7 @@ export class HookAggregator { // Extract additionalContext from various hook types if ( 'additionalContext' in specific && - // eslint-disable-next-line no-restricted-syntax + typeof specific['additionalContext'] === 'string' ) { contexts.push(specific['additionalContext']); diff --git a/packages/core/src/hooks/hookEventHandler.test.ts b/packages/core/src/hooks/hookEventHandler.test.ts index 9e93850101..edae449fe5 100644 --- a/packages/core/src/hooks/hookEventHandler.test.ts +++ b/packages/core/src/hooks/hookEventHandler.test.ts @@ -16,6 +16,7 @@ import { SessionStartSource, HookEventName, HookType, + DefaultHookOutput, type HookConfig, type HookExecutionResult, } from './types.js'; @@ -891,4 +892,100 @@ describe('HookEventHandler', () => { ); }); }); + + describe('refreshContext handling', () => { + it('should call updateSystemInstructionIfInitialized when refreshContext is true', async () => { + const hookConfig = { + type: HookType.Command, + command: './after-agent.sh', + } as HookConfig; + const mockResults: HookExecutionResult[] = [ + { + success: true, + duration: 100, + hookConfig, + eventName: HookEventName.AfterAgent, + output: { refreshContext: true }, + }, + ]; + const mockAggregated = { + success: true, + finalOutput: new DefaultHookOutput({ refreshContext: true }), + allOutputs: mockResults.map((r) => r.output!), + errors: [], + totalDuration: 100, + }; + + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue({ + hookConfigs: [hookConfig], + sequential: false, + eventName: HookEventName.AfterAgent, + }); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue( + mockResults, + ); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + mockAggregated, + ); + + const updateSpy = vi.fn(); + (mockConfig as unknown as Record)[ + 'updateSystemInstructionIfInitialized' + ] = updateSpy; + + await hookEventHandler.fireAfterAgentEvent( + 'test prompt', + 'test response', + ); + + expect(updateSpy).toHaveBeenCalled(); + }); + + it('should not call updateSystemInstructionIfInitialized when refreshContext is not set', async () => { + const hookConfig = { + type: HookType.Command, + command: './after-agent.sh', + } as HookConfig; + const mockResults: HookExecutionResult[] = [ + { + success: true, + duration: 100, + hookConfig, + eventName: HookEventName.AfterAgent, + output: {}, + }, + ]; + const mockAggregated = { + success: true, + finalOutput: new DefaultHookOutput({}), + allOutputs: mockResults.map((r) => r.output!), + errors: [], + totalDuration: 100, + }; + + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue({ + hookConfigs: [hookConfig], + sequential: false, + eventName: HookEventName.AfterAgent, + }); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue( + mockResults, + ); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + mockAggregated, + ); + + const updateSpy = vi.fn(); + (mockConfig as unknown as Record)[ + 'updateSystemInstructionIfInitialized' + ] = updateSpy; + + await hookEventHandler.fireAfterAgentEvent( + 'test prompt', + 'test response', + ); + + expect(updateSpy).not.toHaveBeenCalled(); + }); + }); }); diff --git a/packages/core/src/hooks/hookEventHandler.ts b/packages/core/src/hooks/hookEventHandler.ts index a092bed334..bbf1aedcd1 100644 --- a/packages/core/src/hooks/hookEventHandler.ts +++ b/packages/core/src/hooks/hookEventHandler.ts @@ -29,6 +29,7 @@ import { type PreCompressTrigger, type HookExecutionResult, type McpToolContext, + type IdleInput, } from './types.js'; import { defaultHookTranslator } from './hookTranslator.js'; import type { @@ -204,16 +205,30 @@ export class HookEventHandler { */ async firePreCompressEvent( trigger: PreCompressTrigger, + history: Array<{ role: string; parts: Array<{ text?: string }> }>, ): Promise { const input: PreCompressInput = { ...this.createBaseInput(HookEventName.PreCompress), trigger, + history, }; const context: HookEventContext = { trigger }; return this.executeHooks(HookEventName.PreCompress, input, context); } + /** + * Fire an Idle event + */ + async fireIdleEvent(idleSeconds: number): Promise { + const input: IdleInput = { + ...this.createBaseInput(HookEventName.Idle), + idle_seconds: idleSeconds, + }; + + return this.executeHooks(HookEventName.Idle, input); + } + /** * Fire a BeforeModel event * Called by handleHookExecutionRequest - executes hooks directly @@ -493,6 +508,11 @@ export class HookEventHandler { // This is just logging the request centrally } + // Handle refreshContext - reload the system instruction to pick up context file changes + if (aggregated.finalOutput.refreshContext) { + this.context.config.updateSystemInstructionIfInitialized(); + } + // Other common fields like decision/reason are handled by specific hook output classes } diff --git a/packages/core/src/hooks/hookRegistry.ts b/packages/core/src/hooks/hookRegistry.ts index 1dad67bad5..1de6438c61 100644 --- a/packages/core/src/hooks/hookRegistry.ts +++ b/packages/core/src/hooks/hookRegistry.ts @@ -26,6 +26,7 @@ export interface HookRegistryEntry { matcher?: string; sequential?: boolean; enabled: boolean; + idleTimeout?: number; } /** @@ -279,6 +280,7 @@ please review the project settings (.gemini/settings.json) and remove them.`; matcher: definition.matcher, sequential: definition.sequential, enabled: !isDisabled, + idleTimeout: definition.idleTimeout, }); } else { // Invalid hooks are logged and discarded here, they won't reach HookRunner diff --git a/packages/core/src/hooks/hookSystem.ts b/packages/core/src/hooks/hookSystem.ts index f748665985..96f0c7522f 100644 --- a/packages/core/src/hooks/hookSystem.ts +++ b/packages/core/src/hooks/hookSystem.ts @@ -232,8 +232,15 @@ export class HookSystem { async firePreCompressEvent( trigger: PreCompressTrigger, + history: Array<{ role: string; parts: Array<{ text?: string }> }>, ): Promise { - return this.hookEventHandler.firePreCompressEvent(trigger); + return this.hookEventHandler.firePreCompressEvent(trigger, history); + } + + async fireIdleEvent( + idleSeconds: number, + ): Promise { + return this.hookEventHandler.fireIdleEvent(idleSeconds); } async fireBeforeAgentEvent( diff --git a/packages/core/src/hooks/types.test.ts b/packages/core/src/hooks/types.test.ts index ab809cbec7..0b95e505a1 100644 --- a/packages/core/src/hooks/types.test.ts +++ b/packages/core/src/hooks/types.test.ts @@ -57,6 +57,7 @@ describe('Hook Types', () => { 'BeforeModel', 'AfterModel', 'BeforeToolSelection', + 'Idle', ]; for (const event of expectedEvents) { @@ -112,6 +113,7 @@ describe('Hook Output Classes', () => { systemMessage: 'test system message', decision: 'block' as HookDecision, reason: 'test reason', + refreshContext: true, hookSpecificOutput: { key: 'value' }, }; const output = new DefaultHookOutput(data); @@ -121,6 +123,7 @@ describe('Hook Output Classes', () => { expect(output.systemMessage).toBe(data.systemMessage); expect(output.decision).toBe(data.decision); expect(output.reason).toBe(data.reason); + expect(output.refreshContext).toBe(true); expect(output.hookSpecificOutput).toEqual(data.hookSpecificOutput); }); diff --git a/packages/core/src/hooks/types.ts b/packages/core/src/hooks/types.ts index 9c6217ffa4..d70074a74e 100644 --- a/packages/core/src/hooks/types.ts +++ b/packages/core/src/hooks/types.ts @@ -43,6 +43,7 @@ export enum HookEventName { BeforeModel = 'BeforeModel', AfterModel = 'AfterModel', BeforeToolSelection = 'BeforeToolSelection', + Idle = 'Idle', } /** @@ -104,6 +105,8 @@ export interface HookDefinition { matcher?: string; sequential?: boolean; hooks: HookConfig[]; + /** Seconds before the Idle hook fires. Only meaningful for Idle hooks. */ + idleTimeout?: number; } /** @@ -147,6 +150,8 @@ export interface HookOutput { systemMessage?: string; decision?: HookDecision; reason?: string; + /** When true, refreshes the system instruction after hook execution to pick up context file changes (e.g. GEMINI.md). */ + refreshContext?: boolean; hookSpecificOutput?: Record; } @@ -184,6 +189,7 @@ export class DefaultHookOutput implements HookOutput { systemMessage?: string; decision?: HookDecision; reason?: string; + refreshContext?: boolean; hookSpecificOutput?: Record; constructor(data: Partial = {}) { @@ -193,6 +199,7 @@ export class DefaultHookOutput implements HookOutput { this.systemMessage = data.systemMessage; this.decision = data.decision; this.reason = data.reason; + this.refreshContext = data.refreshContext; this.hookSpecificOutput = data.hookSpecificOutput; } @@ -642,14 +649,37 @@ export enum PreCompressTrigger { */ export interface PreCompressInput extends HookInput { trigger: PreCompressTrigger; + history: Array<{ role: string; parts: Array<{ text?: string }> }>; } - /** * PreCompress hook output */ export interface PreCompressOutput { suppressOutput?: boolean; systemMessage?: string; + hookSpecificOutput?: { + hookEventName: 'PreCompress'; + newHistory?: Array<{ role: string; parts: Array<{ text?: string }> }>; + }; +} + +/** + * Idle hook input + */ +export interface IdleInput extends HookInput { + idle_seconds: number; +} + +/** + * Idle hook output + */ +export interface IdleOutput { + suppressOutput?: boolean; + systemMessage?: string; + hookSpecificOutput?: { + hookEventName: 'Idle'; + prompt?: string; + }; } /** diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 47412dd73c..d05e1bab19 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -125,6 +125,7 @@ export * from './services/fileDiscoveryService.js'; export * from './services/gitService.js'; export * from './services/FolderTrustDiscoveryService.js'; export * from './services/chatRecordingService.js'; +export * from './services/work-scheduler.js'; export * from './services/fileSystemService.js'; export * from './services/sessionSummaryUtils.js'; export * from './services/contextManager.js'; diff --git a/packages/core/src/prompts/promptProvider.test.ts b/packages/core/src/prompts/promptProvider.test.ts index c2253a9b57..32e9fc628d 100644 --- a/packages/core/src/prompts/promptProvider.test.ts +++ b/packages/core/src/prompts/promptProvider.test.ts @@ -71,6 +71,7 @@ describe('PromptProvider', () => { getApprovedPlanPath: vi.fn().mockReturnValue(undefined), getApprovalMode: vi.fn(), isTrackerEnabled: vi.fn().mockReturnValue(false), + getIsForeverMode: vi.fn().mockReturnValue(false), } as unknown as Config; }); diff --git a/packages/core/src/services/chatCompressionService.test.ts b/packages/core/src/services/chatCompressionService.test.ts index c4f26dedc0..9b40410c38 100644 --- a/packages/core/src/services/chatCompressionService.test.ts +++ b/packages/core/src/services/chatCompressionService.test.ts @@ -186,7 +186,7 @@ describe('ChatCompressionService', () => { }), getEnableHooks: vi.fn().mockReturnValue(false), getMessageBus: vi.fn().mockReturnValue(undefined), - getHookSystem: () => undefined, + getHookSystem: vi.fn().mockReturnValue(undefined), getNextCompressionTruncationId: vi.fn().mockReturnValue(1), getTruncateToolOutputThreshold: vi.fn().mockReturnValue(40000), storage: { @@ -897,4 +897,151 @@ describe('ChatCompressionService', () => { ); }); }); + + describe('PreCompress hook replacement', () => { + it('should use hook-provided newHistory and skip built-in compression', async () => { + const history: Content[] = [ + { role: 'user', parts: [{ text: 'msg1' }] }, + { role: 'model', parts: [{ text: 'msg2' }] }, + { role: 'user', parts: [{ text: 'msg3' }] }, + { role: 'model', parts: [{ text: 'msg4' }] }, + ]; + vi.mocked(mockChat.getHistory).mockReturnValue(history); + vi.mocked(mockChat.getLastPromptTokenCount).mockReturnValue(600000); + vi.mocked(tokenLimit).mockReturnValue(1_000_000); + + const hookReplacementHistory = [ + { + role: 'user', + parts: [{ text: 'Archive summary: topics discussed...' }], + }, + { + role: 'model', + parts: [{ text: 'Understood, continuing from archive.' }], + }, + ]; + + const mockHookSystem = { + firePreCompressEvent: vi.fn().mockResolvedValue({ + success: true, + finalOutput: { + hookSpecificOutput: { + hookEventName: 'PreCompress', + newHistory: hookReplacementHistory, + }, + }, + allOutputs: [], + errors: [], + totalDuration: 100, + }), + }; + + vi.mocked(mockConfig.getHookSystem).mockReturnValue( + mockHookSystem as unknown as ReturnType, + ); + + const result = await service.compress( + mockChat, + mockPromptId, + true, + mockModel, + mockConfig, + false, + ); + + expect(result.info.compressionStatus).toBe( + CompressionStatus.HOOK_REPLACED, + ); + expect(result.newHistory).not.toBeNull(); + expect(result.newHistory!.length).toBe(2); + expect(result.newHistory![0].parts![0].text).toBe( + 'Archive summary: topics discussed...', + ); + // Built-in LLM compression should NOT have been called + expect( + mockConfig.getBaseLlmClient().generateContent, + ).not.toHaveBeenCalled(); + }); + + it('should proceed with built-in compression when hook returns no newHistory', async () => { + const history: Content[] = [ + { role: 'user', parts: [{ text: 'msg1' }] }, + { role: 'model', parts: [{ text: 'msg2' }] }, + { role: 'user', parts: [{ text: 'msg3' }] }, + { role: 'model', parts: [{ text: 'msg4' }] }, + ]; + vi.mocked(mockChat.getHistory).mockReturnValue(history); + vi.mocked(mockChat.getLastPromptTokenCount).mockReturnValue(600000); + vi.mocked(tokenLimit).mockReturnValue(1_000_000); + + const mockHookSystem = { + firePreCompressEvent: vi.fn().mockResolvedValue({ + success: true, + finalOutput: { + systemMessage: 'Compression starting...', + }, + allOutputs: [], + errors: [], + totalDuration: 50, + }), + }; + + vi.mocked(mockConfig.getHookSystem).mockReturnValue( + mockHookSystem as unknown as ReturnType, + ); + + const result = await service.compress( + mockChat, + mockPromptId, + true, + mockModel, + mockConfig, + false, + ); + + // Should fall through to normal compression + expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); + expect(mockConfig.getBaseLlmClient().generateContent).toHaveBeenCalled(); + }); + + it('should pass history to the hook', async () => { + const history: Content[] = [ + { role: 'user', parts: [{ text: 'hello' }] }, + { role: 'model', parts: [{ text: 'world' }] }, + ]; + vi.mocked(mockChat.getHistory).mockReturnValue(history); + vi.mocked(mockChat.getLastPromptTokenCount).mockReturnValue(600000); + vi.mocked(tokenLimit).mockReturnValue(1_000_000); + + const mockHookSystem = { + firePreCompressEvent: vi.fn().mockResolvedValue({ + success: true, + allOutputs: [], + errors: [], + totalDuration: 10, + }), + }; + + vi.mocked(mockConfig.getHookSystem).mockReturnValue( + mockHookSystem as unknown as ReturnType, + ); + + await service.compress( + mockChat, + mockPromptId, + true, + mockModel, + mockConfig, + false, + ); + + expect(mockHookSystem.firePreCompressEvent).toHaveBeenCalledWith( + 'manual', + [ + { role: 'user', parts: [{ text: 'hello' }] }, + { role: 'model', parts: [{ text: 'world' }] }, + ], + ); + }); + }); }); diff --git a/packages/core/src/services/chatCompressionService.ts b/packages/core/src/services/chatCompressionService.ts index a1f9c12f2c..8d6fc64ac3 100644 --- a/packages/core/src/services/chatCompressionService.ts +++ b/packages/core/src/services/chatCompressionService.ts @@ -156,13 +156,13 @@ async function truncateHistoryToBudget( } else if (responseObj && typeof responseObj === 'object') { if ( 'output' in responseObj && - // eslint-disable-next-line no-restricted-syntax + typeof responseObj['output'] === 'string' ) { contentStr = responseObj['output']; } else if ( 'content' in responseObj && - // eslint-disable-next-line no-restricted-syntax + typeof responseObj['content'] === 'string' ) { contentStr = responseObj['content']; @@ -254,11 +254,6 @@ export class ChatCompressionService { }; } - // Fire PreCompress hook before compression - // This fires for both manual and auto compression attempts - const trigger = force ? PreCompressTrigger.Manual : PreCompressTrigger.Auto; - await config.getHookSystem()?.firePreCompressEvent(trigger); - const originalTokenCount = chat.getLastPromptTokenCount(); // Don't compress if not forced and we are under the limit. @@ -278,6 +273,63 @@ export class ChatCompressionService { } } + // Fire PreCompress hook — only when compression will actually proceed + const trigger = force ? PreCompressTrigger.Manual : PreCompressTrigger.Auto; + + // Serialize history for the hook: strip non-text parts to keep payload manageable + const curatedForHook = curatedHistory.map((c) => ({ + role: c.role ?? 'user', + parts: (c.parts ?? []) + .filter((p): p is { text: string } => typeof p.text === 'string') + .map((p) => ({ text: p.text })), + })); + + const hookResult = await config + .getHookSystem() + ?.firePreCompressEvent(trigger, curatedForHook); + + // If a hook provided replacement history, use it and skip built-in compression + const hookNewHistory = + hookResult?.finalOutput?.hookSpecificOutput?.['newHistory']; + if (Array.isArray(hookNewHistory) && hookNewHistory.length > 0) { + // Convert hook output back to Content[] + const replacementHistory: Content[] = hookNewHistory.map( + (entry: { role?: string; parts?: Array<{ text?: string }> }) => { + const role = + entry.role === 'model' || entry.role === 'user' + ? entry.role + : 'user'; + return { + role, + parts: (entry.parts ?? []).map((p: { text?: string }) => ({ + text: p.text ?? '', + })), + }; + }, + ); + + const newTokenCount = estimateTokenCountSync( + replacementHistory.flatMap((c) => c.parts || []), + ); + + logChatCompression( + config, + makeChatCompressionEvent({ + tokens_before: originalTokenCount, + tokens_after: newTokenCount, + }), + ); + + return { + newHistory: replacementHistory, + info: { + originalTokenCount, + newTokenCount, + compressionStatus: CompressionStatus.HOOK_REPLACED, + }, + }; + } + // Apply token-based truncation to the entire history before splitting. // This ensures that even the "to compress" portion is within safe limits for the summarization model. const truncatedHistory = await truncateHistoryToBudget( diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 2591d90bb4..631e500643 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -20,6 +20,7 @@ import type { import { debugLogger } from '../utils/debugLogger.js'; import type { ToolResultDisplay } from '../tools/tools.js'; import type { AgentLoopContext } from '../config/agent-loop-context.js'; +import type { SerializedScheduledItem } from './work-scheduler.js'; export const SESSION_FILE_PREFIX = 'session-'; @@ -104,6 +105,10 @@ export interface ConversationRecord { directories?: string[]; /** The kind of conversation (main agent or subagent) */ kind?: 'main' | 'subagent'; + /** Index into messages[] after the last compression, used to skip pre-compressed messages on resume */ + lastCompressionIndex?: number; + /** Pending scheduled work items persisted for session resume */ + scheduledWork?: SerializedScheduledItem[]; } /** @@ -567,6 +572,23 @@ export class ChatRecordingService { } } + /** + * Records pending scheduled work items to the session file. + * Called when the work schedule changes so items survive session restart. + */ + recordScheduledWork(items: SerializedScheduledItem[]): void { + if (!this.conversationFile) return; + + try { + this.updateConversation((conversation) => { + conversation.scheduledWork = items.length > 0 ? items : undefined; + }); + } catch (error) { + debugLogger.error('Error saving scheduled work to chat history.', error); + // Don't throw - we want graceful degradation + } + } + /** * Gets the current conversation data (for summary generation). */ @@ -727,6 +749,17 @@ export class ChatRecordingService { } } + /** + * Stamps the current end of the messages array so that future session + * resumes can skip the pre-compression portion of the history. + */ + recordCompressionPoint(): void { + if (!this.conversationFile) return; + this.updateConversation((conversation) => { + conversation.lastCompressionIndex = conversation.messages.length; + }); + } + /** * Rewinds the conversation to the state just before the specified message ID. * All messages from (and including) the specified ID onwards are removed. @@ -759,37 +792,39 @@ export class ChatRecordingService { updateMessagesFromHistory(history: readonly Content[]): void { if (!this.conversationFile) return; + // Build the partsMap before touching the file — skip I/O entirely when + // there are no tool results to sync. + const partsMap = new Map(); + for (const content of history) { + if (content.role === 'user' && content.parts) { + // Find all unique call IDs in this message + const callIds = content.parts + .map((p) => p.functionResponse?.id) + .filter((id): id is string => !!id); + + if (callIds.length === 0) continue; + + // Use the first ID as a seed to capture any "leading" non-ID parts + // in this specific content block. + let currentCallId = callIds[0]; + for (const part of content.parts) { + if (part.functionResponse?.id) { + currentCallId = part.functionResponse.id; + } + + if (!partsMap.has(currentCallId)) { + partsMap.set(currentCallId, []); + } + partsMap.get(currentCallId)!.push(part); + } + } + } + + // No tool results to update — skip file I/O entirely. + if (partsMap.size === 0) return; + try { this.updateConversation((conversation) => { - // Create a map of tool results from the API history for quick lookup by call ID. - // We store the full list of parts associated with each tool call ID to preserve - // multi-modal data and proper trajectory structure. - const partsMap = new Map(); - for (const content of history) { - if (content.role === 'user' && content.parts) { - // Find all unique call IDs in this message - const callIds = content.parts - .map((p) => p.functionResponse?.id) - .filter((id): id is string => !!id); - - if (callIds.length === 0) continue; - - // Use the first ID as a seed to capture any "leading" non-ID parts - // in this specific content block. - let currentCallId = callIds[0]; - for (const part of content.parts) { - if (part.functionResponse?.id) { - currentCallId = part.functionResponse.id; - } - - if (!partsMap.has(currentCallId)) { - partsMap.set(currentCallId, []); - } - partsMap.get(currentCallId)!.push(part); - } - } - } - // Update the conversation records tool results if they've changed. for (const message of conversation.messages) { if (message.type === 'gemini' && message.toolCalls) { diff --git a/packages/core/src/services/work-scheduler.test.ts b/packages/core/src/services/work-scheduler.test.ts new file mode 100644 index 0000000000..b50777d377 --- /dev/null +++ b/packages/core/src/services/work-scheduler.test.ts @@ -0,0 +1,498 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { WorkScheduler } from './work-scheduler.js'; +import type { SerializedScheduledItem } from './work-scheduler.js'; + +describe('WorkScheduler', () => { + let scheduler: WorkScheduler; + + beforeEach(() => { + vi.useFakeTimers(); + scheduler = new WorkScheduler(); + }); + + afterEach(() => { + scheduler.dispose(); + vi.useRealTimers(); + }); + + describe('add', () => { + it('should add an item and return it with pending status', () => { + const fireAt = new Date(Date.now() + 60_000); + const item = scheduler.add('do something', fireAt); + + expect(item.id).toBeDefined(); + expect(item.prompt).toBe('do something'); + expect(item.fireAt).toEqual(fireAt); + expect(item.status).toBe('pending'); + expect(item.createdAt).toBeInstanceOf(Date); + }); + + it('should keep items sorted by fireAt', () => { + const now = Date.now(); + scheduler.add('third', new Date(now + 300_000)); + scheduler.add('first', new Date(now + 60_000)); + scheduler.add('second', new Date(now + 120_000)); + + const pending = scheduler.getPendingItems(); + expect(pending).toHaveLength(3); + expect(pending[0].prompt).toBe('first'); + expect(pending[1].prompt).toBe('second'); + expect(pending[2].prompt).toBe('third'); + }); + + it('should emit changed event on add', () => { + const changedSpy = vi.fn(); + scheduler.on('changed', changedSpy); + + scheduler.add('test', new Date(Date.now() + 60_000)); + + expect(changedSpy).toHaveBeenCalledTimes(1); + }); + + it('should fire past-due items immediately on add', () => { + const fireSpy = vi.fn(); + scheduler.on('fire', fireSpy); + + const pastDate = new Date(Date.now() - 10_000); + const item = scheduler.add('overdue task', pastDate); + + expect(fireSpy).toHaveBeenCalledWith('overdue task'); + expect(item.status).toBe('fired'); + }); + + it('should rearm timer when adding an item sooner than the current next', () => { + const fireSpy = vi.fn(); + scheduler.on('fire', fireSpy); + + scheduler.addRelative('ten min', 10); + scheduler.addRelative('five min', 5); + + vi.advanceTimersByTime(5 * 60_000); + expect(fireSpy).toHaveBeenCalledTimes(1); + expect(fireSpy).toHaveBeenCalledWith('five min'); + + vi.advanceTimersByTime(5 * 60_000); + expect(fireSpy).toHaveBeenCalledTimes(2); + expect(fireSpy).toHaveBeenCalledWith('ten min'); + }); + }); + + describe('addRelative', () => { + it('should add an item relative to current time', () => { + const now = Date.now(); + const item = scheduler.addRelative('in 5 minutes', 5); + + expect(item.fireAt.getTime()).toBe(now + 5 * 60_000); + expect(item.status).toBe('pending'); + }); + + it('should appear in pending items', () => { + scheduler.addRelative('relative item', 10); + + const pending = scheduler.getPendingItems(); + expect(pending).toHaveLength(1); + expect(pending[0].prompt).toBe('relative item'); + }); + }); + + describe('cancel', () => { + it('should cancel a pending item and return true', () => { + const item = scheduler.add('cancel me', new Date(Date.now() + 60_000)); + + const result = scheduler.cancel(item.id); + + expect(result).toBe(true); + expect(scheduler.getPendingItems()).toHaveLength(0); + }); + + it('should return false for a non-existent id', () => { + const result = scheduler.cancel('non-existent-id'); + + expect(result).toBe(false); + }); + + it('should return false for an already cancelled item', () => { + const item = scheduler.add('cancel me', new Date(Date.now() + 60_000)); + scheduler.cancel(item.id); + + const result = scheduler.cancel(item.id); + + expect(result).toBe(false); + }); + + it('should emit changed event on successful cancel', () => { + const item = scheduler.add('cancel me', new Date(Date.now() + 60_000)); + const changedSpy = vi.fn(); + scheduler.on('changed', changedSpy); + + scheduler.cancel(item.id); + + expect(changedSpy).toHaveBeenCalledTimes(1); + }); + + it('should not emit changed event on failed cancel', () => { + const changedSpy = vi.fn(); + scheduler.on('changed', changedSpy); + + scheduler.cancel('non-existent-id'); + + expect(changedSpy).not.toHaveBeenCalled(); + }); + + it('should rearm timer to next item when cancelling the currently armed item', () => { + const fireSpy = vi.fn(); + scheduler.on('fire', fireSpy); + + const first = scheduler.addRelative('five min', 5); + scheduler.addRelative('ten min', 10); + + scheduler.cancel(first.id); + + vi.advanceTimersByTime(5 * 60_000); + expect(fireSpy).not.toHaveBeenCalled(); + + vi.advanceTimersByTime(5 * 60_000); + expect(fireSpy).toHaveBeenCalledTimes(1); + expect(fireSpy).toHaveBeenCalledWith('ten min'); + }); + }); + + describe('getPendingItems', () => { + it('should return empty array when no items', () => { + expect(scheduler.getPendingItems()).toHaveLength(0); + }); + + it('should exclude cancelled and fired items', () => { + const item1 = scheduler.add('keep', new Date(Date.now() + 60_000)); + scheduler.add('cancel', new Date(Date.now() + 120_000)); + scheduler.add('also keep', new Date(Date.now() + 180_000)); + + scheduler.cancel(scheduler.getPendingItems()[1].id); + + const pending = scheduler.getPendingItems(); + expect(pending).toHaveLength(2); + expect(pending[0].id).toBe(item1.id); + expect(pending[1].prompt).toBe('also keep'); + }); + }); + + describe('getNextPending', () => { + it('should return undefined when no items', () => { + expect(scheduler.getNextPending()).toBeUndefined(); + }); + + it('should return the soonest pending item', () => { + const now = Date.now(); + scheduler.add('later', new Date(now + 120_000)); + scheduler.add('sooner', new Date(now + 60_000)); + + const next = scheduler.getNextPending(); + expect(next?.prompt).toBe('sooner'); + }); + }); + + describe('fire event', () => { + it('should emit fire event when timer expires', () => { + const fireSpy = vi.fn(); + scheduler.on('fire', fireSpy); + + scheduler.addRelative('fire me', 5); + + expect(fireSpy).not.toHaveBeenCalled(); + + vi.advanceTimersByTime(5 * 60_000); + + expect(fireSpy).toHaveBeenCalledWith('fire me'); + }); + + it('should emit changed event when item fires via timer', () => { + scheduler.addRelative('fire me', 5); + const changedSpy = vi.fn(); + scheduler.on('changed', changedSpy); + + vi.advanceTimersByTime(5 * 60_000); + + expect(changedSpy).toHaveBeenCalled(); + }); + + it('should fire multiple items in sequence as their timers expire', () => { + const fireSpy = vi.fn(); + scheduler.on('fire', fireSpy); + + scheduler.addRelative('first', 1); + scheduler.addRelative('second', 2); + + vi.advanceTimersByTime(60_000); + expect(fireSpy).toHaveBeenCalledTimes(1); + expect(fireSpy).toHaveBeenCalledWith('first'); + + vi.advanceTimersByTime(60_000); + expect(fireSpy).toHaveBeenCalledTimes(2); + expect(fireSpy).toHaveBeenCalledWith('second'); + }); + + it('should not fire cancelled items', () => { + const fireSpy = vi.fn(); + scheduler.on('fire', fireSpy); + + const item = scheduler.addRelative('cancel me', 5); + scheduler.cancel(item.id); + + vi.advanceTimersByTime(5 * 60_000); + + expect(fireSpy).not.toHaveBeenCalled(); + }); + }); + + describe('serialize', () => { + it('should serialize pending items to ISO string format', () => { + const now = Date.now(); + scheduler.add('item 1', new Date(now + 60_000)); + scheduler.add('item 2', new Date(now + 120_000)); + + const serialized = scheduler.serialize(); + + expect(serialized).toHaveLength(2); + expect(serialized[0].prompt).toBe('item 1'); + expect(serialized[0].fireAt).toBe(new Date(now + 60_000).toISOString()); + expect(serialized[0].id).toBeDefined(); + expect(serialized[0].createdAt).toBeDefined(); + expect(serialized[1].prompt).toBe('item 2'); + }); + + it('should not include cancelled items', () => { + const item = scheduler.add('cancel me', new Date(Date.now() + 60_000)); + scheduler.add('keep me', new Date(Date.now() + 120_000)); + scheduler.cancel(item.id); + + const serialized = scheduler.serialize(); + + expect(serialized).toHaveLength(1); + expect(serialized[0].prompt).toBe('keep me'); + }); + + it('should return empty array when no pending items', () => { + expect(scheduler.serialize()).toEqual([]); + }); + + it('should exclude fired items', () => { + scheduler.addRelative('will fire', 5); + + vi.advanceTimersByTime(5 * 60_000); + + const serialized = scheduler.serialize(); + expect(serialized).toEqual([]); + }); + }); + + describe('restore', () => { + it('should restore items from serialized data', () => { + const now = Date.now(); + const serialized: SerializedScheduledItem[] = [ + { + id: 'restored-1', + prompt: 'restored item', + fireAt: new Date(now + 60_000).toISOString(), + createdAt: new Date(now - 10_000).toISOString(), + }, + ]; + + scheduler.restore(serialized); + + const pending = scheduler.getPendingItems(); + expect(pending).toHaveLength(1); + expect(pending[0].id).toBe('restored-1'); + expect(pending[0].prompt).toBe('restored item'); + expect(pending[0].status).toBe('pending'); + }); + + it('should emit changed event on restore', () => { + const changedSpy = vi.fn(); + scheduler.on('changed', changedSpy); + + scheduler.restore([ + { + id: 'id-1', + prompt: 'test', + fireAt: new Date(Date.now() + 60_000).toISOString(), + createdAt: new Date().toISOString(), + }, + ]); + + expect(changedSpy).toHaveBeenCalledTimes(1); + }); + + it('should fire past-due items immediately on restore', () => { + const fireSpy = vi.fn(); + scheduler.on('fire', fireSpy); + + const serialized: SerializedScheduledItem[] = [ + { + id: 'past-due-1', + prompt: 'overdue restored', + fireAt: new Date(Date.now() - 30_000).toISOString(), + createdAt: new Date(Date.now() - 120_000).toISOString(), + }, + ]; + + scheduler.restore(serialized); + + expect(fireSpy).toHaveBeenCalledWith('overdue restored'); + }); + + it('should handle round-trip serialize and restore', () => { + const now = Date.now(); + scheduler.add('round trip 1', new Date(now + 60_000)); + scheduler.add('round trip 2', new Date(now + 120_000)); + + const serialized = scheduler.serialize(); + + const newScheduler = new WorkScheduler(); + newScheduler.restore(serialized); + + const pending = newScheduler.getPendingItems(); + expect(pending).toHaveLength(2); + expect(pending[0].prompt).toBe('round trip 1'); + expect(pending[1].prompt).toBe('round trip 2'); + + newScheduler.dispose(); + }); + + it('should sort restored items by fireAt', () => { + const now = Date.now(); + const serialized: SerializedScheduledItem[] = [ + { + id: 'later', + prompt: 'later', + fireAt: new Date(now + 120_000).toISOString(), + createdAt: new Date().toISOString(), + }, + { + id: 'sooner', + prompt: 'sooner', + fireAt: new Date(now + 60_000).toISOString(), + createdAt: new Date().toISOString(), + }, + ]; + + scheduler.restore(serialized); + + const pending = scheduler.getPendingItems(); + expect(pending[0].prompt).toBe('sooner'); + expect(pending[1].prompt).toBe('later'); + }); + + it('should handle mix of past-due and future items', () => { + const fireSpy = vi.fn(); + scheduler.on('fire', fireSpy); + + const now = Date.now(); + const serialized: SerializedScheduledItem[] = [ + { + id: 'past-due', + prompt: 'overdue', + fireAt: new Date(now - 30_000).toISOString(), + createdAt: new Date(now - 120_000).toISOString(), + }, + { + id: 'future', + prompt: 'upcoming', + fireAt: new Date(now + 5 * 60_000).toISOString(), + createdAt: new Date(now - 60_000).toISOString(), + }, + ]; + + scheduler.restore(serialized); + + expect(fireSpy).toHaveBeenCalledTimes(1); + expect(fireSpy).toHaveBeenCalledWith('overdue'); + + const pending = scheduler.getPendingItems(); + expect(pending).toHaveLength(1); + expect(pending[0].prompt).toBe('upcoming'); + + vi.advanceTimersByTime(5 * 60_000); + expect(fireSpy).toHaveBeenCalledTimes(2); + expect(fireSpy).toHaveBeenCalledWith('upcoming'); + }); + }); + + describe('formatScheduleSummary', () => { + it('should include current time and timezone', () => { + const summary = scheduler.formatScheduleSummary(); + + expect(summary).toContain('Current time:'); + const tz = Intl.DateTimeFormat().resolvedOptions().timeZone; + expect(summary).toContain(tz); + }); + + it('should show no scheduled items message when empty', () => { + const summary = scheduler.formatScheduleSummary(); + + expect(summary).toContain('No scheduled items.'); + }); + + it('should list all pending items', () => { + scheduler.add('first task', new Date(Date.now() + 60_000)); + scheduler.add('second task', new Date(Date.now() + 120_000)); + + const summary = scheduler.formatScheduleSummary(); + + expect(summary).toContain('Active schedule:'); + expect(summary).toContain('first task'); + expect(summary).toContain('second task'); + expect(summary).toContain('1.'); + expect(summary).toContain('2.'); + }); + + it('should show truncated item IDs', () => { + const item = scheduler.add('test', new Date(Date.now() + 60_000)); + + const summary = scheduler.formatScheduleSummary(); + + expect(summary).toContain(`[${item.id.slice(0, 8)}]`); + }); + + it('should not contain ISO UTC format (toISOString)', () => { + scheduler.add('check format', new Date(Date.now() + 60_000)); + + const summary = scheduler.formatScheduleSummary(); + + expect(summary).not.toMatch(/\.\d{3}Z/); + }); + }); + + describe('dispose', () => { + it('should prevent items from firing after dispose', () => { + const fireSpy = vi.fn(); + scheduler.on('fire', fireSpy); + + scheduler.addRelative('should not fire', 5); + scheduler.dispose(); + + vi.advanceTimersByTime(5 * 60_000); + + expect(fireSpy).not.toHaveBeenCalled(); + }); + + it('should be safe to call dispose multiple times', () => { + scheduler.dispose(); + scheduler.dispose(); + // No error thrown + }); + + it('should not throw when adding items after dispose', () => { + scheduler.dispose(); + + expect(() => { + scheduler.add('after dispose', new Date(Date.now() + 60_000)); + }).not.toThrow(); + }); + }); +}); diff --git a/packages/core/src/services/work-scheduler.ts b/packages/core/src/services/work-scheduler.ts new file mode 100644 index 0000000000..2baccc84dc --- /dev/null +++ b/packages/core/src/services/work-scheduler.ts @@ -0,0 +1,223 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { randomUUID } from 'node:crypto'; +import { EventEmitter } from 'node:events'; + +/** + * A single scheduled work item. + */ +export interface ScheduledItem { + id: string; + prompt: string; + fireAt: Date; + createdAt: Date; + status: 'pending' | 'fired' | 'cancelled'; +} + +/** + * Serializable representation of a scheduled item for persistence. + */ +export interface SerializedScheduledItem { + id: string; + prompt: string; + fireAt: string; + createdAt: string; +} + +export interface WorkSchedulerEvents { + fire: [prompt: string]; + changed: []; +} + +/** + * Manages a time-based list of scheduled work items. + * Emits 'fire' when a scheduled item's time arrives, with the prompt text. + * Emits 'changed' whenever the schedule is mutated. + */ +export class WorkScheduler extends EventEmitter { + private items: ScheduledItem[] = []; + private timer: ReturnType | null = null; + + /** + * Add a scheduled item at an absolute time. + * @returns The created ScheduledItem. + */ + add(prompt: string, fireAt: Date): ScheduledItem { + const item: ScheduledItem = { + id: randomUUID(), + prompt, + fireAt, + createdAt: new Date(), + status: 'pending', + }; + this.items.push(item); + this.items.sort((a, b) => a.fireAt.getTime() - b.fireAt.getTime()); + this.rearm(); + this.emit('changed'); + return item; + } + + /** + * Add a scheduled item using a relative delay in minutes. + * @returns The created ScheduledItem. + */ + addRelative(prompt: string, inMinutes: number): ScheduledItem { + const fireAt = new Date(Date.now() + inMinutes * 60 * 1000); + return this.add(prompt, fireAt); + } + + /** + * Cancel a scheduled item by ID. + * @returns true if the item was found and cancelled, false otherwise. + */ + cancel(id: string): boolean { + const item = this.items.find((i) => i.id === id && i.status === 'pending'); + if (!item) { + return false; + } + item.status = 'cancelled'; + this.rearm(); + this.emit('changed'); + return true; + } + + /** + * Get all pending items, sorted by fireAt. + */ + getPendingItems(): readonly ScheduledItem[] { + return this.items.filter((i) => i.status === 'pending'); + } + + /** + * Get the next pending item (soonest fireAt). + */ + getNextPending(): ScheduledItem | undefined { + return this.items.find((i) => i.status === 'pending'); + } + + /** + * Serialize pending items for session persistence. + */ + serialize(): SerializedScheduledItem[] { + return this.getPendingItems().map((item) => ({ + id: item.id, + prompt: item.prompt, + fireAt: item.fireAt.toISOString(), + createdAt: item.createdAt.toISOString(), + })); + } + + /** + * Restore scheduled items from serialized data (e.g. session resume). + * Items whose fireAt is in the past fire immediately (queued sequentially). + * Items in the future get timers re-armed. + */ + restore(serialized: SerializedScheduledItem[]): void { + for (const s of serialized) { + const item: ScheduledItem = { + id: s.id, + prompt: s.prompt, + fireAt: new Date(s.fireAt), + createdAt: new Date(s.createdAt), + status: 'pending', + }; + this.items.push(item); + } + this.items.sort((a, b) => a.fireAt.getTime() - b.fireAt.getTime()); + this.rearm(); + this.emit('changed'); + } + + /** + * Format a human-readable summary of the current schedule, including current time. + */ + formatScheduleSummary(): string { + const now = new Date(); + const pending = this.getPendingItems(); + const tz = Intl.DateTimeFormat().resolvedOptions().timeZone; + + const localTime = now.toLocaleString([], { + year: 'numeric', + month: '2-digit', + day: '2-digit', + hour: '2-digit', + minute: '2-digit', + second: '2-digit', + }); + const lines: string[] = [`Current time: ${localTime} (${tz})`]; + + if (pending.length === 0) { + lines.push('No scheduled items.'); + } else { + lines.push(''); + lines.push('Active schedule:'); + for (let i = 0; i < pending.length; i++) { + const item = pending[i]; + const diffMs = item.fireAt.getTime() - now.getTime(); + const diffMins = Math.max(0, Math.ceil(diffMs / 60000)); + const itemTime = item.fireAt.toLocaleTimeString([], { + hour: '2-digit', + minute: '2-digit', + }); + lines.push( + ` ${i + 1}. [${item.id.slice(0, 8)}] ${itemTime} (in ${diffMins}m) — "${item.prompt}"`, + ); + } + } + return lines.join('\n'); + } + + /** + * Stop all timers. Call on cleanup/shutdown. + */ + dispose(): void { + if (this.timer) { + clearTimeout(this.timer); + this.timer = null; + } + } + + /** + * Re-arm the internal timer to point at the next pending item. + */ + private rearm(): void { + if (this.timer) { + clearTimeout(this.timer); + this.timer = null; + } + + // Fire all past-due items immediately + const now = Date.now(); + const pastDue = this.items.filter( + (i) => i.status === 'pending' && i.fireAt.getTime() <= now, + ); + for (const item of pastDue) { + item.status = 'fired'; + this.emit('fire', item.prompt); + } + + // Find next future pending item + const next = this.items.find( + (i) => i.status === 'pending' && i.fireAt.getTime() > now, + ); + if (!next) { + return; + } + + const delayMs = Math.max(0, next.fireAt.getTime() - Date.now()); + this.timer = setTimeout(() => { + this.timer = null; + if (next.status === 'pending') { + next.status = 'fired'; + this.emit('fire', next.prompt); + this.emit('changed'); + } + // Re-arm for the next item after this one + this.rearm(); + }, delayMs); + } +} diff --git a/packages/core/src/tools/schedule-work.test.ts b/packages/core/src/tools/schedule-work.test.ts new file mode 100644 index 0000000000..4bbcb736f8 --- /dev/null +++ b/packages/core/src/tools/schedule-work.test.ts @@ -0,0 +1,308 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { + ScheduleWorkTool, + ScheduleWorkInvocation, + type ScheduleWorkParams, +} from './schedule-work.js'; +import { WorkScheduler } from '../services/work-scheduler.js'; +import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; + +describe('ScheduleWorkTool', () => { + let tool: ScheduleWorkTool; + let scheduler: WorkScheduler; + let mockMessageBus: MessageBus; + const signal = new AbortController().signal; + + beforeEach(() => { + vi.useFakeTimers(); + mockMessageBus = createMockMessageBus(); + scheduler = new WorkScheduler(); + tool = new ScheduleWorkTool(mockMessageBus, scheduler); + }); + + afterEach(() => { + scheduler.dispose(); + vi.useRealTimers(); + }); + + describe('validation', () => { + it('should reject add without prompt', () => { + const params: ScheduleWorkParams = { + action: 'add', + inMinutes: 5, + }; + const result = tool.validateToolParams(params); + expect(result).toContain('"prompt" is required'); + }); + + it('should reject add with empty prompt', () => { + const params: ScheduleWorkParams = { + action: 'add', + prompt: ' ', + inMinutes: 5, + }; + const result = tool.validateToolParams(params); + expect(result).toContain('"prompt" is required'); + }); + + it('should reject add without at or inMinutes', () => { + const params: ScheduleWorkParams = { + action: 'add', + prompt: 'do something', + }; + const result = tool.validateToolParams(params); + expect(result).toContain('One of "at" or "inMinutes" is required'); + }); + + it('should reject add with both at and inMinutes', () => { + const params: ScheduleWorkParams = { + action: 'add', + prompt: 'do something', + at: '2025-06-15T14:00:00', + inMinutes: 5, + }; + const result = tool.validateToolParams(params); + expect(result).toContain('"at" and "inMinutes" are mutually exclusive'); + }); + + it('should reject inMinutes <= 0', () => { + const params: ScheduleWorkParams = { + action: 'add', + prompt: 'do something', + inMinutes: 0, + }; + const result = tool.validateToolParams(params); + expect(result).toContain('"inMinutes" must be greater than 0'); + }); + + it('should reject negative inMinutes', () => { + const params: ScheduleWorkParams = { + action: 'add', + prompt: 'do something', + inMinutes: -5, + }; + const result = tool.validateToolParams(params); + expect(result).toContain('"inMinutes" must be greater than 0'); + }); + + it('should reject invalid date format for at', () => { + const params: ScheduleWorkParams = { + action: 'add', + prompt: 'do something', + at: 'not-a-date', + }; + const result = tool.validateToolParams(params); + expect(result).toContain('Invalid date format'); + }); + + it('should reject cancel without id', () => { + const params: ScheduleWorkParams = { + action: 'cancel', + }; + const result = tool.validateToolParams(params); + expect(result).toContain('"id" is required'); + }); + + it('should reject cancel with empty id', () => { + const params: ScheduleWorkParams = { + action: 'cancel', + id: ' ', + }; + const result = tool.validateToolParams(params); + expect(result).toContain('"id" is required'); + }); + + it('should accept valid add with inMinutes', () => { + const params: ScheduleWorkParams = { + action: 'add', + prompt: 'do something', + inMinutes: 5, + }; + const result = tool.validateToolParams(params); + expect(result).toBeNull(); + }); + + it('should accept valid add with at', () => { + const params: ScheduleWorkParams = { + action: 'add', + prompt: 'do something', + at: '2025-06-15T14:00:00', + }; + const result = tool.validateToolParams(params); + expect(result).toBeNull(); + }); + + it('should accept valid cancel with id', () => { + const params: ScheduleWorkParams = { + action: 'cancel', + id: 'some-id', + }; + const result = tool.validateToolParams(params); + expect(result).toBeNull(); + }); + }); + + describe('execute - add with inMinutes', () => { + it('should call scheduler.addRelative and return confirmation', async () => { + const addRelativeSpy = vi.spyOn(scheduler, 'addRelative'); + + const result = await tool.buildAndExecute( + { action: 'add', prompt: 'run tests', inMinutes: 10 }, + signal, + ); + + expect(addRelativeSpy).toHaveBeenCalledWith('run tests', 10); + expect(result.llmContent).toContain('Scheduled item'); + expect(result.llmContent).toContain('Will fire at'); + expect(result.llmContent).toContain('Current time:'); + expect(result.returnDisplay).toContain('run tests'); + }); + }); + + describe('execute - add with at', () => { + it('should call scheduler.add with parsed Date and return confirmation', async () => { + const addSpy = vi.spyOn(scheduler, 'add'); + + const result = await tool.buildAndExecute( + { action: 'add', prompt: 'deploy', at: '2025-06-15T14:00:00' }, + signal, + ); + + expect(addSpy).toHaveBeenCalledWith( + 'deploy', + new Date('2025-06-15T14:00:00'), + ); + expect(result.llmContent).toContain('Scheduled item'); + expect(result.returnDisplay).toContain('deploy'); + }); + + it('should handle scheduling with an at time in the past', async () => { + const pastDate = new Date(Date.now() - 60_000); + const fireSpy = vi.fn(); + scheduler.on('fire', fireSpy); + + const result = await tool.buildAndExecute( + { action: 'add', prompt: 'overdue task', at: pastDate.toISOString() }, + signal, + ); + + expect(result.llmContent).toContain('Scheduled item'); + expect(fireSpy).toHaveBeenCalledWith('overdue task'); + }); + }); + + describe('execute - cancel', () => { + it('should cancel a matching item and return confirmation', async () => { + const item = scheduler.addRelative('cancel me', 10); + + const result = await tool.buildAndExecute( + { action: 'cancel', id: item.id }, + signal, + ); + + expect(result.llmContent).toContain('Cancelled item'); + expect(result.llmContent).toContain('cancel me'); + expect(result.error).toBeUndefined(); + expect(scheduler.getPendingItems()).toHaveLength(0); + }); + + it('should cancel by ID prefix', async () => { + const item = scheduler.addRelative('prefix cancel', 10); + const prefix = item.id.slice(0, 8); + + const result = await tool.buildAndExecute( + { action: 'cancel', id: prefix }, + signal, + ); + + expect(result.llmContent).toContain('Cancelled item'); + expect(result.llmContent).toContain('prefix cancel'); + expect(scheduler.getPendingItems()).toHaveLength(0); + }); + + it('should return error when item not found', async () => { + const result = await tool.buildAndExecute( + { action: 'cancel', id: 'nonexistent' }, + signal, + ); + + expect(result.llmContent).toContain('No pending item found'); + expect(result.error).toBeDefined(); + expect(result.error?.message).toContain('No pending item found'); + }); + + it('should cancel the first match when prefix matches multiple items', async () => { + const item1 = scheduler.addRelative('first item', 10); + scheduler.addRelative('second item', 20); + + const prefix = item1.id.slice(0, 4); + + const result = await tool.buildAndExecute( + { action: 'cancel', id: prefix }, + signal, + ); + + expect(result.llmContent).toContain('Cancelled item'); + expect(scheduler.getPendingItems()).toHaveLength(1); + }); + }); + + describe('execute - unknown action', () => { + it('should return error for unknown action', async () => { + const invocation = new ScheduleWorkInvocation( + { action: 'unknown' as ScheduleWorkParams['action'] }, + mockMessageBus, + 'schedule_work', + 'Schedule Work', + scheduler, + ); + + const result = await invocation.execute(signal); + + expect(result.llmContent).toContain('Unknown action'); + expect(result.error).toBeDefined(); + expect(result.error?.message).toContain('Unknown action'); + }); + }); + + describe('getDescription', () => { + it('should return correct description for add with inMinutes', () => { + const invocation = tool.build({ + action: 'add', + prompt: 'run tests', + inMinutes: 10, + }); + expect(invocation.getDescription()).toBe( + 'Scheduling work to fire in 10 minutes.', + ); + }); + + it('should return correct description for add with at', () => { + const invocation = tool.build({ + action: 'add', + prompt: 'deploy', + at: '2025-06-15T14:00:00', + }); + expect(invocation.getDescription()).toBe( + 'Scheduling work to fire at 2025-06-15T14:00:00.', + ); + }); + + it('should return correct description for cancel', () => { + const invocation = tool.build({ + action: 'cancel', + id: 'abc123', + }); + expect(invocation.getDescription()).toBe( + 'Cancelling scheduled item abc123.', + ); + }); + }); +}); diff --git a/packages/core/src/tools/schedule-work.ts b/packages/core/src/tools/schedule-work.ts new file mode 100644 index 0000000000..82de80aead --- /dev/null +++ b/packages/core/src/tools/schedule-work.ts @@ -0,0 +1,210 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + BaseDeclarativeTool, + BaseToolInvocation, + type ToolResult, + Kind, +} from './tools.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; +import { SCHEDULE_WORK_TOOL_NAME } from './tool-names.js'; +import type { WorkScheduler } from '../services/work-scheduler.js'; + +export interface ScheduleWorkParams { + action: 'add' | 'cancel'; + prompt?: string; + at?: string; + inMinutes?: number; + id?: string; +} + +export class ScheduleWorkTool extends BaseDeclarativeTool< + ScheduleWorkParams, + ToolResult +> { + private readonly scheduler: WorkScheduler; + + constructor(messageBus: MessageBus, scheduler: WorkScheduler) { + super( + SCHEDULE_WORK_TOOL_NAME, + 'Schedule Work', + 'Manage a scheduled work list. Schedule prompts to be automatically injected at specific times (queued if the agent is busy). The current schedule is always visible in the context prepended to each message. The "at" parameter is interpreted as the system\'s local timezone.', + Kind.Communicate, + { + type: 'object', + required: ['action'], + properties: { + action: { + type: 'string', + enum: ['add', 'cancel'], + description: + 'Action to perform: "add" to schedule a new item, "cancel" to remove a scheduled item.', + }, + prompt: { + type: 'string', + description: + 'The prompt to inject when the scheduled time arrives. Required for "add".', + }, + at: { + type: 'string', + description: + 'Absolute local time to fire the prompt, in ISO 8601 format without timezone offset (e.g. "2025-01-15T14:30:00"). Interpreted as the system\'s local timezone. Mutually exclusive with inMinutes. One of "at" or "inMinutes" is required for "add".', + }, + inMinutes: { + type: 'number', + description: + 'Minutes from now to fire the prompt. Mutually exclusive with "at". One of "at" or "inMinutes" is required for "add".', + }, + id: { + type: 'string', + description: + 'ID of the scheduled item to cancel. Required for "cancel". The current schedule with IDs is visible in the context prepended to each message.', + }, + }, + }, + messageBus, + ); + this.scheduler = scheduler; + } + + protected override validateToolParamValues( + params: ScheduleWorkParams, + ): string | null { + if (params.action === 'add') { + if (!params.prompt || params.prompt.trim() === '') { + return '"prompt" is required for action "add".'; + } + if (params.at == null && params.inMinutes == null) { + return 'One of "at" or "inMinutes" is required for action "add".'; + } + if (params.at != null && params.inMinutes != null) { + return '"at" and "inMinutes" are mutually exclusive.'; + } + if (params.inMinutes != null && params.inMinutes <= 0) { + return '"inMinutes" must be greater than 0.'; + } + if (params.at != null) { + const parsed = new Date(params.at); + if (isNaN(parsed.getTime())) { + return `Invalid date format for "at": "${params.at}". Use ISO 8601 format (e.g. "2025-01-15T14:30:00").`; + } + } + } + if (params.action === 'cancel') { + if (!params.id || params.id.trim() === '') { + return '"id" is required for action "cancel".'; + } + } + return null; + } + + protected createInvocation( + params: ScheduleWorkParams, + messageBus: MessageBus, + toolName: string, + toolDisplayName: string, + ): ScheduleWorkInvocation { + return new ScheduleWorkInvocation( + params, + messageBus, + toolName, + toolDisplayName, + this.scheduler, + ); + } +} + +export class ScheduleWorkInvocation extends BaseToolInvocation< + ScheduleWorkParams, + ToolResult +> { + private readonly scheduler: WorkScheduler; + + constructor( + params: ScheduleWorkParams, + messageBus: MessageBus, + toolName: string, + toolDisplayName: string, + scheduler: WorkScheduler, + ) { + super(params, messageBus, toolName, toolDisplayName); + this.scheduler = scheduler; + } + + getDescription(): string { + switch (this.params.action) { + case 'add': + if (this.params.inMinutes) { + return `Scheduling work to fire in ${this.params.inMinutes} minutes.`; + } + return `Scheduling work to fire at ${this.params.at}.`; + case 'cancel': + return `Cancelling scheduled item ${this.params.id}.`; + default: + return 'Managing scheduled work.'; + } + } + + async execute(_signal: AbortSignal): Promise { + switch (this.params.action) { + case 'add': + return this.handleAdd(); + case 'cancel': + return this.handleCancel(); + default: + return { + llmContent: `Unknown action: "${String(this.params.action)}". Use "add" or "cancel".`, + returnDisplay: 'Unknown action.', + error: { message: 'Unknown action.' }, + }; + } + } + + private handleAdd(): ToolResult { + const prompt = this.params.prompt!; + let item; + + if (this.params.inMinutes != null) { + item = this.scheduler.addRelative(prompt, this.params.inMinutes); + } else { + // Parse "at" as local time + const fireAt = new Date(this.params.at!); + item = this.scheduler.add(prompt, fireAt); + } + + const summary = this.scheduler.formatScheduleSummary(); + return { + llmContent: `Scheduled item [${item.id.slice(0, 8)}] created. Will fire at ${item.fireAt.toISOString()}.\n\n${summary}`, + returnDisplay: `Scheduled: "${prompt}" at ${item.fireAt.toLocaleTimeString()}`, + }; + } + + private handleCancel(): ToolResult { + const id = this.params.id!; + // Allow matching by full ID or prefix + const pending = this.scheduler.getPendingItems(); + const match = pending.find( + (item) => item.id === id || item.id.startsWith(id), + ); + + if (!match) { + const summary = this.scheduler.formatScheduleSummary(); + return { + llmContent: `No pending item found with ID "${id}".\n\n${summary}`, + returnDisplay: `Item not found: ${id}`, + error: { message: `No pending item found with ID "${id}".` }, + }; + } + + this.scheduler.cancel(match.id); + const summary = this.scheduler.formatScheduleSummary(); + return { + llmContent: `Cancelled item [${match.id.slice(0, 8)}] — "${match.prompt}".\n\n${summary}`, + returnDisplay: `Cancelled: "${match.prompt}"`, + }; + } +} diff --git a/packages/core/src/tools/tool-names.ts b/packages/core/src/tools/tool-names.ts index e818881662..8ee1203c99 100644 --- a/packages/core/src/tools/tool-names.ts +++ b/packages/core/src/tools/tool-names.ts @@ -151,6 +151,7 @@ export { }; export const LS_TOOL_NAME_LEGACY = 'list_directory'; // Just to be safe if anything used the old exported name directly +export const SCHEDULE_WORK_TOOL_NAME = 'schedule_work'; export const EDIT_TOOL_NAMES = new Set([EDIT_TOOL_NAME, WRITE_FILE_TOOL_NAME]); @@ -251,6 +252,7 @@ export const ALL_BUILTIN_TOOL_NAMES = [ GET_INTERNAL_DOCS_TOOL_NAME, ENTER_PLAN_MODE_TOOL_NAME, EXIT_PLAN_MODE_TOOL_NAME, + SCHEDULE_WORK_TOOL_NAME, ] as const; /** diff --git a/packages/core/src/utils/sessionUtils.ts b/packages/core/src/utils/sessionUtils.ts index 4803dd4f07..4361ada10c 100644 --- a/packages/core/src/utils/sessionUtils.ts +++ b/packages/core/src/utils/sessionUtils.ts @@ -29,10 +29,16 @@ function ensurePartArray(content: PartListUnion): Part[] { */ export function convertSessionToClientHistory( messages: ConversationRecord['messages'], + startIndex?: number, ): Array<{ role: 'user' | 'model'; parts: Part[] }> { const clientHistory: Array<{ role: 'user' | 'model'; parts: Part[] }> = []; - for (const msg of messages) { + const slice = + startIndex != null && startIndex > 0 + ? messages.slice(startIndex) + : messages; + + for (const msg of slice) { if (msg.type === 'info' || msg.type === 'error' || msg.type === 'warning') { continue; }