mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-05-15 06:12:50 -07:00
fix(core): ensure complete_task tool calls are recorded in chat history (#24437)
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -8,12 +8,10 @@ import { type AgentLoopContext } from '../config/agent-loop-context.js';
|
||||
import { reportError } from '../utils/errorReporting.js';
|
||||
import { GeminiChat, StreamEventType } from '../core/geminiChat.js';
|
||||
import {
|
||||
Type,
|
||||
type Content,
|
||||
type Part,
|
||||
type FunctionCall,
|
||||
type FunctionDeclaration,
|
||||
type Schema,
|
||||
} from '@google/genai';
|
||||
import { ToolRegistry } from '../tools/tool-registry.js';
|
||||
import { PromptRegistry } from '../prompts/prompt-registry.js';
|
||||
@@ -64,7 +62,6 @@ import { DEFAULT_GEMINI_MODEL, isAutoModel } from '../config/models.js';
|
||||
import type { RoutingContext } from '../routing/routingStrategy.js';
|
||||
import { parseThought } from '../utils/thoughtUtils.js';
|
||||
import { type z } from 'zod';
|
||||
import { zodToJsonSchema } from 'zod-to-json-schema';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
import { getModelConfigAlias } from './registry.js';
|
||||
import { getVersion } from '../utils/version.js';
|
||||
@@ -76,11 +73,12 @@ import {
|
||||
formatBackgroundCompletionForModel,
|
||||
} from '../utils/fastAckHelper.js';
|
||||
import type { InjectionSource } from '../config/injectionService.js';
|
||||
import { CompleteTaskTool } from '../tools/complete-task.js';
|
||||
import { COMPLETE_TASK_TOOL_NAME } from '../tools/definitions/base-declarations.js';
|
||||
|
||||
/** A callback function to report on agent activity. */
|
||||
export type ActivityCallback = (activity: SubagentActivityEvent) => void;
|
||||
|
||||
const TASK_COMPLETE_TOOL_NAME = 'complete_task';
|
||||
const GRACE_PERIOD_MS = 60 * 1000; // 1 min
|
||||
|
||||
/** The possible outcomes of a single agent turn. */
|
||||
@@ -256,6 +254,15 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
|
||||
agentToolRegistry.sortTools();
|
||||
|
||||
// Register the mandatory completion tool for this agent.
|
||||
agentToolRegistry.registerTool(
|
||||
new CompleteTaskTool(
|
||||
subagentMessageBus,
|
||||
definition.outputConfig,
|
||||
definition.processOutput,
|
||||
),
|
||||
);
|
||||
|
||||
// Get the parent tool call ID from context
|
||||
const toolContext = getToolCallContext();
|
||||
const parentCallId = toolContext?.callId;
|
||||
@@ -341,7 +348,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
// If the model stops calling tools without calling complete_task, it's an error.
|
||||
if (functionCalls.length === 0) {
|
||||
this.emitActivity('ERROR', {
|
||||
error: `Agent stopped calling tools but did not call '${TASK_COMPLETE_TOOL_NAME}' to finalize the session.`,
|
||||
error: `Agent stopped calling tools but did not call '${COMPLETE_TASK_TOOL_NAME}' to finalize the session.`,
|
||||
context: 'protocol_violation',
|
||||
errorType: SubagentActivityErrorType.GENERIC,
|
||||
});
|
||||
@@ -409,7 +416,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
default:
|
||||
throw new Error(`Unknown terminate reason: ${reason}`);
|
||||
}
|
||||
return `${explanation} You have one final chance to complete the task with a short grace period. You MUST call \`${TASK_COMPLETE_TOOL_NAME}\` immediately with your best answer and explain that your investigation was interrupted. Do not call any other tools.`;
|
||||
return `${explanation} You have one final chance to complete the task with a short grace period. You MUST call \`${COMPLETE_TASK_TOOL_NAME}\` immediately with your best answer and explain that your investigation was interrupted. Do not call any other tools.`;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -720,7 +727,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
// The finalResult was already set by executeTurn, but we re-emit just in case.
|
||||
finalResult =
|
||||
finalResult ||
|
||||
`Agent stopped calling tools but did not call '${TASK_COMPLETE_TOOL_NAME}'.`;
|
||||
`Agent stopped calling tools but did not call '${COMPLETE_TASK_TOOL_NAME}'.`;
|
||||
this.emitActivity('ERROR', {
|
||||
error: finalResult,
|
||||
context: 'protocol_violation',
|
||||
@@ -1031,23 +1038,42 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
aborted: boolean;
|
||||
}> {
|
||||
const allowedToolNames = new Set(this.toolRegistry.getAllToolNames());
|
||||
// Always allow the completion tool
|
||||
allowedToolNames.add(TASK_COMPLETE_TOOL_NAME);
|
||||
|
||||
let submittedOutput: string | null = null;
|
||||
let taskCompleted = false;
|
||||
let aborted = false;
|
||||
|
||||
// We'll separate complete_task from other tools
|
||||
const toolRequests: ToolCallRequestInfo[] = [];
|
||||
// Map to keep track of tool name by callId for activity emission
|
||||
const toolNameMap = new Map<string, string>();
|
||||
// Synchronous results (like complete_task or unauthorized calls)
|
||||
// Synchronous results (like unauthorized calls)
|
||||
const syncResults = new Map<string, Part>();
|
||||
|
||||
for (const [index, functionCall] of functionCalls.entries()) {
|
||||
const callId = functionCall.id ?? `${promptId}-${index}`;
|
||||
const args = functionCall.args ?? {};
|
||||
const { args, error: parseError } = this.parseToolArguments(functionCall);
|
||||
|
||||
if (parseError) {
|
||||
debugLogger.warn(`[LocalAgentExecutor] ${parseError}`);
|
||||
|
||||
syncResults.set(callId, {
|
||||
functionResponse: {
|
||||
name: functionCall.name,
|
||||
id: callId,
|
||||
response: { error: parseError },
|
||||
},
|
||||
});
|
||||
|
||||
this.emitActivity('ERROR', {
|
||||
context: 'tool_call',
|
||||
name: functionCall.name,
|
||||
callId,
|
||||
error: parseError,
|
||||
errorType: SubagentActivityErrorType.GENERIC,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
const toolName = functionCall.name as string;
|
||||
|
||||
@@ -1073,144 +1099,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
callId,
|
||||
});
|
||||
|
||||
if (toolName === TASK_COMPLETE_TOOL_NAME) {
|
||||
if (taskCompleted) {
|
||||
const error =
|
||||
'Task already marked complete in this turn. Ignoring duplicate call.';
|
||||
syncResults.set(callId, {
|
||||
functionResponse: {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
response: { error },
|
||||
id: callId,
|
||||
},
|
||||
});
|
||||
this.emitActivity('ERROR', {
|
||||
context: 'tool_call',
|
||||
name: toolName,
|
||||
error,
|
||||
errorType: SubagentActivityErrorType.GENERIC,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
const { outputConfig } = this.definition;
|
||||
taskCompleted = true; // Signal completion regardless of output presence
|
||||
|
||||
if (outputConfig) {
|
||||
const outputName = outputConfig.outputName;
|
||||
if (args[outputName] !== undefined) {
|
||||
const outputValue = args[outputName];
|
||||
const validationResult = outputConfig.schema.safeParse(outputValue);
|
||||
|
||||
if (!validationResult.success) {
|
||||
taskCompleted = false; // Validation failed, revoke completion
|
||||
const error = `Output validation failed: ${JSON.stringify(validationResult.error.flatten())}`;
|
||||
syncResults.set(callId, {
|
||||
functionResponse: {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
response: { error },
|
||||
id: callId,
|
||||
},
|
||||
});
|
||||
this.emitActivity('ERROR', {
|
||||
context: 'tool_call',
|
||||
name: toolName,
|
||||
error,
|
||||
errorType: SubagentActivityErrorType.GENERIC,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
|
||||
const validatedOutput = validationResult.data;
|
||||
if (this.definition.processOutput) {
|
||||
submittedOutput = this.definition.processOutput(validatedOutput);
|
||||
} else {
|
||||
submittedOutput =
|
||||
typeof outputValue === 'string'
|
||||
? outputValue
|
||||
: JSON.stringify(outputValue, null, 2);
|
||||
}
|
||||
syncResults.set(callId, {
|
||||
functionResponse: {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
response: { result: 'Output submitted and task completed.' },
|
||||
id: callId,
|
||||
},
|
||||
});
|
||||
this.emitActivity('TOOL_CALL_END', {
|
||||
name: toolName,
|
||||
id: callId,
|
||||
output: 'Output submitted and task completed.',
|
||||
});
|
||||
} else {
|
||||
// Failed to provide required output.
|
||||
taskCompleted = false; // Revoke completion status
|
||||
const error = `Missing required argument '${outputName}' for completion.`;
|
||||
syncResults.set(callId, {
|
||||
functionResponse: {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
response: { error },
|
||||
id: callId,
|
||||
},
|
||||
});
|
||||
this.emitActivity('ERROR', {
|
||||
context: 'tool_call',
|
||||
name: toolName,
|
||||
callId,
|
||||
error,
|
||||
errorType: SubagentActivityErrorType.GENERIC,
|
||||
});
|
||||
}
|
||||
} else {
|
||||
// No outputConfig - use default 'result' parameter
|
||||
const resultArg = args['result'];
|
||||
if (
|
||||
resultArg !== undefined &&
|
||||
resultArg !== null &&
|
||||
resultArg !== ''
|
||||
) {
|
||||
submittedOutput =
|
||||
typeof resultArg === 'string'
|
||||
? resultArg
|
||||
: JSON.stringify(resultArg, null, 2);
|
||||
syncResults.set(callId, {
|
||||
functionResponse: {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
response: { status: 'Result submitted and task completed.' },
|
||||
id: callId,
|
||||
},
|
||||
});
|
||||
this.emitActivity('TOOL_CALL_END', {
|
||||
name: toolName,
|
||||
id: callId,
|
||||
output: 'Result submitted and task completed.',
|
||||
});
|
||||
} else {
|
||||
// No result provided - this is an error for agents expected to return results
|
||||
taskCompleted = false; // Revoke completion
|
||||
const error =
|
||||
'Missing required "result" argument. You must provide your findings when calling complete_task.';
|
||||
syncResults.set(callId, {
|
||||
functionResponse: {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
response: { error },
|
||||
id: callId,
|
||||
},
|
||||
});
|
||||
this.emitActivity('ERROR', {
|
||||
context: 'tool_call',
|
||||
name: toolName,
|
||||
callId,
|
||||
error,
|
||||
errorType: SubagentActivityErrorType.GENERIC,
|
||||
});
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Handle standard tools
|
||||
// Handle unauthorized tools
|
||||
if (!allowedToolNames.has(toolName)) {
|
||||
const error = createUnauthorizedToolError(toolName);
|
||||
debugLogger.warn(`[LocalAgentExecutor] Blocked call: ${error}`);
|
||||
@@ -1274,6 +1163,22 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
output: call.response.resultDisplay,
|
||||
data: call.response.data,
|
||||
});
|
||||
|
||||
// Check if this was a completion tool call
|
||||
const isCompletionTool =
|
||||
call.request.name === COMPLETE_TASK_TOOL_NAME;
|
||||
const data = call.response.data;
|
||||
if (
|
||||
isCompletionTool &&
|
||||
!taskCompleted &&
|
||||
data?.['taskCompleted'] === true
|
||||
) {
|
||||
taskCompleted = true;
|
||||
const output = data['submittedOutput'];
|
||||
if (typeof output === 'string') {
|
||||
submittedOutput = output;
|
||||
}
|
||||
}
|
||||
} else if (call.status === 'error') {
|
||||
this.emitActivity('ERROR', {
|
||||
context: 'tool_call',
|
||||
@@ -1287,7 +1192,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
call.outcome === ToolConfirmationOutcome.Cancel;
|
||||
|
||||
if (isSoftRejection) {
|
||||
const error = `${SUBAGENT_REJECTED_ERROR_PREFIX} Please acknowledge this, rethink your strategy, and try a different approach. If you cannot proceed without the rejected operation, summarize the issue and use \`${TASK_COMPLETE_TOOL_NAME}\` to report your findings and the blocker.`;
|
||||
const error = `${SUBAGENT_REJECTED_ERROR_PREFIX} Please acknowledge this, rethink your strategy, and try a different approach. If you cannot proceed without the rejected operation, summarize the issue and use \`${COMPLETE_TASK_TOOL_NAME}\` to report your findings and the blocker.`;
|
||||
this.emitActivity('ERROR', {
|
||||
context: 'tool_call',
|
||||
name: toolName,
|
||||
@@ -1358,7 +1263,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
*/
|
||||
private prepareToolsList(): FunctionDeclaration[] {
|
||||
const toolsList: FunctionDeclaration[] = [];
|
||||
const { toolConfig, outputConfig } = this.definition;
|
||||
const { toolConfig } = this.definition;
|
||||
|
||||
if (toolConfig) {
|
||||
for (const toolRef of toolConfig.tools) {
|
||||
@@ -1375,43 +1280,6 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
),
|
||||
);
|
||||
|
||||
// Always inject complete_task.
|
||||
// Configure its schema based on whether output is expected.
|
||||
const completeTool: FunctionDeclaration = {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
description: outputConfig
|
||||
? 'Call this tool to submit your final answer and complete the task. This is the ONLY way to finish.'
|
||||
: 'Call this tool to submit your final findings and complete the task. This is the ONLY way to finish.',
|
||||
parameters: {
|
||||
type: Type.OBJECT,
|
||||
properties: {},
|
||||
required: [],
|
||||
},
|
||||
};
|
||||
|
||||
if (outputConfig) {
|
||||
const jsonSchema = zodToJsonSchema(outputConfig.schema);
|
||||
const {
|
||||
$schema: _$schema,
|
||||
definitions: _definitions,
|
||||
...schema
|
||||
} = jsonSchema;
|
||||
completeTool.parameters!.properties![outputConfig.outputName] =
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
schema as Schema;
|
||||
completeTool.parameters!.required!.push(outputConfig.outputName);
|
||||
} else {
|
||||
completeTool.parameters!.properties!['result'] = {
|
||||
type: Type.STRING,
|
||||
description:
|
||||
'Your final results or findings to return to the orchestrator. ' +
|
||||
'Ensure this is comprehensive and follows any formatting requested in your instructions.',
|
||||
};
|
||||
completeTool.parameters!.required!.push('result');
|
||||
}
|
||||
|
||||
toolsList.push(completeTool);
|
||||
|
||||
return toolsList;
|
||||
}
|
||||
|
||||
@@ -1445,15 +1313,15 @@ Important Rules:
|
||||
|
||||
if (this.definition.outputConfig) {
|
||||
finalPrompt += `
|
||||
* When you have completed your task, you MUST call the \`${TASK_COMPLETE_TOOL_NAME}\` tool with your structured output.
|
||||
* Do not call any other tools in the same turn as \`${TASK_COMPLETE_TOOL_NAME}\`.
|
||||
* When you have completed your task, you MUST call the \`${COMPLETE_TASK_TOOL_NAME}\` tool with your structured output.
|
||||
* Do not call any other tools in the same turn as \`${COMPLETE_TASK_TOOL_NAME}\`.
|
||||
* This is the ONLY way to complete your mission. If you stop calling tools without calling this, you have failed.`;
|
||||
} else {
|
||||
finalPrompt += `
|
||||
* When you have completed your task, you MUST call the \`${TASK_COMPLETE_TOOL_NAME}\` tool.
|
||||
* When you have completed your task, you MUST call the \`${COMPLETE_TASK_TOOL_NAME}\` tool.
|
||||
* You MUST include your final findings in the "result" parameter. This is how you return the necessary results for the task to be marked complete.
|
||||
* Ensure your findings are comprehensive and follow any specific formatting requirements provided in your instructions.
|
||||
* Do not call any other tools in the same turn as \`${TASK_COMPLETE_TOOL_NAME}\`.
|
||||
* Do not call any other tools in the same turn as \`${COMPLETE_TASK_TOOL_NAME}\`.
|
||||
* This is the ONLY way to complete your mission. If you stop calling tools without calling this, you have failed.`;
|
||||
}
|
||||
|
||||
@@ -1524,4 +1392,31 @@ Important Rules:
|
||||
}
|
||||
return chars.slice(0, 197).join('') + '...';
|
||||
}
|
||||
|
||||
/**
|
||||
* Parses the arguments for a tool call, handling both JSON strings and objects.
|
||||
*/
|
||||
private parseToolArguments(functionCall: FunctionCall): {
|
||||
args: Record<string, unknown>;
|
||||
error?: string;
|
||||
} {
|
||||
const args: Record<string, unknown> = {};
|
||||
if (typeof functionCall.args === 'string') {
|
||||
try {
|
||||
const parsed: unknown = JSON.parse(functionCall.args);
|
||||
if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) {
|
||||
Object.assign(args, parsed);
|
||||
}
|
||||
return { args };
|
||||
} catch (_) {
|
||||
return {
|
||||
args: {},
|
||||
error: `Failed to parse JSON arguments for tool "${functionCall.name}": ${functionCall.args}. Ensure you provide a valid JSON object.`,
|
||||
};
|
||||
}
|
||||
} else if (functionCall.args) {
|
||||
return { args: functionCall.args };
|
||||
}
|
||||
return { args: {} };
|
||||
}
|
||||
}
|
||||
|
||||
@@ -61,4 +61,10 @@ priority = 50
|
||||
[[rule]]
|
||||
toolName = "update_topic"
|
||||
decision = "allow"
|
||||
priority = 50
|
||||
|
||||
# Core agent lifecycle tool
|
||||
[[rule]]
|
||||
toolName = "complete_task"
|
||||
decision = "allow"
|
||||
priority = 50
|
||||
@@ -0,0 +1,160 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { CompleteTaskTool } from './complete-task.js';
|
||||
import { type MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import { z } from 'zod';
|
||||
|
||||
describe('CompleteTaskTool', () => {
|
||||
let mockMessageBus: MessageBus;
|
||||
|
||||
beforeEach(() => {
|
||||
mockMessageBus = {
|
||||
publish: vi.fn().mockResolvedValue(undefined),
|
||||
subscribe: vi.fn(),
|
||||
unsubscribe: vi.fn(),
|
||||
} as unknown as MessageBus;
|
||||
});
|
||||
|
||||
describe('Default Configuration (no outputConfig)', () => {
|
||||
let tool: CompleteTaskTool;
|
||||
|
||||
beforeEach(() => {
|
||||
tool = new CompleteTaskTool(mockMessageBus);
|
||||
});
|
||||
|
||||
it('should have correct metadata', () => {
|
||||
expect(tool.name).toBe('complete_task');
|
||||
expect(tool.displayName).toBe('Complete Task');
|
||||
});
|
||||
|
||||
it('should generate correct schema', () => {
|
||||
const schema = tool.getSchema();
|
||||
const parameters = schema.parametersJsonSchema as Record<string, unknown>;
|
||||
const properties = parameters['properties'] as Record<string, unknown>;
|
||||
|
||||
expect(properties).toHaveProperty('result');
|
||||
expect(parameters['required']).toContain('result');
|
||||
|
||||
const resultProp = properties['result'] as Record<string, unknown>;
|
||||
expect(resultProp['type']).toBe('string');
|
||||
});
|
||||
|
||||
it('should validate successfully with result', () => {
|
||||
const result = tool.validateToolParams({ result: 'Task done' });
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
it('should fail validation if result is missing', () => {
|
||||
const result = tool.validateToolParams({});
|
||||
expect(result).toContain("must have required property 'result'");
|
||||
});
|
||||
|
||||
it('should fail validation if result is only whitespace', () => {
|
||||
const result = tool.validateToolParams({ result: ' ' });
|
||||
expect(result).toContain(
|
||||
'Missing required "result" argument. You must provide your findings when calling complete_task.',
|
||||
);
|
||||
});
|
||||
|
||||
it('should execute and return correct data', async () => {
|
||||
const invocation = tool.build({ result: 'Success message' });
|
||||
const result = await invocation.execute(new AbortController().signal);
|
||||
|
||||
expect(result.data).toEqual({
|
||||
taskCompleted: true,
|
||||
submittedOutput: 'Success message',
|
||||
});
|
||||
expect(result.returnDisplay).toBe('Result submitted and task completed.');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Structured Configuration (with outputConfig)', () => {
|
||||
const schema = z.object({
|
||||
report: z.string(),
|
||||
score: z.number(),
|
||||
});
|
||||
const outputConfig = {
|
||||
outputName: 'my_output',
|
||||
description: 'The final report',
|
||||
schema,
|
||||
};
|
||||
let tool: CompleteTaskTool<typeof schema>;
|
||||
|
||||
beforeEach(() => {
|
||||
tool = new CompleteTaskTool(mockMessageBus, outputConfig);
|
||||
});
|
||||
|
||||
it('should generate schema based on outputConfig', () => {
|
||||
const toolSchema = tool.getSchema();
|
||||
|
||||
expect(toolSchema.parametersJsonSchema).toHaveProperty(
|
||||
'properties.my_output',
|
||||
);
|
||||
expect(toolSchema.parametersJsonSchema).toHaveProperty(
|
||||
'properties.my_output.type',
|
||||
'object',
|
||||
);
|
||||
expect(toolSchema.parametersJsonSchema).toHaveProperty(
|
||||
'properties.my_output.properties.report',
|
||||
);
|
||||
expect(toolSchema.parametersJsonSchema).toHaveProperty(
|
||||
'properties.my_output.properties.score',
|
||||
);
|
||||
expect(toolSchema.parametersJsonSchema).toHaveProperty(
|
||||
'required',
|
||||
expect.arrayContaining(['my_output']),
|
||||
);
|
||||
});
|
||||
|
||||
it('should validate successfully with correct structure', () => {
|
||||
const result = tool.validateToolParams({
|
||||
my_output: { report: 'All good', score: 100 },
|
||||
});
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
it('should fail validation if output is missing', () => {
|
||||
const result = tool.validateToolParams({});
|
||||
expect(result).toContain("must have required property 'my_output'");
|
||||
});
|
||||
|
||||
it('should fail validation if schema mismatch', () => {
|
||||
const result = tool.validateToolParams({
|
||||
my_output: { report: 'All good', score: 'not a number' },
|
||||
});
|
||||
expect(result).toContain('must be number');
|
||||
});
|
||||
|
||||
it('should execute and return structured data', async () => {
|
||||
const outputValue = { report: 'Final findings', score: 42 };
|
||||
const invocation = tool.build({ my_output: outputValue });
|
||||
const result = await invocation.execute(new AbortController().signal);
|
||||
|
||||
expect(result.data?.['taskCompleted']).toBe(true);
|
||||
expect(result.data?.['submittedOutput']).toBe(
|
||||
JSON.stringify(outputValue, null, 2),
|
||||
);
|
||||
});
|
||||
|
||||
it('should use processOutput if provided', async () => {
|
||||
const processOutput = (val: z.infer<typeof schema>) =>
|
||||
`Score was ${val.score}`;
|
||||
const toolWithProcess = new CompleteTaskTool(
|
||||
mockMessageBus,
|
||||
outputConfig,
|
||||
processOutput,
|
||||
);
|
||||
|
||||
const outputValue = { report: 'Final findings', score: 42 };
|
||||
const invocation = toolWithProcess.build({ my_output: outputValue });
|
||||
const result = await invocation.execute(new AbortController().signal);
|
||||
|
||||
expect(result.data?.['submittedOutput']).toBe('Score was 42');
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,179 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
BaseDeclarativeTool,
|
||||
BaseToolInvocation,
|
||||
type ToolResult,
|
||||
Kind,
|
||||
} from './tools.js';
|
||||
import {
|
||||
COMPLETE_TASK_TOOL_NAME,
|
||||
COMPLETE_TASK_DISPLAY_NAME,
|
||||
} from './definitions/base-declarations.js';
|
||||
import { type OutputConfig } from '../agents/types.js';
|
||||
import { type z } from 'zod';
|
||||
import { type MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import { zodToJsonSchema } from 'zod-to-json-schema';
|
||||
|
||||
/**
|
||||
* Tool for signaling task completion and optionally returning structured output.
|
||||
* This tool is specifically designed for use in subagent loops.
|
||||
*/
|
||||
export class CompleteTaskTool<
|
||||
TOutput extends z.ZodTypeAny = z.ZodTypeAny,
|
||||
> extends BaseDeclarativeTool<Record<string, unknown>, ToolResult> {
|
||||
static readonly Name = COMPLETE_TASK_TOOL_NAME;
|
||||
|
||||
constructor(
|
||||
messageBus: MessageBus,
|
||||
private readonly outputConfig?: OutputConfig<TOutput>,
|
||||
private readonly processOutput?: (output: z.infer<TOutput>) => string,
|
||||
) {
|
||||
super(
|
||||
CompleteTaskTool.Name,
|
||||
COMPLETE_TASK_DISPLAY_NAME,
|
||||
outputConfig
|
||||
? 'Call this tool to submit your final answer and complete the task. This is the ONLY way to finish.'
|
||||
: 'Call this tool to submit your final findings and complete the task. This is the ONLY way to finish.',
|
||||
Kind.Other,
|
||||
CompleteTaskTool.buildParameterSchema(outputConfig),
|
||||
messageBus,
|
||||
);
|
||||
}
|
||||
|
||||
private static buildParameterSchema(
|
||||
outputConfig?: OutputConfig<z.ZodTypeAny>,
|
||||
): unknown {
|
||||
if (outputConfig) {
|
||||
const jsonSchema = zodToJsonSchema(outputConfig.schema);
|
||||
const {
|
||||
$schema: _$schema,
|
||||
definitions: _definitions,
|
||||
...schema
|
||||
} = jsonSchema;
|
||||
return {
|
||||
type: 'object',
|
||||
properties: {
|
||||
[outputConfig.outputName]: schema,
|
||||
},
|
||||
required: [outputConfig.outputName],
|
||||
};
|
||||
}
|
||||
return {
|
||||
type: 'object',
|
||||
properties: {
|
||||
result: {
|
||||
type: 'string',
|
||||
description:
|
||||
'Your final results or findings to return to the orchestrator. ' +
|
||||
'Ensure this is comprehensive and follows any formatting requested in your instructions.',
|
||||
},
|
||||
},
|
||||
required: ['result'],
|
||||
};
|
||||
}
|
||||
|
||||
protected override validateToolParamValues(
|
||||
params: Record<string, unknown>,
|
||||
): string | null {
|
||||
if (this.outputConfig) {
|
||||
const outputName = this.outputConfig.outputName;
|
||||
if (params[outputName] === undefined) {
|
||||
return `Missing required argument '${outputName}' for completion.`;
|
||||
}
|
||||
|
||||
const validationResult = this.outputConfig.schema.safeParse(
|
||||
params[outputName],
|
||||
);
|
||||
if (!validationResult.success) {
|
||||
return `Output validation failed: ${JSON.stringify(validationResult.error.flatten())}`;
|
||||
}
|
||||
} else {
|
||||
const resultArg = params['result'];
|
||||
if (
|
||||
resultArg === undefined ||
|
||||
resultArg === null ||
|
||||
(typeof resultArg === 'string' && resultArg.trim() === '')
|
||||
) {
|
||||
return 'Missing required "result" argument. You must provide your findings when calling complete_task.';
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
protected createInvocation(
|
||||
params: Record<string, unknown>,
|
||||
messageBus: MessageBus,
|
||||
toolName: string,
|
||||
toolDisplayName: string,
|
||||
): CompleteTaskInvocation<TOutput> {
|
||||
return new CompleteTaskInvocation(
|
||||
params,
|
||||
messageBus,
|
||||
toolName,
|
||||
toolDisplayName,
|
||||
this.outputConfig,
|
||||
this.processOutput,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
export class CompleteTaskInvocation<
|
||||
TOutput extends z.ZodTypeAny = z.ZodTypeAny,
|
||||
> extends BaseToolInvocation<Record<string, unknown>, ToolResult> {
|
||||
constructor(
|
||||
params: Record<string, unknown>,
|
||||
messageBus: MessageBus,
|
||||
toolName: string,
|
||||
toolDisplayName: string,
|
||||
private readonly outputConfig?: OutputConfig<TOutput>,
|
||||
private readonly processOutput?: (output: z.infer<TOutput>) => string,
|
||||
) {
|
||||
super(params, messageBus, toolName, toolDisplayName);
|
||||
}
|
||||
|
||||
getDescription(): string {
|
||||
return 'Completing task and submitting results.';
|
||||
}
|
||||
|
||||
async execute(_signal: AbortSignal): Promise<ToolResult> {
|
||||
let submittedOutput: string | null = null;
|
||||
let outputValue: unknown;
|
||||
|
||||
if (this.outputConfig) {
|
||||
outputValue = this.params[this.outputConfig.outputName];
|
||||
if (this.processOutput) {
|
||||
// We validated the params in validateToolParamValues, so safe to cast
|
||||
submittedOutput = this.processOutput(outputValue as z.infer<TOutput>);
|
||||
} else {
|
||||
submittedOutput =
|
||||
typeof outputValue === 'string'
|
||||
? outputValue
|
||||
: JSON.stringify(outputValue, null, 2);
|
||||
}
|
||||
} else {
|
||||
outputValue = this.params['result'];
|
||||
submittedOutput =
|
||||
typeof outputValue === 'string'
|
||||
? outputValue
|
||||
: JSON.stringify(outputValue, null, 2);
|
||||
}
|
||||
|
||||
const returnDisplay = this.outputConfig
|
||||
? 'Output submitted and task completed.'
|
||||
: 'Result submitted and task completed.';
|
||||
|
||||
return {
|
||||
llmContent: returnDisplay,
|
||||
returnDisplay,
|
||||
data: {
|
||||
taskCompleted: true,
|
||||
submittedOutput,
|
||||
},
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -133,3 +133,7 @@ export const UPDATE_TOPIC_DISPLAY_NAME = 'Update Topic Context';
|
||||
export const TOPIC_PARAM_TITLE = 'title';
|
||||
export const TOPIC_PARAM_SUMMARY = 'summary';
|
||||
export const TOPIC_PARAM_STRATEGIC_INTENT = 'strategic_intent';
|
||||
|
||||
// -- complete_task --
|
||||
export const COMPLETE_TASK_TOOL_NAME = 'complete_task';
|
||||
export const COMPLETE_TASK_DISPLAY_NAME = 'Complete Task';
|
||||
|
||||
@@ -41,6 +41,8 @@ export {
|
||||
ENTER_PLAN_MODE_TOOL_NAME,
|
||||
UPDATE_TOPIC_TOOL_NAME,
|
||||
UPDATE_TOPIC_DISPLAY_NAME,
|
||||
COMPLETE_TASK_TOOL_NAME,
|
||||
COMPLETE_TASK_DISPLAY_NAME,
|
||||
// Shared parameter names
|
||||
PARAM_FILE_PATH,
|
||||
PARAM_DIR_PATH,
|
||||
|
||||
@@ -77,6 +77,8 @@ import {
|
||||
SKILL_PARAM_NAME,
|
||||
UPDATE_TOPIC_TOOL_NAME,
|
||||
UPDATE_TOPIC_DISPLAY_NAME,
|
||||
COMPLETE_TASK_TOOL_NAME,
|
||||
COMPLETE_TASK_DISPLAY_NAME,
|
||||
TOPIC_PARAM_TITLE,
|
||||
TOPIC_PARAM_SUMMARY,
|
||||
TOPIC_PARAM_STRATEGIC_INTENT,
|
||||
@@ -102,6 +104,8 @@ export {
|
||||
ENTER_PLAN_MODE_TOOL_NAME,
|
||||
UPDATE_TOPIC_TOOL_NAME,
|
||||
UPDATE_TOPIC_DISPLAY_NAME,
|
||||
COMPLETE_TASK_TOOL_NAME,
|
||||
COMPLETE_TASK_DISPLAY_NAME,
|
||||
// Shared parameter names
|
||||
PARAM_FILE_PATH,
|
||||
PARAM_DIR_PATH,
|
||||
@@ -264,6 +268,7 @@ export const ALL_BUILTIN_TOOL_NAMES = [
|
||||
ENTER_PLAN_MODE_TOOL_NAME,
|
||||
EXIT_PLAN_MODE_TOOL_NAME,
|
||||
UPDATE_TOPIC_TOOL_NAME,
|
||||
COMPLETE_TASK_TOOL_NAME,
|
||||
] as const;
|
||||
|
||||
/**
|
||||
|
||||
Reference in New Issue
Block a user