Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/guard-background-questions.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@moonshot-ai/kimi-code": patch
---

Prevent AskUserQuestion from starting background tasks when task controls are unavailable.
Original file line number Diff line number Diff line change
@@ -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 {
Expand All @@ -23,6 +22,7 @@ import type {
QuestionResult,
} from '#/session/question/question';
import {
AskUserQuestionInputSchema,
AskUserQuestionInputSchemaWithBackground,
IAskUserQuestionTool,
questionUniquenessError,
Expand All @@ -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<string, unknown>;

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<string, unknown> {
return this.allowBackground() ? PARAMETERS_WITH_BACKGROUND : PARAMETERS_FOREGROUND_ONLY;
}

resolveExecution(args: AskUserQuestionInput): ToolExecution {
Expand All @@ -67,6 +80,10 @@ export class AskUserQuestionTool implements IAskUserQuestionTool {
args: AskUserQuestionInput,
{ toolCallId, signal, turnId, trace }: ExecutableToolContext,
): Promise<ExecutableToolResult> {
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 };
Expand All @@ -79,8 +96,12 @@ export class AskUserQuestionTool implements IAskUserQuestionTool {
return this.executeQuestion(args, { toolCallId, turnId, signal, trace });
}

private inputSchema(): z.ZodType<AskUserQuestionInput> {
return AskUserQuestionInputSchemaWithBackground;
private allowBackground(): boolean {
return (
this.toolPolicy.isToolActive('TaskList') &&
this.toolPolicy.isToolActive('TaskOutput') &&
this.toolPolicy.isToolActive('TaskStop')
);
}

private executeInBackground(
Expand Down
168 changes: 154 additions & 14 deletions packages/agent-core-v2/test/agent/questionTools/tools/ask-user.test.ts
Original file line number Diff line number Diff line change
@@ -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<AskUserQuestionInput['questions'][number]> = {},
Expand All @@ -41,13 +51,14 @@ function input(

function makeTool(
options: {
readonly activeTaskTools?: ReadonlySet<string>;
readonly request?: (
req: QuestionRequest,
requestOptions?: { readonly signal?: AbortSignal },
) => Promise<QuestionResult>;
} = {},
): {
readonly tool: AskUserQuestionTool;
readonly tool: IAskUserQuestionTool;
readonly request: ReturnType<typeof vi.fn>;
readonly telemetryTrack: ReturnType<typeof vi.fn>;
readonly registerTask: ReturnType<typeof vi.fn>;
Expand All @@ -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();

Expand Down Expand Up @@ -167,14 +209,112 @@ 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 } };
};

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<string, unknown> };

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 () => {
Expand Down
3 changes: 2 additions & 1 deletion packages/agent-core/src/agent/tool/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -828,7 +828,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 }),
Comment thread
pvzheroes125 marked this conversation as resolved.
Outdated
new b.TodoListTool(this.toolStore),
new b.TaskListTool(background),
new b.TaskOutputTool(background),
Expand Down
24 changes: 21 additions & 3 deletions packages/agent-core/src/tools/builtin/collaboration/ask-user.ts
Original file line number Diff line number Diff line change
Expand Up @@ -128,15 +128,27 @@ 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.';

// ── Implementation ───────────────────────────────────────────────────

export class AskUserQuestionTool implements BuiltinTool<AskUserQuestionInput> {
readonly name = 'AskUserQuestion' as const;
readonly description: string;
readonly parameters: Record<string, unknown>;

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.`;
private readonly allowBackground: boolean;

constructor(
private readonly agent: Agent,
options?: { allowBackground?: boolean },
) {
this.allowBackground = options?.allowBackground ?? true;
this.description = `${DESCRIPTION}${this.allowBackground ? BACKGROUND_DESCRIPTION : ''}`;
this.parameters = toInputJsonSchema(this.inputSchema());
}

Expand All @@ -160,6 +172,10 @@ export class AskUserQuestionTool implements BuiltinTool<AskUserQuestionInput> {
turnId,
}: ExecutableToolContext,
): Promise<ExecutableToolResult> {
if (args.background === true && !this.allowBackground) {
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);
Expand All @@ -175,7 +191,9 @@ export class AskUserQuestionTool implements BuiltinTool<AskUserQuestionInput> {
}

private inputSchema(): z.ZodType<AskUserQuestionInput> {
return AskUserQuestionInputSchemaWithBackground;
return this.allowBackground
? AskUserQuestionInputSchemaWithBackground
: AskUserQuestionInputSchema;
}

private async executeQuestion(
Expand Down
Loading
Loading