@gram-ai/elements 1.28.0 → 1.30.0

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);
+  });
+});

package/src/lib/tools.test.ts

@@ -0,0 +1,259 @@
+import { describe, expect, it } from "vitest";
+import {
+  convertToModelMessages,
+  isToolUIPart,
+  jsonSchema,
+  lastAssistantMessageIsCompleteWithToolCalls,
+  readUIMessageStream,
+  stepCountIs,
+  streamText,
+  type ToolSet,
+  type UIMessage,
+  type UIMessagePart,
+} from "ai";
+import { MockLanguageModelV2 } from "ai/test";
+
+type MockStream = Extract<
+  NonNullable<
+    NonNullable<
+      ConstructorParameters<typeof MockLanguageModelV2>[0]
+    >["doStream"]
+  >,
+  (...a: never[]) => PromiseLike<{ stream: ReadableStream<unknown> }>
+>;
+type StreamPart =
+  Awaited<ReturnType<MockStream>>["stream"] extends ReadableStream<infer T>
+    ? T
+    : never;
+
+/**
+ * Repro for the assistants-onboarding "Skip bugged state":
+ *
+ * 1. Assistant calls a frontend tool (e.g. `request_environment_secrets`) that
+ *    renders a form with a Skip button.
+ * 2. User clicks Skip. The form calls `draft.resolvePending(toolCallId, { cancelled: true })`.
+ * 3. Expected: tool-result is patched onto the message, the agent continues,
+ *    chat returns to a ready state.
+ * 4. Observed (pre-fix): the chat stayed stuck — the next user message landed
+ *    with an invalid tool sequence and the provider rejected it with a
+ *    "message needing to be sent with role: assistant"-shaped error.
+ *
+ * `streamText` runs without an `execute` for frontend tools: AI-SDK's
+ * `frontendTools()` helper strips execute so client-side logic can take over.
+ * The missing link on main was that the runtime patched in the tool result
+ * but nothing resumed the turn. The fix wires `sendAutomaticallyWhen:
+ * lastAssistantMessageIsCompleteWithToolCalls` into `useChatRuntime`, which
+ * flips that resume on.
+ */
+
+function toolCallChunks(opts: {
+  toolCallId: string;
+  toolName: string;
+  input: string;
+}): StreamPart[] {
+  return [
+    { type: "stream-start", warnings: [] },
+    {
+      type: "response-metadata",
+      id: "resp-1",
+      modelId: "m",
+      timestamp: new Date(0),
+    },
+    { type: "tool-input-start", id: opts.toolCallId, toolName: opts.toolName },
+    { type: "tool-input-delta", id: opts.toolCallId, delta: opts.input },
+    { type: "tool-input-end", id: opts.toolCallId },
+    {
+      type: "tool-call",
+      toolCallId: opts.toolCallId,
+      toolName: opts.toolName,
+      input: opts.input,
+    },
+    {
+      type: "finish",
+      finishReason: "tool-calls",
+      usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 },
+    },
+  ];
+}
+
+function makeStream<T>(chunks: T[]): ReadableStream<T> {
+  return new ReadableStream({
+    start(controller) {
+      for (const c of chunks) controller.enqueue(c);
+      controller.close();
+    },
+  });
+}
+
+async function collectUIMessages(
+  stream: AsyncIterable<UIMessage>,
+): Promise<UIMessage[]> {
+  const out: UIMessage[] = [];
+  for await (const msg of stream) {
+    const idx = out.findIndex((m) => m.id === msg.id);
+    if (idx >= 0) out[idx] = msg;
+    else out.push(msg);
+  }
+  return out;
+}
+
+async function streamToolCallOnly(toolCallId: string): Promise<UIMessage[]> {
+  const toolsNoExecute = {
+    request_environment_secrets: {
+      description: "Ask the user to enter secrets for an env.",
+      inputSchema: jsonSchema({
+        type: "object",
+        properties: {
+          keys: {
+            type: "array",
+            items: {
+              type: "object",
+              properties: { name: { type: "string" } },
+              required: ["name"],
+            },
+          },
+        },
+        required: ["keys"],
+      }),
+    },
+  } as unknown as ToolSet;
+
+  const model = new MockLanguageModelV2({
+    doStream: async () => ({
+      stream: makeStream([
+        ...toolCallChunks({
+          toolCallId,
+          toolName: "request_environment_secrets",
+          input: JSON.stringify({ keys: [{ name: "SLACK_BOT_TOKEN" }] }),
+        }),
+      ]),
+    }),
+  });
+
+  const result = streamText({
+    model,
+    messages: [{ role: "user", content: "Set up Slack" }],
+    tools: toolsNoExecute,
+    stopWhen: stepCountIs(5),
+  });
+  return collectUIMessages(
+    readUIMessageStream({ stream: result.toUIMessageStream() }),
+  );
+}
+
+describe("frontend tool Skip flow (sendAutomaticallyWhen fix)", () => {
+  it("without a tool-result, the message sequence is invalid — this is the bug we are fixing", async () => {
+    // Mirrors the elements flow on main: frontend tool has no execute inside
+    // streamText (it's run client-side by useToolInvocations). If nothing
+    // patches a tool-result onto the message, sending a follow-up user
+    // message produces an invalid sequence.
+    const toolCallId = "call_unresolved";
+    const messages = await streamToolCallOnly(toolCallId);
+
+    const assistant = messages.find((m) => m.role === "assistant")!;
+    const toolParts = (assistant.parts as UIMessagePart<never, never>[]).filter(
+      (p) => isToolUIPart(p),
+    );
+    expect(toolParts).toHaveLength(1);
+    expect((toolParts[0] as unknown as { state: string }).state).toBe(
+      "input-available",
+    );
+
+    const follow: UIMessage[] = [
+      ...messages,
+      {
+        id: "u2",
+        role: "user",
+        parts: [{ type: "text", text: "skip" }],
+      } as unknown as UIMessage,
+    ];
+
+    // Also: `lastAssistantMessageIsCompleteWithToolCalls` must return `false`
+    // here — there is no tool result, so the runtime should NOT auto-resume.
+    expect(lastAssistantMessageIsCompleteWithToolCalls({ messages })).toBe(
+      false,
+    );
+
+    // And the resulting model-message sequence contains a bogus `role: "tool"`
+    // with empty content — the provider will reject this as an invalid tool
+    // message, surfacing to the user as the "needs role: assistant" error.
+    const modelMsgs = convertToModelMessages(follow);
+    const assistantIdx = modelMsgs.findIndex(
+      (m) =>
+        m.role === "assistant" &&
+        Array.isArray(m.content) &&
+        (m.content as Array<{ type: string }>).some(
+          (c) => c.type === "tool-call",
+        ),
+    );
+    expect(assistantIdx).toBeGreaterThanOrEqual(0);
+    expect(modelMsgs[assistantIdx + 1]?.role).toBe("tool");
+    expect(modelMsgs[assistantIdx + 1]?.content).toEqual([]);
+  });
+
+  it("once the tool-result is patched onto the message, sendAutomaticallyWhen fires and the sequence is valid", async () => {
+    // Simulates the full post-fix behaviour: `useToolInvocations` ran execute
+    // client-side and called `addToolResult`, which flips the tool part to
+    // `output-available`. With the result in place:
+    //   - `lastAssistantMessageIsCompleteWithToolCalls` returns true, so the
+    //     runtime re-issues the model turn (this is the 1-line fix).
+    //   - `convertToModelMessages` produces a real `role: "tool"` message
+    //     with the result, which the provider accepts.
+    const toolCallId = "call_resolved";
+    const rawMessages = await streamToolCallOnly(toolCallId);
+
+    // Patch in the tool-output-available state — this is the shape
+    // `chatHelpers.addToolResult` produces under the hood.
+    const patched: UIMessage[] = rawMessages.map((m) => {
+      if (m.role !== "assistant") return m;
+      return {
+        ...m,
+        parts: (m.parts as Array<Record<string, unknown>>).map((p) =>
+          isToolUIPart(p as UIMessagePart<never, never>)
+            ? {
+                ...p,
+                state: "output-available",
+                output: { ok: true, cancelled: true },
+              }
+            : p,
+        ),
+      } as UIMessage;
+    });
+
+    const toolPart = (
+      (
+        patched.find((m) => m.role === "assistant")!.parts as UIMessagePart<
+          never,
+          never
+        >[]
+      ).filter((p) => isToolUIPart(p))[0] as unknown as { state: string }
+    ).state;
+    expect(toolPart).toBe("output-available");
+
+    // Pre-condition for the fix: the runtime auto-resumes the turn.
+    expect(
+      lastAssistantMessageIsCompleteWithToolCalls({ messages: patched }),
+    ).toBe(true);
+
+    // And the sequence handed to the model is well-formed (assistant
+    // tool-call is followed by a real role:"tool" message with a result).
+    const modelMsgs = convertToModelMessages(patched);
+    const assistantIdx = modelMsgs.findIndex(
+      (m) =>
+        m.role === "assistant" &&
+        Array.isArray(m.content) &&
+        (m.content as Array<{ type: string }>).some(
+          (c) => c.type === "tool-call",
+        ),
+    );
+    const next = modelMsgs[assistantIdx + 1];
+    expect(next?.role).toBe("tool");
+    expect(Array.isArray(next?.content) && next.content.length).toBeGreaterThan(
+      0,
+    );
+    const toolResult = (
+      next?.content as Array<{ type: string; output?: { type?: string } }>
+    )[0];
+    expect(toolResult?.type).toBe("tool-result");
+  });
+});

