mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-08-15 03:35:32 +00:00
58898de020
* feat(agent-core): cap AgentSwarm concurrency via env var Add KIMI_CODE_AGENT_SWARM_MAX_CONCURRENCY to limit how many subagents run concurrently during the initial ramp, so large swarms do not trip provider rate limits as easily. Leave it unset to keep the previous uncapped ramp behavior. * chore: drop ignored agent-core from changeset * fix(agent-core): fail fast on invalid AgentSwarm concurrency cap
905 lines
28 KiB
TypeScript
905 lines
28 KiB
TypeScript
import { createControlledPromise } from '@antfu/utils';
|
|
import { APIProviderRateLimitError } from '@moonshot-ai/kosong';
|
|
import { describe, expect, it, vi } from 'vitest';
|
|
|
|
import {
|
|
type QueuedSubagentTask,
|
|
type RunSubagentOptions,
|
|
type SpawnSubagentOptions,
|
|
type SubagentHandle,
|
|
} from '../../src/session/subagent-host';
|
|
import {
|
|
SubagentBatch,
|
|
resolveSwarmMaxConcurrency,
|
|
type SubagentBatchLauncher,
|
|
type SubagentResult,
|
|
type SubagentSuspendedEvent,
|
|
} from '../../src/session/subagent-batch';
|
|
import { userCancellationReason } from '../../src/utils/abort';
|
|
|
|
const signal = new AbortController().signal;
|
|
|
|
describe('SubagentBatch scheduling contract', () => {
|
|
it('normal phase starts five tasks immediately, then one task every 700ms', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const { runBatch, attempts } = createMockBatchRunner();
|
|
const running = runBatch(
|
|
Array.from({ length: 9 }, (_, index) => queuedTask(index + 1)),
|
|
{ signal },
|
|
);
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(699);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(1);
|
|
expect(attempts).toHaveLength(6);
|
|
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(attempts).toHaveLength(7);
|
|
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(attempts).toHaveLength(8);
|
|
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(attempts).toHaveLength(9);
|
|
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(attempts).toHaveLength(9);
|
|
|
|
attempts.forEach((attempt, index) => {
|
|
attempt.outcome.resolve({
|
|
task: attempt.task,
|
|
agentId: `agent-${String(index + 1)}`,
|
|
status: 'completed',
|
|
result: `result ${String(index + 1)}`,
|
|
});
|
|
});
|
|
const results = await running;
|
|
|
|
expect(results).toHaveLength(9);
|
|
expect(results.every((result) => result.status === 'completed')).toBe(true);
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('rate-limit phase starts when the first provider rate limit stops the normal ramp', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const controller = new AbortController();
|
|
const { runBatch, attempts } = createMockBatchRunner();
|
|
const running = runBatch(Array.from({ length: 9 }, (_, index) => queuedTask(index + 1)), {
|
|
signal: controller.signal,
|
|
});
|
|
void running.catch(() => {});
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
attempts.forEach((attempt) => {
|
|
attempt.markReady();
|
|
});
|
|
|
|
attempts[0]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-1' });
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
attempts[1]!.outcome.resolve({
|
|
task: attempts[1]!.task,
|
|
agentId: 'agent-2',
|
|
status: 'completed',
|
|
result: 'completed 2',
|
|
});
|
|
await vi.advanceTimersByTimeAsync(3000);
|
|
expect(attempts).toHaveLength(6);
|
|
expect(attempts[5]!.task.data).toBe(1);
|
|
expect(attempts[5]!.retryAgentId).toBe('agent-1');
|
|
|
|
controller.abort();
|
|
await expect(running).rejects.toThrow();
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('user cancellation returns completed, started, and not-started task results', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const controller = new AbortController();
|
|
const { runBatch, attempts } = createMockBatchRunner();
|
|
const running = runBatch(Array.from({ length: 6 }, (_, index) => queuedTask(index + 1)), {
|
|
signal: controller.signal,
|
|
});
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
attempts[0]!.outcome.resolve({
|
|
task: attempts[0]!.task,
|
|
agentId: 'agent-1',
|
|
status: 'completed',
|
|
result: 'completed 1',
|
|
});
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
|
|
controller.abort(userCancellationReason());
|
|
const results = await running;
|
|
|
|
expect(results.map((result) => ({
|
|
data: result.task.data,
|
|
agentId: result.agentId,
|
|
status: result.status,
|
|
state: result.state,
|
|
result: result.result,
|
|
error: result.error,
|
|
}))).toEqual([
|
|
{
|
|
data: 1,
|
|
agentId: 'agent-1',
|
|
status: 'completed',
|
|
state: undefined,
|
|
result: 'completed 1',
|
|
error: undefined,
|
|
},
|
|
{
|
|
data: 2,
|
|
agentId: 'agent-2',
|
|
status: 'aborted',
|
|
state: 'started',
|
|
result: undefined,
|
|
error: 'The user manually interrupted this subagent batch before this subagent finished.',
|
|
},
|
|
{
|
|
data: 3,
|
|
agentId: 'agent-3',
|
|
status: 'aborted',
|
|
state: 'started',
|
|
result: undefined,
|
|
error: 'The user manually interrupted this subagent batch before this subagent finished.',
|
|
},
|
|
{
|
|
data: 4,
|
|
agentId: 'agent-4',
|
|
status: 'aborted',
|
|
state: 'started',
|
|
result: undefined,
|
|
error: 'The user manually interrupted this subagent batch before this subagent finished.',
|
|
},
|
|
{
|
|
data: 5,
|
|
agentId: 'agent-5',
|
|
status: 'aborted',
|
|
state: 'started',
|
|
result: undefined,
|
|
error: 'The user manually interrupted this subagent batch before this subagent finished.',
|
|
},
|
|
{
|
|
data: 6,
|
|
agentId: undefined,
|
|
status: 'aborted',
|
|
state: 'not_started',
|
|
result: undefined,
|
|
error: 'The user manually interrupted this subagent batch before this subagent was started.',
|
|
},
|
|
]);
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('normal phase keeps processing completions while waiting for the next launch', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const { runBatch, attempts } = createMockBatchRunner();
|
|
const running = runBatch(
|
|
Array.from({ length: 6 }, (_, index) => queuedTask(index + 1)),
|
|
{ signal },
|
|
);
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
attempts[0]!.outcome.resolve({
|
|
task: attempts[0]!.task,
|
|
agentId: 'agent-1',
|
|
status: 'completed',
|
|
result: 'completed 1',
|
|
});
|
|
|
|
await vi.advanceTimersByTimeAsync(699);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(1);
|
|
expect(attempts).toHaveLength(6);
|
|
|
|
attempts.slice(1).forEach((attempt, index) => {
|
|
attempt.outcome.resolve({
|
|
task: attempt.task,
|
|
agentId: `agent-${String(index + 2)}`,
|
|
status: 'completed',
|
|
result: `completed ${String(index + 2)}`,
|
|
});
|
|
});
|
|
await expect(running).resolves.toHaveLength(6);
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('rate-limit phase requeues 429 tasks, emits suspended, and throttles launches', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const controller = new AbortController();
|
|
const onSuspended = vi.fn();
|
|
const { runBatch, attempts } = createMockBatchRunner({ onSuspended });
|
|
const running = runBatch(Array.from({ length: 8 }, (_, index) => queuedTask(index + 1)), {
|
|
signal: controller.signal,
|
|
});
|
|
void running.catch(() => {});
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
attempts.forEach((attempt) => {
|
|
attempt.markReady();
|
|
});
|
|
attempts[0]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-1' });
|
|
attempts[1]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-2' });
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(onSuspended).toHaveBeenCalledTimes(2);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(500);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(2500);
|
|
expect(attempts).toHaveLength(6);
|
|
expect(attempts[5]!.task.data).toBe(2);
|
|
expect(attempts[5]!.retryAgentId).toBe('agent-2');
|
|
|
|
controller.abort();
|
|
await expect(running).rejects.toThrow();
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('fails the only unfinished task on provider rate limit instead of suspending forever', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const onSuspended = vi.fn();
|
|
const { runBatch, attempts } = createMockBatchRunner({ onSuspended });
|
|
const running = runBatch(Array.from({ length: 2 }, (_, index) => queuedTask(index + 1)), {
|
|
signal,
|
|
});
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(2);
|
|
attempts.forEach((attempt) => {
|
|
attempt.markReady();
|
|
});
|
|
|
|
attempts[0]!.outcome.resolve({
|
|
task: attempts[0]!.task,
|
|
agentId: 'agent-1',
|
|
status: 'completed',
|
|
result: 'completed 1',
|
|
});
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
|
|
attempts[1]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-2' });
|
|
await expect(running).resolves.toMatchObject([
|
|
{
|
|
task: { data: 1 },
|
|
agentId: 'agent-1',
|
|
status: 'completed',
|
|
result: 'completed 1',
|
|
},
|
|
{
|
|
task: { data: 2 },
|
|
agentId: 'agent-2',
|
|
status: 'failed',
|
|
state: 'started',
|
|
error: 'Rate limited',
|
|
},
|
|
]);
|
|
expect(onSuspended).not.toHaveBeenCalled();
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('rate-limit capacity blocks launches while active attempts fill all slots', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const controller = new AbortController();
|
|
const { runBatch, attempts } = createMockBatchRunner();
|
|
const running = runBatch(Array.from({ length: 12 }, (_, index) => queuedTask(index + 1)), {
|
|
signal: controller.signal,
|
|
});
|
|
void running.catch(() => {});
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
attempts.slice(0, 5).forEach((attempt) => {
|
|
attempt.markReady();
|
|
});
|
|
|
|
for (let count = 6; count <= 12; count += 1) {
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(attempts).toHaveLength(count);
|
|
attempts[count - 1]!.markReady();
|
|
}
|
|
|
|
attempts.slice(0, 12).forEach((attempt) => {
|
|
attempt.markReady();
|
|
});
|
|
|
|
for (let index = 0; index < 1; index += 1) {
|
|
attempts[index]!.outcome.resolve({
|
|
type: 'rate_limited',
|
|
agentId: `agent-${String(index + 1)}`,
|
|
});
|
|
}
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
|
|
await vi.advanceTimersByTimeAsync(3000);
|
|
expect(attempts).toHaveLength(12);
|
|
|
|
controller.abort();
|
|
await expect(running).rejects.toThrow();
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('rate-limit recovery adds one capacity slot after three quiet minutes with queued work', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const controller = new AbortController();
|
|
const { runBatch, attempts } = createMockBatchRunner();
|
|
const running = runBatch(Array.from({ length: 6 }, (_, index) => queuedTask(index + 1)), {
|
|
signal: controller.signal,
|
|
});
|
|
void running.catch(() => {});
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
attempts.forEach((attempt) => {
|
|
attempt.markReady();
|
|
});
|
|
|
|
attempts[0]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-1' });
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(2000);
|
|
attempts[1]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-2' });
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(2000);
|
|
attempts[2]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-3' });
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(2000);
|
|
attempts[3]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-4' });
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(179_999);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(1);
|
|
expect(attempts).toHaveLength(6);
|
|
expect(attempts[5]!.task.data).toBe(4);
|
|
expect(attempts[5]!.retryAgentId).toBe('agent-4');
|
|
|
|
controller.abort();
|
|
await expect(running).rejects.toThrow();
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('rate-limit phase keeps launches bounded after repeated 429s', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const controller = new AbortController();
|
|
const { runBatch, attempts } = createMockBatchRunner();
|
|
const running = runBatch(Array.from({ length: 8 }, (_, index) => queuedTask(index + 1)), {
|
|
signal: controller.signal,
|
|
});
|
|
void running.catch(() => {});
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
attempts.forEach((attempt) => {
|
|
attempt.markReady();
|
|
});
|
|
|
|
for (let index = 0; index < 3; index += 1) {
|
|
attempts[index]!.outcome.resolve({
|
|
type: 'rate_limited',
|
|
agentId: `agent-${String(index + 1)}`,
|
|
});
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
}
|
|
|
|
await vi.advanceTimersByTimeAsync(3000);
|
|
expect(attempts).toHaveLength(6);
|
|
expect(attempts[5]!.task.data).toBe(3);
|
|
expect(attempts[5]!.retryAgentId).toBe('agent-3');
|
|
|
|
await vi.advanceTimersByTimeAsync(3000);
|
|
expect(attempts).toHaveLength(7);
|
|
expect(attempts[6]!.task.data).toBe(2);
|
|
expect(attempts[6]!.retryAgentId).toBe('agent-2');
|
|
|
|
controller.abort();
|
|
await expect(running).rejects.toThrow();
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('rate-limit phase schedules another launch after starting while capacity remains', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const controller = new AbortController();
|
|
const { runBatch, attempts } = createMockBatchRunner();
|
|
const running = runBatch(Array.from({ length: 8 }, (_, index) => queuedTask(index + 1)), {
|
|
signal: controller.signal,
|
|
});
|
|
void running.catch(() => {});
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
attempts.forEach((attempt) => {
|
|
attempt.markReady();
|
|
});
|
|
|
|
attempts[0]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-1' });
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
attempts[1]!.outcome.resolve({
|
|
task: attempts[1]!.task,
|
|
agentId: 'agent-2',
|
|
status: 'completed',
|
|
result: 'completed 2',
|
|
});
|
|
attempts[2]!.outcome.resolve({
|
|
task: attempts[2]!.task,
|
|
agentId: 'agent-3',
|
|
status: 'completed',
|
|
result: 'completed 3',
|
|
});
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(2_999);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(1);
|
|
expect(attempts).toHaveLength(6);
|
|
expect(attempts[5]!.task.data).toBe(1);
|
|
expect(attempts[5]!.retryAgentId).toBe('agent-1');
|
|
|
|
await vi.advanceTimersByTimeAsync(2_999);
|
|
expect(attempts).toHaveLength(6);
|
|
|
|
await vi.advanceTimersByTimeAsync(1);
|
|
expect(attempts).toHaveLength(7);
|
|
expect(attempts[6]!.task.data).toBe(6);
|
|
expect(attempts[6]!.retryAgentId).toBeUndefined();
|
|
|
|
controller.abort();
|
|
await expect(running).rejects.toThrow();
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('task timeout fails only that task', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const { runBatch, attempts } = createMockBatchRunner();
|
|
const running = runBatch([{ ...queuedTask(1), timeout: 10_000 }], { signal });
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
attempts[0]!.markReady();
|
|
|
|
await vi.advanceTimersByTimeAsync(9999);
|
|
expect(attempts).toHaveLength(1);
|
|
|
|
await vi.advanceTimersByTimeAsync(1);
|
|
await expect(running).resolves.toMatchObject([
|
|
{
|
|
task: { data: 1 },
|
|
agentId: 'agent-1',
|
|
status: 'failed',
|
|
state: 'started',
|
|
error: 'Subagent timed out.',
|
|
},
|
|
]);
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('does not spend task timeout while the task is queued', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
let settled = false;
|
|
const { runBatch, attempts } = createMockBatchRunner();
|
|
const running = runBatch(
|
|
[
|
|
...Array.from({ length: 5 }, (_, index) => queuedTask(index + 1)),
|
|
{ ...queuedTask(6), timeout: 1000 },
|
|
],
|
|
{ signal },
|
|
);
|
|
void running.finally(() => {
|
|
settled = true;
|
|
});
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(699);
|
|
expect(attempts).toHaveLength(5);
|
|
|
|
await vi.advanceTimersByTimeAsync(1);
|
|
expect(attempts).toHaveLength(6);
|
|
|
|
await vi.advanceTimersByTimeAsync(999);
|
|
expect(settled).toBe(false);
|
|
|
|
attempts.slice(0, 5).forEach((attempt, index) => {
|
|
attempt.outcome.resolve({
|
|
task: attempt.task,
|
|
agentId: `agent-${String(index + 1)}`,
|
|
status: 'completed',
|
|
result: `completed ${String(index + 1)}`,
|
|
});
|
|
});
|
|
await vi.advanceTimersByTimeAsync(1);
|
|
|
|
await expect(running).resolves.toMatchObject([
|
|
{ task: { data: 1 }, status: 'completed' },
|
|
{ task: { data: 2 }, status: 'completed' },
|
|
{ task: { data: 3 }, status: 'completed' },
|
|
{ task: { data: 4 }, status: 'completed' },
|
|
{ task: { data: 5 }, status: 'completed' },
|
|
{
|
|
task: { data: 6 },
|
|
agentId: 'agent-6',
|
|
status: 'failed',
|
|
state: 'started',
|
|
error: 'Subagent timed out.',
|
|
},
|
|
]);
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
|
|
it('rate-limit phase continues launching after rate-limited attempts settle', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const controller = new AbortController();
|
|
const { runBatch, attempts } = createMockBatchRunner({
|
|
readyDelay: (attemptIndex) => (attemptIndex >= 7 ? 100 : undefined),
|
|
});
|
|
|
|
const running = runBatch(
|
|
Array.from({ length: 9 }, (_, index) => queuedTask(index + 1)),
|
|
{ signal: controller.signal },
|
|
);
|
|
void running.catch(() => {});
|
|
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
attempts.slice(0, 5).forEach((attempt) => {
|
|
attempt.markReady();
|
|
});
|
|
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(attempts).toHaveLength(6);
|
|
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(attempts).toHaveLength(7);
|
|
|
|
attempts[5]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-6' });
|
|
attempts[6]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-7' });
|
|
attempts[0]!.outcome.resolve({
|
|
task: attempts[0]!.task,
|
|
agentId: 'agent-1',
|
|
status: 'completed',
|
|
result: 'completed 1',
|
|
});
|
|
attempts[1]!.outcome.resolve({
|
|
task: attempts[1]!.task,
|
|
agentId: 'agent-2',
|
|
status: 'completed',
|
|
result: 'completed 2',
|
|
});
|
|
await vi.advanceTimersByTimeAsync(12_000);
|
|
expect(attempts).toHaveLength(8);
|
|
expect(attempts[7]!.task.data).toBe(7);
|
|
expect(attempts[7]!.retryAgentId).toBe('agent-7');
|
|
|
|
controller.abort();
|
|
await expect(running).rejects.toThrow();
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
});
|
|
|
|
describe('resolveSwarmMaxConcurrency', () => {
|
|
it('returns undefined when the variable is unset', () => {
|
|
expect(resolveSwarmMaxConcurrency({})).toBeUndefined();
|
|
});
|
|
|
|
it('returns undefined for empty or whitespace-only values', () => {
|
|
expect(
|
|
resolveSwarmMaxConcurrency({ KIMI_CODE_AGENT_SWARM_MAX_CONCURRENCY: '' }),
|
|
).toBeUndefined();
|
|
expect(
|
|
resolveSwarmMaxConcurrency({ KIMI_CODE_AGENT_SWARM_MAX_CONCURRENCY: ' ' }),
|
|
).toBeUndefined();
|
|
});
|
|
|
|
it('throws for non-positive, non-integer, or non-numeric values', () => {
|
|
for (const raw of ['0', '-1', '2.5', 'abc']) {
|
|
expect(() =>
|
|
resolveSwarmMaxConcurrency({ KIMI_CODE_AGENT_SWARM_MAX_CONCURRENCY: raw }),
|
|
).toThrow(/KIMI_CODE_AGENT_SWARM_MAX_CONCURRENCY.*positive integer/);
|
|
}
|
|
});
|
|
|
|
it('returns the integer for a positive integer value', () => {
|
|
expect(resolveSwarmMaxConcurrency({ KIMI_CODE_AGENT_SWARM_MAX_CONCURRENCY: '3' })).toBe(3);
|
|
expect(resolveSwarmMaxConcurrency({ KIMI_CODE_AGENT_SWARM_MAX_CONCURRENCY: ' 8 ' })).toBe(8);
|
|
});
|
|
});
|
|
|
|
describe('SubagentBatch max concurrency cap', () => {
|
|
it('caps in-flight tasks at maxConcurrency during the normal phase', async () => {
|
|
vi.useFakeTimers();
|
|
try {
|
|
const { runBatch, attempts } = createMockBatchRunner({ maxConcurrency: 3 });
|
|
const running = runBatch(Array.from({ length: 9 }, (_, index) => queuedTask(index + 1)), {
|
|
signal,
|
|
});
|
|
const resolved = new Set<number>();
|
|
const resolveOne = (index: number) => {
|
|
const attempt = attempts[index]!;
|
|
resolved.add(index);
|
|
attempt.outcome.resolve({
|
|
task: attempt.task,
|
|
agentId: `agent-${String(index + 1)}`,
|
|
status: 'completed',
|
|
result: `result ${String(index + 1)}`,
|
|
});
|
|
};
|
|
const inFlight = () => attempts.length - resolved.size;
|
|
|
|
// Initial burst is capped at 3 instead of the default 5.
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(3);
|
|
expect(inFlight()).toBe(3);
|
|
|
|
// The 700ms ramp tick does not exceed the cap while all slots are occupied.
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(attempts).toHaveLength(3);
|
|
|
|
// Freeing one slot refills it without exceeding the cap.
|
|
resolveOne(0);
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(4);
|
|
expect(inFlight()).toBeLessThanOrEqual(3);
|
|
|
|
resolveOne(1);
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
expect(inFlight()).toBeLessThanOrEqual(3);
|
|
|
|
// Once the initial burst budget (5) is exhausted, further launches wait for
|
|
// the 700ms ramp tick, but the in-flight count still never exceeds the cap.
|
|
resolveOne(2);
|
|
await vi.advanceTimersByTimeAsync(0);
|
|
expect(attempts).toHaveLength(5);
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(attempts).toHaveLength(6);
|
|
expect(inFlight()).toBeLessThanOrEqual(3);
|
|
|
|
// Drain the remaining attempts.
|
|
for (let index = 3; index < 9; index += 1) {
|
|
resolveOne(index);
|
|
await vi.advanceTimersByTimeAsync(700);
|
|
expect(inFlight()).toBeLessThanOrEqual(3);
|
|
}
|
|
|
|
const results = await running;
|
|
expect(results).toHaveLength(9);
|
|
expect(results.every((result) => result.status === 'completed')).toBe(true);
|
|
} finally {
|
|
vi.useRealTimers();
|
|
}
|
|
});
|
|
});
|
|
|
|
type MockAttemptOutcome<T> =
|
|
| SubagentResult<T>
|
|
| {
|
|
readonly type: 'rate_limited';
|
|
readonly agentId: string;
|
|
};
|
|
|
|
type MockAttemptRecord = {
|
|
readonly task: QueuedSubagentTask<number>;
|
|
readonly retryAgentId?: string;
|
|
readonly markReady: () => void;
|
|
readonly outcome: ReturnType<typeof createControlledPromise<MockAttemptOutcome<number>>>;
|
|
};
|
|
|
|
type MockBatchRunnerOptions = {
|
|
readonly onSuspended?: (event: SubagentSuspendedEvent) => void;
|
|
readonly readyDelay?: (attemptIndex: number) => number | undefined;
|
|
readonly maxConcurrency?: number;
|
|
};
|
|
|
|
function createMockBatchRunner(
|
|
options: MockBatchRunnerOptions = {},
|
|
): {
|
|
readonly runBatch: <T>(
|
|
tasks: readonly QueuedSubagentTask<T>[],
|
|
options?: { readonly signal?: AbortSignal },
|
|
) => Promise<Array<SubagentResult<T>>>;
|
|
readonly attempts: MockAttemptRecord[];
|
|
} {
|
|
const attempts: MockAttemptRecord[] = [];
|
|
let activeTasks: readonly QueuedSubagentTask<unknown>[] = [];
|
|
|
|
const createHandle = <T,>(
|
|
runOptions: RunSubagentOptions,
|
|
agentId: string,
|
|
profileName: string,
|
|
resumed: boolean,
|
|
retryAgentId?: string,
|
|
): SubagentHandle => {
|
|
const task = findMockTask<T>(activeTasks, runOptions);
|
|
const outcome = createControlledPromise<MockAttemptOutcome<T>>();
|
|
const markReady = () => {
|
|
runOptions.onReady?.();
|
|
};
|
|
const attemptIndex = attempts.length;
|
|
attempts.push({
|
|
task: task as unknown as QueuedSubagentTask<number>,
|
|
retryAgentId,
|
|
markReady,
|
|
outcome: outcome as unknown as MockAttemptRecord['outcome'],
|
|
});
|
|
|
|
const delay = options.readyDelay?.(attemptIndex);
|
|
if (delay !== undefined) setTimeout(markReady, delay);
|
|
|
|
return {
|
|
agentId,
|
|
profileName,
|
|
resumed,
|
|
completion: completionFromMockOutcome(outcome, runOptions.signal),
|
|
};
|
|
};
|
|
|
|
const host = {
|
|
spawn: async (spawnOptions: SpawnSubagentOptions) => {
|
|
const task = findMockTask(activeTasks, spawnOptions);
|
|
return createHandle(
|
|
spawnOptions,
|
|
mockAgentId(task, attempts.length),
|
|
spawnOptions.profileName,
|
|
false,
|
|
);
|
|
},
|
|
resume: async (agentId: string, runOptions: RunSubagentOptions) =>
|
|
createHandle(runOptions, agentId, 'subagent', true),
|
|
retry: async (agentId: string, runOptions: RunSubagentOptions) =>
|
|
createHandle(runOptions, agentId, 'subagent', true, agentId),
|
|
suspended: (event: SubagentSuspendedEvent) => {
|
|
options.onSuspended?.(event);
|
|
},
|
|
} satisfies SubagentBatchLauncher;
|
|
|
|
return {
|
|
runBatch: <T,>(
|
|
tasks: readonly QueuedSubagentTask<T>[],
|
|
runOptions?: { readonly signal?: AbortSignal },
|
|
) => {
|
|
activeTasks = tasks.map((task) => ({
|
|
...task,
|
|
signal: task.signal ?? runOptions?.signal,
|
|
}));
|
|
return new SubagentBatch(host, activeTasks as readonly QueuedSubagentTask<T>[], {
|
|
maxConcurrency: options.maxConcurrency,
|
|
}).run();
|
|
},
|
|
attempts,
|
|
};
|
|
}
|
|
|
|
function findMockTask<T>(
|
|
tasks: readonly QueuedSubagentTask<unknown>[],
|
|
options: RunSubagentOptions,
|
|
): QueuedSubagentTask<T> {
|
|
const task = tasks.find(
|
|
(candidate) =>
|
|
candidate.prompt === options.prompt &&
|
|
candidate.parentToolCallId === options.parentToolCallId,
|
|
);
|
|
if (task === undefined) {
|
|
throw new Error(`No mock queued task for prompt "${options.prompt}"`);
|
|
}
|
|
return task as QueuedSubagentTask<T>;
|
|
}
|
|
|
|
function mockAgentId(task: QueuedSubagentTask<unknown>, attemptIndex: number): string {
|
|
if (typeof task.data === 'number') return `agent-${String(task.data)}`;
|
|
return `agent-${String(attemptIndex + 1)}`;
|
|
}
|
|
|
|
function completionFromMockOutcome<T>(
|
|
outcome: ReturnType<typeof createControlledPromise<MockAttemptOutcome<T>>>,
|
|
signal: AbortSignal,
|
|
): SubagentHandle['completion'] {
|
|
return new Promise((resolve, reject) => {
|
|
const abort = () => {
|
|
reject(signal.reason ?? new Error('Aborted'));
|
|
};
|
|
signal.addEventListener('abort', abort, { once: true });
|
|
outcome.then(
|
|
(result) => {
|
|
signal.removeEventListener('abort', abort);
|
|
if (isMockRateLimitOutcome(result)) {
|
|
reject(new APIProviderRateLimitError('Rate limited', result.agentId));
|
|
return;
|
|
}
|
|
if (result.status === 'completed') {
|
|
resolve({ result: result.result ?? '', usage: result.usage });
|
|
return;
|
|
}
|
|
reject(new Error(result.error ?? result.status));
|
|
},
|
|
(error: unknown) => {
|
|
signal.removeEventListener('abort', abort);
|
|
reject(error);
|
|
},
|
|
);
|
|
});
|
|
}
|
|
|
|
function isMockRateLimitOutcome<T>(
|
|
outcome: MockAttemptOutcome<T>,
|
|
): outcome is Extract<MockAttemptOutcome<T>, { readonly type: 'rate_limited' }> {
|
|
return 'type' in outcome && outcome.type === 'rate_limited';
|
|
}
|
|
|
|
function queuedTask(index: number): QueuedSubagentTask<number> {
|
|
return {
|
|
kind: 'spawn',
|
|
data: index,
|
|
profileName: 'coder',
|
|
parentToolCallId: 'call_swarm',
|
|
prompt: `Review item-${String(index)}`,
|
|
description: `Review #${String(index)}`,
|
|
runInBackground: false,
|
|
};
|
|
}
|