mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-07-10 01:39:25 +00:00
feat(agent-core): add Swarm tool wired to SwarmCoordinator with recursion guard
This commit is contained in:
parent
985fd5c6f6
commit
9c309b19ef
5 changed files with 142 additions and 0 deletions
|
|
@ -393,6 +393,9 @@ export class ToolManager {
|
|||
log: this.agent.log,
|
||||
},
|
||||
),
|
||||
this.agent.subagentHost &&
|
||||
this.agent.type !== 'sub' &&
|
||||
new b.SwarmTool(this.agent.subagentHost, { log: this.agent.log }),
|
||||
toolServices?.webSearcher && new b.WebSearchTool(toolServices.webSearcher),
|
||||
toolServices?.urlFetcher && new b.FetchURLTool(toolServices.urlFetcher),
|
||||
]
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ tools:
|
|||
- Skill
|
||||
- WebSearch
|
||||
- Agent
|
||||
- Swarm
|
||||
- FetchURL
|
||||
- AskUserQuestion
|
||||
- EnterPlanMode
|
||||
|
|
|
|||
90
packages/agent-core/src/tools/builtin/collaboration/swarm.ts
Normal file
90
packages/agent-core/src/tools/builtin/collaboration/swarm.ts
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
/**
|
||||
* SwarmTool — collaboration tool that runs a task as a self-directed agent
|
||||
* swarm.
|
||||
*
|
||||
* Like {@link AgentTool}, this is a "collaboration tool": it uses
|
||||
* `SessionSubagentHost` (injected via the constructor) to create in-process
|
||||
* subagents. The {@link SwarmCoordinator} dynamically decomposes the task into
|
||||
* parallel role-specialized workers, then synthesizes their outputs into one
|
||||
* answer.
|
||||
*
|
||||
* Workers are spawned with an ad-hoc `profileOverride`, and the tool is
|
||||
* registered only on non-sub agents so a swarm worker can never launch another
|
||||
* swarm (recursion guard).
|
||||
*/
|
||||
|
||||
import { z } from 'zod';
|
||||
|
||||
import type { BuiltinTool } from '../../../agent/tool';
|
||||
import type { Logger } from '../../../logging';
|
||||
import type { ExecutableToolContext, ExecutableToolResult, ToolExecution } from '../../../loop/types';
|
||||
import type { SessionSubagentHost } from '../../../session/subagent-host';
|
||||
import { toInputJsonSchema } from '../../support/input-schema';
|
||||
import { SwarmCoordinator } from '../../../agent/swarm/coordinator';
|
||||
|
||||
export const SwarmToolInputSchema = z.object({
|
||||
task: z.string().describe('The high-level task to decompose and run as a parallel agent swarm.'),
|
||||
});
|
||||
|
||||
export type SwarmToolInput = z.infer<typeof SwarmToolInputSchema>;
|
||||
|
||||
const SWARM_DESCRIPTION =
|
||||
'Run a task as a self-directed agent swarm: dynamically decompose it into parallel ' +
|
||||
'role-specialized subagents, then synthesize their outputs into one answer. ' +
|
||||
'Use for broad, parallelizable tasks (research, multi-file analysis). ' +
|
||||
'Subagents run in isolated contexts and cannot themselves launch swarms.';
|
||||
|
||||
const DEFAULT_MAX_CONCURRENCY = 4;
|
||||
|
||||
export class SwarmTool implements BuiltinTool<SwarmToolInput> {
|
||||
readonly name: string = 'Swarm';
|
||||
readonly description: string = SWARM_DESCRIPTION;
|
||||
readonly parameters: Record<string, unknown> = toInputJsonSchema(SwarmToolInputSchema);
|
||||
private readonly log: Logger | undefined;
|
||||
|
||||
constructor(
|
||||
private readonly subagentHost: SessionSubagentHost,
|
||||
options?: { log?: Logger },
|
||||
) {
|
||||
this.log = options?.log;
|
||||
}
|
||||
|
||||
resolveExecution(args: SwarmToolInput): ToolExecution {
|
||||
return {
|
||||
description: `Running swarm: ${args.task.slice(0, 60)}`,
|
||||
approvalRule: 'Swarm',
|
||||
execute: (ctx) => this.execution(args, ctx),
|
||||
};
|
||||
}
|
||||
|
||||
private async execution(
|
||||
args: SwarmToolInput,
|
||||
ctx: ExecutableToolContext,
|
||||
): Promise<ExecutableToolResult> {
|
||||
const coordinator = new SwarmCoordinator({
|
||||
signal: ctx.signal,
|
||||
maxConcurrency: DEFAULT_MAX_CONCURRENCY,
|
||||
onProgress: (text) => ctx.onUpdate?.({ kind: 'status', text }),
|
||||
spawnSubagent: async ({ profileName, systemPrompt, tools, prompt, description, signal }) => {
|
||||
const handle = await this.subagentHost.spawn(profileName, {
|
||||
parentToolCallId: ctx.toolCallId,
|
||||
prompt,
|
||||
description,
|
||||
runInBackground: false,
|
||||
signal,
|
||||
profileOverride: { systemPrompt, tools },
|
||||
});
|
||||
return handle.completion;
|
||||
},
|
||||
});
|
||||
|
||||
try {
|
||||
const output = await coordinator.run(args.task);
|
||||
return { output };
|
||||
} catch (err) {
|
||||
const message = err instanceof Error ? err.message : String(err);
|
||||
this.log?.error(`swarm failed: ${message}`);
|
||||
return { output: `Swarm failed: ${message}`, isError: true };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -8,6 +8,7 @@ export * from '../cron/cron-list';
|
|||
export * from './collaboration/agent';
|
||||
export * from './collaboration/ask-user';
|
||||
export * from './collaboration/skill-tool';
|
||||
export * from './collaboration/swarm';
|
||||
export * from './file/edit';
|
||||
export * from './file/glob';
|
||||
export * from './file/grep';
|
||||
|
|
|
|||
47
packages/agent-core/test/swarm/swarm-tool.test.ts
Normal file
47
packages/agent-core/test/swarm/swarm-tool.test.ts
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
import { describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { SwarmTool } from '../../src/tools/builtin/collaboration/swarm';
|
||||
import type { SessionSubagentHost } from '../../src/session/subagent-host';
|
||||
|
||||
const PLAN_JSON = JSON.stringify({
|
||||
subtasks: [{ role: 'R', systemPrompt: 'sp', prompt: 'p' }],
|
||||
});
|
||||
|
||||
function fakeHost(): SessionSubagentHost {
|
||||
const spawn = vi.fn(async (profileName: string) => {
|
||||
const result =
|
||||
profileName === 'swarm-planner'
|
||||
? PLAN_JSON
|
||||
: profileName === 'swarm-synthesizer'
|
||||
? 'FINAL'
|
||||
: 'worker-out';
|
||||
return { agentId: 'a', profileName, resumed: false, completion: Promise.resolve({ result }) };
|
||||
});
|
||||
return { spawn } as unknown as SessionSubagentHost;
|
||||
}
|
||||
|
||||
describe('SwarmTool', () => {
|
||||
it('exposes a task parameter and an approval rule', () => {
|
||||
const tool = new SwarmTool(fakeHost());
|
||||
expect(tool.name).toBe('Swarm');
|
||||
const exec = tool.resolveExecution({ task: 'hello' });
|
||||
expect('approvalRule' in exec && exec.approvalRule).toBe('Swarm');
|
||||
});
|
||||
|
||||
it('runs the coordinator and returns the synthesized output', async () => {
|
||||
const tool = new SwarmTool(fakeHost());
|
||||
const exec = tool.resolveExecution({ task: 'do it' });
|
||||
if (!('execute' in exec)) throw new Error('expected runnable execution');
|
||||
const updates: string[] = [];
|
||||
const result = await exec.execute({
|
||||
turnId: 't1',
|
||||
toolCallId: 'tc1',
|
||||
signal: new AbortController().signal,
|
||||
onUpdate: (u) => {
|
||||
if (u.text !== undefined) updates.push(u.text);
|
||||
},
|
||||
});
|
||||
expect('output' in result && result.output).toBe('FINAL');
|
||||
expect(updates.length).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue