mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-08-18 05:05:28 +00:00
516 lines
16 KiB
TypeScript
516 lines
16 KiB
TypeScript
import { ChatProviderError } from '#/errors';
|
|
import type { Message, StreamedMessagePart, ToolCall } from '#/message';
|
|
import type { ChatProvider, StreamedMessage, ThinkingEffort } from '#/provider';
|
|
import { step } from './fixtures/step';
|
|
import type { Tool } from '#/tool';
|
|
import { toolOk } from './fixtures/simple-toolset';
|
|
import type { ToolResult, Toolset } from './fixtures/simple-toolset';
|
|
import type { TokenUsage } from '#/usage';
|
|
import { describe, expect, it } from 'vitest';
|
|
function createMockStream(
|
|
parts: StreamedMessagePart[],
|
|
opts?: { id?: string; usage?: TokenUsage },
|
|
): StreamedMessage {
|
|
return {
|
|
get id(): string | null {
|
|
return opts?.id ?? null;
|
|
},
|
|
get usage(): TokenUsage | null {
|
|
return opts?.usage ?? null;
|
|
},
|
|
finishReason: null,
|
|
rawFinishReason: null,
|
|
async *[Symbol.asyncIterator](): AsyncIterator<StreamedMessagePart> {
|
|
for (const part of parts) {
|
|
yield part;
|
|
}
|
|
},
|
|
};
|
|
}
|
|
|
|
function createMockProvider(stream: StreamedMessage): ChatProvider {
|
|
return {
|
|
name: 'mock',
|
|
modelName: 'mock-model',
|
|
thinkingEffort: null,
|
|
generate(_systemPrompt: string, _tools: Tool[], _history: Message[]): Promise<StreamedMessage> {
|
|
return Promise.resolve(stream);
|
|
},
|
|
withThinking(_effort: ThinkingEffort): ChatProvider {
|
|
return this;
|
|
},
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Create a mock toolset that returns a simple success result for any tool call.
|
|
*/
|
|
function createSimpleMockToolset(
|
|
handler?: (toolCall: ToolCall) => Promise<ToolResult> | ToolResult,
|
|
): Toolset {
|
|
return {
|
|
tools: [
|
|
{
|
|
name: 'plus',
|
|
description: 'Add two numbers',
|
|
parameters: {
|
|
type: 'object',
|
|
properties: {
|
|
a: { type: 'integer' },
|
|
b: { type: 'integer' },
|
|
},
|
|
},
|
|
},
|
|
],
|
|
handle:
|
|
handler ??
|
|
((toolCall: ToolCall): ToolResult => ({
|
|
toolCallId: toolCall.id,
|
|
returnValue: toolOk({ output: 'mock-result' }),
|
|
})),
|
|
};
|
|
}
|
|
describe('step()', () => {
|
|
it('returns a StepResult with message, toolCalls, and toolResults', async () => {
|
|
const plusToolCall: ToolCall = {
|
|
type: 'function',
|
|
id: 'plus#123',
|
|
name: 'plus', arguments: '{"a": 1, "b": 2}',
|
|
};
|
|
const stream = createMockStream([{ type: 'text', text: 'Hello, world!' }, plusToolCall]);
|
|
const provider = createMockProvider(stream);
|
|
|
|
const toolset = createSimpleMockToolset(
|
|
(toolCall: ToolCall): ToolResult => ({
|
|
toolCallId: toolCall.id,
|
|
returnValue: toolOk({ output: '3' }),
|
|
}),
|
|
);
|
|
|
|
const result = await step(provider, '', toolset, []);
|
|
|
|
expect(result.message.content).toEqual([{ type: 'text', text: 'Hello, world!' }]);
|
|
expect(result.toolCalls).toHaveLength(1);
|
|
expect(result.toolCalls[0]).toEqual(plusToolCall);
|
|
|
|
const toolResults = await result.toolResults();
|
|
expect(toolResults).toHaveLength(1);
|
|
expect(toolResults[0]!.toolCallId).toBe('plus#123');
|
|
expect(toolResults[0]!.returnValue.output).toBe('3');
|
|
});
|
|
|
|
it('fires onMessagePart and onToolResult callbacks', async () => {
|
|
const plusToolCall: ToolCall = {
|
|
type: 'function',
|
|
id: 'plus#123',
|
|
name: 'plus', arguments: '{"a": 1, "b": 2}',
|
|
};
|
|
const stream = createMockStream([{ type: 'text', text: 'Hello, world!' }, plusToolCall]);
|
|
const provider = createMockProvider(stream);
|
|
|
|
const toolset = createSimpleMockToolset(
|
|
(toolCall: ToolCall): ToolResult => ({
|
|
toolCallId: toolCall.id,
|
|
returnValue: toolOk({ output: '3' }),
|
|
}),
|
|
);
|
|
|
|
const outputParts: StreamedMessagePart[] = [];
|
|
const collectedToolResults: ToolResult[] = [];
|
|
|
|
const stepResult = await step(provider, '', toolset, [], {
|
|
onMessagePart(part: StreamedMessagePart): void {
|
|
outputParts.push(part);
|
|
},
|
|
onToolResult(toolResult: ToolResult): void {
|
|
collectedToolResults.push(toolResult);
|
|
},
|
|
});
|
|
|
|
// onMessagePart should receive each raw streamed part.
|
|
expect(outputParts).toHaveLength(2);
|
|
expect(outputParts[0]).toEqual({ type: 'text', text: 'Hello, world!' });
|
|
|
|
// Await results so onToolResult fires.
|
|
const toolResults = await stepResult.toolResults();
|
|
|
|
expect(toolResults).toHaveLength(1);
|
|
expect(collectedToolResults).toHaveLength(1);
|
|
expect(collectedToolResults[0]!.toolCallId).toBe('plus#123');
|
|
});
|
|
|
|
it('returns empty toolResults when no tool calls', async () => {
|
|
const stream = createMockStream([{ type: 'text', text: 'No tools needed.' }]);
|
|
const provider = createMockProvider(stream);
|
|
const toolset = createSimpleMockToolset();
|
|
|
|
const stepResult = await step(provider, '', toolset, []);
|
|
const toolResults = await stepResult.toolResults();
|
|
|
|
expect(stepResult.toolCalls).toHaveLength(0);
|
|
expect(toolResults).toHaveLength(0);
|
|
});
|
|
|
|
it('preserves stream id and usage', async () => {
|
|
const usage: TokenUsage = {
|
|
inputOther: 100,
|
|
output: 50,
|
|
inputCacheRead: 200,
|
|
inputCacheCreation: 10,
|
|
};
|
|
const stream = createMockStream([{ type: 'text', text: 'hi' }], {
|
|
id: 'msg-123',
|
|
usage,
|
|
});
|
|
const provider = createMockProvider(stream);
|
|
const toolset = createSimpleMockToolset();
|
|
|
|
const stepResult = await step(provider, '', toolset, []);
|
|
|
|
expect(stepResult.id).toBe('msg-123');
|
|
expect(stepResult.usage).toEqual(usage);
|
|
});
|
|
|
|
it('handles multiple tool calls', async () => {
|
|
const tc1: ToolCall = {
|
|
type: 'function',
|
|
id: 'call-1',
|
|
name: 'plus', arguments: '{"a":1,"b":2}',
|
|
};
|
|
const tc2: ToolCall = {
|
|
type: 'function',
|
|
id: 'call-2',
|
|
name: 'plus', arguments: '{"a":3,"b":4}',
|
|
};
|
|
const stream = createMockStream([tc1, tc2]);
|
|
const provider = createMockProvider(stream);
|
|
|
|
const toolset = createSimpleMockToolset(
|
|
(toolCall: ToolCall): ToolResult => ({
|
|
toolCallId: toolCall.id,
|
|
returnValue: toolOk({ output: `result-${toolCall.id}` }),
|
|
}),
|
|
);
|
|
|
|
const stepResult = await step(provider, '', toolset, []);
|
|
const toolResults = await stepResult.toolResults();
|
|
|
|
expect(stepResult.toolCalls).toHaveLength(2);
|
|
expect(toolResults).toHaveLength(2);
|
|
expect(toolResults[0]!.toolCallId).toBe('call-1');
|
|
expect(toolResults[1]!.toolCallId).toBe('call-2');
|
|
});
|
|
it('does not dispatch tool calls when provider throws mid-stream', async () => {
|
|
// With deferred dispatch, onToolCall only fires *after* the stream
|
|
// has drained. A mid-stream provider error therefore leaves zero
|
|
// pending tool futures — toolset.handle must never be called.
|
|
//
|
|
// This is the desired behaviour for parallel tool call streams
|
|
// where arguments can interleave; dispatching eagerly would hand
|
|
// tools half-parsed JSON. See generate.test.ts for the full
|
|
// rationale.
|
|
let toolHandleCalls = 0;
|
|
|
|
const tc1: ToolCall = {
|
|
type: 'function',
|
|
id: 'call-first',
|
|
name: 'slow', arguments: '{}',
|
|
};
|
|
const tc2: ToolCall = {
|
|
type: 'function',
|
|
id: 'call-second',
|
|
name: 'slow', arguments: '{}',
|
|
};
|
|
|
|
const throwingStream: StreamedMessage = {
|
|
get id(): string | null {
|
|
return null;
|
|
},
|
|
get usage(): TokenUsage | null {
|
|
return null;
|
|
},
|
|
finishReason: null,
|
|
rawFinishReason: null,
|
|
async *[Symbol.asyncIterator](): AsyncIterator<StreamedMessagePart> {
|
|
yield tc1;
|
|
yield tc2;
|
|
throw new ChatProviderError('provider blew up mid-stream');
|
|
},
|
|
};
|
|
|
|
const provider: ChatProvider = {
|
|
name: 'mock-throwing',
|
|
modelName: 'mock-model',
|
|
thinkingEffort: null,
|
|
generate(): Promise<StreamedMessage> {
|
|
return Promise.resolve(throwingStream);
|
|
},
|
|
withThinking(_effort: ThinkingEffort): ChatProvider {
|
|
return this;
|
|
},
|
|
};
|
|
|
|
const toolset: Toolset = {
|
|
tools: [],
|
|
handle(toolCall: ToolCall): Promise<ToolResult> {
|
|
toolHandleCalls++;
|
|
return Promise.resolve({
|
|
toolCallId: toolCall.id,
|
|
returnValue: toolOk({ output: 'should-never-run' }),
|
|
});
|
|
},
|
|
};
|
|
|
|
await expect(step(provider, '', toolset, [])).rejects.toThrow(ChatProviderError);
|
|
|
|
// No tools were dispatched because onToolCall is deferred until
|
|
// after the stream completes successfully.
|
|
expect(toolHandleCalls).toBe(0);
|
|
});
|
|
|
|
it('mid-stream provider error does not leak unhandled rejections', async () => {
|
|
const tc1: ToolCall = {
|
|
type: 'function',
|
|
id: 'call-rejected',
|
|
name: 'boom', arguments: '{}',
|
|
};
|
|
const tc2: ToolCall = {
|
|
type: 'function',
|
|
id: 'call-next',
|
|
name: 'boom', arguments: '{}',
|
|
};
|
|
|
|
const throwingStream: StreamedMessage = {
|
|
get id(): string | null {
|
|
return null;
|
|
},
|
|
get usage(): TokenUsage | null {
|
|
return null;
|
|
},
|
|
finishReason: null,
|
|
rawFinishReason: null,
|
|
async *[Symbol.asyncIterator](): AsyncIterator<StreamedMessagePart> {
|
|
yield tc1;
|
|
yield tc2;
|
|
throw new ChatProviderError('kaboom');
|
|
},
|
|
};
|
|
|
|
const provider: ChatProvider = {
|
|
name: 'mock-throwing',
|
|
modelName: 'mock-model',
|
|
thinkingEffort: null,
|
|
generate(): Promise<StreamedMessage> {
|
|
return Promise.resolve(throwingStream);
|
|
},
|
|
withThinking(_effort: ThinkingEffort): ChatProvider {
|
|
return this;
|
|
},
|
|
};
|
|
|
|
// Toolset whose handle() returns a rejecting promise. With deferred
|
|
// dispatch, handle() is never even invoked on a mid-stream error,
|
|
// so this never produces a rejected promise. The listener check
|
|
// below still guards against any future regression where step()
|
|
// accidentally drops a reference to a rejected promise.
|
|
const toolset: Toolset = {
|
|
tools: [],
|
|
handle(): Promise<ToolResult> {
|
|
return Promise.reject(new Error('tool failed'));
|
|
},
|
|
};
|
|
|
|
const unhandled: unknown[] = [];
|
|
const listener = (reason: unknown): void => {
|
|
unhandled.push(reason);
|
|
};
|
|
process.on('unhandledRejection', listener);
|
|
try {
|
|
await expect(step(provider, '', toolset, [])).rejects.toThrow(ChatProviderError);
|
|
// Give the event loop a few ticks for any unhandled rejections to fire.
|
|
await new Promise<void>((resolve) => {
|
|
setTimeout(resolve, 10);
|
|
});
|
|
} finally {
|
|
process.off('unhandledRejection', listener);
|
|
}
|
|
|
|
expect(unhandled).toEqual([]);
|
|
});
|
|
|
|
it('does not emit unhandledRejection when toolResults() is not awaited', async () => {
|
|
const stream = createMockStream([
|
|
{
|
|
type: 'function',
|
|
id: 'call-rejected',
|
|
name: 'plus', arguments: '{"a":1,"b":2}',
|
|
},
|
|
]);
|
|
const provider = createMockProvider(stream);
|
|
|
|
const toolset: Toolset = {
|
|
tools: [],
|
|
handle(): Promise<ToolResult> {
|
|
return Promise.reject(new Error('tool failed'));
|
|
},
|
|
};
|
|
|
|
const unhandled: unknown[] = [];
|
|
const listener = (reason: unknown): void => {
|
|
unhandled.push(reason);
|
|
};
|
|
|
|
process.on('unhandledRejection', listener);
|
|
try {
|
|
await step(provider, '', toolset, []);
|
|
await new Promise<void>((resolve) => {
|
|
setTimeout(resolve, 100);
|
|
});
|
|
} finally {
|
|
process.off('unhandledRejection', listener);
|
|
}
|
|
|
|
expect(unhandled).toEqual([]);
|
|
});
|
|
|
|
it('dispatches tool calls after the stream completes (deferred dispatch)', async () => {
|
|
// Positive regression: with interleaved parallel tool call streams,
|
|
// step() must still dispatch every tool exactly once, and each
|
|
// toolset.handle() invocation must receive fully-assembled
|
|
// arguments.
|
|
const seenArgs: Array<{ id: string; args: string | null }> = [];
|
|
|
|
const stream = createMockStream([
|
|
{
|
|
type: 'function',
|
|
id: 'tc-a',
|
|
name: 'plus', arguments: null,
|
|
_streamIndex: 0,
|
|
},
|
|
{
|
|
type: 'function',
|
|
id: 'tc-b',
|
|
name: 'plus', arguments: null,
|
|
_streamIndex: 1,
|
|
},
|
|
{ type: 'tool_call_part', argumentsPart: '{"a":1,', index: 0 },
|
|
{ type: 'tool_call_part', argumentsPart: '{"a":3,', index: 1 },
|
|
{ type: 'tool_call_part', argumentsPart: '"b":2}', index: 0 },
|
|
{ type: 'tool_call_part', argumentsPart: '"b":4}', index: 1 },
|
|
]);
|
|
const provider = createMockProvider(stream);
|
|
|
|
const toolset: Toolset = {
|
|
tools: [],
|
|
handle(toolCall: ToolCall): Promise<ToolResult> {
|
|
seenArgs.push({ id: toolCall.id, args: toolCall.arguments });
|
|
return Promise.resolve({
|
|
toolCallId: toolCall.id,
|
|
returnValue: toolOk({ output: `ok-${toolCall.id}` }),
|
|
});
|
|
},
|
|
};
|
|
|
|
const stepResult = await step(provider, '', toolset, []);
|
|
const toolResults = await stepResult.toolResults();
|
|
|
|
// Every handle() invocation saw the fully-assembled arguments for
|
|
// its own tool call — no partial JSON.
|
|
expect(seenArgs).toEqual([
|
|
{ id: 'tc-a', args: '{"a":1,"b":2}' },
|
|
{ id: 'tc-b', args: '{"a":3,"b":4}' },
|
|
]);
|
|
expect(toolResults).toHaveLength(2);
|
|
expect(toolResults[0]!.toolCallId).toBe('tc-a');
|
|
expect(toolResults[1]!.toolCallId).toBe('tc-b');
|
|
});
|
|
|
|
it('handles sync toolset.handle() return', async () => {
|
|
const tc: ToolCall = {
|
|
type: 'function',
|
|
id: 'call-sync',
|
|
name: 'plus', arguments: '{}',
|
|
};
|
|
const stream = createMockStream([tc]);
|
|
const provider = createMockProvider(stream);
|
|
|
|
// Toolset that returns a sync ToolResult (not a Promise).
|
|
const toolset: Toolset = {
|
|
tools: [],
|
|
handle(toolCall: ToolCall): ToolResult {
|
|
return {
|
|
toolCallId: toolCall.id,
|
|
returnValue: toolOk({ output: 'sync-result' }),
|
|
};
|
|
},
|
|
};
|
|
|
|
const stepResult = await step(provider, '', toolset, []);
|
|
const toolResults = await stepResult.toolResults();
|
|
|
|
expect(toolResults).toHaveLength(1);
|
|
expect(toolResults[0]!.returnValue.output).toBe('sync-result');
|
|
});
|
|
|
|
describe('finishReason propagation', () => {
|
|
function streamWithFinish(
|
|
parts: StreamedMessagePart[],
|
|
finishReason:
|
|
| 'completed'
|
|
| 'truncated'
|
|
| 'tool_calls'
|
|
| 'filtered'
|
|
| 'paused'
|
|
| 'other'
|
|
| null,
|
|
rawFinishReason: string | null,
|
|
): StreamedMessage {
|
|
return {
|
|
id: 'mock',
|
|
usage: null,
|
|
finishReason,
|
|
rawFinishReason,
|
|
async *[Symbol.asyncIterator](): AsyncIterator<StreamedMessagePart> {
|
|
for (const part of parts) {
|
|
yield part;
|
|
}
|
|
},
|
|
};
|
|
}
|
|
|
|
it('copies stream.finishReason onto the StepResult', async () => {
|
|
const stream = streamWithFinish([{ type: 'text', text: 'hi' }], 'truncated', 'length');
|
|
const provider = createMockProvider(stream);
|
|
const toolset = createSimpleMockToolset();
|
|
const result = await step(provider, '', toolset, []);
|
|
expect(result.finishReason).toBe('truncated');
|
|
expect(result.rawFinishReason).toBe('length');
|
|
await result.toolResults();
|
|
});
|
|
|
|
it('copies null finishReason onto the StepResult', async () => {
|
|
const stream = streamWithFinish([{ type: 'text', text: 'hi' }], null, null);
|
|
const provider = createMockProvider(stream);
|
|
const toolset = createSimpleMockToolset();
|
|
const result = await step(provider, '', toolset, []);
|
|
expect(result.finishReason).toBeNull();
|
|
expect(result.rawFinishReason).toBeNull();
|
|
await result.toolResults();
|
|
});
|
|
|
|
it('propagates tool_calls finishReason when the model requested tools', async () => {
|
|
const tc: ToolCall = {
|
|
type: 'function',
|
|
id: 'tc-propagate',
|
|
name: 'plus', arguments: '{"a":1,"b":2}',
|
|
};
|
|
const stream = streamWithFinish([tc], 'tool_calls', 'tool_calls');
|
|
const provider = createMockProvider(stream);
|
|
const toolset = createSimpleMockToolset();
|
|
const result = await step(provider, '', toolset, []);
|
|
expect(result.finishReason).toBe('tool_calls');
|
|
expect(result.rawFinishReason).toBe('tool_calls');
|
|
await result.toolResults();
|
|
});
|
|
});
|
|
});
|