@gram-ai/elements 1.28.0 → 1.29.0

package/src/contexts/ElementsProvider.tsx

@@ -15,9 +15,11 @@ import {
   setFrontendToolApprovalConfig,
   toAISDKTools,
   wrapToolsWithApproval,
+  wrapToolsWithByteCap,
   type ApprovalHelpers,
   type FrontendTool,
 } from "@/lib/tools";
+import { compactForModel } from "@/lib/contextCompaction";
 import { cn } from "@/lib/utils";
 import { recommended } from "@/plugins";
 import { ElementsConfig, Model } from "@/types";
@@ -37,6 +39,7 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import {
   convertToModelMessages,
   createUIMessageStream,
+  lastAssistantMessageIsCompleteWithToolCalls,
   LanguageModel,
   smoothStream,
   stepCountIs,
@@ -366,12 +369,19 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
         } as ToolSet;
 
         // Wrap tools that require approval
-        const tools = wrapToolsWithApproval(
+        const approvedTools = wrapToolsWithApproval(
           combinedTools,
           config.tools?.toolsRequiringApproval,
           getApprovalHelpers(),
         );
 
+        // Cap oversized tool results so one greedy tool call (e.g. a wide log
+        // search) can't fill the context window in a single step.
+        const tools = wrapToolsWithByteCap(
+          approvedTools,
+          config.tools?.maxOutputBytes,
+        );
+
         // Stream the response
         const modelToUse = config.languageModel
           ? config.languageModel
@@ -387,7 +397,29 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
           const nonSystemMessages = cleanedMessages.filter(
             (m) => m.role !== "system",
           );
-          const modelMessages = convertToModelMessages(nonSystemMessages);
+          const rawModelMessages = convertToModelMessages(nonSystemMessages);
+
+          // Auto-compact older turns if the estimated input is approaching
+          // the model's context window. System prompt + last few turns are
+          // always preserved. No-op when the conversation is small.
+          const compaction = config.contextCompaction?.disabled
+            ? {
+                messages: rawModelMessages,
+                droppedCount: 0,
+                estimatedTokensBefore: 0,
+                estimatedTokensAfter: 0,
+              }
+            : compactForModel(rawModelMessages, model, {
+                maxTokens: config.contextCompaction?.maxTokens,
+                compactAtFraction: config.contextCompaction?.compactAtFraction,
+                keepRecent: config.contextCompaction?.keepRecent,
+              });
+          if (compaction.droppedCount > 0) {
+            console.warn(
+              `[elements] compacted ${compaction.droppedCount} older turn(s) from ${compaction.estimatedTokensBefore} → ${compaction.estimatedTokensAfter} est. tokens (model ${model})`,
+            );
+          }
+          const modelMessages = compaction.messages;
 
           const result = streamText({
             system: systemPrompt,
@@ -456,6 +488,11 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
     [
       config.languageModel,
       config.tools?.toolsRequiringApproval,
+      config.tools?.maxOutputBytes,
+      config.contextCompaction?.disabled,
+      config.contextCompaction?.maxTokens,
+      config.contextCompaction?.compactAtFraction,
+      config.contextCompaction?.keepRecent,
       model,
       systemPrompt,
       mcpTools,
@@ -606,9 +643,14 @@ const ElementsProviderWithHistory = ({
   });
   const initialThreadId = contextValue?.config.history?.initialThreadId;
 
-  // Hook factory for creating the base chat runtime
+  // Without `sendAutomaticallyWhen`, client-side frontend tools leave the turn
+  // half-finished: the tool-result is patched in but the agent never resumes,
+  // so the next user message lands on top of an unresolved tool-call sequence.
   const useChatRuntimeHook = useCallback(() => {
-    return useChatRuntime({ transport });
+    return useChatRuntime({
+      transport,
+      sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
+    });
   }, [transport]);
 
   const runtime = useRemoteThreadListRuntime({
@@ -691,7 +733,10 @@ const ElementsProviderWithoutHistory = ({
   executableTools,
   currentChatId,
 }: ElementsProviderWithoutHistoryProps) => {
-  const runtime = useChatRuntime({ transport });
+  const runtime = useChatRuntime({
+    transport,
+    sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
+  });
 
   // Populate runtimeRef so transport can access thread context
   useEffect(() => {

package/src/lib/contextCompaction.test.ts

@@ -0,0 +1,201 @@
+import { describe, expect, it } from "vitest";
+import type { ModelMessage } from "ai";
+import {
+  compactBySlidingWindow,
+  compactForModel,
+  DEFAULT_CONTEXT_LIMIT,
+  estimateTokens,
+  getModelContextLimit,
+} from "./contextCompaction";
+
+function msg(
+  role: "system" | "user" | "assistant" | "tool",
+  content: string,
+): ModelMessage {
+  return { role, content } as ModelMessage;
+}
+
+describe("estimateTokens", () => {
+  it("returns roughly chars/4", () => {
+    const messages = [msg("user", "a".repeat(400))];
+    const n = estimateTokens(messages);
+    // Actual output is JSON-wrapped so it's slightly larger than 100
+    expect(n).toBeGreaterThan(100);
+    expect(n).toBeLessThan(200);
+  });
+
+  it("grows with message count", () => {
+    const one = estimateTokens([msg("user", "hello")]);
+    const many = estimateTokens(
+      Array.from({ length: 100 }, () => msg("user", "hello")),
+    );
+    expect(many).toBeGreaterThan(one * 50);
+  });
+});
+
+describe("getModelContextLimit", () => {
+  it("returns known mapping for Sonnet 4.6", () => {
+    expect(getModelContextLimit("anthropic/claude-sonnet-4.6")).toBe(1_000_000);
+  });
+
+  it("returns known mapping for Claude 4 (non-1M)", () => {
+    expect(getModelContextLimit("anthropic/claude-sonnet-4")).toBe(200_000);
+  });
+
+  it("returns DEFAULT_CONTEXT_LIMIT for unknown models", () => {
+    expect(getModelContextLimit("acme/very-new-model")).toBe(
+      DEFAULT_CONTEXT_LIMIT,
+    );
+  });
+});
+
+describe("compactBySlidingWindow", () => {
+  it("no-ops when under the limit", () => {
+    const messages = [msg("user", "hi"), msg("assistant", "hello")];
+    const result = compactBySlidingWindow(messages, 1_000_000);
+    expect(result.droppedCount).toBe(0);
+    expect(result.messages).toBe(messages);
+  });
+
+  it("drops oldest non-system turns to fit", () => {
+    // 10 bulky messages, tiny limit → forces dropping
+    const messages: ModelMessage[] = [];
+    for (let i = 0; i < 10; i++) {
+      messages.push(msg("user", `query-${i} ` + "x".repeat(400)));
+      messages.push(msg("assistant", `reply-${i} ` + "y".repeat(400)));
+    }
+    const maxTokens = 500;
+    const result = compactBySlidingWindow(messages, maxTokens, 4);
+    expect(result.droppedCount).toBeGreaterThan(0);
+    expect(result.estimatedTokensAfter).toBeLessThanOrEqual(
+      result.estimatedTokensBefore,
+    );
+    // Last 4 are preserved verbatim
+    const tail = result.messages.slice(-4);
+    expect(tail[tail.length - 1]).toEqual(messages[messages.length - 1]);
+    // Marker prepended
+    const markerPresent = result.messages.some(
+      (m) => typeof m.content === "string" && m.content.includes("omitted"),
+    );
+    expect(markerPresent).toBe(true);
+  });
+
+  it("always preserves system messages", () => {
+    const messages: ModelMessage[] = [
+      msg("system", "sys " + "s".repeat(1000)),
+      ...Array.from({ length: 20 }, (_, i) =>
+        msg("user", `q-${i} ` + "x".repeat(500)),
+      ),
+    ];
+    const result = compactBySlidingWindow(messages, 300, 2);
+    expect(result.droppedCount).toBeGreaterThan(0);
+    expect(result.messages[0]!.role).toBe("system");
+  });
+
+  it("preserves at least keepRecent messages even if over limit", () => {
+    const messages = Array.from({ length: 10 }, (_, i) =>
+      msg("user", "x".repeat(1000) + `-${i}`),
+    );
+    const result = compactBySlidingWindow(messages, 10, 3);
+    // keepRecent preserved even though we can't get under the limit
+    expect(result.messages.length).toBeGreaterThanOrEqual(3);
+    // Last 3 are intact
+    const tail = result.messages.slice(-3);
+    expect(tail).toEqual(messages.slice(-3));
+  });
+});
+
+describe("compactBySlidingWindow — tool message pairing", () => {
+  it("never leaves a tool message at the head of the retained window", () => {
+    // Scenario from Devin: dropping oldest-first could split an
+    // assistant(tool_calls) → tool pair, leaving an orphan tool at the
+    // head of the retained set. Providers reject this with a 400.
+    const messages: ModelMessage[] = [
+      msg("user", "q1 " + "x".repeat(400)),
+      msg("assistant", "a1-with-tool-call " + "x".repeat(400)),
+      msg("tool", "t1-result " + "x".repeat(400)),
+      msg("assistant", "a1-final " + "x".repeat(400)),
+      msg("user", "q2 " + "x".repeat(400)),
+      msg("assistant", "a2-with-tool-call " + "x".repeat(400)),
+      msg("tool", "t2-result " + "x".repeat(400)),
+      msg("assistant", "a2-final " + "x".repeat(400)),
+    ];
+
+    const result = compactBySlidingWindow(messages, 400, 4);
+    expect(result.droppedCount).toBeGreaterThan(0);
+
+    // The retained non-system messages should never start with a tool.
+    const nonSystem = result.messages.filter((m) => m.role !== "system");
+    // Skip the synthetic assistant marker if present.
+    const firstReal = nonSystem.find(
+      (m) =>
+        !(
+          m.role === "assistant" &&
+          typeof m.content === "string" &&
+          m.content.includes("omitted")
+        ),
+    );
+    expect(firstReal?.role).not.toBe("tool");
+  });
+
+  it("drops an assistant+tool pair atomically (not one without the other)", () => {
+    const messages: ModelMessage[] = [
+      msg("user", "old"),
+      msg("assistant", "calling tool"),
+      msg("tool", "result " + "x".repeat(2000)),
+      msg("user", "recent " + "x".repeat(200)),
+      msg("assistant", "recent reply " + "x".repeat(200)),
+    ];
+    const result = compactBySlidingWindow(messages, 300, 2);
+    // If the group was dropped atomically, both the assistant and its tool
+    // are gone together. If the bug was still present, we'd see the tool
+    // message lingering alone.
+    const nonSystem = result.messages.filter((m) => m.role !== "system");
+    const hasLoneTool = nonSystem.some(
+      (m, i) =>
+        m.role === "tool" && (i === 0 || nonSystem[i - 1]!.role === "user"),
+    );
+    expect(hasLoneTool).toBe(false);
+  });
+
+  it("does not split a tool group when aligning the recent window", () => {
+    // keepRecent=3 would cut mid-group with naive slicing. Grouping should
+    // expand the recent window to keep the assistant+tools together.
+    const messages: ModelMessage[] = [
+      msg("user", "old " + "x".repeat(1000)),
+      msg("assistant", "calling 2 tools"),
+      msg("tool", "result1"),
+      msg("tool", "result2"),
+      msg("assistant", "final"),
+    ];
+    const result = compactBySlidingWindow(messages, 200, 3);
+    // If the first tool was the "recent" cut-off, we'd see a tool at
+    // the head of retained — but grouping should have pulled the
+    // assistant with it.
+    const kept = result.messages.filter((m) => m.role !== "system");
+    const firstTool = kept.findIndex((m) => m.role === "tool");
+    if (firstTool !== -1) {
+      expect(kept[firstTool - 1]?.role).toMatch(/assistant|tool/);
+    }
+  });
+});
+
+describe("compactForModel", () => {
+  it("uses 70% of the nominal ceiling by default", () => {
+    const small = [msg("user", "hi")];
+    const result = compactForModel(small, "anthropic/claude-sonnet-4.6");
+    expect(result.droppedCount).toBe(0);
+    expect(result.messages).toBe(small);
+  });
+
+  it("honors explicit maxTokens override", () => {
+    const messages = Array.from({ length: 30 }, (_, i) =>
+      msg("user", "x".repeat(500) + `-${i}`),
+    );
+    const result = compactForModel(messages, "anthropic/claude-sonnet-4.6", {
+      maxTokens: 2000,
+      keepRecent: 2,
+    });
+    expect(result.droppedCount).toBeGreaterThan(0);
+  });
+});

package/src/lib/contextCompaction.ts

@@ -0,0 +1,211 @@
+import type { ModelMessage } from "ai";
+import { MODELS } from "./models";
+
+type KnownModelId = (typeof MODELS)[number];
+
+/**
+ * Fraction-of-limit at which compaction kicks in. Below this, messages pass
+ * through untouched; above this, oldest non-system turns are dropped until
+ * the estimated token count is back under the threshold.
+ */
+export const DEFAULT_COMPACTION_FRACTION = 0.7;
+
+/**
+ * Number of most-recent messages preserved verbatim, even if the conversation
+ * is already over the limit. Ensures the assistant always has the latest turn
+ * and its immediate predecessor.
+ */
+export const DEFAULT_KEEP_RECENT = 4;
+
+/**
+ * Conservative fallback when we encounter a model we haven't mapped — big
+ * enough to be useful for unknown models, small enough to still trigger
+ * compaction before hitting upstream 400s.
+ */
+export const DEFAULT_CONTEXT_LIMIT = 200_000;
+
+/**
+ * Known input-token ceilings per model (nominal upstream maximum). Keyed by
+ * MODELS so TypeScript catches drift — adding a model id here that isn't in
+ * MODELS, or misspelling an id, is a compile error. Coverage is intentionally
+ * partial: models without an explicit entry fall back to DEFAULT_CONTEXT_LIMIT.
+ */
+const MODEL_CONTEXT_LIMITS: Partial<Record<KnownModelId, number>> = {
+  // Anthropic (1M tier where available, else 200K)
+  "anthropic/claude-opus-4.6": 1_000_000,
+  "anthropic/claude-opus-4.5": 1_000_000,
+  "anthropic/claude-opus-4.1": 200_000,
+  "anthropic/claude-sonnet-4.6": 1_000_000,
+  "anthropic/claude-sonnet-4.5": 1_000_000,
+  "anthropic/claude-sonnet-4": 200_000,
+  "anthropic/claude-haiku-4.5": 200_000,
+
+  // OpenAI
+  "openai/gpt-5.4": 400_000,
+  "openai/gpt-5.4-mini": 400_000,
+  "openai/gpt-5.1": 400_000,
+  "openai/gpt-5.1-codex": 400_000,
+  "openai/gpt-5": 400_000,
+  "openai/gpt-4.1": 1_000_000,
+  "openai/o4-mini": 200_000,
+  "openai/o3": 200_000,
+
+  // Google
+  "google/gemini-3.1-pro-preview": 1_000_000,
+  "google/gemini-2.5-pro": 1_000_000,
+  "google/gemini-2.5-flash": 1_000_000,
+
+  // Others
+  "deepseek/deepseek-r1": 128_000,
+  "deepseek/deepseek-v3.2": 128_000,
+  "meta-llama/llama-4-maverick": 1_000_000,
+  "x-ai/grok-4": 256_000,
+  "qwen/qwen3-coder": 256_000,
+  "moonshotai/kimi-k2.5": 128_000,
+  "mistralai/mistral-medium-3.1": 128_000,
+  "mistralai/codestral-2508": 256_000,
+  "mistralai/devstral-small": 128_000,
+};
+
+/**
+ * Returns the input-token ceiling for a given OpenRouter model id, or
+ * DEFAULT_CONTEXT_LIMIT if unknown.
+ */
+export function getModelContextLimit(modelId: string): number {
+  return MODEL_CONTEXT_LIMITS[modelId as KnownModelId] ?? DEFAULT_CONTEXT_LIMIT;
+}
+
+/**
+ * Rough input-token estimate using a chars/4 heuristic on the JSON serialized
+ * conversation. Tokens-per-char varies by model and content, but a chars/4
+ * heuristic matches OpenAI's rule-of-thumb within ~15% for English prose and
+ * is deterministic + zero-cost — good enough to trigger compaction.
+ */
+export function estimateTokens(messages: ModelMessage[]): number {
+  const serialized = JSON.stringify(messages);
+  return Math.ceil(serialized.length / 4);
+}
+
+export interface CompactionResult {
+  messages: ModelMessage[];
+  droppedCount: number;
+  estimatedTokensBefore: number;
+  estimatedTokensAfter: number;
+}
+
+/**
+ * Drops oldest non-system messages until the estimated token count is under
+ * maxTokens. Always preserves the last `keepRecent` messages and any system
+ * role messages. When any messages are dropped, prepends a synthetic assistant
+ * note so the model knows earlier context was elided.
+ */
+export function compactBySlidingWindow(
+  messages: ModelMessage[],
+  maxTokens: number,
+  keepRecent: number = DEFAULT_KEEP_RECENT,
+): CompactionResult {
+  const estimatedTokensBefore = estimateTokens(messages);
+
+  if (estimatedTokensBefore <= maxTokens || messages.length <= keepRecent) {
+    return {
+      messages,
+      droppedCount: 0,
+      estimatedTokensBefore,
+      estimatedTokensAfter: estimatedTokensBefore,
+    };
+  }
+
+  const systemMessages = messages.filter((m) => m.role === "system");
+  const nonSystem = messages.filter((m) => m.role !== "system");
+
+  // Group consecutive `tool` messages with the assistant message that
+  // precedes them. OpenAI-compatible providers require every tool-result
+  // message to be immediately preceded by the assistant message holding its
+  // tool_calls — splitting these produces an invalid conversation that
+  // providers reject with a 400. Grouping ensures we drop or keep the
+  // full assistant+tools unit atomically.
+  const groups: ModelMessage[][] = [];
+  for (const m of nonSystem) {
+    if (m.role === "tool" && groups.length > 0) {
+      groups[groups.length - 1]!.push(m);
+    } else {
+      groups.push([m]);
+    }
+  }
+
+  // Reserve the trailing groups that together contain at least `keepRecent`
+  // messages. Using groups (not raw messages) keeps assistant+tool pairs
+  // intact at the boundary between retained and dropped.
+  let recentMsgCount = 0;
+  let recentStart = groups.length;
+  while (recentStart > 0 && recentMsgCount < keepRecent) {
+    recentStart -= 1;
+    recentMsgCount += groups[recentStart]!.length;
+  }
+  const recentGroups = groups.slice(recentStart);
+  const droppableGroups = groups.slice(0, recentStart);
+
+  let droppedCount = 0;
+  let working = [
+    ...systemMessages,
+    ...droppableGroups.flat(),
+    ...recentGroups.flat(),
+  ];
+
+  while (droppableGroups.length > 0 && estimateTokens(working) > maxTokens) {
+    const droppedGroup = droppableGroups.shift()!;
+    droppedCount += droppedGroup.length;
+    working = [
+      ...systemMessages,
+      ...droppableGroups.flat(),
+      ...recentGroups.flat(),
+    ];
+  }
+
+  if (droppedCount > 0) {
+    const marker: ModelMessage = {
+      role: "assistant",
+      content: `[${droppedCount} earlier message${
+        droppedCount === 1 ? "" : "s"
+      } omitted to stay under context length. If the user asks about them, say you no longer have that context and suggest they restate the relevant details.]`,
+    };
+    working = [
+      ...systemMessages,
+      marker,
+      ...droppableGroups.flat(),
+      ...recentGroups.flat(),
+    ];
+  }
+
+  return {
+    messages: working,
+    droppedCount,
+    estimatedTokensBefore,
+    estimatedTokensAfter: estimateTokens(working),
+  };
+}
+
+export interface CompactionOptions {
+  /** Override the model's nominal input ceiling. */
+  maxTokens?: number;
+  /** Fraction of maxTokens at which compaction kicks in. */
+  compactAtFraction?: number;
+  /** Most-recent messages preserved verbatim. */
+  keepRecent?: number;
+}
+
+/**
+ * Convenience wrapper that picks the model ceiling, applies compactAtFraction,
+ * and runs compactBySlidingWindow. Returns the (possibly unchanged) messages
+ * plus diagnostics.
+ */
+export function compactForModel(
+  messages: ModelMessage[],
+  modelId: string,
+  opts: CompactionOptions = {},
+): CompactionResult {
+  const ceiling = opts.maxTokens ?? getModelContextLimit(modelId);
+  const fraction = opts.compactAtFraction ?? DEFAULT_COMPACTION_FRACTION;
+  const limit = Math.floor(ceiling * fraction);
+  return compactBySlidingWindow(messages, limit, opts.keepRecent);
+}

package/src/lib/errorTracking.config.ts

@@ -7,10 +7,12 @@
  * - VITE_DATADOG_APPLICATION_ID
  * - VITE_DATADOG_CLIENT_TOKEN
  * - VITE_DATADOG_SITE (optional, defaults to datadoghq.com)
+ * - VITE_DATADOG_ENV (optional, defaults to prod)
  */
 export const DATADOG_CONFIG = {
   applicationId: import.meta.env.VITE_DATADOG_APPLICATION_ID ?? "",
   clientToken: import.meta.env.VITE_DATADOG_CLIENT_TOKEN ?? "",
   site: import.meta.env.VITE_DATADOG_SITE ?? "datadoghq.com",
+  env: import.meta.env.VITE_DATADOG_ENV ?? "prod",
   service: "gram-elements",
 } as const;

package/src/lib/errorTracking.ts

@@ -44,7 +44,7 @@ export function initErrorTracking(config: ErrorTrackingConfig = {}): void {
       clientToken: DATADOG_CONFIG.clientToken,
       site: DATADOG_CONFIG.site,
       service: DATADOG_CONFIG.service,
-      env: process.env.NODE_ENV || "production",
+      env: DATADOG_CONFIG.env,
       sessionSampleRate: 100,
       sessionReplaySampleRate: 100,
       trackUserInteractions: true,

package/src/lib/tools.byte-cap.test.ts

@@ -0,0 +1,132 @@
+import { describe, expect, it, vi } from "vitest";
+import {
+  capToolResultBytes,
+  truncateTextToByteCap,
+  wrapToolsWithByteCap,
+} from "./tools";
+import type { ToolSet } from "ai";
+
+describe("truncateTextToByteCap", () => {
+  it("returns original when under cap", () => {
+    expect(truncateTextToByteCap("hello world", 100)).toBe("hello world");
+  });
+
+  it("truncates with head + tail + notice when over cap", () => {
+    const text = "a".repeat(1000) + "-MIDDLE-" + "b".repeat(1000);
+    const out = truncateTextToByteCap(text, 200);
+    expect(out.length).toBeLessThan(text.length);
+    expect(out).toContain("tool output truncated");
+    expect(out.startsWith("a")).toBe(true);
+    expect(out.endsWith("b")).toBe(true);
+    expect(out).not.toContain("MIDDLE");
+  });
+
+  it("output always stays at or under maxBytes (notice included in budget)", () => {
+    // Regression test — earlier version appended the notice *without*
+    // reserving budget for it, so the output overshot maxBytes by ~100.
+    for (const maxBytes of [256, 512, 1024, 4096]) {
+      const text = "x".repeat(50_000);
+      const out = truncateTextToByteCap(text, maxBytes);
+      const outBytes = new TextEncoder().encode(out).byteLength;
+      expect(outBytes).toBeLessThanOrEqual(maxBytes);
+    }
+  });
+
+  it("passes through when maxBytes <= 0 (disabled)", () => {
+    const text = "x".repeat(10_000);
+    expect(truncateTextToByteCap(text, 0)).toBe(text);
+    expect(truncateTextToByteCap(text, -1)).toBe(text);
+  });
+
+  it("handles multibyte UTF-8 without crashing", () => {
+    const text = "🎉".repeat(500);
+    const out = truncateTextToByteCap(text, 200);
+    expect(out).toContain("tool output truncated");
+    expect(new TextEncoder().encode(out).byteLength).toBeGreaterThan(0);
+  });
+});
+
+describe("capToolResultBytes", () => {
+  it("truncates plain string results", () => {
+    const out = capToolResultBytes("x".repeat(5_000), 100);
+    expect(typeof out).toBe("string");
+    expect(out).not.toBe("x".repeat(5_000));
+    expect(out).toContain("tool output truncated");
+  });
+
+  it("truncates text chunks inside MCP-shaped results", () => {
+    const result = {
+      content: [
+        { type: "text", text: "short" },
+        { type: "text", text: "big".repeat(5_000) },
+      ],
+      isError: false,
+    };
+    const out = capToolResultBytes(result, 100) as typeof result;
+    expect(out.content[0]).toEqual({ type: "text", text: "short" });
+    expect((out.content[1] as { text: string }).text).toContain(
+      "tool output truncated",
+    );
+    expect(out.isError).toBe(false);
+  });
+
+  it("leaves non-text chunks alone", () => {
+    const result = {
+      content: [
+        { type: "image", data: "x".repeat(5_000), mimeType: "image/png" },
+      ],
+    };
+    const out = capToolResultBytes(result, 100) as typeof result;
+    expect(out.content[0]).toEqual(result.content[0]);
+  });
+
+  it("preserves isError flag", () => {
+    const result = {
+      content: [{ type: "text", text: "tool blew up: " + "x".repeat(5_000) }],
+      isError: true,
+    };
+    const out = capToolResultBytes(result, 100) as typeof result;
+    expect(out.isError).toBe(true);
+  });
+
+  it("passes unknown shapes through", () => {
+    expect(capToolResultBytes(42, 100)).toBe(42);
+    expect(capToolResultBytes(null, 100)).toBe(null);
+    expect(capToolResultBytes({ foo: "bar" }, 100)).toEqual({ foo: "bar" });
+  });
+});
+
+describe("wrapToolsWithByteCap", () => {
+  it("is a no-op when maxBytes is undefined/0", () => {
+    const execute = vi.fn().mockResolvedValue("anything");
+    const tools: ToolSet = {
+      t: { description: "", inputSchema: { type: "object" }, execute } as never,
+    };
+    expect(wrapToolsWithByteCap(tools, undefined)).toBe(tools);
+    expect(wrapToolsWithByteCap(tools, 0)).toBe(tools);
+  });
+
+  it("wraps execute and truncates oversized result", async () => {
+    const execute = vi.fn().mockResolvedValue({
+      content: [{ type: "text", text: "z".repeat(10_000) }],
+    });
+    const tools: ToolSet = {
+      t: { description: "", inputSchema: { type: "object" }, execute } as never,
+    };
+    const wrapped = wrapToolsWithByteCap(tools, 256);
+    const wrappedExecute = wrapped.t.execute!;
+    const out = (await wrappedExecute({}, { toolCallId: "id" } as never)) as {
+      content: Array<{ text: string }>;
+    };
+    expect(out.content[0]!.text).toContain("tool output truncated");
+    expect(execute).toHaveBeenCalledOnce();
+  });
+
+  it("leaves tools without execute alone", () => {
+    const tools: ToolSet = {
+      t: { description: "", inputSchema: { type: "object" } } as never,
+    };
+    const wrapped = wrapToolsWithByteCap(tools, 256);
+    expect(wrapped.t).toBe(tools.t);
+  });
+});