package/src/lib/tools.ts

@@ -156,6 +156,128 @@ export interface ApprovalHelpers {
   whitelistTool: (toolName: string) => void;
 }
 
+/**
+ * Default head/tail split (bytes) when a tool result exceeds the cap. Head keeps
+ * early context (e.g. the preamble of a log query); tail keeps the most recent
+ * lines, which are usually the most relevant.
+ */
+const BYTE_CAP_HEAD_FRACTION = 0.9;
+
+/**
+ * Truncates a single string to maxBytes using a head + tail preserving strategy
+ * when it exceeds the cap. Returns the original string when under the cap.
+ */
+export function truncateTextToByteCap(text: string, maxBytes: number): string {
+  if (maxBytes <= 0) return text;
+  const original = text;
+  // Work in UTF-8 bytes to match what OpenRouter counts.
+  const encoded = new TextEncoder().encode(original);
+  if (encoded.byteLength <= maxBytes) return original;
+
+  // Reserve room for the notice up-front so final output stays under maxBytes.
+  // Without this deduction, output would be head + notice + tail ≈ maxBytes
+  // + ~100 bytes, which silently overshoots the cap.
+  const notice = `\n\n[…tool output truncated from ${encoded.byteLength} bytes to ${maxBytes}; ask a narrower question to see more…]\n\n`;
+  const noticeBytes = new TextEncoder().encode(notice).byteLength;
+  const availableBytes = Math.max(0, maxBytes - noticeBytes);
+
+  const headBytes = Math.max(
+    0,
+    Math.floor(availableBytes * BYTE_CAP_HEAD_FRACTION),
+  );
+  const tailBytes = Math.max(0, availableBytes - headBytes);
+  const decoder = new TextDecoder("utf-8", { fatal: false });
+  const head = decoder.decode(encoded.slice(0, headBytes));
+  const tail =
+    tailBytes > 0
+      ? decoder.decode(encoded.slice(encoded.byteLength - tailBytes))
+      : "";
+
+  return tail ? `${head}${notice}${tail}` : `${head}${notice}`;
+}
+
+/**
+ * Walks the shape returned by MCP/AI-SDK tool executors and truncates any
+ * over-sized text payload in place. Handles:
+ *   - plain strings
+ *   - { content: Array<{ type, text?, ... }>, isError? }
+ * Other shapes pass through untouched.
+ */
+export function capToolResultBytes(result: unknown, maxBytes: number): unknown {
+  if (maxBytes <= 0) return result;
+
+  if (typeof result === "string") {
+    return truncateTextToByteCap(result, maxBytes);
+  }
+
+  if (result && typeof result === "object" && "content" in result) {
+    const r = result as {
+      content?: unknown;
+      isError?: boolean;
+      [k: string]: unknown;
+    };
+    if (Array.isArray(r.content)) {
+      const cappedContent = r.content.map((chunk) => {
+        if (
+          chunk &&
+          typeof chunk === "object" &&
+          (chunk as { type?: unknown }).type === "text" &&
+          typeof (chunk as { text?: unknown }).text === "string"
+        ) {
+          return {
+            ...(chunk as Record<string, unknown>),
+            text: truncateTextToByteCap(
+              (chunk as { text: string }).text,
+              maxBytes,
+            ),
+          };
+        }
+        return chunk;
+      });
+      return { ...r, content: cappedContent };
+    }
+  }
+
+  return result;
+}
+
+/**
+ * Wraps tools so that oversized results are truncated before they reach the
+ * conversation history. Tools whose result fits under the cap pass through
+ * untouched. Composes cleanly before or after wrapToolsWithApproval.
+ */
+export function wrapToolsWithByteCap(
+  tools: ToolSet,
+  maxBytes: number | undefined,
+): ToolSet {
+  if (!maxBytes || maxBytes <= 0) {
+    return tools;
+  }
+
+  return Object.fromEntries(
+    Object.entries(tools).map(([name, tool]) => {
+      const originalExecute = tool.execute;
+      if (!originalExecute) {
+        return [name, tool];
+      }
+
+      return [
+        name,
+        {
+          ...tool,
+          execute: async (args: unknown, options?: ToolCallOptions) => {
+            const result = await originalExecute(
+              args,
+              options as Parameters<typeof originalExecute>[1],
+            );
+            return capToolResultBytes(result, maxBytes);
+          },
+        },
+      ];
+    }),
+  ) as ToolSet;
+}
+
 /**
  * Wraps tools with approval logic based on the approval config.
  */

