mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-07-10 01:39:25 +00:00
refactor: onToolResult, LLMRequestSource
This commit is contained in:
parent
8ba75bbc02
commit
ab64d48154
14 changed files with 326 additions and 116 deletions
|
|
@ -396,7 +396,7 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull
|
|||
{
|
||||
messages,
|
||||
maxOutputSize: compactionMaxOutputSize,
|
||||
requestLogFields: { requestKind: 'full_compaction' },
|
||||
source: { type: 'operation', requestKind: 'full_compaction' },
|
||||
},
|
||||
undefined,
|
||||
signal,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import { createDecorator } from '#/_base/di';
|
||||
import type { UsageRecordContext } from '#/agent/usage';
|
||||
import type {
|
||||
FinishReason,
|
||||
Message,
|
||||
|
|
@ -11,6 +10,19 @@ import type { LogContext } from '#/app/log';
|
|||
|
||||
export type LLMRequestLogFields = Readonly<LogContext>;
|
||||
|
||||
export type LLMRequestSource =
|
||||
| {
|
||||
readonly type: 'turn';
|
||||
readonly turnId: number;
|
||||
readonly step?: number;
|
||||
readonly logFields?: LLMRequestLogFields;
|
||||
}
|
||||
| {
|
||||
readonly type: 'operation';
|
||||
readonly requestKind?: string;
|
||||
readonly logFields?: LLMRequestLogFields;
|
||||
};
|
||||
|
||||
export interface LLMRequestRetryContext {
|
||||
readonly failedAttempt: number;
|
||||
readonly nextAttempt: number;
|
||||
|
|
@ -59,7 +71,7 @@ export interface LLMRequestParams {
|
|||
messages: Message[];
|
||||
tools: readonly Tool[];
|
||||
signal: AbortSignal;
|
||||
requestLogFields?: LLMRequestLogFields;
|
||||
source?: LLMRequestSource;
|
||||
}
|
||||
|
||||
export interface LLMRequestFinish {
|
||||
|
|
@ -81,8 +93,7 @@ export interface LLMRequestOverrides {
|
|||
messages?: readonly Message[];
|
||||
tools?: readonly Tool[];
|
||||
systemPrompt?: string;
|
||||
requestLogFields?: LLMRequestLogFields;
|
||||
usageContext?: UsageRecordContext;
|
||||
source?: LLMRequestSource;
|
||||
maxOutputSize?: number;
|
||||
retry?: LLMRequestRetryOptions;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ import type {
|
|||
LLMRequestLogFields,
|
||||
LLMRequestOverrides,
|
||||
LLMRequestPartHandler,
|
||||
LLMRequestSource,
|
||||
LLMStreamTiming,
|
||||
} from './index';
|
||||
import { IAgentLLMRequesterService } from './llmRequester';
|
||||
|
|
@ -72,8 +73,8 @@ interface ResolvedLLMRequest {
|
|||
readonly systemPrompt: string;
|
||||
readonly tools: readonly Tool[];
|
||||
readonly messages: Message[];
|
||||
readonly requestLogFields: LLMRequestOverrides['requestLogFields'];
|
||||
readonly usageContext: LLMRequestOverrides['usageContext'];
|
||||
readonly source: LLMRequestSource | undefined;
|
||||
readonly logFields: LLMRequestLogFields;
|
||||
}
|
||||
|
||||
interface LLMRequestLogInput {
|
||||
|
|
@ -165,7 +166,8 @@ export class AgentLLMRequesterService implements IAgentLLMRequesterService {
|
|||
): Promise<LLMRequestFinish> {
|
||||
signal?.throwIfAborted();
|
||||
const request = this.resolveRequest(
|
||||
requestOverridesForAttempt(overrides, attempt, maxAttempts),
|
||||
overrides,
|
||||
attempt === 1 ? undefined : { attempt: `${String(attempt)}/${String(maxAttempts)}` },
|
||||
);
|
||||
return await this.runRequest(request, onPart, signal);
|
||||
}
|
||||
|
|
@ -179,7 +181,7 @@ export class AgentLLMRequesterService implements IAgentLLMRequesterService {
|
|||
): void {
|
||||
if (isAbortError(error) || signal?.aborted === true) return;
|
||||
const payload: LogContext = {
|
||||
...overrides.requestLogFields,
|
||||
...logFieldsForSource(overrides.source),
|
||||
attempt: `${String(attempt)}/${String(maxAttempts)}`,
|
||||
model: this.profile.data().modelAlias ?? 'unknown',
|
||||
...retryErrorFields(error),
|
||||
|
|
@ -217,7 +219,7 @@ export class AgentLLMRequesterService implements IAgentLLMRequesterService {
|
|||
systemPrompt: request.systemPrompt,
|
||||
tools: request.tools,
|
||||
messages: request.messages,
|
||||
fields: request.requestLogFields,
|
||||
fields: request.logFields,
|
||||
});
|
||||
|
||||
const input = {
|
||||
|
|
@ -256,9 +258,9 @@ export class AgentLLMRequesterService implements IAgentLLMRequesterService {
|
|||
}
|
||||
|
||||
const usageModel = request.modelAlias;
|
||||
this.usage.record(usageModel, usage, request.usageContext);
|
||||
this.usage.record(usageModel, usage, request.source);
|
||||
this.contextSize.measured(request.messages, [message], usage);
|
||||
this.logResponse(request.requestLogFields, usage, timing);
|
||||
this.logResponse(request.logFields, usage, timing);
|
||||
|
||||
return {
|
||||
message,
|
||||
|
|
@ -271,7 +273,10 @@ export class AgentLLMRequesterService implements IAgentLLMRequesterService {
|
|||
};
|
||||
}
|
||||
|
||||
private resolveRequest(overrides: LLMRequestOverrides): ResolvedLLMRequest {
|
||||
private resolveRequest(
|
||||
overrides: LLMRequestOverrides,
|
||||
extraLogFields?: LLMRequestLogFields,
|
||||
): ResolvedLLMRequest {
|
||||
const resolved = this.profile.resolveModelContext();
|
||||
let model = this.profile.getProvider();
|
||||
model = applyCompletionBudget({
|
||||
|
|
@ -293,13 +298,13 @@ export class AgentLLMRequesterService implements IAgentLLMRequesterService {
|
|||
systemPrompt: overrides.systemPrompt ?? this.profile.getSystemPrompt(),
|
||||
tools: [...(overrides.tools ?? this.defaultTools())],
|
||||
messages: [...messages],
|
||||
requestLogFields: overrides.requestLogFields,
|
||||
usageContext: overrides.usageContext,
|
||||
source: overrides.source,
|
||||
logFields: logFieldsForSource(overrides.source, extraLogFields),
|
||||
};
|
||||
}
|
||||
|
||||
private logRequest(input: LLMRequestLogInput): void {
|
||||
const requestLogFields: LLMRequestLogFields = input.fields ?? {};
|
||||
const logFields: LLMRequestLogFields = input.fields ?? {};
|
||||
const config = {
|
||||
provider: input.protocol,
|
||||
model: input.modelName,
|
||||
|
|
@ -315,11 +320,11 @@ export class AgentLLMRequesterService implements IAgentLLMRequesterService {
|
|||
});
|
||||
if (signature !== this.lastConfigLogSignature) {
|
||||
this.lastConfigLogSignature = signature;
|
||||
this.log.info('llm config', { ...requestLogFields, ...config });
|
||||
this.log.info('llm config', { ...logFields, ...config });
|
||||
}
|
||||
|
||||
const partialMessageCount = input.messages.filter((message) => message.partial === true).length;
|
||||
const requestFields: LogContext = { ...requestLogFields };
|
||||
const requestFields: LogContext = { ...logFields };
|
||||
if (partialMessageCount > 0) requestFields['partialMessageCount'] = partialMessageCount;
|
||||
this.log.info('llm request', requestFields);
|
||||
}
|
||||
|
|
@ -357,21 +362,28 @@ export class AgentLLMRequesterService implements IAgentLLMRequesterService {
|
|||
}
|
||||
}
|
||||
|
||||
function requestOverridesForAttempt(
|
||||
overrides: LLMRequestOverrides,
|
||||
attempt: number,
|
||||
maxAttempts: number,
|
||||
): LLMRequestOverrides {
|
||||
if (attempt === 1) {
|
||||
return overrides;
|
||||
function logFieldsForSource(
|
||||
source: LLMRequestSource | undefined,
|
||||
extraFields?: LLMRequestLogFields,
|
||||
): LLMRequestLogFields {
|
||||
switch (source?.type) {
|
||||
case 'turn':
|
||||
return {
|
||||
...source.logFields,
|
||||
...(source.step === undefined
|
||||
? {}
|
||||
: { turnStep: `${String(source.turnId)}.${String(source.step)}` }),
|
||||
...extraFields,
|
||||
};
|
||||
case 'operation':
|
||||
return {
|
||||
...source.logFields,
|
||||
...(source.requestKind === undefined ? {} : { requestKind: source.requestKind }),
|
||||
...extraFields,
|
||||
};
|
||||
default:
|
||||
return extraFields ?? {};
|
||||
}
|
||||
return {
|
||||
...overrides,
|
||||
requestLogFields: {
|
||||
...overrides.requestLogFields,
|
||||
attempt: `${String(attempt)}/${String(maxAttempts)}`,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function toolSignature(tools: readonly Tool[]) {
|
||||
|
|
|
|||
|
|
@ -148,14 +148,13 @@ export class AgentLoopService implements IAgentLoopService {
|
|||
signal.throwIfAborted();
|
||||
|
||||
const stepUuid = randomUUID();
|
||||
const turnStep = `${turnId}.${String(currentStep)}`;
|
||||
|
||||
this.record.signal({ type: 'turn.step.started', turnId, step: currentStep, stepId: stepUuid });
|
||||
|
||||
const emitStreamPart = this.createStreamPartHandler(turnId);
|
||||
const response = await this.llmRequester.request(
|
||||
{
|
||||
requestLogFields: { turnStep },
|
||||
source: { type: 'turn', turnId, step: currentStep },
|
||||
retry: {
|
||||
maxAttempts: this.config.get<LoopControl>(LOOP_CONTROL_SECTION)?.maxRetriesPerStep,
|
||||
onRetry: (retry) => {
|
||||
|
|
@ -174,7 +173,6 @@ export class AgentLoopService implements IAgentLoopService {
|
|||
});
|
||||
},
|
||||
},
|
||||
usageContext: { type: 'turn', turnId },
|
||||
},
|
||||
emitStreamPart,
|
||||
signal,
|
||||
|
|
@ -193,18 +191,20 @@ export class AgentLoopService implements IAgentLoopService {
|
|||
let finishReason: FinishReason = providerFinishReason ?? 'completed';
|
||||
const hasToolCalls = message.toolCalls.length > 0;
|
||||
if (hasToolCalls) {
|
||||
const toolResults = await this.toolExecutor.execute(response.message.toolCalls, {
|
||||
let stopTurn = false;
|
||||
for await (const toolResult of this.toolExecutor.execute(response.message.toolCalls, {
|
||||
signal,
|
||||
turnId,
|
||||
onToolResult: (toolCallId, result) => {
|
||||
this.append({
|
||||
...createToolMessage(toolCallId, toolResultOutputForModel(result)),
|
||||
role: 'tool',
|
||||
isError: result.isError,
|
||||
});
|
||||
},
|
||||
});
|
||||
if (toolResults.some((r) => r.stopTurn === true)) {
|
||||
})) {
|
||||
const { result } = toolResult;
|
||||
this.append({
|
||||
...createToolMessage(toolResult.toolCallId, toolResultOutputForModel(result)),
|
||||
role: 'tool',
|
||||
isError: result.isError,
|
||||
});
|
||||
if (result.stopTurn === true) stopTurn = true;
|
||||
}
|
||||
if (stopTurn) {
|
||||
finishReason = 'completed';
|
||||
} else {
|
||||
finishReason = 'tool_calls';
|
||||
|
|
|
|||
|
|
@ -10,13 +10,18 @@ import type { OrderedHookSlot } from '#/hooks';
|
|||
export interface ToolExecutorExecuteOptions {
|
||||
readonly signal: AbortSignal;
|
||||
readonly turnId: number;
|
||||
readonly onToolResult?: (toolCallId: string, result: ToolResult) => void | Promise<void>;
|
||||
}
|
||||
|
||||
export interface ToolExecutionResult {
|
||||
readonly toolCallId: string;
|
||||
readonly toolName: string;
|
||||
readonly result: ToolResult;
|
||||
}
|
||||
|
||||
export interface IAgentToolExecutorService {
|
||||
readonly _serviceBrand: undefined;
|
||||
|
||||
execute(calls: ToolCall[], options: ToolExecutorExecuteOptions): Promise<ToolResult[]>;
|
||||
execute(calls: ToolCall[], options: ToolExecutorExecuteOptions): AsyncIterable<ToolExecutionResult>;
|
||||
|
||||
readonly hooks: {
|
||||
readonly onWillExecuteTool: OrderedHookSlot<ToolWillExecuteContext>;
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ import { ITelemetryService } from '#/app/telemetry';
|
|||
import { OrderedHookSlot } from '#/hooks';
|
||||
import {
|
||||
IAgentToolExecutorService,
|
||||
type ToolExecutionResult,
|
||||
type ToolExecutorExecuteOptions,
|
||||
} from './toolExecutor';
|
||||
import { ToolScheduler } from './toolScheduler';
|
||||
|
|
@ -47,10 +48,30 @@ export interface ToolExecutionTask {
|
|||
}
|
||||
|
||||
interface TimedToolResult {
|
||||
readonly index: number;
|
||||
readonly result: ToolResult;
|
||||
readonly durationMs: number;
|
||||
}
|
||||
|
||||
type SettledTimedToolResult =
|
||||
| { readonly status: 'fulfilled'; readonly value: TimedToolResult }
|
||||
| { readonly status: 'rejected'; readonly index: number; readonly reason: unknown };
|
||||
|
||||
type SettledToolExecutionResult =
|
||||
| { readonly status: 'fulfilled'; readonly value: ToolExecutionResult }
|
||||
| { readonly status: 'rejected'; readonly reason: unknown };
|
||||
|
||||
type ToolExecutionResultPromise = Promise<SettledToolExecutionResult>;
|
||||
|
||||
type ToolExecutionStreamEvent =
|
||||
| { readonly type: 'timed'; readonly result: IteratorResult<TimedToolResult> }
|
||||
| { readonly type: 'timedRejected'; readonly reason: unknown }
|
||||
| {
|
||||
readonly type: 'finalized';
|
||||
readonly promise: ToolExecutionResultPromise;
|
||||
readonly settled: SettledToolExecutionResult;
|
||||
};
|
||||
|
||||
export class AgentToolExecutorService implements IAgentToolExecutorService {
|
||||
declare readonly _serviceBrand: undefined;
|
||||
readonly hooks = {
|
||||
|
|
@ -65,11 +86,11 @@ export class AgentToolExecutorService implements IAgentToolExecutorService {
|
|||
@ILogService private readonly log?: ILogService,
|
||||
) {}
|
||||
|
||||
async execute(
|
||||
async *execute(
|
||||
calls: ToolCall[],
|
||||
options: ToolExecutorExecuteOptions,
|
||||
): Promise<ToolResult[]> {
|
||||
if (calls.length === 0) return [];
|
||||
): AsyncIterable<ToolExecutionResult> {
|
||||
if (calls.length === 0) return;
|
||||
|
||||
const preflighted = calls.map((call) => preflightToolCall(this.toolRegistry, call, this.log));
|
||||
const preparedTasks: Array<{
|
||||
|
|
@ -96,25 +117,86 @@ export class AgentToolExecutorService implements IAgentToolExecutorService {
|
|||
}
|
||||
}
|
||||
|
||||
const timedResults = await this.executeBatch(
|
||||
const timedResults = this.executeBatch(
|
||||
preparedTasks.map(({ task }) => task),
|
||||
options.signal,
|
||||
);
|
||||
)[Symbol.asyncIterator]();
|
||||
let nextTimed: Promise<IteratorResult<TimedToolResult>> | undefined = timedResults.next();
|
||||
const finalizations = new Set<ToolExecutionResultPromise>();
|
||||
|
||||
const results: ToolResult[] = [];
|
||||
for (let index = 0; index < preparedTasks.length; index += 1) {
|
||||
const prepared = preparedTasks[index]!;
|
||||
const { call } = prepared;
|
||||
const timedResult = timedResults[index]!;
|
||||
const rawResult = timedResult.result;
|
||||
const finalized = await this.finalizeToolResult(call, rawResult, options);
|
||||
results.push(finalized);
|
||||
try {
|
||||
while (nextTimed !== undefined || finalizations.size > 0) {
|
||||
const candidates: Array<Promise<ToolExecutionStreamEvent>> = [];
|
||||
if (nextTimed !== undefined) {
|
||||
candidates.push(
|
||||
nextTimed.then(
|
||||
(result): ToolExecutionStreamEvent => ({ type: 'timed', result }),
|
||||
(reason): ToolExecutionStreamEvent => ({ type: 'timedRejected', reason }),
|
||||
),
|
||||
);
|
||||
}
|
||||
for (const promise of finalizations) {
|
||||
candidates.push(
|
||||
promise.then((settled): ToolExecutionStreamEvent => ({
|
||||
type: 'finalized',
|
||||
promise,
|
||||
settled,
|
||||
})),
|
||||
);
|
||||
}
|
||||
|
||||
await this.dispatchToolResult(call, finalized, options);
|
||||
this.trackToolCall(call, finalized, timedResult.durationMs, options.turnId);
|
||||
const event = await Promise.race(candidates);
|
||||
if (event.type === 'timedRejected') {
|
||||
throw event.reason;
|
||||
}
|
||||
if (event.type === 'timed') {
|
||||
if (event.result.done === true) {
|
||||
nextTimed = undefined;
|
||||
continue;
|
||||
}
|
||||
|
||||
const finalization = this.finalizeTimedResult(
|
||||
preparedTasks[event.result.value.index]!,
|
||||
event.result.value,
|
||||
options,
|
||||
).then(
|
||||
(value): SettledToolExecutionResult => ({ status: 'fulfilled', value }),
|
||||
(reason): SettledToolExecutionResult => ({ status: 'rejected', reason }),
|
||||
);
|
||||
finalizations.add(finalization);
|
||||
nextTimed = timedResults.next();
|
||||
continue;
|
||||
}
|
||||
|
||||
finalizations.delete(event.promise);
|
||||
if (event.settled.status === 'rejected') throw event.settled.reason;
|
||||
yield event.settled.value;
|
||||
}
|
||||
} finally {
|
||||
await timedResults.return?.();
|
||||
await Promise.allSettled(finalizations);
|
||||
}
|
||||
}
|
||||
|
||||
return results;
|
||||
private async finalizeTimedResult(
|
||||
prepared: {
|
||||
readonly call: PreflightedToolCall;
|
||||
},
|
||||
timedResult: TimedToolResult,
|
||||
options: ToolExecutorExecuteOptions,
|
||||
): Promise<ToolExecutionResult> {
|
||||
const { call } = prepared;
|
||||
const rawResult = timedResult.result;
|
||||
const finalized = await this.finalizeToolResult(call, rawResult, options);
|
||||
|
||||
this.dispatchToolResult(call, finalized, options);
|
||||
this.trackToolCall(call, finalized, timedResult.durationMs, options.turnId);
|
||||
|
||||
return {
|
||||
toolCallId: call.toolCall.id,
|
||||
toolName: call.toolName,
|
||||
result: finalized,
|
||||
};
|
||||
}
|
||||
|
||||
private trackToolCall(
|
||||
|
|
@ -247,30 +329,49 @@ export class AgentToolExecutorService implements IAgentToolExecutorService {
|
|||
return makeResolvedTask(makeErrorToolResult(call, call.args, output));
|
||||
}
|
||||
|
||||
private async executeBatch(
|
||||
private async *executeBatch(
|
||||
tasks: ToolExecutionTask[],
|
||||
signal: AbortSignal,
|
||||
): Promise<TimedToolResult[]> {
|
||||
): AsyncIterable<TimedToolResult> {
|
||||
const scheduler = new ToolScheduler<TimedToolResult>();
|
||||
const pendingResults = tasks.map((task) =>
|
||||
scheduler.add({
|
||||
const allResults: Array<Promise<TimedToolResult>> = [];
|
||||
const pendingResults = new Map<number, Promise<SettledTimedToolResult>>();
|
||||
|
||||
for (let index = 0; index < tasks.length; index += 1) {
|
||||
const task = tasks[index]!;
|
||||
const pendingResult = scheduler.add({
|
||||
accesses: task.accesses,
|
||||
start: async () => {
|
||||
const startedAt = Date.now();
|
||||
return {
|
||||
result: task.execute(signal).then((result) => ({
|
||||
index,
|
||||
result,
|
||||
durationMs: Math.max(0, Date.now() - startedAt),
|
||||
})),
|
||||
};
|
||||
},
|
||||
}),
|
||||
);
|
||||
});
|
||||
allResults.push(pendingResult);
|
||||
pendingResults.set(
|
||||
index,
|
||||
pendingResult.then(
|
||||
(value): SettledTimedToolResult => ({ status: 'fulfilled', value }),
|
||||
(reason): SettledTimedToolResult => ({ status: 'rejected', index, reason }),
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
return await Promise.all(pendingResults);
|
||||
while (pendingResults.size > 0) {
|
||||
const settled = await Promise.race(pendingResults.values());
|
||||
const index = settled.status === 'fulfilled' ? settled.value.index : settled.index;
|
||||
pendingResults.delete(index);
|
||||
if (settled.status === 'rejected') throw settled.reason;
|
||||
yield settled.value;
|
||||
}
|
||||
} finally {
|
||||
await Promise.allSettled(pendingResults);
|
||||
await Promise.allSettled(allResults);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -346,12 +447,11 @@ export class AgentToolExecutorService implements IAgentToolExecutorService {
|
|||
});
|
||||
}
|
||||
|
||||
private async dispatchToolResult(
|
||||
private dispatchToolResult(
|
||||
call: PreflightedToolCall,
|
||||
result: ToolResult,
|
||||
options: ToolExecutorExecuteOptions,
|
||||
): Promise<void> {
|
||||
await options.onToolResult?.(call.toolCall.id, result);
|
||||
): void {
|
||||
this.record.signal({
|
||||
type: 'tool.result',
|
||||
turnId: options.turnId,
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
* The scheduler owns only execution ordering:
|
||||
* - tasks with non-conflicting resource accesses may overlap
|
||||
* - tasks with conflicting resource accesses wait for the conflicting active tasks
|
||||
* - drained results are handed back in provider order
|
||||
* - callers decide whether to drain results in provider order or completion order
|
||||
*
|
||||
* Validation, hooks, event construction, and result finalization stay in
|
||||
* `toolExecutorService.ts`.
|
||||
|
|
|
|||
|
|
@ -1,11 +1,7 @@
|
|||
import type { LLMRequestSource } from '#/agent/llmRequester/llmRequester';
|
||||
import type { TokenUsage } from '#/app/llmProtocol';
|
||||
|
||||
import { createDecorator } from "#/_base/di";
|
||||
|
||||
export interface UsageRecordContext {
|
||||
readonly type: 'turn';
|
||||
readonly turnId: number;
|
||||
}
|
||||
import { createDecorator } from '#/_base/di';
|
||||
|
||||
export interface UsageStatus {
|
||||
readonly byModel?: Record<string, TokenUsage>;
|
||||
|
|
@ -15,7 +11,7 @@ export interface UsageStatus {
|
|||
|
||||
export interface IAgentUsageService {
|
||||
readonly _serviceBrand: undefined;
|
||||
record(model: string, usage: TokenUsage, context?: UsageRecordContext): void;
|
||||
record(model: string, usage: TokenUsage, source?: LLMRequestSource): void;
|
||||
status(): UsageStatus;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,12 +1,11 @@
|
|||
import {
|
||||
addUsage,
|
||||
type TokenUsage } from '#/app/llmProtocol';
|
||||
import { addUsage, type TokenUsage } from '#/app/llmProtocol';
|
||||
import { Disposable } from '#/_base/di/lifecycle';
|
||||
import { InstantiationType } from '#/_base/di/extensions';
|
||||
import { LifecycleScope, registerScopedService } from '#/_base/di/scope';
|
||||
import { Disposable } from '#/_base/di/lifecycle';
|
||||
|
||||
import type { LLMRequestSource } from '#/agent/llmRequester/llmRequester';
|
||||
import { IAgentRecordService } from '#/agent/record';
|
||||
import type { UsageRecordContext, UsageStatus } from './usage';
|
||||
import type { UsageStatus } from './usage';
|
||||
import { IAgentUsageService } from './usage';
|
||||
|
||||
declare module '#/agent/wireRecord' {
|
||||
|
|
@ -14,7 +13,7 @@ declare module '#/agent/wireRecord' {
|
|||
'usage.record': {
|
||||
model: string;
|
||||
usage: TokenUsage;
|
||||
context?: UsageRecordContext;
|
||||
context?: LLMRequestSource;
|
||||
};
|
||||
}
|
||||
}
|
||||
|
|
@ -36,14 +35,14 @@ export class AgentUsageService extends Disposable implements IAgentUsageService
|
|||
);
|
||||
}
|
||||
|
||||
record(model: string, usage: TokenUsage, context?: UsageRecordContext): void {
|
||||
record(model: string, usage: TokenUsage, source?: LLMRequestSource): void {
|
||||
this.records.append({
|
||||
type: 'usage.record',
|
||||
model,
|
||||
usage,
|
||||
context,
|
||||
context: source,
|
||||
});
|
||||
this.apply(model, usage, context);
|
||||
this.apply(model, usage, source);
|
||||
this.publishChanged();
|
||||
}
|
||||
|
||||
|
|
@ -58,13 +57,13 @@ export class AgentUsageService extends Disposable implements IAgentUsageService
|
|||
};
|
||||
}
|
||||
|
||||
private apply(model: string, usage: TokenUsage, context: UsageRecordContext | undefined): void {
|
||||
private apply(model: string, usage: TokenUsage, source: LLMRequestSource | undefined): void {
|
||||
const current = this.byModel[model];
|
||||
this.byModel[model] = current === undefined ? copyUsage(usage) : addUsage(current, usage);
|
||||
|
||||
if (context?.type === 'turn') {
|
||||
if (this.currentTurnId !== context.turnId) {
|
||||
this.currentTurnId = context.turnId;
|
||||
if (source?.type === 'turn') {
|
||||
if (this.currentTurnId !== source.turnId) {
|
||||
this.currentTurnId = source.turnId;
|
||||
this.currentTurn = copyUsage(usage);
|
||||
} else {
|
||||
this.currentTurn =
|
||||
|
|
|
|||
|
|
@ -209,7 +209,11 @@ describe('LLMRequester service migration coverage', () => {
|
|||
|
||||
await expect(
|
||||
llmRequester.request({
|
||||
requestLogFields: { requestKind: 'direct_test', turnStep: '0.1' },
|
||||
source: {
|
||||
type: 'operation',
|
||||
requestKind: 'direct_test',
|
||||
logFields: { turnStep: '0.1' },
|
||||
},
|
||||
retry: { maxAttempts: 1 },
|
||||
}),
|
||||
).rejects.toMatchObject({ message: 'temporary provider failure' });
|
||||
|
|
@ -323,7 +327,13 @@ describe('LLMRequester service migration coverage', () => {
|
|||
it('logs successful LLM responses with caller-provided request fields', async () => {
|
||||
await collectLLMRequest((onPart) =>
|
||||
llmRequester.request(
|
||||
{ requestLogFields: { requestKind: 'direct_test', turnStep: '0.1' } },
|
||||
{
|
||||
source: {
|
||||
type: 'operation',
|
||||
requestKind: 'direct_test',
|
||||
logFields: { turnStep: '0.1' },
|
||||
},
|
||||
},
|
||||
onPart,
|
||||
),
|
||||
);
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import {
|
|||
ExitPlanModeTool,
|
||||
type ExitPlanModeInput,
|
||||
} from '#/agent/plan/tools/exit-plan-mode';
|
||||
import type { ToolResult } from '#/agent/tool';
|
||||
import type { ITelemetryService } from '#/app/telemetry';
|
||||
import { IAgentToolExecutorService } from '#/agent/toolExecutor';
|
||||
|
||||
|
|
@ -243,10 +244,13 @@ describe('AgentPlanService EnterPlanMode telemetry', () => {
|
|||
arguments: '{}',
|
||||
};
|
||||
|
||||
const result = await toolExecutor.execute([call], {
|
||||
const result: ToolResult[] = [];
|
||||
for await (const item of toolExecutor.execute([call], {
|
||||
turnId: 1,
|
||||
signal: new AbortController().signal,
|
||||
});
|
||||
})) {
|
||||
result.push(item.result);
|
||||
}
|
||||
|
||||
expect(result[0]?.isError).toBeFalsy();
|
||||
expect(result[0]?.output).toContain('Plan mode is now active');
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import { ITelemetryService } from '#/app/telemetry';
|
|||
import type { ToolCall } from '#/app/llmProtocol';
|
||||
import { IAgentLoopService } from '#/agent/loop';
|
||||
import { IAgentRecordService } from '#/agent/record';
|
||||
import type { ExecutableTool, ExecutableToolContext, ToolExecution } from '#/agent/tool';
|
||||
import type { ExecutableTool, ExecutableToolContext, ToolExecution, ToolResult } from '#/agent/tool';
|
||||
import {
|
||||
IAgentToolDedupeService,
|
||||
AgentToolDedupeService,
|
||||
|
|
@ -181,13 +181,16 @@ describe('AgentToolDedupeService', () => {
|
|||
step: 1,
|
||||
signal: new AbortController().signal,
|
||||
});
|
||||
const results = await executor.execute(
|
||||
const results: ToolResult[] = [];
|
||||
for await (const item of executor.execute(
|
||||
[
|
||||
toolCall('call_1', 'Echo', { text: 'same' }),
|
||||
toolCall('call_2', 'Echo', { text: 'same' }),
|
||||
],
|
||||
{ turnId: 3, signal: new AbortController().signal },
|
||||
);
|
||||
)) {
|
||||
results.push(item.result);
|
||||
}
|
||||
|
||||
expect(tool.calls).toHaveLength(1);
|
||||
expect(results.map((result) => result.output)).toEqual(['same', 'same']);
|
||||
|
|
|
|||
|
|
@ -279,17 +279,69 @@ describe('AgentToolExecutorService', () => {
|
|||
toolCall('call_second', 'second', {}),
|
||||
]);
|
||||
|
||||
expect(results).toEqual([
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results).toEqual(expect.arrayContaining([
|
||||
expect.objectContaining({ output: 'first result', stopBatchAfterThis: true }),
|
||||
expect.objectContaining({
|
||||
output: 'Tool skipped because a previous tool call stopped the turn.',
|
||||
isError: true,
|
||||
}),
|
||||
]);
|
||||
]));
|
||||
expect(first.calls).toHaveLength(1);
|
||||
expect(second.calls).toEqual([]);
|
||||
});
|
||||
|
||||
it('yields independent tool results as each call finishes', async () => {
|
||||
const slowRelease = deferred();
|
||||
const fastRelease = deferred();
|
||||
const slowStarted = deferred();
|
||||
const fastStarted = deferred();
|
||||
const firstYielded = deferred();
|
||||
const slow = new TestTool('slow', {
|
||||
accesses: ToolAccesses.readFile('/repo/slow.txt'),
|
||||
execute: async () => {
|
||||
slowStarted.resolve();
|
||||
await slowRelease.promise;
|
||||
return { output: 'slow' };
|
||||
},
|
||||
});
|
||||
const fast = new TestTool('fast', {
|
||||
accesses: ToolAccesses.readFile('/repo/fast.txt'),
|
||||
execute: async () => {
|
||||
fastStarted.resolve();
|
||||
await fastRelease.promise;
|
||||
return { output: 'fast' };
|
||||
},
|
||||
});
|
||||
registry.register(slow);
|
||||
registry.register(fast);
|
||||
|
||||
const yielded: string[] = [];
|
||||
const execution = (async () => {
|
||||
for await (const item of executor.execute(
|
||||
[
|
||||
toolCall('call_slow', 'slow', {}),
|
||||
toolCall('call_fast', 'fast', {}),
|
||||
],
|
||||
{ turnId: 0, signal: new AbortController().signal },
|
||||
)) {
|
||||
yielded.push(String(item.result.output));
|
||||
if (yielded.length === 1) firstYielded.resolve();
|
||||
}
|
||||
})();
|
||||
|
||||
await Promise.all([slowStarted.promise, fastStarted.promise]);
|
||||
fastRelease.resolve();
|
||||
await firstYielded.promise;
|
||||
|
||||
expect(yielded).toEqual(['fast']);
|
||||
|
||||
slowRelease.resolve();
|
||||
await execution;
|
||||
|
||||
expect(yielded).toEqual(['fast', 'slow']);
|
||||
});
|
||||
|
||||
it('writes resolveExecution description and display onto tool.call.started events', async () => {
|
||||
const tool = new TestTool('display', {
|
||||
description: 'Prepared display description',
|
||||
|
|
@ -417,10 +469,12 @@ describe('AgentToolExecutorService', () => {
|
|||
controller.abort();
|
||||
await execution;
|
||||
|
||||
expect(pairedToolCallIds()).toEqual({
|
||||
calls: ['call_first', 'call_second', 'call_third'],
|
||||
results: ['call_first', 'call_second', 'call_third'],
|
||||
});
|
||||
const paired = pairedToolCallIds();
|
||||
expect(paired.calls).toEqual(['call_first', 'call_second', 'call_third']);
|
||||
expect(paired.results).toHaveLength(3);
|
||||
expect(paired.results).toEqual(
|
||||
expect.arrayContaining(['call_first', 'call_second', 'call_third']),
|
||||
);
|
||||
});
|
||||
|
||||
it('preserves media-only image output with a text companion', async () => {
|
||||
|
|
@ -535,14 +589,16 @@ describe('parseToolCallArguments', () => {
|
|||
});
|
||||
});
|
||||
|
||||
function execute(calls: ToolCall[], signal?: AbortSignal): Promise<ToolResult[]> {
|
||||
return executor.execute(calls, {
|
||||
async function execute(calls: ToolCall[], signal?: AbortSignal): Promise<ToolResult[]> {
|
||||
const results: ToolResult[] = [];
|
||||
for await (const item of executor.execute(calls, {
|
||||
turnId: 0,
|
||||
signal: signal ?? new AbortController().signal,
|
||||
onToolResult: (toolCallId, result) => {
|
||||
events.push({ type: 'tool.result', toolCallId, result });
|
||||
},
|
||||
});
|
||||
})) {
|
||||
results.push(item.result);
|
||||
events.push({ type: 'tool.result', toolCallId: item.toolCallId, result: item.result });
|
||||
}
|
||||
return results;
|
||||
}
|
||||
|
||||
function toolCall(id: string, name: string, args: unknown): ToolCall {
|
||||
|
|
@ -579,6 +635,20 @@ function pairedToolCallIds(): { readonly calls: string[]; readonly results: stri
|
|||
};
|
||||
}
|
||||
|
||||
function deferred<T = void>(): {
|
||||
readonly promise: Promise<T>;
|
||||
readonly resolve: (value: T | PromiseLike<T>) => void;
|
||||
readonly reject: (reason?: unknown) => void;
|
||||
} {
|
||||
let resolve!: (value: T | PromiseLike<T>) => void;
|
||||
let reject!: (reason?: unknown) => void;
|
||||
const promise = new Promise<T>((resolvePromise, rejectPromise) => {
|
||||
resolve = resolvePromise;
|
||||
reject = rejectPromise;
|
||||
});
|
||||
return { promise, resolve, reject };
|
||||
}
|
||||
|
||||
class TestTool implements ExecutableTool<Record<string, unknown>> {
|
||||
readonly description = 'Test tool.';
|
||||
readonly parameters: Record<string, unknown>;
|
||||
|
|
|
|||
|
|
@ -111,12 +111,12 @@ export function stubLoopWithHooks(): IAgentLoopService {
|
|||
* An `IAgentToolExecutorService` stub whose tool-execution hooks (`onWillExecuteTool` /
|
||||
* `onDidExecuteTool`) are real `OrderedHookSlot`s, so services that register
|
||||
* gate hooks in their constructor (AgentPermissionGate, AgentMcpService, …) can be built
|
||||
* in tests. `execute` resolves to an empty batch by default.
|
||||
* in tests. `execute` yields an empty batch by default.
|
||||
*/
|
||||
export function stubToolExecutor(): IAgentToolExecutorService {
|
||||
return {
|
||||
_serviceBrand: undefined,
|
||||
execute: async () => [],
|
||||
execute: async function* () {},
|
||||
hooks: createHooks([
|
||||
'onWillExecuteTool',
|
||||
'onDidExecuteTool',
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue