diff --git a/.changeset/guard-background-questions.md b/.changeset/guard-background-questions.md new file mode 100644 index 0000000000..de33a761ee --- /dev/null +++ b/.changeset/guard-background-questions.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Prevent AskUserQuestion from starting background tasks when task controls are unavailable. diff --git a/packages/agent-core-v2/src/agent/tools/ask-user-question/askUserQuestionTool.ts b/packages/agent-core-v2/src/agent/tools/ask-user-question/askUserQuestionTool.ts index 47129f7eea..38375003cd 100644 --- a/packages/agent-core-v2/src/agent/tools/ask-user-question/askUserQuestionTool.ts +++ b/packages/agent-core-v2/src/agent/tools/ask-user-question/askUserQuestionTool.ts @@ -1,11 +1,10 @@ -import { z } from 'zod'; - import { CoreErrors } from '#/_base/errors/codes'; import { Error2 } from '#/_base/errors/errors'; import { toInputJsonSchema } from '#/tool/input-schema'; import { isAbortError } from '#/_base/utils/abort'; import { IAgentTaskService } from '#/agent/task/task'; import { IAgentScopeContext } from '#/agent/scopeContext/scopeContext'; +import { IAgentToolPolicyService } from '#/agent/toolPolicy/toolPolicy'; import { ITelemetryService } from '#/app/telemetry/telemetry'; import type { QuestionAnsweredEvent, QuestionDismissedEvent } from '#/app/telemetry/events'; import type { @@ -23,6 +22,7 @@ import type { QuestionResult, } from '#/session/question/question'; import { + AskUserQuestionInputSchema, AskUserQuestionInputSchemaWithBackground, IAskUserQuestionTool, questionUniquenessError, @@ -36,20 +36,33 @@ const QUESTION_DISMISSED_MESSAGE = 'User dismissed the question without answerin const QUESTION_UNSUPPORTED_FAILURE_MESSAGE = 'The connected client does not support interactive questions. Do NOT call this tool again. Ask the user directly in your text response instead.'; +const BACKGROUND_DESCRIPTION = + '- Set background=true when you can keep working without the answer. This starts a background question task and returns a task_id immediately. The answer arrives automatically in a later turn — you do not need to poll, sleep, or check on it. Continue with other work; never fabricate or predict the answer.'; + +const BACKGROUND_UNAVAILABLE_MESSAGE = + 'Background questions are not available for this agent because TaskList, TaskOutput, and TaskStop are not enabled.'; + +const PARAMETERS_WITH_BACKGROUND = toInputJsonSchema(AskUserQuestionInputSchemaWithBackground); +const PARAMETERS_FOREGROUND_ONLY = toInputJsonSchema(AskUserQuestionInputSchema); + export class AskUserQuestionTool implements IAskUserQuestionTool { declare readonly _serviceBrand: undefined; readonly name = 'AskUserQuestion' as const; - readonly description: string; - readonly parameters: Record; constructor( @ISessionQuestionService private readonly question: ISessionQuestionService, @ITelemetryService private readonly telemetry: ITelemetryService, @IAgentTaskService private readonly tasks: IAgentTaskService, @IAgentScopeContext private readonly scopeContext: IAgentScopeContext, - ) { - this.description = `${DESCRIPTION}- Set background=true when you can keep working without the answer. This starts a background question task and returns a task_id immediately. The answer arrives automatically in a later turn — you do not need to poll, sleep, or check on it. Continue with other work; never fabricate or predict the answer.`; - this.parameters = toInputJsonSchema(this.inputSchema()); + @IAgentToolPolicyService private readonly toolPolicy: IAgentToolPolicyService, + ) {} + + get description(): string { + return `${DESCRIPTION}${this.allowBackground() ? BACKGROUND_DESCRIPTION : ''}`; + } + + get parameters(): Record { + return this.allowBackground() ? PARAMETERS_WITH_BACKGROUND : PARAMETERS_FOREGROUND_ONLY; } resolveExecution(args: AskUserQuestionInput): ToolExecution { @@ -67,6 +80,10 @@ export class AskUserQuestionTool implements IAskUserQuestionTool { args: AskUserQuestionInput, { toolCallId, signal, turnId, trace }: ExecutableToolContext, ): Promise { + if (args.background === true && !this.allowBackground()) { + return { isError: true, output: BACKGROUND_UNAVAILABLE_MESSAGE }; + } + const uniquenessError = questionUniquenessError(args.questions); if (uniquenessError !== null) { return { isError: true, output: uniquenessError }; @@ -79,8 +96,12 @@ export class AskUserQuestionTool implements IAskUserQuestionTool { return this.executeQuestion(args, { toolCallId, turnId, signal, trace }); } - private inputSchema(): z.ZodType { - return AskUserQuestionInputSchemaWithBackground; + private allowBackground(): boolean { + return ( + this.toolPolicy.isToolActive('TaskList') && + this.toolPolicy.isToolActive('TaskOutput') && + this.toolPolicy.isToolActive('TaskStop') + ); } private executeInBackground( diff --git a/packages/agent-core-v2/test/agent/questionTools/tools/ask-user.test.ts b/packages/agent-core-v2/test/agent/questionTools/tools/ask-user.test.ts index a4ebcc7b37..cbfabe6c86 100644 --- a/packages/agent-core-v2/test/agent/questionTools/tools/ask-user.test.ts +++ b/packages/agent-core-v2/test/agent/questionTools/tools/ask-user.test.ts @@ -1,24 +1,34 @@ -import { describe, expect, it, vi } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { DisposableStore } from '#/_base/di/lifecycle'; +import { createServices } from '#/_base/di/test'; import { CoreErrors } from '#/_base/errors/codes'; import { Error2 } from '#/_base/errors/errors'; import { AskUserQuestionInputSchema, + IAskUserQuestionTool, type AskUserQuestionInput, } from '#/agent/tools/ask-user-question/ask-user-question'; import { AskUserQuestionTool } from '#/agent/tools/ask-user-question/askUserQuestionTool'; import { ITelemetryService } from '#/app/telemetry/telemetry'; import { IAgentTaskService } from '#/agent/task/task'; import { IAgentScopeContext } from '#/agent/scopeContext/scopeContext'; -import type { +import { IAgentToolPolicyService } from '#/agent/toolPolicy/toolPolicy'; +import { ISessionQuestionService, - QuestionRequest, - QuestionResult, + type QuestionRequest, + type QuestionResult, } from '#/session/question/question'; -import type { QuestionBackgroundTask } from '#/agent/tools/ask-user-question/question-background-task'; +import type { + QuestionBackgroundTask, + QuestionTaskInfo, +} from '#/agent/tools/ask-user-question/question-background-task'; import { executeTool } from '../../../tools/fixtures/execute-tool'; const signal = new AbortController().signal; +const TASK_TOOLS = new Set(['TaskList', 'TaskOutput', 'TaskStop']); + +let disposables: DisposableStore; function input( overrides: Partial = {}, @@ -41,13 +51,14 @@ function input( function makeTool( options: { + readonly activeTaskTools?: ReadonlySet; readonly request?: ( req: QuestionRequest, requestOptions?: { readonly signal?: AbortSignal }, ) => Promise; } = {}, ): { - readonly tool: AskUserQuestionTool; + readonly tool: IAskUserQuestionTool; readonly request: ReturnType; readonly telemetryTrack: ReturnType; readonly registerTask: ReturnType; @@ -56,23 +67,54 @@ function makeTool( } { const request = vi.fn(options.request ?? (async () => ({ Postgres: true }) as QuestionResult)); const telemetryTrack = vi.fn(); - const question = { request } as unknown as ISessionQuestionService; - const telemetry = { track2: telemetryTrack } as unknown as ITelemetryService; let lastTask: QuestionBackgroundTask | undefined; const registerTask = vi.fn((task: QuestionBackgroundTask) => { lastTask = task; return 'q_test_task_id'; }); - const getTask = vi.fn((id: string) => - id === 'q_test_task_id' ? { status: 'running' } : undefined, + const getTask = vi.fn( + (id: string): QuestionTaskInfo | undefined => + id === 'q_test_task_id' + ? { + taskId: id, + description: 'Which database?', + status: 'running', + detached: true, + startedAt: 0, + endedAt: null, + kind: 'question', + questionCount: 1, + toolCallId: 'call_bg', + } + : undefined, ); - const tasks = { registerTask, getTask } as unknown as IAgentTaskService; - const scopeContext = { agentId: 'main' } as unknown as IAgentScopeContext; - const tool = new AskUserQuestionTool(question, telemetry, tasks, scopeContext); + const activeTaskTools = options.activeTaskTools ?? TASK_TOOLS; + const ix = createServices(disposables, { + additionalServices: (reg) => { + reg.definePartialInstance(ISessionQuestionService, { request }); + reg.definePartialInstance(ITelemetryService, { track2: telemetryTrack }); + reg.definePartialInstance(IAgentTaskService, { registerTask, getTask }); + reg.definePartialInstance(IAgentScopeContext, { agentId: 'main' }); + reg.definePartialInstance(IAgentToolPolicyService, { + isToolActive: (name: string) => activeTaskTools.has(name), + }); + reg.define(IAskUserQuestionTool, AskUserQuestionTool); + }, + strict: true, + }); + const tool = ix.get(IAskUserQuestionTool); return { tool, request, telemetryTrack, registerTask, getTask, lastRegisteredTask: () => lastTask }; } describe('AskUserQuestionTool', () => { + beforeEach(() => { + disposables = new DisposableStore(); + }); + + afterEach(() => { + disposables.dispose(); + }); + it('exposes current metadata and schema', () => { const { tool } = makeTool(); @@ -167,7 +209,7 @@ describe('AskUserQuestionTool', () => { expect(request).toHaveBeenCalledOnce(); }); - it('builds the v1-aligned schema including an optional background flag', () => { + it('exposes background mode when all task controls are active', () => { const { tool } = makeTool(); const params = tool.parameters as { properties: { background?: { type?: string; default?: boolean } }; @@ -175,6 +217,104 @@ describe('AskUserQuestionTool', () => { expect(params.properties.background?.type).toBe('boolean'); expect(params.properties.background?.default).toBe(false); + expect(tool.description).toContain('background=true'); + expect(tool.description).toContain('task_id'); + }); + + it('hides and rejects background mode after a task control becomes inactive', async () => { + const activeTaskTools = new Set(TASK_TOOLS); + const { tool, request, registerTask } = makeTool({ + activeTaskTools, + }); + + expect(tool.parameters).toHaveProperty('properties.background'); + activeTaskTools.delete('TaskStop'); + + const params = tool.parameters as { properties: Record }; + + expect(params.properties).not.toHaveProperty('background'); + expect(tool.description.toLowerCase()).not.toContain('background'); + expect(tool.description).not.toContain('task_id'); + expect(tool.description).not.toContain('TaskOutput'); + + const result = await executeTool(tool, { + turnId: 0, + toolCallId: 'call_bg_disabled', + args: { ...input(), background: true }, + signal, + }); + + expect(result).toEqual({ + isError: true, + output: + 'Background questions are not available for this agent because TaskList, TaskOutput, and TaskStop are not enabled.', + }); + expect(registerTask).not.toHaveBeenCalled(); + expect(request).not.toHaveBeenCalled(); + }); + + it('preserves foreground answers when background mode is unavailable', async () => { + const { tool, request } = makeTool({ activeTaskTools: new Set() }); + + const result = await executeTool(tool, { + turnId: 0, + toolCallId: 'call_fg_disabled', + args: input(), + signal, + }); + + expect(result).toEqual({ + isError: false, + output: JSON.stringify({ answers: { Postgres: true } }), + }); + expect(request).toHaveBeenCalledOnce(); + }); + + it('preserves foreground dismissal when background mode is unavailable', async () => { + const { tool } = makeTool({ + activeTaskTools: new Set(), + request: async () => null, + }); + + const result = await executeTool(tool, { + turnId: 0, + toolCallId: 'call_fg_dismissed', + args: input(), + signal, + }); + + expect(result).toEqual({ + isError: false, + output: JSON.stringify({ + answers: {}, + note: 'User dismissed the question without answering.', + }), + }); + }); + + it('preserves foreground errors when background mode is unavailable', async () => { + const { tool } = makeTool({ + activeTaskTools: new Set(), + request: async () => { + throw new Error2( + CoreErrors.codes.NOT_IMPLEMENTED, + 'Client does not support questions', + ); + }, + }); + + const result = await executeTool(tool, { + turnId: 0, + toolCallId: 'call_fg_unsupported', + args: input(), + signal, + }); + + expect(result).toEqual({ + isError: true, + output: + 'The connected client does not support interactive questions. Do NOT call this tool again. Ask the user directly in your text response instead.', + }); }); it('dispatches questions through the session question service', async () => { diff --git a/packages/agent-core/src/agent/tool/index.ts b/packages/agent-core/src/agent/tool/index.ts index e5874018e8..5dfe3f171a 100644 --- a/packages/agent-core/src/agent/tool/index.ts +++ b/packages/agent-core/src/agent/tool/index.ts @@ -787,10 +787,11 @@ export class ToolManager { }, this.agent.skills?.registry.getSkillRoots() ?? [], ); - const allowBackground = + const canRunInBackground = () => this.isExactToolEnabled('TaskList') && this.isExactToolEnabled('TaskOutput') && this.isExactToolEnabled('TaskStop'); + const allowBackground = canRunInBackground(); const goalToolsEnabled = this.agent.type === 'main'; this.builtinTools = new Map( [ @@ -828,7 +829,8 @@ export class ToolManager { goalToolsEnabled && new b.GetGoalTool(this.agent), goalToolsEnabled && new b.SetGoalBudgetTool(this.agent), goalToolsEnabled && new b.UpdateGoalTool(this.agent), - this.agent.rpc?.requestQuestion && new b.AskUserQuestionTool(this.agent), + this.agent.rpc?.requestQuestion && + new b.AskUserQuestionTool(this.agent, { allowBackground: canRunInBackground }), new b.TodoListTool(this.toolStore), new b.TaskListTool(background), new b.TaskOutputTool(background), diff --git a/packages/agent-core/src/tools/builtin/collaboration/ask-user.ts b/packages/agent-core/src/tools/builtin/collaboration/ask-user.ts index e42adb58ef..2759a231da 100644 --- a/packages/agent-core/src/tools/builtin/collaboration/ask-user.ts +++ b/packages/agent-core/src/tools/builtin/collaboration/ask-user.ts @@ -128,16 +128,39 @@ const QUESTION_DISMISSED_MESSAGE = 'User dismissed the question without answerin const QUESTION_UNSUPPORTED_FAILURE_MESSAGE = 'The connected client does not support interactive questions. Do NOT call this tool again. Ask the user directly in your text response instead.'; +const BACKGROUND_DESCRIPTION = + '- Set background=true when you can keep working without the answer. This starts a background question task and returns a task_id immediately. The answer arrives automatically in a later turn — you do not need to poll, sleep, or check on it. Continue with other work; never fabricate or predict the answer.'; + +const BACKGROUND_UNAVAILABLE_MESSAGE = + 'Background questions are not available for this agent because TaskList, TaskOutput, and TaskStop are not enabled.'; + +const PARAMETERS_WITH_BACKGROUND = toInputJsonSchema(AskUserQuestionInputSchemaWithBackground); +const PARAMETERS_FOREGROUND_ONLY = toInputJsonSchema(AskUserQuestionInputSchema); + // ── Implementation ─────────────────────────────────────────────────── export class AskUserQuestionTool implements BuiltinTool { readonly name = 'AskUserQuestion' as const; - readonly description: string; - readonly parameters: Record; - constructor(private readonly agent: Agent) { - this.description = `${DESCRIPTION}- Set background=true when you can keep working without the answer. This starts a background question task and returns a task_id immediately. The answer arrives automatically in a later turn — you do not need to poll, sleep, or check on it. Continue with other work; never fabricate or predict the answer.`; - this.parameters = toInputJsonSchema(this.inputSchema()); + private readonly canRunInBackground: () => boolean; + + constructor( + private readonly agent: Agent, + options?: { allowBackground?: boolean | (() => boolean) }, + ) { + const allowBackground = options?.allowBackground ?? true; + this.canRunInBackground = + typeof allowBackground === 'function' ? allowBackground : () => allowBackground; + } + + get description(): string { + return `${DESCRIPTION}${this.canRunInBackground() ? BACKGROUND_DESCRIPTION : ''}`; + } + + get parameters(): Record { + return this.canRunInBackground() + ? PARAMETERS_WITH_BACKGROUND + : PARAMETERS_FOREGROUND_ONLY; } resolveExecution(args: AskUserQuestionInput): ToolExecution { @@ -160,6 +183,10 @@ export class AskUserQuestionTool implements BuiltinTool { turnId, }: ExecutableToolContext, ): Promise { + if (args.background === true && !this.canRunInBackground()) { + return { isError: true, output: BACKGROUND_UNAVAILABLE_MESSAGE }; + } + // AJV (the runtime arg validator) cannot express the uniqueness refine, // so enforce it here before any UI interaction or task registration. const uniquenessError = questionUniquenessError(args.questions); @@ -174,10 +201,6 @@ export class AskUserQuestionTool implements BuiltinTool { return this.executeQuestion(args, { toolCallId, turnId, signal, traceId }); } - private inputSchema(): z.ZodType { - return AskUserQuestionInputSchemaWithBackground; - } - private async executeQuestion( args: AskUserQuestionInput, { diff --git a/packages/agent-core/test/agent/tool.test.ts b/packages/agent-core/test/agent/tool.test.ts index 5accc6527c..a0d885d14c 100644 --- a/packages/agent-core/test/agent/tool.test.ts +++ b/packages/agent-core/test/agent/tool.test.ts @@ -326,6 +326,56 @@ describe('Agent tools', () => { expect(subagentHost.spawn).not.toHaveBeenCalled(); }); + it('rechecks AskUserQuestion background mode after the task policy changes', async () => { + const ctx = testAgent(); + ctx.configure(); + ctx.agent.tools.setActiveTools([ + 'AskUserQuestion', + 'TaskList', + 'TaskOutput', + 'TaskStop', + ]); + + const retainedTool = ctx.agent.tools.loopTools.find( + (tool) => tool.name === 'AskUserQuestion', + ); + expect(retainedTool).toBeDefined(); + expect(retainedTool!.parameters).toHaveProperty('properties.background'); + expect(retainedTool!.description).toContain('background=true'); + + const registerTask = vi.spyOn(ctx.agent.background, 'registerTask'); + ctx.agent.tools.setActiveTools(['AskUserQuestion', 'TaskList', 'TaskOutput']); + + expect(retainedTool!.parameters).not.toHaveProperty('properties.background'); + expect(retainedTool!.description.toLowerCase()).not.toContain('background'); + await expect( + executeTool(retainedTool!, { + turnId: '0', + toolCallId: 'call_question', + args: { + background: true, + questions: [ + { + question: 'Which database?', + header: 'Storage', + options: [ + { label: 'Postgres', description: 'Relational storage' }, + { label: 'SQLite', description: 'Embedded storage' }, + ], + multi_select: false, + }, + ], + }, + signal, + }), + ).resolves.toEqual({ + isError: true, + output: + 'Background questions are not available for this agent because TaskList, TaskOutput, and TaskStop are not enabled.', + }); + expect(registerTask).not.toHaveBeenCalled(); + }); + it('removes denied exact tool names from the active set', () => { const ctx = testAgent(); ctx.configure(); diff --git a/packages/agent-core/test/tools/ask-user.test.ts b/packages/agent-core/test/tools/ask-user.test.ts index 1727019e67..1f8fe653e2 100644 --- a/packages/agent-core/test/tools/ask-user.test.ts +++ b/packages/agent-core/test/tools/ask-user.test.ts @@ -35,6 +35,7 @@ function input( function makeTool( options: { + readonly allowBackground?: boolean; readonly mode?: PermissionMode; readonly requestQuestion?: ( request: QuestionRequest, @@ -58,7 +59,11 @@ function makeTool( rpc: { requestQuestion }, telemetry: { track: telemetryTrack }, } as unknown as Agent; - return { tool: new AskUserQuestionTool(agent), requestQuestion, telemetryTrack }; + return { + tool: new AskUserQuestionTool(agent, { allowBackground: options.allowBackground }), + requestQuestion, + telemetryTrack, + }; } describe('AskUserQuestionTool', () => { @@ -185,16 +190,111 @@ describe('AskUserQuestionTool', () => { expect(requestQuestion).not.toHaveBeenCalled(); }); - it('always builds the background-question schema', () => { + it('keeps the background schema and description when background questions are allowed', () => { const agent = { rpc: { requestQuestion: vi.fn() }, telemetry: { track: vi.fn() }, background: createBackgroundManager().manager, } as unknown as Agent; - const tool = new AskUserQuestionTool(agent); + const tool = new AskUserQuestionTool(agent, { allowBackground: true }); expect(JSON.stringify(tool.parameters)).toContain('background'); + expect(tool.description).toContain('background=true'); + expect(tool.description).toContain('task_id'); + }); + + it('hides and rejects background questions when background is not allowed', async () => { + const { manager } = createBackgroundManager(); + const registerTask = vi.spyOn(manager, 'registerTask'); + const requestQuestion = vi.fn(); + const agent = { + rpc: { requestQuestion }, + telemetry: { track: vi.fn() }, + background: manager, + } as unknown as Agent; + const tool = new AskUserQuestionTool(agent, { allowBackground: false }); + + expect(JSON.stringify(tool.parameters)).not.toContain('background'); + expect(tool.description.toLowerCase()).not.toContain('background'); + expect(tool.description).not.toContain('task_id'); + expect(tool.description).not.toContain('TaskOutput'); + + const result = await executeTool(tool, { + turnId: '0', + toolCallId: 'call_bg_disabled', + args: { ...input(), background: true }, + signal, + }); + + expect(result).toEqual({ + isError: true, + output: + 'Background questions are not available for this agent because TaskList, TaskOutput, and TaskStop are not enabled.', + }); + expect(registerTask).not.toHaveBeenCalled(); + expect(requestQuestion).not.toHaveBeenCalled(); + }); + + it('preserves foreground answers when background questions are disabled', async () => { + const { tool, requestQuestion } = makeTool({ allowBackground: false }); + + const result = await executeTool(tool, { + turnId: '0', + toolCallId: 'call_fg_disabled', + args: input(), + signal, + }); + + expect(result).toEqual({ + isError: false, + output: JSON.stringify({ answers: { Postgres: true } }), + }); + expect(requestQuestion).toHaveBeenCalledOnce(); + }); + + it('preserves foreground dismissal semantics when background questions are disabled', async () => { + const { tool } = makeTool({ + allowBackground: false, + requestQuestion: async () => null, + }); + + const result = await executeTool(tool, { + turnId: '0', + toolCallId: 'call_fg_dismissed', + args: input(), + signal, + }); + + expect(result).toEqual({ + isError: false, + output: JSON.stringify({ + answers: {}, + note: 'User dismissed the question without answering.', + }), + }); + }); + + it('preserves foreground error semantics when background questions are disabled', async () => { + const { tool } = makeTool({ + allowBackground: false, + requestQuestion: async () => { + throw new KimiError(ErrorCodes.NOT_IMPLEMENTED, 'Client does not support questions'); + }, + }); + + const result = await executeTool(tool, { + turnId: '0', + toolCallId: 'call_fg_unsupported', + args: input(), + signal, + }); + + expect(result).toEqual({ + isError: true, + output: + 'The connected client does not support interactive questions. Do NOT call this tool again. Ask the user directly in your text response instead.', + }); }); it.each(['manual', 'yolo'] as const)(