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
This commit is contained in:
Sandy Tao
2026-03-17 21:34:50 -07:00
parent a5a461c234
commit 1a6f730439
51 changed files with 2737 additions and 55 deletions
+4 -1
View File
@@ -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();
+8
View File
@@ -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,
});
}
+12
View File
@@ -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',
+453
View File
@@ -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<string, A2ATask>();
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<string, unknown>;
}
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<string> {
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<string, unknown>,
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<string, unknown>,
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<ExternalListenerResult> {
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);
});
});
}
+1
View File
@@ -513,6 +513,7 @@ describe('gemini.tsx main function kitty protocol', () => {
rawOutput: undefined,
acceptRawOutputRisk: undefined,
isCommand: undefined,
forever: undefined,
});
await act(async () => {
+20
View File
@@ -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());
}
+1
View File
@@ -222,6 +222,7 @@ export async function runNonInteractive({
await geminiClient.resumeChat(
convertSessionToClientHistory(
resumedSessionData.conversation.messages,
resumedSessionData.conversation.lastCompressionIndex,
),
resumedSessionData,
);
@@ -173,6 +173,14 @@ export const createMockConfig = (overrides: Partial<Config> = {}): 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;
+1
View File
@@ -82,6 +82,7 @@ describe('App', () => {
updateItem: vi.fn(),
clearItems: vi.fn(),
loadHistory: vi.fn(),
pruneItems: vi.fn(),
},
history: [],
pendingHistoryItems: [],
+87
View File
@@ -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<number | null>(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,
],
);
@@ -61,6 +61,10 @@ vi.mock('./StatusDisplay.js', () => ({
StatusDisplay: () => <Text>StatusDisplay</Text>,
}));
vi.mock('./ScheduledWorkDisplay.js', () => ({
ScheduledWorkDisplay: () => null,
}));
vi.mock('./ToastDisplay.js', () => ({
ToastDisplay: () => <Text>ToastDisplay</Text>,
shouldShowToast: (uiState: UIState) =>
@@ -202,6 +206,7 @@ const createMockUIState = (overrides: Partial<UIState> = {}): UIState =>
activeHooks: [],
isBackgroundShellVisible: false,
embeddedShellFocused: false,
a2aListenerPort: null,
quota: {
userTier: undefined,
stats: undefined,
@@ -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 && <ShortcutsHelp />}
{showUiDetails && <HorizontalLine />}
<ScheduledWorkDisplay />
{showUiDetails && (
<Box
justifyContent={
@@ -0,0 +1,66 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type React from 'react';
import { useState, useEffect } from 'react';
import { Box, Text } from 'ink';
import { theme } from '../semantic-colors.js';
import { useConfig } from '../contexts/ConfigContext.js';
import type { ScheduledItem } from '@google/gemini-cli-core';
/**
* Displays all pending scheduled work items above the context summary.
* Only renders when there are pending items.
*/
export const ScheduledWorkDisplay: React.FC = () => {
const config = useConfig();
const [items, setItems] = useState<readonly ScheduledItem[]>([]);
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 (
<Box flexDirection="column" paddingX={1}>
<Text color={theme.text.secondary}>
Scheduled work ({items.length}):
</Text>
{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 (
<Text key={item.id} color={theme.text.secondary}>
{' '}
{timeStr} (in {diffMins}m) {truncatedPrompt}
</Text>
);
})}
</Box>
);
};
@@ -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();
});
});
@@ -28,6 +28,12 @@ export const StatusDisplay: React.FC<StatusDisplayProps> = ({
return <Text color={theme.status.error}>|_|</Text>;
}
if (uiState.a2aListenerPort !== null) {
return (
<Text color={theme.text.accent}> A2A :{uiState.a2aListenerPort}</Text>
);
}
if (
uiState.activeHooks.length > 0 &&
settings.merged.hooksConfig.notifications
@@ -223,6 +223,7 @@ export interface UIState {
showIsExpandableHint: boolean;
hintMode: boolean;
hintBuffer: string;
a2aListenerPort: number | null;
transientMessage: {
text: string;
type: TransientMessageType;
@@ -57,6 +57,7 @@ describe('handleCreditsFlow', () => {
updateItem: vi.fn(),
clearItems: vi.fn(),
loadHistory: vi.fn(),
pruneItems: vi.fn(),
};
isDialogPending = { current: false };
mockSetOverageMenuRequest = vi.fn();
+79 -2
View File
@@ -220,6 +220,7 @@ export const useGeminiStream = (
terminalHeight: number,
isShellFocused?: boolean,
consumeUserHint?: () => string | null,
pruneItems?: () => void,
) => {
const [initError, setInitError] = useState<string | null>(null);
const [retryStatus, setRetryStatus] = useState<RetryAttemptPayload | null>(
@@ -256,6 +257,7 @@ export const useGeminiStream = (
const [_isFirstToolInGroup, isFirstToolInGroupRef, setIsFirstToolInGroup] =
useStateAndRef<boolean>(true);
const processedMemoryToolsRef = useRef<Set<string>>(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<ReturnType<typeof setTimeout> | 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,
@@ -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);
});
});
});
+28 -1
View File
@@ -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],
);
}
@@ -91,6 +91,7 @@ describe('useIncludeDirsTrust', () => {
updateItem: vi.fn(),
clearItems: vi.fn(),
loadHistory: vi.fn(),
pruneItems: vi.fn(),
};
mockSetCustomDialog = vi.fn();
});
@@ -85,6 +85,7 @@ describe('useQuotaAndFallback', () => {
updateItem: vi.fn(),
clearItems: vi.fn(),
loadHistory: vi.fn(),
pruneItems: vi.fn(),
};
mockSetModelSwitchedFromQuotaError = vi.fn();
mockOnShowAuthSelection = vi.fn();
@@ -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) {
@@ -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);
});
});
});
+17 -1
View File
@@ -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,
);
}
+4
View File
@@ -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<AppEvents>();
+13 -1
View File
@@ -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) {
+33 -6
View File
@@ -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<boolean | undefined> {
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()) {
+3
View File
@@ -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,
+7
View File
@@ -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();
+2
View File
@@ -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;
},
+3
View File
@@ -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 {
@@ -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();
});
});
});
+6 -1
View File
@@ -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']);
@@ -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<string, unknown>)[
'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<string, unknown>)[
'updateSystemInstructionIfInitialized'
] = updateSpy;
await hookEventHandler.fireAfterAgentEvent(
'test prompt',
'test response',
);
expect(updateSpy).not.toHaveBeenCalled();
});
});
});
@@ -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<AggregatedHookResult> {
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<AggregatedHookResult> {
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
}
+2
View File
@@ -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
+8 -1
View File
@@ -232,8 +232,15 @@ export class HookSystem {
async firePreCompressEvent(
trigger: PreCompressTrigger,
history: Array<{ role: string; parts: Array<{ text?: string }> }>,
): Promise<AggregatedHookResult | undefined> {
return this.hookEventHandler.firePreCompressEvent(trigger);
return this.hookEventHandler.firePreCompressEvent(trigger, history);
}
async fireIdleEvent(
idleSeconds: number,
): Promise<AggregatedHookResult | undefined> {
return this.hookEventHandler.fireIdleEvent(idleSeconds);
}
async fireBeforeAgentEvent(
+3
View File
@@ -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);
});
+31 -1
View File
@@ -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<string, unknown>;
}
@@ -184,6 +189,7 @@ export class DefaultHookOutput implements HookOutput {
systemMessage?: string;
decision?: HookDecision;
reason?: string;
refreshContext?: boolean;
hookSpecificOutput?: Record<string, unknown>;
constructor(data: Partial<HookOutput> = {}) {
@@ -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;
};
}
/**
+1
View File
@@ -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';
@@ -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;
});
@@ -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<Config['getHookSystem']>,
);
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<Config['getHookSystem']>,
);
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<Config['getHookSystem']>,
);
await service.compress(
mockChat,
mockPromptId,
true,
mockModel,
mockConfig,
false,
);
expect(mockHookSystem.firePreCompressEvent).toHaveBeenCalledWith(
'manual',
[
{ role: 'user', parts: [{ text: 'hello' }] },
{ role: 'model', parts: [{ text: 'world' }] },
],
);
});
});
});
@@ -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(
@@ -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<string, Part[]>();
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<string, Part[]>();
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) {
@@ -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();
});
});
});
@@ -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<WorkSchedulerEvents> {
private items: ScheduledItem[] = [];
private timer: ReturnType<typeof setTimeout> | 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);
}
}
@@ -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.',
);
});
});
});
+210
View File
@@ -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<ToolResult> {
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}"`,
};
}
}
+2
View File
@@ -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;
/**
+7 -1
View File
@@ -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;
}