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

packages/core/src/tools/complete-task.test.ts

@@ -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');
+    });
+  });
+});

packages/core/src/tools/complete-task.ts

@@ -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,
+      },
+    };
+  }
+}

packages/core/src/tools/definitions/base-declarations.ts

@@ -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';

packages/core/src/tools/definitions/coreTools.ts

@@ -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,

packages/core/src/tools/tool-names.ts

@@ -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;
 
 /**