@gram-ai/elements 1.28.0 → 1.29.0

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/types/index.ts

@@ -268,6 +268,13 @@ export interface ElementsConfig {
    */
   tools?: ToolsConfig;
 
+  /**
+   * Configuration for automatic conversation compaction when the estimated
+   * input size approaches the model's context window. Defaults are safe for
+   * all models; override per-page to tighten or disable.
+   */
+  contextCompaction?: ContextCompactionConfig;
+
   /**
    * Configuration for chat history and thread persistence.
    * When enabled, conversations are saved and the thread list is shown.
@@ -690,6 +697,54 @@ export interface ToolsConfig {
    * }
    */
   toolsToInclude?: ToolsFilter;
+
+  /**
+   * Maximum UTF-8 byte size for any single tool call's result. Results larger
+   * than this are truncated with a head+tail preserving strategy and a notice
+   * suffix before being added to the conversation. Prevents one greedy tool
+   * call (e.g. a wide log search) from filling the model's context window.
+   *
+   * Omit or set to 0 to disable.
+   *
+   * @example
+   * tools: {
+   *   maxOutputBytes: 50_000, // ~12.5K tokens per tool call
+   * }
+   */
+  maxOutputBytes?: number;
+}
+
+/**
+ * Configuration for automatic compaction of older conversation turns when the
+ * estimated input size approaches the model's context window. Prevents
+ * upstream 400 "prompt too long" errors without losing the system prompt or
+ * the most recent turns.
+ */
+export interface ContextCompactionConfig {
+  /**
+   * Hard ceiling (in tokens) for the outbound request. Overrides the built-in
+   * per-model map. Use this when you know your upstream provider enforces a
+   * smaller limit than the model's nominal maximum.
+   */
+  maxTokens?: number;
+
+  /**
+   * Fraction of the model ceiling at which compaction kicks in. Defaults to
+   * 0.7 — leaves room for the assistant's response and some slack for the
+   * chars/4 token heuristic's error.
+   */
+  compactAtFraction?: number;
+
+  /**
+   * Number of most-recent messages preserved verbatim during compaction.
+   * Defaults to 4 (covers the current turn + its immediate predecessor).
+   */
+  keepRecent?: number;
+
+  /**
+   * Disable compaction entirely. Useful in tests and for opting out per-page.
+   */
+  disabled?: boolean;
 }
 
 export interface WelcomeConfig {