fix(mastra): vendor sync partial-JSON parser, drop @ai-sdk/ui-utils

integrations/mastra/typescript/package.json

@@ -57,10 +57,9 @@
   },
   "dependencies": {
     "@ag-ui/a2ui-toolkit": "workspace:*",
-    "@ai-sdk/ui-utils": "^1.1.19",
     "fast-json-patch": "^3.1.1",
     "rxjs": "7.8.1",
-    "zod": "^3.25.76"
+    "zod": "^3.25.76 || ^4.0.0"
   },
   "peerDependencies": {
     "@ag-ui/core": ">=1.0.0",

integrations/mastra/typescript/src/__tests__/state-delta.test.ts

@@ -74,6 +74,43 @@ function streamingWorkingMemoryChunks(
   ];
 }
 
+// Like streamingWorkingMemoryChunks, but streams the given RAW working-memory
+// JSON text, so a test can send keys (such as `__proto__`) that a JS object
+// literal would not keep as own properties. The final `tool-call` carries
+// `finalMemory`, so only the streamed deltas go through parsePartialJson.
+function streamingRawWorkingMemoryChunks(
+  toolCallId: string,
+  memoryJson: string,
+  finalMemory: Record<string, any>,
+  numChunks = 8,
+) {
+  const argsText = JSON.stringify({ memory: memoryJson });
+  const size = Math.max(1, Math.ceil(argsText.length / numChunks));
+  const deltas: any[] = [];
+  for (let i = 0; i < argsText.length; i += size) {
+    deltas.push({
+      type: "tool-call-delta",
+      payload: { toolCallId, argsTextDelta: argsText.slice(i, i + size) },
+    });
+  }
+  return [
+    {
+      type: "tool-call-input-streaming-start",
+      payload: { toolCallId, toolName: "updateWorkingMemory" },
+    },
+    ...deltas,
+    { type: "tool-call-input-streaming-end", payload: { toolCallId } },
+    {
+      type: "tool-call",
+      payload: {
+        toolCallId,
+        toolName: "updateWorkingMemory",
+        args: { memory: finalMemory },
+      },
+    },
+  ];
+}
+
 // Reconstruct the client-side state after applying, in order, the seeded base
 // (input.state) followed by every STATE_DELTA / STATE_SNAPSHOT the run emitted.
 function reconstructState(
@@ -260,6 +297,54 @@ describe("Mastra working-memory updates -> AG-UI STATE_DELTA", () => {
         expect(types).toContain(EventType.RUN_FINISHED);
       });
 
+      // One delta puts the whole parsed update into the first snapshot; many
+      // deltas take the incremental merge path.
+      it.each([1, 12])(
+        "never streams a __proto__ key from working-memory args into shared state (%i deltas)",
+        async (numChunks) => {
+          const memoryJson =
+            '{"recipe":{"title":"Soup","__proto__":{"polluted":true}},' +
+            '"__proto__":{"polluted":true}}';
+          const agent = makeAgent({
+            streamChunks: [
+              ...streamingRawWorkingMemoryChunks(
+                "wm-1",
+                memoryJson,
+                { recipe: { title: "Soup" } },
+                numChunks,
+              ),
+              { type: "finish", payload: { finishReason: "stop" } },
+            ],
+          });
+
+          const events = await collectEvents(agent, makeInput());
+          const stateEvents = events.filter(
+            (e) =>
+              e.type === EventType.STATE_SNAPSHOT ||
+              e.type === EventType.STATE_DELTA,
+          );
+          expect(stateEvents.length).toBeGreaterThan(0);
+          expect(JSON.stringify(stateEvents)).not.toContain("__proto__");
+          expect(JSON.stringify(stateEvents)).not.toContain("polluted");
+          // A merged `__proto__` key would become the object's prototype, which
+          // JSON hides, so also check the emitted objects for inherited values.
+          for (const e of stateEvents) {
+            const payload =
+              e.type === EventType.STATE_SNAPSHOT
+                ? [(e as StateSnapshotEvent).snapshot]
+                : (e as StateDeltaEvent).delta.map((op: any) => op.value);
+            for (const obj of payload) {
+              if (obj && typeof obj === "object") {
+                expect((obj as any).polluted).toBeUndefined();
+                expect((obj as any).recipe?.polluted).toBeUndefined();
+              }
+            }
+          }
+          expect(({} as Record<string, unknown>).polluted).toBeUndefined();
+          expect(events.map((e) => e.type)).toContain(EventType.RUN_FINISHED);
+        },
+      );
+
       it("streams state (snapshot + deltas) BEFORE RUN_FINISHED, not only at run end", async () => {
         const agent = makeAgent({
           streamChunks: [

integrations/mastra/typescript/src/mastra.ts

@@ -35,7 +35,11 @@ import type { Agent as LocalMastraAgent } from "@mastra/core/agent";
 import { RequestContext } from "@mastra/core/request-context";
 import { randomUUID } from "@ag-ui/client";
 import jsonpatch from "fast-json-patch";
-import { parsePartialJson } from "@ai-sdk/ui-utils";
+// Vendored synchronous copy of @ai-sdk/ui-utils' parsePartialJson, as Mastra did in
+// https://github.com/mastra-ai/mastra/pull/27297. Fixes
+// https://github.com/ag-ui-protocol/ag-ui/issues/3096 and
+// https://github.com/ag-ui-protocol/ag-ui/issues/2874.
+import { parsePartialJson } from "./utils/ai-sdk-v4";
 import { Observable, type Subscriber } from "rxjs";
 import { MastraClient } from "@mastra/client-js";
 import {

integrations/mastra/typescript/src/utils/ai-sdk-v4/ai-sdk-v4.test.ts

@@ -0,0 +1,156 @@
+import { parsePartialJson } from ".";
+import { fixJson } from "./fix-json";
+
+// The vendored parser replaces @ai-sdk/ui-utils' parsePartialJson (#3096, #2874),
+// following Mastra's copy in https://github.com/mastra-ai/mastra/pull/27297.
+// The first case is Mastra's own test; the rest pin the shapes the working-memory
+// streaming bridge in mastra.ts feeds it.
+
+describe("parsePartialJson (vendored)", () => {
+  it("reports each parse state (Mastra's test case)", () => {
+    expect(parsePartialJson(undefined)).toEqual({
+      value: undefined,
+      state: "undefined-input",
+    });
+    expect(parsePartialJson('{"a":1}')).toEqual({
+      value: { a: 1 },
+      state: "successful-parse",
+    });
+    expect(parsePartialJson('{"a":[1,"tw')).toEqual({
+      value: { a: [1, "tw"] },
+      state: "repaired-parse",
+    });
+    expect(parsePartialJson("}{")).toEqual({
+      value: undefined,
+      state: "failed-parse",
+    });
+  });
+
+  it("is synchronous", () => {
+    expect(parsePartialJson('{"a":1}')).not.toBeInstanceOf(Promise);
+  });
+
+  it("repairs the streamed working-memory envelope and its nested JSON string", () => {
+    const outer = parsePartialJson(
+      '{"memory":"{\\"recipe\\":{\\"title\\":\\"Sp',
+    );
+    expect(outer.state).toBe("repaired-parse");
+    const memory = (outer.value as { memory: string }).memory;
+    expect(parsePartialJson(memory).value).toEqual({
+      recipe: { title: "Sp" },
+    });
+  });
+
+  it("repairs dangling keys, literals and numbers", () => {
+    expect(parsePartialJson('{"memory":').value).toEqual({});
+    expect(parsePartialJson('{"done":tr').value).toEqual({ done: true });
+    expect(parsePartialJson('{"a":1,"b":').value).toEqual({ a: 1 });
+  });
+});
+
+// @ai-sdk/ui-utils parsed through @ai-sdk/provider-utils' safeParseJSON, which
+// uses secure-json-parse with its defaults: a leading BOM is stripped, and any
+// `__proto__` key or `constructor` object with its own `prototype` is rejected.
+// The vendored parser keeps those rules; a rejection is a failed parse.
+describe("parsePartialJson (vendored) keeps secure-json-parse's rules", () => {
+  const failed = { value: undefined, state: "failed-parse" };
+
+  it("rejects a __proto__ key in complete input", () => {
+    expect(parsePartialJson('{"__proto__":{"polluted":true}}')).toEqual(failed);
+    expect(parsePartialJson('{"a":1,"__proto__":null}')).toEqual(failed);
+  });
+
+  it("rejects a __proto__ key in input that needs repair", () => {
+    expect(parsePartialJson('{"__proto__":{"polluted":tr')).toEqual(failed);
+    expect(parsePartialJson('{"a":"x","__proto__":{')).toEqual(failed);
+  });
+
+  it("rejects a nested __proto__ key", () => {
+    expect(
+      parsePartialJson('{"recipe":{"steps":[{"__proto__":{"x":1}}]}}'),
+    ).toEqual(failed);
+    expect(parsePartialJson('{"recipe":{"meta":{"__proto__":{"x":1')).toEqual(
+      failed,
+    );
+  });
+
+  it("rejects a Unicode-escaped __proto__ key", () => {
+    expect(
+      parsePartialJson('{"\\u005f\\u005fproto\\u005f\\u005f":{"x":1}}'),
+    ).toEqual(failed);
+    expect(parsePartialJson('{"\\u005f_proto__":{"x":1')).toEqual(failed);
+  });
+
+  it("rejects a constructor object with its own prototype", () => {
+    expect(parsePartialJson('{"constructor":{"prototype":{"x":1}}}')).toEqual(
+      failed,
+    );
+    expect(parsePartialJson('{"a":{"constructor":{"prototype":{"x":1')).toEqual(
+      failed,
+    );
+    expect(
+      parsePartialJson('{"\\u0063onstructor":{"prototyp\\u0065":null}}'),
+    ).toEqual(failed);
+  });
+
+  it("does not pollute Object.prototype", () => {
+    parsePartialJson('{"__proto__":{"polluted":true}}');
+    parsePartialJson('{"__proto__":{"polluted":tr');
+    expect(({} as Record<string, unknown>).polluted).toBeUndefined();
+  });
+
+  it("keeps ordinary constructor and prototype fields valid", () => {
+    expect(parsePartialJson('{"constructor":"Recipe"}')).toEqual({
+      value: { constructor: "Recipe" },
+      state: "successful-parse",
+    });
+    expect(parsePartialJson('{"constructor":{"name":"Recipe"}}')).toEqual({
+      value: { constructor: { name: "Recipe" } },
+      state: "successful-parse",
+    });
+    // secure-json-parse 2.7.0 failed this one by accident (hasOwnProperty on
+    // null throws a TypeError); a null constructor carries no prototype.
+    expect(parsePartialJson('{"constructor":null}')).toEqual({
+      value: { constructor: null },
+      state: "successful-parse",
+    });
+    expect(parsePartialJson('{"prototype":{"constructor":"x"}}')).toEqual({
+      value: { prototype: { constructor: "x" } },
+      state: "successful-parse",
+    });
+    expect(parsePartialJson('{"proto":1,"__proto":2,"proto__":3}')).toEqual({
+      value: { proto: 1, __proto: 2, proto__: 3 },
+      state: "successful-parse",
+    });
+    expect(parsePartialJson('{"constructor":{"name":"Rec')).toEqual({
+      value: { constructor: { name: "Rec" } },
+      state: "repaired-parse",
+    });
+  });
+
+  it("keeps __proto__ and constructor as string values valid", () => {
+    expect(parsePartialJson('{"key":"__proto__","t":"constructor"}')).toEqual({
+      value: { key: "__proto__", t: "constructor" },
+      state: "successful-parse",
+    });
+  });
+
+  it("strips a leading BOM", () => {
+    expect(parsePartialJson('\uFEFF{"a":1}')).toEqual({
+      value: { a: 1 },
+      state: "successful-parse",
+    });
+    expect(parsePartialJson('\uFEFF{"a":[1,"tw')).toEqual({
+      value: { a: [1, "tw"] },
+      state: "repaired-parse",
+    });
+  });
+});
+
+describe("fixJson (vendored)", () => {
+  it("closes open strings, arrays and objects", () => {
+    expect(JSON.parse(fixJson('{"a":"he'))).toEqual({ a: "he" });
+    expect(JSON.parse(fixJson('{"a":['))).toEqual({ a: [] });
+    expect(JSON.parse(fixJson('{"a":"x\\"y'))).toEqual({ a: 'x"y' });
+  });
+});