diff options
| -rw-r--r-- | packages/cli/src/config/config.ts | 1 | ||||
| -rw-r--r-- | packages/cli/src/config/settings.ts | 1 | ||||
| -rw-r--r-- | packages/cli/src/services/CommandService.test.ts | 1 | ||||
| -rw-r--r-- | packages/cli/src/services/CommandService.ts | 2 | ||||
| -rw-r--r-- | packages/cli/src/ui/App.test.tsx | 3 | ||||
| -rw-r--r-- | packages/cli/src/ui/commands/backgroundCommand.ts | 262 | ||||
| -rw-r--r-- | packages/core/src/background/backgroundAgent.ts | 126 | ||||
| -rw-r--r-- | packages/core/src/background/backgroundManager.ts | 40 | ||||
| -rw-r--r-- | packages/core/src/background/types.ts | 107 | ||||
| -rw-r--r-- | packages/core/src/config/config.ts | 16 | ||||
| -rw-r--r-- | packages/core/src/index.ts | 3 | ||||
| -rw-r--r-- | packages/core/src/tools/mcp-tool.ts | 4 |
12 files changed, 563 insertions, 3 deletions
diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 2d33daa3..84ca6815 100644 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -382,6 +382,7 @@ export async function loadCliConfig( toolCallCommand: settings.toolCallCommand, mcpServerCommand: settings.mcpServerCommand, mcpServers, + backgroundAgents: settings.backgroundAgents, userMemory: memoryContent, geminiMdFileCount: fileCount, approvalMode: argv.yolo || false ? ApprovalMode.YOLO : ApprovalMode.DEFAULT, diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts index 604e89dc..2da2e592 100644 --- a/packages/cli/src/config/settings.ts +++ b/packages/cli/src/config/settings.ts @@ -64,6 +64,7 @@ export interface Settings { toolCallCommand?: string; mcpServerCommand?: string; mcpServers?: Record<string, MCPServerConfig>; + backgroundAgents?: Record<string, MCPServerConfig>; allowMCPServers?: string[]; excludeMCPServers?: string[]; showMemoryUsage?: boolean; diff --git a/packages/cli/src/services/CommandService.test.ts b/packages/cli/src/services/CommandService.test.ts index d03bf988..5945e3f6 100644 --- a/packages/cli/src/services/CommandService.test.ts +++ b/packages/cli/src/services/CommandService.test.ts @@ -96,6 +96,7 @@ describe('CommandService', () => { mockConfig = { getIdeMode: vi.fn(), getCheckpointingEnabled: vi.fn(), + getBackgroundAgentManager: vi.fn(), } as unknown as Mocked<Config>; vi.mocked(ideCommand).mockReturnValue(null); vi.mocked(restoreCommand).mockReturnValue(null); diff --git a/packages/cli/src/services/CommandService.ts b/packages/cli/src/services/CommandService.ts index def8cfcc..9db4e9e6 100644 --- a/packages/cli/src/services/CommandService.ts +++ b/packages/cli/src/services/CommandService.ts @@ -7,6 +7,7 @@ import { Config } from '@google/gemini-cli-core'; import { SlashCommand } from '../ui/commands/types.js'; import { memoryCommand } from '../ui/commands/memoryCommand.js'; +import { backgroundCommand } from '../ui/commands/backgroundCommand.js'; import { helpCommand } from '../ui/commands/helpCommand.js'; import { clearCommand } from '../ui/commands/clearCommand.js'; import { corgiCommand } from '../ui/commands/corgiCommand.js'; @@ -33,6 +34,7 @@ const loadBuiltInCommands = async ( const allCommands = [ aboutCommand, authCommand, + backgroundCommand(config), bugCommand, chatCommand, clearCommand, diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index e03c80ae..24d30f74 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -9,6 +9,7 @@ import { render } from 'ink-testing-library'; import { AppWrapper as App } from './App.js'; import { Config as ServerConfig, + BackgroundAgentManager, MCPServerConfig, ApprovalMode, ToolRegistry, @@ -51,6 +52,7 @@ interface MockServerConfig { getSandbox: Mock<() => SandboxConfig | undefined>; getTargetDir: Mock<() => string>; getToolRegistry: Mock<() => ToolRegistry>; // Use imported ToolRegistry type + getBackgroundAgentManager: Mock<() => BackgroundAgentManager>; getDebugMode: Mock<() => boolean>; getQuestion: Mock<() => string | undefined>; getFullContext: Mock<() => boolean>; @@ -117,6 +119,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { getSandbox: vi.fn(() => opts.sandbox), getTargetDir: vi.fn(() => opts.targetDir || '/test/dir'), getToolRegistry: vi.fn(() => ({}) as ToolRegistry), // Simple mock + getBackgroundAgentManager: vi.fn(() => new BackgroundAgentManager([])), getDebugMode: vi.fn(() => opts.debugMode || false), getQuestion: vi.fn(() => opts.question), getFullContext: vi.fn(() => opts.fullContext ?? false), diff --git a/packages/cli/src/ui/commands/backgroundCommand.ts b/packages/cli/src/ui/commands/backgroundCommand.ts new file mode 100644 index 00000000..cbeb18df --- /dev/null +++ b/packages/cli/src/ui/commands/backgroundCommand.ts @@ -0,0 +1,262 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { SlashCommand, CommandContext } from './types.js'; +import { + Config, + BackgroundAgentMessage, + partListUnionToString, +} from '@google/gemini-cli-core'; + +const MAX_STATUS_MESSAGE_LENGTH = 100; + +function toMessageString(message?: BackgroundAgentMessage): string { + return partListUnionToString(message?.parts ?? []).trim(); +} + +function toOneliner(input: string, maxlength: number) { + let output = input.replace(/\r?\n|\r/g, ' '); + if (output.length > maxlength) { + output = output.substring(0, maxlength) + '...'; + } + return output; +} + +function getActiveAgent(context: CommandContext) { + const agent = + context.services.config?.getBackgroundAgentManager()?.activeAgent; + if (!agent) { + throw Error('There is no active background agent.'); + } + return agent; +} + +function addClientHistory(context: CommandContext, text: string) { + context.services.config!.getGeminiClient().addHistory({ + role: 'user', + parts: [{ text }], + }); + + context.services.config!.getGeminiClient().addHistory({ + role: 'model', + parts: [{ text: 'Got it.' }], + }); +} + +const startSubcommand: SlashCommand = { + name: 'start', + description: + 'Start a new task with the provided prompt. Usage: /bg start <prompt>', + action: async (context, args) => { + if (!args || args.trim() === '') { + return { + type: 'message', + messageType: 'error', + content: 'The `start` command requires a prompt.', + }; + } + + const agent = getActiveAgent(context); + const task = await agent.startTask(args); + + addClientHistory( + context, + `I started a background task with id '${task.id}' and prompt:\n${args}`, + ); + + return { + type: 'message', + messageType: 'info', + content: `Started background task with id '${task.id}' and prompt:\n${args}`, + }; + }, +}; + +const stopSubcommand: SlashCommand = { + name: 'stop', + description: 'Stops a running task. Usage: /bg stop <task_id>', + action: async (context, args) => { + if (!args || args.trim() === '') { + return { + type: 'message', + messageType: 'error', + content: 'The `stop` command requires a task id.', + }; + } + const agent = getActiveAgent(context); + await agent.cancelTask(args); + addClientHistory(context, `I canceled the background task with id ${args}`); + return { + type: 'message', + messageType: 'info', + content: `Stopped background task with id ${args}.`, + }; + }, +}; + +const listSubcommand: SlashCommand = { + name: 'list', + description: 'List all tasks', + action: async (context, args) => { + if (args && args.trim() !== '') { + return { + type: 'message', + messageType: 'error', + content: 'The `list` command takes no arguments.', + }; + } + + const agent = getActiveAgent(context); + const tasks = await agent.listTasks(); + let content: string; + if (tasks.length === 0) { + content = 'No background tasks found.'; + } else { + const taskList = tasks + .map((task) => { + const shortStatus = toOneliner( + toMessageString(task.status.message), + MAX_STATUS_MESSAGE_LENGTH, + ); + return ` - ${task.id}: (${task.status.state}) ${shortStatus}`; + }) + .join('\n'); + content = `Background tasks:\n${taskList}`; + } + return { + type: 'message', + messageType: 'info', + content, + }; + }, +}; + +const getSubcommand: SlashCommand = { + name: 'get', + description: 'View a task. Usage: /bg get <task_id>', + action: async (context, args) => { + if (!args || args.trim() === '') { + return { + type: 'message', + messageType: 'error', + content: 'The `get` command requires a task id.', + }; + } + const agent = getActiveAgent(context); + const task = await agent.getTask(args); + const content = `Task Details for ${task.id}: +Status: (${task.status.state}) ${toMessageString(task.status.message)}}`; + + return { + type: 'message', + messageType: 'info', + content, + }; + }, +}; + +const logsSubcommand: SlashCommand = { + name: 'logs', + description: "View a task's recent logs. Usage: /bg log <task_id>", + action: async (context, args) => { + if (!args || args.trim() === '') { + return { + type: 'message', + messageType: 'error', + content: 'The `log` command requires a task id.', + }; + } + const agent = getActiveAgent(context); + const task = await agent.getTask(args, 5); + const contents = [ + `Task logs for ${task.id}. status: (${task.status.state})`, + ]; + (task.history ?? []).forEach((message) => { + contents.push(toMessageString(message)); + }); + return { + type: 'message', + messageType: 'info', + content: contents.join('\n\n'), + }; + }, +}; + +const messageSubcommand: SlashCommand = { + name: 'message', + description: + 'Send a message to a task. Usage: /bg message <task_id> <message>', + action: async (context, args) => { + if (!args || args.trim() === '' || !args.trim().includes(' ')) { + return { + type: 'message', + messageType: 'error', + content: 'The `message` command requires a task id and a message.', + }; + } + + const firstSpaceIndex = args.indexOf(' '); + const id = args.substring(0, firstSpaceIndex); + const message = args.substring(firstSpaceIndex + 1); + + const agent = getActiveAgent(context); + await agent.messageTask(id, message); + addClientHistory( + context, + `I sent a message to the background task with id '${id}':\n${message}`, + ); + + return { + type: 'message', + messageType: 'info', + content: `Sent a message to the background task with id '${id}':\n${message}`, + }; + }, +}; + +const deleteSubcommand: SlashCommand = { + name: 'delete', + description: 'Deletes a task. Usage: /bg delete <task_id>', + action: async (context, args) => { + if (!args) { + return { + type: 'message', + messageType: 'error', + content: 'The `delete` command requires a task id.', + }; + } + const agent = getActiveAgent(context); + await agent.deleteTask(args); + addClientHistory(context, `I deleted the background task with id ${args}`); + return { + type: 'message', + messageType: 'info', + content: `Task ${args} deleted.`, + }; + }, +}; + +export const backgroundCommand = ( + config: Config | null, +): SlashCommand | null => { + if (!config?.getBackgroundAgentManager()?.activeAgent) { + return null; + } + return { + name: 'background', + altName: 'bg', + description: "Commands for managing the background agent's tasks", + subCommands: [ + startSubcommand, + stopSubcommand, + listSubcommand, + getSubcommand, + logsSubcommand, + messageSubcommand, + deleteSubcommand, + ], + }; +}; diff --git a/packages/core/src/background/backgroundAgent.ts b/packages/core/src/background/backgroundAgent.ts new file mode 100644 index 00000000..008010cd --- /dev/null +++ b/packages/core/src/background/backgroundAgent.ts @@ -0,0 +1,126 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { MCPServerConfig } from '../config/config.js'; +import { connectToMcpServer, discoverTools } from '../tools/mcp-client.js'; +import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; +import { + BackgroundAgentTasksResponseSchema, + BackgroundAgentTaskResponseSchema, + BackgroundAgentTask, +} from './types.js'; + +export async function loadBackgroundAgent( + name: string, + config: MCPServerConfig, + debugMode: boolean, +): Promise<BackgroundAgent> { + const server = await connectToMcpServer(name, config, debugMode); + try { + const tools = await discoverTools(name, config, server); + return new BackgroundAgent(name, tools); + } catch (error) { + await server.close(); + throw error; + } +} + +export class BackgroundAgent { + readonly startTaskTool: DiscoveredMCPTool; + readonly getTaskTool: DiscoveredMCPTool; + readonly listTasksTool: DiscoveredMCPTool; + readonly messageTaskTool: DiscoveredMCPTool; + readonly deleteTaskTool: DiscoveredMCPTool; + readonly cancelTaskTool: DiscoveredMCPTool; + + constructor( + readonly serverName: string, + tools: DiscoveredMCPTool[], + ) { + const getToolOrFail = (name: string): DiscoveredMCPTool => { + for (const tool of tools) { + if (tool.serverToolName === name) { + return tool; + } + } + throw new Error(`missing expected tool: ${name}`); + }; + + this.startTaskTool = getToolOrFail('startTask'); + this.getTaskTool = getToolOrFail('getTask'); + this.listTasksTool = getToolOrFail('listTasks'); + this.messageTaskTool = getToolOrFail('messageTask'); + this.deleteTaskTool = getToolOrFail('deleteTask'); + this.cancelTaskTool = getToolOrFail('cancelTask'); + } + + async startTask(prompt: string): Promise<BackgroundAgentTask> { + const resp = await this.callTool(this.startTaskTool, { + prompt: { + role: 'user', + parts: [{ text: prompt }], + }, + }); + const taskResp = await BackgroundAgentTaskResponseSchema.parseAsync(resp); + return taskResp.structuredContent; + } + + async getTask( + id: string, + historyLength?: number, + ): Promise<BackgroundAgentTask> { + const resp = await this.callTool(this.getTaskTool, { + id, + historyLength, + }); + const taskResp = await BackgroundAgentTaskResponseSchema.parseAsync(resp); + return taskResp.structuredContent; + } + + async listTasks(): Promise<BackgroundAgentTask[]> { + const resp = await this.callTool(this.listTasksTool, {}); + const tasksResp = await BackgroundAgentTasksResponseSchema.parseAsync(resp); + return tasksResp.structuredContent; + } + + async messageTask(id: string, message: string) { + await this.callTool(this.messageTaskTool, { + id, + message: { + role: 'user', + parts: [{ text: message }], + }, + }); + } + + async deleteTask(id: string) { + await this.callTool(this.deleteTaskTool, { id }); + } + + async cancelTask(id: string) { + await this.callTool(this.cancelTaskTool, { id }); + } + + private async callTool( + tool: DiscoveredMCPTool, + params: Record<string, unknown>, + ): Promise<Record<string, unknown>> { + const { llmContent: parts } = await tool.execute(params); + if ( + !Array.isArray(parts) || + parts.length !== 1 || + typeof parts[0] !== 'object' || + parts[0]?.functionResponse?.response === undefined + ) { + throw new Error('Expected exactly one part with a functionResponse'); + } + const resp = parts[0].functionResponse.response; + if ('isError' in resp && resp.isError) { + throw new Error(`Error calling ${tool.displayName}: ${resp}`); + } + return resp; + } +} diff --git a/packages/core/src/background/backgroundManager.ts b/packages/core/src/background/backgroundManager.ts new file mode 100644 index 00000000..a3ec526c --- /dev/null +++ b/packages/core/src/background/backgroundManager.ts @@ -0,0 +1,40 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { MCPServerConfig } from '../config/config.js'; +import { BackgroundAgent, loadBackgroundAgent } from './backgroundAgent.js'; + +export async function loadBackgroundAgentManager( + backgroundAgentConfigs: Record<string, MCPServerConfig> | undefined, + debugMode: boolean, +): Promise<BackgroundAgentManager> { + const agents = await Promise.all( + Object.entries(backgroundAgentConfigs ?? {}).map(([name, config]) => + loadBackgroundAgent(name, config, debugMode).catch((error) => { + console.error(`Error loading background agent '${name}': ${error}`); + return null; + }), + ), + ).then((agents) => agents.filter((agent) => agent !== null)); + return new BackgroundAgentManager(agents); +} + +export class BackgroundAgentManager { + // The active agent. May be empty if none are confgured. + activeAgent?: BackgroundAgent; + + constructor(readonly backgroundAgents: BackgroundAgent[]) { + if (backgroundAgents.length !== 0) { + this.activeAgent = backgroundAgents[0]; + } + } + + setActiveAgentByName(name: string) { + this.activeAgent = this.backgroundAgents.find( + (agent) => agent.serverName === name, + ); + } +} diff --git a/packages/core/src/background/types.ts b/packages/core/src/background/types.ts new file mode 100644 index 00000000..60927af6 --- /dev/null +++ b/packages/core/src/background/types.ts @@ -0,0 +1,107 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { z } from 'zod'; +import { Outcome, Language, FunctionResponseScheduling } from '@google/genai'; + +// Should conform to Part in @google/genai +export const PartSchema = z.object({ + videoMetadata: z + .object({ + fps: z.number().optional(), + endOffset: z.string().optional(), + startOffset: z.string().optional(), + }) + .optional(), + thought: z.boolean().optional(), + inlineData: z + .object({ + displayName: z.string().optional(), + data: z.string(), + mimeType: z.string(), + }) + .optional(), + fileData: z + .object({ + displayName: z.string().optional(), + fileUri: z.string(), + mimeType: z.string(), + }) + .optional(), + thoughtSignature: z.string().optional(), + codeExecutionResult: z + .object({ + outcome: z.nativeEnum(Outcome).optional(), + output: z.string().optional(), + }) + .optional(), + executableCode: z + .object({ + code: z.string().optional(), + language: z.nativeEnum(Language).optional(), + }) + .optional(), + functionCall: z + .object({ + id: z.string().optional(), + args: z.record(z.unknown()).optional(), + name: z.string(), + }) + .optional(), + functionResponse: z + .object({ + willContinue: z.boolean().optional(), + scheduling: z.nativeEnum(FunctionResponseScheduling).optional(), + id: z.string().optional(), + name: z.string(), + response: z.record(z.unknown()).optional(), + }) + .optional(), + text: z.string().optional(), +}); + +export const BackgroundAgentMessageSchema = z.object({ + role: z.enum(['user', 'agent']).describe('The role of the sender.'), + parts: z.array(PartSchema).describe('The parts of the message.'), +}); + +export const BackgroundAgentTaskStatusSchema = z.object({ + state: z.enum([ + 'submitted', + 'working', + 'input-required', + 'completed', + 'failed', + ]), + message: BackgroundAgentMessageSchema.describe( + 'Message describing the state of the task.', + ).optional(), +}); + +export const BackgroundAgentTaskSchema = z.object({ + id: z.string().describe('The id of the task. Must match `[a-zA-Z0-9.-_]+`'), + status: BackgroundAgentTaskStatusSchema.describe( + 'The current status of the task.', + ), + history: z + .array(BackgroundAgentMessageSchema) + .describe('Recent history of messages associated with this task') + .optional(), +}); + +export type BackgroundAgentMessage = z.infer< + typeof BackgroundAgentMessageSchema +>; + +export type BackgroundAgentTask = z.infer<typeof BackgroundAgentTaskSchema>; + +export const BackgroundAgentTaskResponseSchema = z.object({ + structuredContent: BackgroundAgentTaskSchema, +}); + +export const BackgroundAgentTasksResponseSchema = z.object({ + structuredContent: z.array(BackgroundAgentTaskSchema), +}); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index f81b3e32..5d02f269 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -45,6 +45,10 @@ import { DEFAULT_GEMINI_FLASH_MODEL, } from './models.js'; import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js'; +import { + BackgroundAgentManager, + loadBackgroundAgentManager, +} from '../background/backgroundManager.js'; export enum ApprovalMode { DEFAULT = 'default', @@ -127,6 +131,7 @@ export interface ConfigParameters { toolCallCommand?: string; mcpServerCommand?: string; mcpServers?: Record<string, MCPServerConfig>; + backgroundAgents?: Record<string, MCPServerConfig>; userMemory?: string; geminiMdFileCount?: number; approvalMode?: ApprovalMode; @@ -158,6 +163,7 @@ export interface ConfigParameters { export class Config { private toolRegistry!: ToolRegistry; + private backgroundAgentManager?: BackgroundAgentManager; private readonly sessionId: string; private contentGeneratorConfig!: ContentGeneratorConfig; private readonly embeddingModel: string; @@ -172,6 +178,7 @@ export class Config { private readonly toolCallCommand: string | undefined; private readonly mcpServerCommand: string | undefined; private readonly mcpServers: Record<string, MCPServerConfig> | undefined; + private readonly backgroundAgents?: Record<string, MCPServerConfig>; private userMemory: string; private geminiMdFileCount: number; private approvalMode: ApprovalMode; @@ -224,6 +231,7 @@ export class Config { this.toolCallCommand = params.toolCallCommand; this.mcpServerCommand = params.mcpServerCommand; this.mcpServers = params.mcpServers; + this.backgroundAgents = params.backgroundAgents; this.userMemory = params.userMemory ?? ''; this.geminiMdFileCount = params.geminiMdFileCount ?? 0; this.approvalMode = params.approvalMode ?? ApprovalMode.DEFAULT; @@ -281,6 +289,10 @@ export class Config { if (this.getCheckpointingEnabled()) { await this.getGitService(); } + this.backgroundAgentManager = await loadBackgroundAgentManager( + this.backgroundAgents, + this.debugMode, + ); this.toolRegistry = await this.createToolRegistry(); } @@ -406,6 +418,10 @@ export class Config { return this.mcpServers; } + getBackgroundAgentManager(): BackgroundAgentManager | undefined { + return this.backgroundAgentManager; + } + getUserMemory(): string { return this.userMemory; } diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 0aab6106..2e85deff 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -24,6 +24,9 @@ export * from './code_assist/oauth2.js'; export * from './code_assist/server.js'; export * from './code_assist/types.js'; +export * from './background/types.js'; +export * from './background/backgroundManager.js'; + // Export utilities export * from './utils/paths.js'; export * from './utils/schemaValidator.js'; diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index 9916d7f9..2cb124ed 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -113,9 +113,7 @@ export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> { args: params, }, ]; - - const responseParts: Part[] = await this.mcpTool.callTool(functionCalls); - + const responseParts = await this.mcpTool.callTool(functionCalls); return { llmContent: responseParts, returnDisplay: getStringifiedResultForDisplay(responseParts), |