package/src/plugins/chart/index.ts

@@ -5,6 +5,7 @@ import { ChartRenderer } from "./component";
  * This plugin renders charts using json-render format.
  */
 export const chart: Plugin = {
+  id: "chart",
   language: "chart",
   prompt: `WHEN TO USE CHARTS:
 Create charts to visualize numerical data when it helps users understand patterns, trends, or comparisons. Use the 'chart' code block format.

package/src/plugins/chart/ui/bar-chart.tsx

@@ -58,6 +58,14 @@ export const BarChart: FC<BarChartProps> = ({
 }) => {
   const isHorizontal = layout === "horizontal";
 
+  // Estimate YAxis width from longest label to prevent clipping
+  const yAxisWidth = isHorizontal
+    ? Math.max(
+        80,
+        Math.min(200, Math.max(...data.map((d) => d.label.length)) * 7 + 16),
+      )
+    : undefined;
+
   return (
     <div className={cn("flex flex-col gap-2", className)}>
       {title && (
@@ -87,7 +95,7 @@ export const BarChart: FC<BarChartProps> = ({
                 <YAxis
                   dataKey="label"
                   type="category"
-                  width={80}
+                  width={yAxisWidth}
                   tick={{ fill: "var(--foreground)", fontSize: 12 }}
                   axisLine={{ stroke: "var(--border)" }}
                   tickLine={{ stroke: "var(--border)" }}

package/src/plugins/generative-ui/index.ts

@@ -6,6 +6,7 @@ import { GenerativeUIRenderer } from "./component";
  * Use the language identifier 'ui' or 'json-render' in code blocks.
  */
 export const generativeUI: Plugin = {
+  id: "generative-ui",
   language: "ui",
   prompt: `Render structured data as visual UI using \`\`\`ui code blocks with valid JSON.
 

package/src/plugins/index.test.ts

@@ -0,0 +1,62 @@
+import { describe, it, expect } from "vitest";
+import { recommended, type Plugin } from "./index";
+import { chart } from "./chart";
+import { generativeUI } from "./generative-ui";
+
+describe("recommended", () => {
+  it("contains chart and generative-ui plugins", () => {
+    expect(recommended).toHaveLength(2);
+    expect(recommended).toContain(chart);
+    expect(recommended).toContain(generativeUI);
+  });
+
+  it("is usable as a plain array", () => {
+    const ids = recommended.map((p) => p.id);
+    expect(ids).toEqual(["chart", "generative-ui"]);
+  });
+});
+
+describe("recommended.except", () => {
+  it("excludes a plugin by id", () => {
+    const result = recommended.except("generative-ui");
+    expect(result).toHaveLength(1);
+    expect(result[0]).toBe(chart);
+  });
+
+  it("excludes multiple plugins", () => {
+    const result = recommended.except("chart", "generative-ui");
+    expect(result).toHaveLength(0);
+  });
+
+  it("returns all plugins when no ids match", () => {
+    const result = recommended.except("nonexistent");
+    expect(result).toHaveLength(2);
+  });
+
+  it("returns all plugins when called with no arguments", () => {
+    const result = recommended.except();
+    expect(result).toHaveLength(2);
+  });
+
+  it("is assignable to Plugin[] for backwards compatibility", () => {
+    // Consumers may have typed variables as Plugin[]
+    const plugins: Plugin[] = recommended.except("chart");
+    expect(plugins).toHaveLength(1);
+
+    // Spread into a plain array
+    const spread: Plugin[] = [...recommended];
+    expect(spread).toHaveLength(2);
+
+    // Nullish coalescing fallback (common pattern in ElementsProvider)
+    const config: { plugins?: Plugin[] } = {};
+    const resolved = config.plugins ?? recommended;
+    expect(resolved).toHaveLength(2);
+  });
+
+  it("does not match on language when id is set", () => {
+    // generative-ui has id="generative-ui" and language="ui"
+    // Excluding by language value "ui" should not match because id takes precedence
+    const result = recommended.except("ui");
+    expect(result).toHaveLength(2);
+  });
+});

package/src/plugins/index.ts

@@ -2,7 +2,20 @@ import type { Plugin } from "@/types/plugins";
 import { chart } from "./chart";
 import { generativeUI } from "./generative-ui";
 
-export const recommended: Plugin[] = [chart, generativeUI];
+export type PluginList = Plugin[] & {
+  except(...ids: string[]): Plugin[];
+};
+
+function createPluginList(plugins: Plugin[]): PluginList {
+  const arr = [...plugins] as PluginList;
+  arr.except = (...ids: string[]) => {
+    const excluded = new Set(ids);
+    return arr.filter((p) => !excluded.has(p.id ?? p.language));
+  };
+  return arr;
+}
+
+export const recommended: PluginList = createPluginList([chart, generativeUI]);
 export { chart } from "./chart";
 export { generativeUI } from "./generative-ui";