abbyswag/qwen-code
1
0
Fork 0
mirror of https://github.com/QwenLM/qwen-code.git synced 2025-12-19 09:33:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: subagent feature wip

Browse Source
This commit is contained in:
tanzhenxin
2025-09-10 13:41:28 +08:00
parent 549f296eb5
commit 6b09aee32b
30 changed files with 329 additions and 239 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
packages/core/src/core/coreToolScheduler.test.ts
View File

@@ -22,6 +22,7 @@ import {
Config,
Kind,
ApprovalMode,
ToolResultDisplay,
ToolRegistry,
} from '../index.js';
import { Part, PartListUnion } from '@google/genai';
@@ -633,6 +634,135 @@ describe('CoreToolScheduler YOLO mode', () => {
});
});
describe('CoreToolScheduler cancellation during executing with live output', () => {
it('sets status to cancelled and preserves last output', async () => {
class StreamingInvocation extends BaseToolInvocation<
{ id: string },
ToolResult
> {
getDescription(): string {
return `Streaming tool ${this.params.id}`;
}
async execute(
signal: AbortSignal,
updateOutput?: (output: ToolResultDisplay) => void,
): Promise<ToolResult> {
updateOutput?.('hello');
// Wait until aborted to emulate a long-running task
await new Promise<void>((resolve) => {
if (signal.aborted) return resolve();
const onAbort = () => {
signal.removeEventListener('abort', onAbort);
resolve();
};
signal.addEventListener('abort', onAbort, { once: true });
});
// Return a normal (non-error) result; scheduler should still mark cancelled
return { llmContent: 'done', returnDisplay: 'done' };
}
}
class StreamingTool extends BaseDeclarativeTool<
{ id: string },
ToolResult
> {
constructor() {
super(
'stream-tool',
'Stream Tool',
'Emits live output and waits for abort',
Kind.Other,
{
type: 'object',
properties: { id: { type: 'string' } },
required: ['id'],
},
true,
true,
);
}
protected createInvocation(params: { id: string }) {
return new StreamingInvocation(params);
}
}
const tool = new StreamingTool();
const mockToolRegistry = {
getTool: () => tool,
getFunctionDeclarations: () => [],
tools: new Map(),
discovery: {},
registerTool: () => {},
getToolByName: () => tool,
getToolByDisplayName: () => tool,
getTools: () => [],
discoverTools: async () => {},
getAllTools: () => [],
getToolsByServer: () => [],
} as unknown as ToolRegistry;
const onAllToolCallsComplete = vi.fn();
const onToolCallsUpdate = vi.fn();
const mockConfig = {
getSessionId: () => 'test-session-id',
getUsageStatisticsEnabled: () => true,
getDebugMode: () => false,
getApprovalMode: () => ApprovalMode.DEFAULT,
getContentGeneratorConfig: () => ({
model: 'test-model',
authType: 'oauth-personal',
}),
} as unknown as Config;
const scheduler = new CoreToolScheduler({
config: mockConfig,
toolRegistry: mockToolRegistry,
onAllToolCallsComplete,
onToolCallsUpdate,
getPreferredEditor: () => 'vscode',
onEditorClose: vi.fn(),
});
const abortController = new AbortController();
const request = {
callId: '1',
name: 'stream-tool',
args: { id: 'x' },
isClientInitiated: true,
prompt_id: 'prompt-stream',
};
const schedulePromise = scheduler.schedule(
[request],
abortController.signal,
);
// Wait until executing
await vi.waitFor(() => {
const calls = onToolCallsUpdate.mock.calls;
const last = calls[calls.length - 1]?.[0][0] as ToolCall | undefined;
expect(last?.status).toBe('executing');
});
// Now abort
abortController.abort();
await schedulePromise;
await vi.waitFor(() => {
expect(onAllToolCallsComplete).toHaveBeenCalled();
});
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('cancelled');
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const cancelled: any = completedCalls[0];
expect(cancelled.response.resultDisplay).toBe('hello');
});
});
describe('CoreToolScheduler request queueing', () => {
it('should queue a request if another is running', async () => {
let resolveFirstCall: (result: ToolResult) => void;

32
packages/core/src/core/coreToolScheduler.ts
View File

@@ -374,6 +374,13 @@ export class CoreToolScheduler {
newContent: waitingCall.confirmationDetails.newContent,
};
}
} else if (currentCall.status === 'executing') {
// If the tool was streaming live output, preserve the latest
// output so the UI can continue to show it after cancellation.
const executingCall = currentCall as ExecutingToolCall;
if (executingCall.liveOutput !== undefined) {
resultDisplay = executingCall.liveOutput;
}
}
return {
@@ -816,20 +823,19 @@ export class CoreToolScheduler {
const invocation = scheduledCall.invocation;
this.setStatusInternal(callId, 'executing');
const liveOutputCallback =
scheduledCall.tool.canUpdateOutput && this.outputUpdateHandler
? (outputChunk: ToolResultDisplay) => {
if (this.outputUpdateHandler) {
this.outputUpdateHandler(callId, outputChunk);
}
this.toolCalls = this.toolCalls.map((tc) =>
tc.request.callId === callId && tc.status === 'executing'
? { ...tc, liveOutput: outputChunk }
: tc,
);
this.notifyToolCallsUpdate();
const liveOutputCallback = scheduledCall.tool.canUpdateOutput
? (outputChunk: ToolResultDisplay) => {
if (this.outputUpdateHandler) {
this.outputUpdateHandler(callId, outputChunk);
}
: undefined;
this.toolCalls = this.toolCalls.map((tc) =>
tc.request.callId === callId && tc.status === 'executing'
? { ...tc, liveOutput: outputChunk }
: tc,
);
this.notifyToolCallsUpdate();
}
: undefined;
invocation
.execute(signal, liveOutputCallback)

1
packages/core/src/subagents/index.ts
View File

@@ -49,7 +49,6 @@ export type {
RunConfig,
ToolConfig,
SubagentTerminateMode,
OutputObject,
} from './types.js';
export { SubAgentScope } from './subagent.js';

22
packages/core/src/subagents/subagent.test.ts
View File

@@ -412,7 +412,7 @@ describe('subagent.ts', () => {
await expect(scope.runNonInteractive(context)).rejects.toThrow(
'Missing context values for the following keys: missing',
);
expect(scope.output.terminate_reason).toBe(SubagentTerminateMode.ERROR);
expect(scope.getTerminateMode()).toBe(SubagentTerminateMode.ERROR);
});
it('should validate that systemPrompt and initialMessages are mutually exclusive', async () => {
@@ -434,7 +434,7 @@ describe('subagent.ts', () => {
await expect(agent.runNonInteractive(context)).rejects.toThrow(
'PromptConfig cannot have both `systemPrompt` and `initialMessages` defined.',
);
expect(agent.output.terminate_reason).toBe(SubagentTerminateMode.ERROR);
expect(agent.getTerminateMode()).toBe(SubagentTerminateMode.ERROR);
});
});
@@ -457,8 +457,7 @@ describe('subagent.ts', () => {
await scope.runNonInteractive(new ContextState());
expect(scope.output.terminate_reason).toBe(SubagentTerminateMode.GOAL);
expect(scope.output.result).toBe('Done.');
expect(scope.getTerminateMode()).toBe(SubagentTerminateMode.GOAL);
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
// Check the initial message
expect(mockSendMessageStream.mock.calls[0][0].message).toEqual([
@@ -482,8 +481,7 @@ describe('subagent.ts', () => {
await scope.runNonInteractive(new ContextState());
expect(scope.output.terminate_reason).toBe(SubagentTerminateMode.GOAL);
expect(scope.output.result).toBe('Done.');
expect(scope.getTerminateMode()).toBe(SubagentTerminateMode.GOAL);
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
});
@@ -549,7 +547,7 @@ describe('subagent.ts', () => {
{ text: 'file1.txt\nfile2.ts' },
]);
expect(scope.output.terminate_reason).toBe(SubagentTerminateMode.GOAL);
expect(scope.getTerminateMode()).toBe(SubagentTerminateMode.GOAL);
});
it('should provide specific tool error responses to the model', async () => {
@@ -645,9 +643,7 @@ describe('subagent.ts', () => {
await scope.runNonInteractive(new ContextState());
expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
expect(scope.output.terminate_reason).toBe(
SubagentTerminateMode.MAX_TURNS,
);
expect(scope.getTerminateMode()).toBe(SubagentTerminateMode.MAX_TURNS);
});
it('should terminate with TIMEOUT if the time limit is reached during an LLM call', async () => {
@@ -690,9 +686,7 @@ describe('subagent.ts', () => {
await runPromise;
expect(scope.output.terminate_reason).toBe(
SubagentTerminateMode.TIMEOUT,
);
expect(scope.getTerminateMode()).toBe(SubagentTerminateMode.TIMEOUT);
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
vi.useRealTimers();
@@ -713,7 +707,7 @@ describe('subagent.ts', () => {
await expect(
scope.runNonInteractive(new ContextState()),
).rejects.toThrow('API Failure');
expect(scope.output.terminate_reason).toBe(SubagentTerminateMode.ERROR);
expect(scope.getTerminateMode()).toBe(SubagentTerminateMode.ERROR);
});
});
});

43
packages/core/src/subagents/subagent.ts
View File

@@ -20,7 +20,6 @@ import {
} from '@google/genai';
import { GeminiChat } from '../core/geminiChat.js';
import {
OutputObject,
SubagentTerminateMode,
PromptConfig,
ModelConfig,
@@ -150,10 +149,6 @@ function templateString(template: string, context: ContextState): string {
* runtime context, and the collection of its outputs.
*/
export class SubAgentScope {
output: OutputObject = {
terminate_reason: SubagentTerminateMode.ERROR,
result: '',
};
executionStats: ExecutionStats = {
startTimeMs: 0,
totalDurationMs: 0,
@@ -179,6 +174,7 @@ export class SubAgentScope {
>();
private eventEmitter?: SubAgentEventEmitter;
private finalText: string = '';
private terminateMode: SubagentTerminateMode = SubagentTerminateMode.ERROR;
private readonly stats = new SubagentStatistics();
private hooks?: SubagentHooks;
private readonly subagentId: string;
@@ -312,14 +308,18 @@ export class SubAgentScope {
const chat = await this.createChatObject(context);
if (!chat) {
this.output.terminate_reason = SubagentTerminateMode.ERROR;
this.terminateMode = SubagentTerminateMode.ERROR;
return;
}
const abortController = new AbortController();
const onAbort = () => abortController.abort();
if (externalSignal) {
if (externalSignal.aborted) abortController.abort();
if (externalSignal.aborted) {
abortController.abort();
this.terminateMode = SubagentTerminateMode.CANCELLED;
return;
}
externalSignal.addEventListener('abort', onAbort, { once: true });
}
const toolRegistry = this.runtimeContext.getToolRegistry();
@@ -381,7 +381,7 @@ export class SubAgentScope {
this.runConfig.max_turns &&
turnCounter >= this.runConfig.max_turns
) {
this.output.terminate_reason = SubagentTerminateMode.MAX_TURNS;
this.terminateMode = SubagentTerminateMode.MAX_TURNS;
break;
}
let durationMin = (Date.now() - startTime) / (1000 * 60);
@@ -389,7 +389,7 @@ export class SubAgentScope {
this.runConfig.max_time_minutes &&
durationMin >= this.runConfig.max_time_minutes
) {
this.output.terminate_reason = SubagentTerminateMode.TIMEOUT;
this.terminateMode = SubagentTerminateMode.TIMEOUT;
break;
}
@@ -418,7 +418,10 @@ export class SubAgentScope {
let lastUsage: GenerateContentResponseUsageMetadata | undefined =
undefined;
for await (const resp of responseStream) {
if (abortController.signal.aborted) return;
if (abortController.signal.aborted) {
this.terminateMode = SubagentTerminateMode.CANCELLED;
return;
}
if (resp.functionCalls) functionCalls.push(...resp.functionCalls);
const content = resp.candidates?.[0]?.content;
const parts = content?.parts || [];
@@ -443,7 +446,7 @@ export class SubAgentScope {
this.runConfig.max_time_minutes &&
durationMin >= this.runConfig.max_time_minutes
) {
this.output.terminate_reason = SubagentTerminateMode.TIMEOUT;
this.terminateMode = SubagentTerminateMode.TIMEOUT;
break;
}
@@ -483,8 +486,7 @@ export class SubAgentScope {
// No tool calls — treat this as the model's final answer.
if (roundText && roundText.trim().length > 0) {
this.finalText = roundText.trim();
this.output.result = this.finalText;
this.output.terminate_reason = SubagentTerminateMode.GOAL;
this.terminateMode = SubagentTerminateMode.GOAL;
break;
}
// Otherwise, nudge the model to finalize a result.
@@ -508,7 +510,7 @@ export class SubAgentScope {
}
} catch (error) {
console.error('Error during subagent execution:', error);
this.output.terminate_reason = SubagentTerminateMode.ERROR;
this.terminateMode = SubagentTerminateMode.ERROR;
this.eventEmitter?.emit(SubAgentEventType.ERROR, {
subagentId: this.subagentId,
error: error instanceof Error ? error.message : String(error),
@@ -529,7 +531,7 @@ export class SubAgentScope {
const summary = this.stats.getSummary(Date.now());
this.eventEmitter?.emit(SubAgentEventType.FINISH, {
subagentId: this.subagentId,
terminate_reason: this.output.terminate_reason,
terminate_reason: this.terminateMode,
timestamp: Date.now(),
rounds: summary.rounds,
totalDurationMs: summary.totalDurationMs,
@@ -541,14 +543,13 @@ export class SubAgentScope {
totalTokens: summary.totalTokens,
} as SubAgentFinishEvent);
// Log telemetry for subagent completion
const completionEvent = new SubagentExecutionEvent(
this.name,
this.output.terminate_reason === SubagentTerminateMode.GOAL
this.terminateMode === SubagentTerminateMode.GOAL
? 'completed'
: 'failed',
{
terminate_reason: this.output.terminate_reason,
terminate_reason: this.terminateMode,
result: this.finalText,
execution_summary: this.stats.formatCompact(
'Subagent execution completed',
@@ -560,7 +561,7 @@ export class SubAgentScope {
await this.hooks?.onStop?.({
subagentId: this.subagentId,
name: this.name,
terminateReason: this.output.terminate_reason,
terminateReason: this.terminateMode,
summary: summary as unknown as Record<string, unknown>,
timestamp: Date.now(),
});
@@ -751,6 +752,10 @@ export class SubAgentScope {
return this.finalText;
}
getTerminateMode(): SubagentTerminateMode {
return this.terminateMode;
}
private async createChatObject(context: ContextState) {
if (!this.promptConfig.systemPrompt && !this.promptConfig.initialMessages) {
throw new Error(

18
packages/core/src/subagents/types.ts
View File

@@ -183,24 +183,10 @@ export enum SubagentTerminateMode {
* Indicates that the subagent's execution terminated because it exceeded the maximum number of turns.
*/
MAX_TURNS = 'MAX_TURNS',
}
/**
* Represents the output structure of a subagent's execution.
* This interface defines the data that a subagent will return upon completion,
* including the final result and the reason for its termination.
*/
export interface OutputObject {
/**
* The final result text returned by the subagent upon completion.
* This contains the direct output from the model's final response.
* Indicates that the subagent's execution was cancelled via an abort signal.
*/
result: string;
/**
* The reason for the subagent's termination, indicating whether it completed
* successfully, timed out, or encountered an error.
*/
terminate_reason: SubagentTerminateMode;
CANCELLED = 'CANCELLED',
}
/**

2
packages/core/src/telemetry/metrics.ts
View File

@@ -294,7 +294,7 @@ export function recordContentRetryFailure(config: Config): void {
export function recordSubagentExecutionMetrics(
config: Config,
subagentName: string,
status: 'started' | 'progress' | 'completed' | 'failed',
status: 'started' | 'completed' | 'failed' | 'cancelled',
terminateReason?: string,
): void {
if (!subagentExecutionCounter || !isMetricsInitialized) return;

4
packages/core/src/telemetry/types.ts
View File

@@ -448,14 +448,14 @@ export class SubagentExecutionEvent implements BaseTelemetryEvent {
'event.name': 'subagent_execution';
'event.timestamp': string;
subagent_name: string;
status: 'started' | 'progress' | 'completed' | 'failed';
status: 'started' | 'completed' | 'failed' | 'cancelled';
terminate_reason?: string;
result?: string;
execution_summary?: string;
constructor(
subagent_name: string,
status: 'started' | 'progress' | 'completed' | 'failed',
status: 'started' | 'completed' | 'failed' | 'cancelled',
options?: {
terminate_reason?: string;
result?: string;

26
packages/core/src/tools/task.test.ts
View File

@@ -258,10 +258,8 @@ describe('TaskTool', () => {
beforeEach(() => {
mockSubagentScope = {
runNonInteractive: vi.fn().mockResolvedValue(undefined),
output: {
result: 'Task completed successfully',
terminate_reason: SubagentTerminateMode.GOAL,
},
result: 'Task completed successfully',
terminateMode: SubagentTerminateMode.GOAL,
getFinalText: vi.fn().mockReturnValue('Task completed successfully'),
formatCompactResult: vi
.fn()
@@ -305,6 +303,7 @@ describe('TaskTool', () => {
successfulToolCalls: 3,
failedToolCalls: 0,
}),
getTerminateMode: vi.fn().mockReturnValue(SubagentTerminateMode.GOAL),
} as unknown as SubAgentScope;
mockContextState = {
@@ -375,25 +374,6 @@ describe('TaskTool', () => {
expect(display.subagentName).toBe('non-existent');
});
it('should handle subagent execution failure', async () => {
mockSubagentScope.output.terminate_reason = SubagentTerminateMode.ERROR;
const params: TaskParams = {
description: 'Search files',
prompt: 'Find all TypeScript files',
subagent_type: 'file-search',
};
const invocation = (
taskTool as TaskToolWithProtectedMethods
).createInvocation(params);
const result = await invocation.execute();
const display = result.returnDisplay as TaskResultDisplay;
expect(display.status).toBe('failed');
expect(display.terminateReason).toBe('ERROR');
});
it('should handle execution errors gracefully', async () => {
vi.mocked(mockSubagentManager.createSubagentScope).mockRejectedValue(
new Error('Creation failed'),

54
packages/core/src/tools/task.ts
View File

@@ -14,7 +14,7 @@ import {
} from './tools.js';
import { Config } from '../config/config.js';
import { SubagentManager } from '../subagents/subagent-manager.js';
import { SubagentConfig } from '../subagents/types.js';
import { SubagentConfig, SubagentTerminateMode } from '../subagents/types.js';
import { ContextState } from '../subagents/subagent.js';
import {
SubAgentEventEmitter,
@@ -409,21 +409,6 @@ class TaskToolInvocation extends BaseToolInvocation<TaskParams, ToolResult> {
// Set up event listeners for real-time updates
this.setupEventListeners(updateOutput);
if (signal) {
signal.addEventListener('abort', () => {
if (this.currentDisplay) {
this.updateDisplay(
{
status: 'failed',
terminateReason: 'CANCELLED',
result: 'Task was cancelled by user',
},
updateOutput,
);
}
});
}
// Send initial display
if (updateOutput) {
updateOutput(this.currentDisplay);
@@ -474,20 +459,31 @@ class TaskToolInvocation extends BaseToolInvocation<TaskParams, ToolResult> {
// Get the results
const finalText = subagentScope.getFinalText();
const terminateReason = subagentScope.output.terminate_reason;
const success = terminateReason === 'GOAL';
const terminateReason = subagentScope.getTerminateMode();
const success = terminateReason === SubagentTerminateMode.GOAL;
const executionSummary = subagentScope.getExecutionSummary();
// Update the final display state
this.updateDisplay(
{
status: success ? 'completed' : 'failed',
terminateReason,
result: finalText,
executionSummary,
},
updateOutput,
);
if (signal?.aborted) {
this.updateDisplay(
{
status: 'cancelled',
terminateReason: 'CANCELLED',
result: finalText || 'Task was cancelled by user',
executionSummary,
},
updateOutput,
);
} else {
this.updateDisplay(
{
status: success ? 'completed' : 'failed',
terminateReason,
result: finalText,
executionSummary,
},
updateOutput,
);
}
return {
llmContent: [{ text: finalText }],
@@ -500,7 +496,7 @@ class TaskToolInvocation extends BaseToolInvocation<TaskParams, ToolResult> {
const errorDisplay: TaskResultDisplay = {
...this.currentDisplay!,
status: 'failed' as const,
status: 'failed',
terminateReason: 'ERROR',
result: `Failed to run subagent: ${errorMessage}`,
};

2
packages/core/src/tools/tools.ts
View File

@@ -428,7 +428,7 @@ export interface TaskResultDisplay {
subagentColor?: string;
taskDescription: string;
taskPrompt: string;
status: 'running' | 'completed' | 'failed';
status: 'running' | 'completed' | 'failed' | 'cancelled';
terminateReason?: string;
result?: string;
executionSummary?: SubagentStatsSummary;