From ae905fb3cffd704dda336fa409a17cd1a5fa3ed4 Mon Sep 17 00:00:00 2001 From: can1357 Date: Sat, 30 May 2026 04:44:40 +0200 Subject: [PATCH] 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. --- packages/agent/CHANGELOG.md | 7 ++ packages/agent/src/agent-loop.ts | 75 +++++++++++- packages/agent/src/agent.ts | 17 +++ packages/agent/src/types.ts | 8 ++ packages/agent/test/agent-loop.test.ts | 115 +++++++++++++++++- packages/coding-agent/CHANGELOG.md | 7 ++ .../coding-agent/src/session/agent-session.ts | 19 +++ .../agent-session-message-pipeline.test.ts | 44 ++++++- 8 files changed, 285 insertions(+), 7 deletions(-) diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 71e2f4558..a375062ca 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -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 diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 75739e5d5..ef3a4d396 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -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 => { + 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; diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 9eb2c52b0..904d1b298 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -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; #metadataResolver?: (provider: string) => Record | 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, diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 17b0f2ff9..acc835f36 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -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). diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index dc4a97365..af0e6e267 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -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 = { + 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 => + 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[] = []; diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index ce8202284..09a4c497b 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -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 diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index ecc875a46..ee24ce201 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -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)); diff --git a/packages/coding-agent/test/agent-session-message-pipeline.test.ts b/packages/coding-agent/test/agent-session-message-pipeline.test.ts index cb81d5cf1..13b287f3b 100644 --- a/packages/coding-agent/test/agent-session-message-pipeline.test.ts +++ b/packages/coding-agent/test/agent-session-message-pipeline.test.ts @@ -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[] = [