feat(agent): added agent tool-call cap enforcement to stream loop
- Added `maxToolCallsPerTurn` support to `AgentOptions` and `AgentLoopConfig`, with Agent getter/setter and serialized state wiring. - Implemented stream-loop cap handling by normalizing bad values and halting after `toolcall_end` reaches the limit. - Added `ANTHROPIC_TOOL_CALL_BATCH_CAP`=8 and wired session cap sync on init, model changes, and restore. - Added tests that truncated a 10-call stream to 8 tool calls, and verified non-Claude models resolve no cap.
This commit is contained in:
@@ -1,6 +1,13 @@
|
||||
# Changelog
|
||||
|
||||
## [Unreleased]
|
||||
### Added
|
||||
|
||||
- Added `maxToolCallsPerTurn` to `AgentLoopConfig`/`AgentOptions`, allowing callers to cut a streamed assistant turn after a completed tool-call batch and execute the runnable partial turn instead of waiting for the provider to yield.
|
||||
|
||||
### Fixed
|
||||
|
||||
- Normalized `maxToolCallsPerTurn` to accept only positive integer limits, with non-finite or non-positive values treated as disabled
|
||||
|
||||
## [15.5.14] - 2026-05-29
|
||||
|
||||
|
||||
@@ -441,6 +441,27 @@ interface StepCounter {
|
||||
count: number;
|
||||
}
|
||||
|
||||
function normalizeMaxToolCallsPerTurn(value: number | undefined): number | undefined {
|
||||
if (value === undefined || !Number.isFinite(value)) return undefined;
|
||||
const normalized = Math.trunc(value);
|
||||
return normalized > 0 ? normalized : undefined;
|
||||
}
|
||||
|
||||
function cloneAssistantMessageForToolCallCap(message: AssistantMessage): AssistantMessage {
|
||||
return {
|
||||
...message,
|
||||
content: message.content.map(block => {
|
||||
if (block.type === "toolCall") {
|
||||
return { ...block, arguments: structuredClone(block.arguments) };
|
||||
}
|
||||
return { ...block };
|
||||
}),
|
||||
stopReason: "toolUse",
|
||||
errorMessage: undefined,
|
||||
errorStatus: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
async function runLoopBody(
|
||||
currentContext: AgentContext,
|
||||
newMessages: AgentMessage[],
|
||||
@@ -712,11 +733,18 @@ async function streamAssistantResponse(
|
||||
const dynamicReasoning = config.getReasoning?.();
|
||||
const harmonyMitigationEnabled = isHarmonyLeakMitigationTarget(config.model);
|
||||
const harmonyAbortController = harmonyMitigationEnabled ? new AbortController() : undefined;
|
||||
const requestSignal = harmonyAbortController
|
||||
? signal
|
||||
? AbortSignal.any([signal, harmonyAbortController.signal])
|
||||
: harmonyAbortController.signal
|
||||
: signal;
|
||||
const maxToolCallsPerTurn = normalizeMaxToolCallsPerTurn(config.maxToolCallsPerTurn);
|
||||
const toolCallCapAbortController = maxToolCallsPerTurn === undefined ? undefined : new AbortController();
|
||||
const requestSignals: AbortSignal[] = [];
|
||||
if (signal) requestSignals.push(signal);
|
||||
if (harmonyAbortController) requestSignals.push(harmonyAbortController.signal);
|
||||
if (toolCallCapAbortController) requestSignals.push(toolCallCapAbortController.signal);
|
||||
const requestSignal =
|
||||
requestSignals.length === 0
|
||||
? undefined
|
||||
: requestSignals.length === 1
|
||||
? requestSignals[0]
|
||||
: AbortSignal.any(requestSignals);
|
||||
const effectiveTemperature =
|
||||
harmonyRetryAttempt > 0 && config.temperature !== undefined ? config.temperature + 0.05 : config.temperature;
|
||||
const effectiveToolChoice = dynamicToolChoice ?? config.toolChoice;
|
||||
@@ -778,6 +806,26 @@ async function streamAssistantResponse(
|
||||
let addedPartial = false;
|
||||
|
||||
const responseIterator = response[Symbol.asyncIterator]();
|
||||
let completedToolCalls = 0;
|
||||
let cappedMessage: AssistantMessage | undefined;
|
||||
let capFinalized = false;
|
||||
|
||||
const finishCappedAssistantMessage = async (): Promise<AssistantMessage | undefined> => {
|
||||
if (!cappedMessage) return undefined;
|
||||
responseIterator.return?.()?.catch(() => {});
|
||||
if (!capFinalized) {
|
||||
if (addedPartial) {
|
||||
context.messages[context.messages.length - 1] = cappedMessage;
|
||||
} else {
|
||||
context.messages.push(cappedMessage);
|
||||
stream.push({ type: "message_start", message: { ...cappedMessage } });
|
||||
}
|
||||
stream.push({ type: "message_end", message: cappedMessage });
|
||||
await finishChat(cappedMessage);
|
||||
capFinalized = true;
|
||||
}
|
||||
return cappedMessage;
|
||||
};
|
||||
|
||||
// Set up a single abort race: register the abort listener once for the whole
|
||||
// stream and reuse the same race promise for every iterator.next() instead of
|
||||
@@ -803,6 +851,10 @@ async function streamAssistantResponse(
|
||||
if (abortRacePromise) {
|
||||
const result = await Promise.race([responseIterator.next(), abortRacePromise]);
|
||||
if (result === ABORTED) {
|
||||
if (toolCallCapAbortController?.signal.aborted) {
|
||||
const capped = await finishCappedAssistantMessage();
|
||||
if (capped) return capped;
|
||||
}
|
||||
responseIterator.return?.()?.catch(() => {});
|
||||
const aborted = emitAbortedAssistantMessage(partialMessage, addedPartial, context, config, stream);
|
||||
await finishChat(aborted);
|
||||
@@ -813,6 +865,10 @@ async function streamAssistantResponse(
|
||||
next = await responseIterator.next();
|
||||
}
|
||||
if (requestSignal?.aborted) {
|
||||
if (toolCallCapAbortController?.signal.aborted) {
|
||||
const capped = await finishCappedAssistantMessage();
|
||||
if (capped) return capped;
|
||||
}
|
||||
const aborted = emitAbortedAssistantMessage(partialMessage, addedPartial, context, config, stream);
|
||||
await finishChat(aborted);
|
||||
return aborted;
|
||||
@@ -853,6 +909,15 @@ async function streamAssistantResponse(
|
||||
assistantMessageEvent: event,
|
||||
message: { ...partialMessage },
|
||||
});
|
||||
if (event.type === "toolcall_end" && maxToolCallsPerTurn !== undefined) {
|
||||
completedToolCalls++;
|
||||
if (completedToolCalls >= maxToolCallsPerTurn) {
|
||||
cappedMessage = cloneAssistantMessageForToolCallCap(partialMessage);
|
||||
toolCallCapAbortController?.abort();
|
||||
const capped = await finishCappedAssistantMessage();
|
||||
if (capped) return capped;
|
||||
}
|
||||
}
|
||||
}
|
||||
break;
|
||||
|
||||
|
||||
@@ -102,6 +102,12 @@ export interface AgentOptions {
|
||||
*/
|
||||
interruptMode?: "immediate" | "wait";
|
||||
|
||||
/**
|
||||
* Maximum completed tool calls to accept from one streamed assistant turn before
|
||||
* executing the batch. Undefined disables batching.
|
||||
*/
|
||||
maxToolCallsPerTurn?: number;
|
||||
|
||||
/**
|
||||
* API format for Kimi Code provider: "openai" or "anthropic" (default: "anthropic")
|
||||
*/
|
||||
@@ -269,6 +275,7 @@ export class Agent {
|
||||
#steeringMode: "all" | "one-at-a-time";
|
||||
#followUpMode: "all" | "one-at-a-time";
|
||||
#interruptMode: "immediate" | "wait";
|
||||
#maxToolCallsPerTurn?: number;
|
||||
#sessionId?: string;
|
||||
#metadata?: Record<string, unknown>;
|
||||
#metadataResolver?: (provider: string) => Record<string, unknown> | undefined;
|
||||
@@ -325,6 +332,7 @@ export class Agent {
|
||||
this.#steeringMode = opts.steeringMode || "one-at-a-time";
|
||||
this.#followUpMode = opts.followUpMode || "one-at-a-time";
|
||||
this.#interruptMode = opts.interruptMode || "immediate";
|
||||
this.#maxToolCallsPerTurn = opts.maxToolCallsPerTurn;
|
||||
this.streamFn = opts.streamFn || streamSimple;
|
||||
this.#sessionId = opts.sessionId;
|
||||
this.#providerSessionState = opts.providerSessionState;
|
||||
@@ -547,6 +555,14 @@ export class Agent {
|
||||
this.#maxRetryDelayMs = value;
|
||||
}
|
||||
|
||||
get maxToolCallsPerTurn(): number | undefined {
|
||||
return this.#maxToolCallsPerTurn;
|
||||
}
|
||||
|
||||
set maxToolCallsPerTurn(value: number | undefined) {
|
||||
this.#maxToolCallsPerTurn = value;
|
||||
}
|
||||
|
||||
get state(): AgentState {
|
||||
return this.#state;
|
||||
}
|
||||
@@ -917,6 +933,7 @@ export class Agent {
|
||||
serviceTier: this.#serviceTier,
|
||||
hideThinkingSummary: this.#hideThinkingSummary,
|
||||
interruptMode: this.#interruptMode,
|
||||
maxToolCallsPerTurn: this.#maxToolCallsPerTurn,
|
||||
sessionId: this.#sessionId,
|
||||
metadata: this.#metadataResolver ? undefined : this.#metadata,
|
||||
metadataResolver: this.#metadataResolver,
|
||||
|
||||
@@ -38,6 +38,14 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
|
||||
*/
|
||||
interruptMode?: "immediate" | "wait";
|
||||
|
||||
/**
|
||||
* Maximum completed tool calls to accept from one streamed assistant turn before
|
||||
* cutting the provider stream and executing that batch. The cap is enforced on
|
||||
* `toolcall_end` so every executed call has complete arguments. Undefined disables
|
||||
* batching.
|
||||
*/
|
||||
maxToolCallsPerTurn?: number;
|
||||
|
||||
/**
|
||||
* Optional session identifier forwarded to LLM providers.
|
||||
* Used by providers that support session-based caching (e.g., OpenAI Codex).
|
||||
|
||||
@@ -7,6 +7,7 @@ import type {
|
||||
AgentMessage,
|
||||
AgentTool,
|
||||
AgentToolContext,
|
||||
StreamFn,
|
||||
ToolCallContext,
|
||||
} from "@oh-my-pi/pi-agent-core/types";
|
||||
import type { AssistantMessage, Message, ToolResultMessage } from "@oh-my-pi/pi-ai";
|
||||
@@ -62,7 +63,7 @@ describe("agentLoop with AgentMessage", () => {
|
||||
tools: [],
|
||||
};
|
||||
const mock = createMockModel();
|
||||
const config: AgentLoopConfig = { model: mock.model, convertToLlm: identityConverter };
|
||||
const config: AgentLoopConfig = { model: mock.model, convertToLlm: identityConverter, maxToolCallsPerTurn: 8 };
|
||||
const controller = new AbortController();
|
||||
// The mock provider would reject without a configured response; we want the
|
||||
// agent's abort path to kick in before any event is emitted. Use a raw stream
|
||||
@@ -273,6 +274,118 @@ describe("agentLoop with AgentMessage", () => {
|
||||
}
|
||||
});
|
||||
|
||||
it("cuts a streamed assistant turn after the configured completed tool-call batch", async () => {
|
||||
const toolSchema = z.object({ value: z.string() });
|
||||
const executed: string[] = [];
|
||||
const tool: AgentTool<typeof toolSchema, { value: string }> = {
|
||||
name: "echo",
|
||||
label: "Echo",
|
||||
description: "Echo tool",
|
||||
parameters: toolSchema,
|
||||
async execute(_toolCallId, params) {
|
||||
executed.push(params.value);
|
||||
return {
|
||||
content: [{ type: "text", text: `echoed: ${params.value}` }],
|
||||
details: { value: params.value },
|
||||
};
|
||||
},
|
||||
};
|
||||
|
||||
const context: AgentContext = { systemPrompt: [""], messages: [], tools: [tool] };
|
||||
const mock = createMockModel();
|
||||
let modelCalls = 0;
|
||||
let firstRequestSignal: AbortSignal | undefined;
|
||||
|
||||
const makeToolCall = (index: number): AssistantMessage["content"][number] => ({
|
||||
type: "toolCall",
|
||||
id: `tool-${index}`,
|
||||
name: "echo",
|
||||
arguments: { value: String(index) },
|
||||
});
|
||||
const makeMessage = (count: number, stopReason: AssistantMessage["stopReason"] = "stop") =>
|
||||
createAssistantMessage(
|
||||
Array.from({ length: count }, (_, index) => makeToolCall(index + 1)),
|
||||
stopReason,
|
||||
);
|
||||
|
||||
const streamFn: StreamFn = (_model, _llmContext, options) => {
|
||||
modelCalls++;
|
||||
const stream = new AssistantMessageEventStream();
|
||||
if (modelCalls > 1) {
|
||||
queueMicrotask(() => {
|
||||
const done = createAssistantMessage([{ type: "text", text: "done" }], "stop");
|
||||
stream.push({ type: "start", partial: done });
|
||||
stream.push({ type: "text_start", contentIndex: 0, partial: done });
|
||||
stream.push({ type: "text_delta", contentIndex: 0, delta: "done", partial: done });
|
||||
stream.push({ type: "text_end", contentIndex: 0, content: "done", partial: done });
|
||||
stream.push({ type: "done", reason: "stop", message: done });
|
||||
});
|
||||
return stream;
|
||||
}
|
||||
|
||||
queueMicrotask(async () => {
|
||||
firstRequestSignal = options?.signal;
|
||||
stream.push({ type: "start", partial: makeMessage(0) });
|
||||
for (let index = 1; index <= 10; index++) {
|
||||
if (options?.signal?.aborted) {
|
||||
const aborted = createAssistantMessage([], "aborted");
|
||||
stream.push({ type: "error", reason: "aborted", error: aborted });
|
||||
return;
|
||||
}
|
||||
const partial = makeMessage(index);
|
||||
const toolCall = partial.content[index - 1];
|
||||
if (!toolCall || toolCall.type !== "toolCall") throw new Error("Expected tool call");
|
||||
stream.push({ type: "toolcall_start", contentIndex: index - 1, partial });
|
||||
stream.push({
|
||||
type: "toolcall_delta",
|
||||
contentIndex: index - 1,
|
||||
delta: JSON.stringify(toolCall.arguments),
|
||||
partial,
|
||||
});
|
||||
stream.push({ type: "toolcall_end", contentIndex: index - 1, toolCall, partial });
|
||||
await Bun.sleep(0);
|
||||
}
|
||||
stream.push({ type: "done", reason: "toolUse", message: makeMessage(10, "toolUse") });
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
maxToolCallsPerTurn: 8,
|
||||
};
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([createUserMessage("echo many")], context, config, undefined, streamFn);
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
expect(executed).toEqual(["1", "2", "3", "4", "5", "6", "7", "8"]);
|
||||
expect(firstRequestSignal?.aborted).toBe(true);
|
||||
expect(modelCalls).toBe(2);
|
||||
|
||||
const batchedTurn = events.find(
|
||||
(event): event is Extract<AgentEvent, { type: "turn_end" }> =>
|
||||
event.type === "turn_end" && event.toolResults.length === 8,
|
||||
);
|
||||
expect(batchedTurn).toBeDefined();
|
||||
if (!batchedTurn || batchedTurn.message.role !== "assistant") return;
|
||||
expect(batchedTurn.message.stopReason).toBe("toolUse");
|
||||
expect(batchedTurn.message.content.filter(block => block.type === "toolCall")).toHaveLength(8);
|
||||
expect(batchedTurn.toolResults.map(result => result.toolCallId).sort()).toEqual([
|
||||
"tool-1",
|
||||
"tool-2",
|
||||
"tool-3",
|
||||
"tool-4",
|
||||
"tool-5",
|
||||
"tool-6",
|
||||
"tool-7",
|
||||
"tool-8",
|
||||
]);
|
||||
});
|
||||
|
||||
it("injects and strips intent when intent tracing is enabled", async () => {
|
||||
const toolSchema = z.object({ value: z.string() });
|
||||
const executedParams: Record<string, unknown>[] = [];
|
||||
|
||||
@@ -1,11 +1,18 @@
|
||||
# Changelog
|
||||
|
||||
## [Unreleased]
|
||||
### Changed
|
||||
|
||||
- Enabled the agent loop's tool-call batch cap for Anthropic Claude sessions, cutting oversized streamed tool-use bursts into runnable batches before continuing the conversation.
|
||||
|
||||
### Removed
|
||||
|
||||
- Removed the `calc` tool (deterministic arithmetic evaluator) and its `calc.enabled` setting. The model can compute via `eval` instead.
|
||||
|
||||
### Fixed
|
||||
|
||||
- Fixed Anthropic Claude tool-call batching to clear and reapply the Claude-specific batch cap whenever the session model changes
|
||||
|
||||
## [15.5.14] - 2026-05-29
|
||||
### Added
|
||||
|
||||
|
||||
@@ -453,6 +453,15 @@ function formatRetryFallbackBaseSelector(selector: RetryFallbackSelector): strin
|
||||
}
|
||||
|
||||
const IRC_REPLY_MAX_BYTES = 4096;
|
||||
export const ANTHROPIC_TOOL_CALL_BATCH_CAP = 4;
|
||||
const CLAUDE_OPUS_4_8_MODEL_ID = /(?:^|[./_-])claude-opus-4[.-]8\b/i;
|
||||
|
||||
export function resolveToolCallBatchCapForModel(model: Model | undefined): number | undefined {
|
||||
if (!model) return undefined;
|
||||
return model.provider === "anthropic" && CLAUDE_OPUS_4_8_MODEL_ID.test(model.id)
|
||||
? ANTHROPIC_TOOL_CALL_BATCH_CAP
|
||||
: undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Collapse degenerate IRC ephemeral replies before they hit the relay.
|
||||
@@ -993,6 +1002,10 @@ export class AgentSession {
|
||||
this.#flushPendingAgentEnd();
|
||||
}
|
||||
|
||||
#syncToolCallBatchCap(model: Model | undefined = this.model): void {
|
||||
this.agent.maxToolCallsPerTurn = resolveToolCallBatchCapForModel(model);
|
||||
}
|
||||
|
||||
#flushPendingAgentEnd(): void {
|
||||
const pending = this.#pendingAgentEndEmit;
|
||||
if (!pending) return;
|
||||
@@ -1097,6 +1110,7 @@ export class AgentSession {
|
||||
this.#agentId = config.agentId;
|
||||
this.#agentRegistry = config.agentRegistry;
|
||||
this.#providerSessionId = config.providerSessionId;
|
||||
this.#syncToolCallBatchCap();
|
||||
this.agent.setAssistantMessageEventInterceptor((message, assistantMessageEvent) => {
|
||||
const event: AgentEvent = {
|
||||
type: "message_update",
|
||||
@@ -6162,6 +6176,7 @@ export class AgentSession {
|
||||
this.#closeProviderSessionsForModelSwitch(currentModel, model);
|
||||
}
|
||||
this.agent.setModel(model);
|
||||
this.#syncToolCallBatchCap(model);
|
||||
|
||||
// Re-evaluate append-only context mode — provider or setting may have changed
|
||||
this.#syncAppendOnlyContext(model);
|
||||
@@ -8214,6 +8229,7 @@ export class AgentSession {
|
||||
this.#setModelWithProviderSessionReset(match);
|
||||
} else {
|
||||
this.agent.setModel(match);
|
||||
this.#syncToolCallBatchCap(match);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -8272,6 +8288,9 @@ export class AgentSession {
|
||||
this.#scheduledHiddenNextTurnGeneration = previousScheduledHiddenNextTurnGeneration;
|
||||
if (previousModel) {
|
||||
this.agent.setModel(previousModel);
|
||||
this.#syncToolCallBatchCap(previousModel);
|
||||
} else {
|
||||
this.#syncToolCallBatchCap(undefined);
|
||||
}
|
||||
this.#thinkingLevel = previousThinkingLevel;
|
||||
this.agent.setThinkingLevel(toReasoningEffort(previousThinkingLevel));
|
||||
|
||||
@@ -9,7 +9,12 @@ import {
|
||||
} from "@oh-my-pi/pi-ai";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
import { AgentSession, type AgentSessionEvent } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
||||
import {
|
||||
AgentSession,
|
||||
type AgentSessionEvent,
|
||||
ANTHROPIC_TOOL_CALL_BATCH_CAP,
|
||||
resolveToolCallBatchCapForModel,
|
||||
} from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
||||
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
||||
import { createAssistantMessage } from "./helpers/agent-session-setup";
|
||||
|
||||
@@ -34,6 +39,43 @@ describe("AgentSession message pipeline", () => {
|
||||
}
|
||||
});
|
||||
|
||||
it("enables the tool-call batch cap only for Anthropic Claude Opus 4.8 models", () => {
|
||||
const baseModel: Model = {
|
||||
id: "gpt-5",
|
||||
name: "GPT-5",
|
||||
api: "openai-responses",
|
||||
provider: "openai",
|
||||
baseUrl: "",
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 200_000,
|
||||
maxTokens: 8_192,
|
||||
};
|
||||
const anthropicOpus48: Model = {
|
||||
...baseModel,
|
||||
id: "claude-opus-4-8",
|
||||
name: "Claude Opus 4.8",
|
||||
api: "anthropic",
|
||||
provider: "anthropic",
|
||||
};
|
||||
|
||||
expect(resolveToolCallBatchCapForModel(anthropicOpus48)).toBe(ANTHROPIC_TOOL_CALL_BATCH_CAP);
|
||||
expect(resolveToolCallBatchCapForModel({ ...anthropicOpus48, id: "claude-opus-4.8" })).toBe(
|
||||
ANTHROPIC_TOOL_CALL_BATCH_CAP,
|
||||
);
|
||||
expect(resolveToolCallBatchCapForModel({ ...anthropicOpus48, id: "claude-opus-4-8-20260530" })).toBe(
|
||||
ANTHROPIC_TOOL_CALL_BATCH_CAP,
|
||||
);
|
||||
expect(resolveToolCallBatchCapForModel({ ...anthropicOpus48, provider: "openrouter" })).toBeUndefined();
|
||||
expect(resolveToolCallBatchCapForModel({ ...anthropicOpus48, id: "claude-sonnet-4-8" })).toBeUndefined();
|
||||
expect(resolveToolCallBatchCapForModel({ ...anthropicOpus48, id: "claude-opus-4-7" })).toBeUndefined();
|
||||
expect(resolveToolCallBatchCapForModel({ ...anthropicOpus48, id: "claude-opus-4-9" })).toBeUndefined();
|
||||
expect(resolveToolCallBatchCapForModel({ ...anthropicOpus48, id: "claude-opus-4-80" })).toBeUndefined();
|
||||
expect(resolveToolCallBatchCapForModel(baseModel)).toBeUndefined();
|
||||
expect(resolveToolCallBatchCapForModel({ ...baseModel, provider: "openai-codex" })).toBeUndefined();
|
||||
});
|
||||
|
||||
it("applies transformContext before convertToLlm", async () => {
|
||||
const inputMessages: AgentMessage[] = [{ role: "user", content: "hello", timestamp: Date.now() }];
|
||||
const transformedMessages: AgentMessage[] = [
|
||||
|
||||
Reference in New Issue
Block a user