diff --git a/packages/agent-core/src/agent-loop.ts b/packages/agent-core/src/agent-loop.ts index 0de9ec85..180924f3 100644 --- a/packages/agent-core/src/agent-loop.ts +++ b/packages/agent-core/src/agent-loop.ts @@ -7,9 +7,16 @@ import { type AssistantMessage, type Context, EventStream, + type ImageContent, + type TextContent, type ToolResultMessage, validateToolArguments, } from "@step-harness/providers"; +import { + truncateStringToBytesFromEnd, + truncateStringToBytesFromStart, + utf8ByteLength, +} from "./harness/utils/truncate.ts"; import { getDefaultStreamFn } from "./stream-fn.ts"; import type { AgentContext, @@ -25,6 +32,82 @@ import type { export type AgentEventSink = (event: AgentEvent) => Promise | void; +/** + * Default cap on the combined byte size of text content blocks in a tool result, + * applied just before the result is turned into a `ToolResultMessage` and persisted + * to session history. Built-in tools already cap their own output well below this + * (see DEFAULT_MAX_BYTES in harness/utils/truncate.ts), so this default leaves + * built-in tool output unaffected; it exists to bound extension/MCP/custom tool + * results, which have no cap of their own otherwise. + */ +export const DEFAULT_MAX_TOOL_RESULT_BYTES = 128 * 1024; // 128KB + +const TOOL_RESULT_HEAD_FRACTION = 0.8; + +function buildToolResultElisionMarker(toolName: string, originalBytes: number, maxBytes: number): string { + return `\n\n[... tool result truncated: "${toolName}" returned ${originalBytes} bytes, exceeding the ${maxBytes}-byte cap. Showing the beginning and end; the middle is elided. Re-call ${toolName} with narrower arguments or pagination to see the rest. ...]\n\n`; +} + +/** + * Caps the combined byte size of text content blocks in a tool result. + * + * This is the single chokepoint that bounds every tool result - built-in, extension, + * MCP, or custom - before it enters session history and gets re-sent on every + * subsequent turn. Only text blocks are measured/trimmed; image blocks pass through + * untouched. + * + * `maxToolResultBytes` semantics: `undefined` applies the default cap + * (`DEFAULT_MAX_TOOL_RESULT_BYTES`); `0` (or any other non-positive value) + * explicitly disables capping; a positive number uses that cap. + */ +function capToolResultContent( + content: (TextContent | ImageContent)[], + toolName: string, + maxToolResultBytes: number | undefined, +): (TextContent | ImageContent)[] { + const effectiveMaxBytes = maxToolResultBytes === undefined ? DEFAULT_MAX_TOOL_RESULT_BYTES : maxToolResultBytes; + if (!effectiveMaxBytes || effectiveMaxBytes <= 0) return content; + + const textIndices: number[] = []; + let totalTextBytes = 0; + for (let i = 0; i < content.length; i++) { + const block = content[i]; + if (block.type === "text") { + textIndices.push(i); + totalTextBytes += utf8ByteLength(block.text); + } + } + if (textIndices.length === 0 || totalTextBytes <= effectiveMaxBytes) { + return content; + } + + const combinedText = textIndices.map((i) => (content[i] as TextContent).text).join("\n"); + const marker = buildToolResultElisionMarker(toolName, totalTextBytes, effectiveMaxBytes); + const markerBytes = utf8ByteLength(marker); + const budget = Math.max(0, effectiveMaxBytes - markerBytes); + const headBudget = Math.ceil(budget * TOOL_RESULT_HEAD_FRACTION); + const tailBudget = budget - headBudget; + + const head = truncateStringToBytesFromStart(combinedText, headBudget); + const tail = tailBudget > 0 ? truncateStringToBytesFromEnd(combinedText, tailBudget) : ""; + const cappedText = head + marker + tail; + + const firstTextIndex = textIndices[0]; + const result: (TextContent | ImageContent)[] = []; + for (let i = 0; i < content.length; i++) { + const block = content[i]; + if (block.type !== "text") { + result.push(block); + continue; + } + if (i === firstTextIndex) { + result.push({ type: "text", text: cappedText }); + } + // Other text blocks are dropped; their content is already folded into cappedText. + } + return result; +} + /** * Start an agent loop with a new prompt message. * The prompt is added to the context and events are emitted for it. @@ -242,7 +325,7 @@ async function runLoop( // them all instead of executing potentially borked calls. const executedToolBatch = message.stopReason === "length" - ? await failToolCallsFromTruncatedMessage(toolCalls, emit) + ? await failToolCallsFromTruncatedMessage(toolCalls, config, emit) : await executeToolCalls(currentContext, message, config, signal, emit); toolResults.push(...executedToolBatch.messages); hasMoreToolCalls = !executedToolBatch.terminate; @@ -417,6 +500,7 @@ async function streamAssistantResponse( */ async function failToolCallsFromTruncatedMessage( toolCalls: AgentToolCall[], + config: AgentLoopConfig, emit: AgentEventSink, ): Promise { const messages: ToolResultMessage[] = []; @@ -435,7 +519,7 @@ async function failToolCallsFromTruncatedMessage( isError: true, }; await emitToolExecutionEnd(finalized, emit); - const toolResultMessage = createToolResultMessage(finalized); + const toolResultMessage = createToolResultMessage(finalized, config.maxToolResultBytes); await emitToolResultMessage(toolResultMessage, emit); messages.push(toolResultMessage); } @@ -507,7 +591,7 @@ async function executeToolCallsSequential( } await emitToolExecutionEnd(finalized, emit); - const toolResultMessage = createToolResultMessage(finalized); + const toolResultMessage = createToolResultMessage(finalized, config.maxToolResultBytes); await emitToolResultMessage(toolResultMessage, emit); finalizedCalls.push(finalized); messages.push(toolResultMessage); @@ -579,7 +663,7 @@ async function executeToolCallsParallel( ); const messages: ToolResultMessage[] = []; for (const finalized of orderedFinalizedCalls) { - const toolResultMessage = createToolResultMessage(finalized); + const toolResultMessage = createToolResultMessage(finalized, config.maxToolResultBytes); await emitToolResultMessage(toolResultMessage, emit); messages.push(toolResultMessage); } @@ -811,14 +895,19 @@ async function emitToolExecutionEnd(finalized: FinalizedToolCallOutcome, emit: A }); } -function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResultMessage { +function createToolResultMessage( + finalized: FinalizedToolCallOutcome, + maxToolResultBytes: number | undefined, +): ToolResultMessage { return { role: "toolResult", toolCallId: finalized.toolCall.id, toolName: finalized.toolCall.name, // Untyped tools (JS extensions) can return results without content; normalize - // so the null never enters session history or provider payloads. - content: finalized.result.content ?? [], + // so the null never enters session history or provider payloads. Cap the + // resulting text so oversized extension/MCP/custom results never enter + // session history or get re-sent on every subsequent turn. + content: capToolResultContent(finalized.result.content ?? [], finalized.toolCall.name, maxToolResultBytes), details: finalized.result.details, usage: finalized.result.usage, ...(finalized.result.addedToolNames?.length ? { addedToolNames: finalized.result.addedToolNames } : {}), diff --git a/packages/agent-core/src/agent.ts b/packages/agent-core/src/agent.ts index 68365590..35df1945 100644 --- a/packages/agent-core/src/agent.ts +++ b/packages/agent-core/src/agent.ts @@ -120,6 +120,8 @@ export interface AgentOptions { transport?: Transport; maxRetryDelayMs?: number; toolExecution?: ToolExecutionMode; + /** See {@link AgentLoopConfig.maxToolResultBytes}. */ + maxToolResultBytes?: number; } class PendingMessageQueue { @@ -212,6 +214,8 @@ export class Agent { public maxRetryDelayMs?: number; /** Tool execution strategy for assistant messages that contain multiple tool calls. */ public toolExecution: ToolExecutionMode; + /** See {@link AgentLoopConfig.maxToolResultBytes}. */ + public maxToolResultBytes?: number; constructor(options: AgentOptions) { // Older compiled consumers may omit options or streamFn even though the current API requires them. @@ -235,6 +239,7 @@ export class Agent { this.transport = runtimeOptions.transport ?? "auto"; this.maxRetryDelayMs = runtimeOptions.maxRetryDelayMs; this.toolExecution = runtimeOptions.toolExecution ?? "parallel"; + this.maxToolResultBytes = runtimeOptions.maxToolResultBytes; } /** @@ -457,6 +462,7 @@ export class Agent { toolExecution: this.toolExecution, beforeToolCall: this.beforeToolCall, afterToolCall: this.afterToolCall, + maxToolResultBytes: this.maxToolResultBytes, shouldStopAfterTurn: shouldStopAfterTurn ? async (context) => await shouldStopAfterTurn(context, this.signal) : undefined, diff --git a/packages/agent-core/src/harness/utils/truncate.ts b/packages/agent-core/src/harness/utils/truncate.ts index 8150458c..4bf64042 100644 --- a/packages/agent-core/src/harness/utils/truncate.ts +++ b/packages/agent-core/src/harness/utils/truncate.ts @@ -51,7 +51,7 @@ interface RuntimeBuffer { const runtimeBuffer = (globalThis as { Buffer?: RuntimeBuffer }).Buffer; const nonAsciiPattern = /[^\x00-\x7f]/; -function utf8ByteLength(content: string): number { +export function utf8ByteLength(content: string): number { if (runtimeBuffer) return runtimeBuffer.byteLength(content, "utf8"); const firstNonAscii = content.search(nonAsciiPattern); @@ -294,11 +294,52 @@ export function truncateTail(content: string, options: TruncationOptions = {}): }; } +/** + * Truncate a string to fit within a byte limit (from the start). + * Handles multi-byte UTF-8 characters correctly; never splits a surrogate pair. + */ +export function truncateStringToBytesFromStart(str: string, maxBytes: number): string { + if (maxBytes <= 0) return ""; + + let outputBytes = 0; + let end = 0; + let needsReplacement = false; + for (let i = 0; i < str.length; ) { + const code = str.charCodeAt(i); + let characterEnd = i + 1; + let characterBytes: number; + let unpairedSurrogate = false; + if (code >= 0xd800 && code <= 0xdbff && i + 1 < str.length) { + const next = str.charCodeAt(i + 1); + if (next >= 0xdc00 && next <= 0xdfff) { + characterEnd = i + 2; + characterBytes = 4; + } else { + characterBytes = 3; + unpairedSurrogate = true; + } + } else if (code >= 0xd800 && code <= 0xdfff) { + characterBytes = 3; + unpairedSurrogate = true; + } else { + characterBytes = code <= 0x7f ? 1 : code <= 0x7ff ? 2 : 3; + } + if (outputBytes + characterBytes > maxBytes) break; + outputBytes += characterBytes; + end = characterEnd; + needsReplacement ||= unpairedSurrogate; + i = characterEnd; + } + + const output = str.slice(0, end); + return needsReplacement ? replaceUnpairedSurrogates(output) : output; +} + /** * Truncate a string to fit within a byte limit (from the end). * Handles multi-byte UTF-8 characters correctly. */ -function truncateStringToBytesFromEnd(str: string, maxBytes: number): string { +export function truncateStringToBytesFromEnd(str: string, maxBytes: number): string { if (maxBytes <= 0) return ""; let outputBytes = 0; diff --git a/packages/agent-core/src/types.ts b/packages/agent-core/src/types.ts index 5db470b4..f67a546f 100644 --- a/packages/agent-core/src/types.ts +++ b/packages/agent-core/src/types.ts @@ -303,6 +303,25 @@ export interface AgentLoopConfig extends SimpleStreamOptions { * The hook receives the agent abort signal and is responsible for honoring it. */ afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; + + /** + * Caps the combined byte size of text content blocks in every tool result, + * applied after `afterToolCall` and just before the result becomes a + * `ToolResultMessage` that enters session history and gets re-sent on every + * subsequent turn. This is the one place that bounds ALL tool results - + * built-in, extension, MCP, or custom - regardless of whether the tool + * itself truncates its own output. + * + * When exceeded, a head (and tail) of the combined text is kept and an + * elision marker is inserted stating the original size, the cap, the tool + * name, and instructing the model to re-call with narrower arguments or + * pagination. Image content blocks are never touched. + * + * `undefined` applies the default cap of 128KB, comfortably above the + * built-in per-tool caps (50KB), so built-in tool output is unaffected. + * Set explicitly to `0` (or any non-positive value) to disable capping. + */ + maxToolResultBytes?: number; } /** diff --git a/packages/agent-core/test/agent-loop.test.ts b/packages/agent-core/test/agent-loop.test.ts index 79957ef1..81cac86e 100644 --- a/packages/agent-core/test/agent-loop.test.ts +++ b/packages/agent-core/test/agent-loop.test.ts @@ -8,7 +8,7 @@ import { } from "@step-harness/providers"; import { Type } from "typebox"; import { describe, expect, it } from "vitest"; -import { agentLoop, agentLoopContinue } from "../src/agent-loop.ts"; +import { agentLoop, agentLoopContinue, DEFAULT_MAX_TOOL_RESULT_BYTES } from "../src/agent-loop.ts"; import { setDefaultStreamFn } from "../src/index.ts"; import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "../src/types.ts"; @@ -1730,3 +1730,109 @@ describe("tool-call markup leak retry", () => { expect(calls()).toBe(2); }); }); + +describe("tool result capping", () => { + const bigToolSchema = Type.Object({}); + + function createBigTool(text: string): AgentTool> { + return { + name: "big", + label: "Big", + description: "Returns a big result", + parameters: bigToolSchema, + async execute() { + return { content: [{ type: "text", text }], details: {} }; + }, + }; + } + + function scriptToolCallThenDone(toolName: string) { + let callIndex = 0; + return () => { + const stream = new MockAssistantStream(); + queueMicrotask(() => { + if (callIndex === 0) { + const message = createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: toolName, arguments: {} }], + "toolUse", + ); + stream.push({ type: "done", reason: "toolUse", message }); + } else { + stream.push({ + type: "done", + reason: "stop", + message: createAssistantMessage([{ type: "text", text: "done" }]), + }); + } + callIndex++; + }); + return stream; + }; + } + + it("caps an oversized tool result text at message_end and preserves isError/details", async () => { + const bigText = "x".repeat(10 * 1024 * 1024); // ~10MB + const tool = createBigTool(bigText); + + const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool] }; + const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter }; + + const events: AgentEvent[] = []; + const stream = agentLoop([createUserMessage("go")], context, config, undefined, scriptToolCallThenDone("big")); + for await (const event of stream) { + events.push(event); + } + + const toolResultEndEvent = events.find((e) => e.type === "message_end" && e.message.role === "toolResult"); + expect(toolResultEndEvent).toBeDefined(); + if (toolResultEndEvent?.type !== "message_end" || toolResultEndEvent.message.role !== "toolResult") { + throw new Error("expected toolResult message_end event"); + } + const message = toolResultEndEvent.message; + const textBlock = message.content.find((c) => c.type === "text"); + expect(textBlock).toBeDefined(); + const text = textBlock && textBlock.type === "text" ? textBlock.text : ""; + + expect(Buffer.byteLength(text, "utf8")).toBeLessThanOrEqual(DEFAULT_MAX_TOOL_RESULT_BYTES); + expect(text).toContain("10485760"); // original byte size mentioned in marker + expect(text.toLowerCase()).toContain("big"); // tool name mentioned + expect(message.isError).toBe(false); + expect(message.details).toEqual({}); + }); + + it("leaves small tool results byte-identical", async () => { + const smallText = "hello world"; + const tool = createBigTool(smallText); + + const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool] }; + const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter }; + + const stream = agentLoop([createUserMessage("go")], context, config, undefined, scriptToolCallThenDone("big")); + const messages = await stream.result(); + const toolResult = messages.find((m) => m.role === "toolResult"); + expect(toolResult?.role === "toolResult" ? toolResult.content : undefined).toEqual([ + { type: "text", text: smallText }, + ]); + }); + + it("honours a custom maxToolResultBytes", async () => { + const text = "a".repeat(5000); + const tool = createBigTool(text); + + const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool] }; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + maxToolResultBytes: 1000, + }; + + const stream = agentLoop([createUserMessage("go")], context, config, undefined, scriptToolCallThenDone("big")); + const messages = await stream.result(); + const toolResult = messages.find((m) => m.role === "toolResult"); + const content = toolResult?.role === "toolResult" ? toolResult.content : []; + const textBlock = content.find((c) => c.type === "text"); + const resultText = textBlock && textBlock.type === "text" ? textBlock.text : ""; + expect(Buffer.byteLength(resultText, "utf8")).toBeLessThanOrEqual(1000); + expect(resultText).toContain("5000"); + }); +}); diff --git a/packages/agent-core/test/harness/truncate.test.ts b/packages/agent-core/test/harness/truncate.test.ts index 42da9d1f..2df2573d 100644 --- a/packages/agent-core/test/harness/truncate.test.ts +++ b/packages/agent-core/test/harness/truncate.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { truncateHead, truncateTail } from "../../src/harness/utils/truncate.ts"; +import { truncateHead, truncateStringToBytesFromStart, truncateTail } from "../../src/harness/utils/truncate.ts"; const encoder = new TextEncoder(); @@ -176,3 +176,51 @@ describe("truncate utilities", () => { } }); }); + +describe("truncateStringToBytesFromStart", () => { + function bufferHead(content: string, maxBytes: number): string { + const bytes = Buffer.from(content, "utf8"); + if (bytes.length <= maxBytes) return content; + let end = maxBytes; + // Back off while sitting on a UTF-8 continuation byte (10xxxxxx). + while (end > 0 && (bytes[end] & 0xc0) === 0x80) end--; + return bytes.subarray(0, end).toString("utf8"); + } + + it("never splits a multi-byte character or an unpaired surrogate, and stays valid UTF-8", () => { + // Mix of ASCII, 2/3/4-byte code points, an emoji (surrogate pair), and lone surrogate halves. + const input = `hello ${"é".repeat(3)}${"中".repeat(3)}${"😀".repeat(5)}𐀀world`; + const totalBytes = Buffer.byteLength(input, "utf8"); + + for (let maxBytes = 0; maxBytes <= totalBytes + 5; maxBytes++) { + const result = truncateStringToBytesFromStart(input, maxBytes); + + // Bounded by the cap. + expect(Buffer.byteLength(result, "utf8")).toBeLessThanOrEqual(maxBytes); + + // Round-trips through UTF-8 without producing/leaving unpaired surrogates + // (an unpaired surrogate would either throw or come back as U+FFFD from Buffer). + const roundTripped = Buffer.from(result, "utf8").toString("utf8"); + expect(roundTripped).toBe(result); + expect(/[\ud800-\udbff](?![\udc00-\udfff])|(? { + const input = "abc😀def中文ghi"; + const totalBytes = Buffer.byteLength(input, "utf8"); + for (let maxBytes = 0; maxBytes <= totalBytes + 2; maxBytes++) { + expect(truncateStringToBytesFromStart(input, maxBytes)).toBe(bufferHead(input, maxBytes)); + } + }); + + it("returns the full string unchanged when it already fits", () => { + const input = "short and sweet"; + expect(truncateStringToBytesFromStart(input, Buffer.byteLength(input, "utf8"))).toBe(input); + }); + + it("returns empty string for a non-positive cap", () => { + expect(truncateStringToBytesFromStart("anything", 0)).toBe(""); + expect(truncateStringToBytesFromStart("anything", -5)).toBe(""); + }); +});