fix(core): ensure complete_task tool calls are recorded in chat history (#24437)

This commit is contained in:
Abhi
2026-04-01 15:53:46 -04:00
committed by GitHub
parent 7cae9a18c1
commit 9054f828c4
8 changed files with 819 additions and 316 deletions
File diff suppressed because it is too large Load Diff
+88 -193
View File
@@ -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: {} };
}
}