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:
can1357
2026-05-30 04:44:40 +02:00
parent 04f9402696
commit ae905fb3cf
8 changed files with 285 additions and 7 deletions
+7
View File
@@ -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
+70 -5
View File
@@ -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;
+17
View File
@@ -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,
+8
View File
@@ -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).
+114 -1
View File
@@ -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>[] = [];
+7
View File
@@ -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[] = [