fix(mastra): support Zod 4 by dropping @ai-sdk/ui-utils

integrations/mastra/typescript/package.json

@@ -57,10 +57,10 @@
   },
   "dependencies": {
     "@ag-ui/a2ui-toolkit": "workspace:*",
-    "@ai-sdk/ui-utils": "^1.1.19",
+    "ai": "^5.0.207",
     "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/mastra.ts

@@ -31,7 +31,10 @@ 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";
+// `parsePartialJson` comes from the current AI SDK (`ai`, zod 3 || 4) instead of
+// `@ai-sdk/ui-utils`, whose zod-3-only peer forced a nested Zod 3 copy into
+// apps on Zod 4. Fixes https://github.com/ag-ui-protocol/ag-ui/issues/2874.
+import { parsePartialJson } from "ai";
 import { Observable, type Subscriber } from "rxjs";
 import { MastraClient } from "@mastra/client-js";
 import {
@@ -237,11 +240,11 @@ function parseWorkingMemoryUpdate(
  * incremental STATE_DELTA as the model writes the update; `undefined` while the
  * text is not yet a usable object (or is markdown-template working memory).
  */
-function parseStreamingWorkingMemoryUpdate(
+async function parseStreamingWorkingMemoryUpdate(
   argsText: string,
-): Record<string, any> | undefined {
+): Promise<Record<string, any> | undefined> {
   if (!argsText) return undefined;
-  const outer = parsePartialJson(argsText);
+  const outer = await parsePartialJson(argsText);
   const outerVal = outer.value;
   if (!outerVal || typeof outerVal !== "object" || Array.isArray(outerVal)) {
     return undefined;
@@ -252,7 +255,7 @@ function parseStreamingWorkingMemoryUpdate(
     return mem as Record<string, any>;
   }
   if (typeof mem === "string") {
-    const inner = parsePartialJson(mem);
+    const inner = await parsePartialJson(mem);
     const innerVal = inner.value;
     if (innerVal && typeof innerVal === "object" && !Array.isArray(innerVal)) {
       return innerVal as Record<string, any>;
@@ -1042,7 +1045,26 @@ export class MastraAgent extends AbstractAgent {
               }),
             );
           }
-          subscriber.complete();
+          // Complete on the NEXT macrotask, not inside the abort event.
+          //
+          // The chunk pipeline is async (`parsePartialJson` from `ai` is async,
+          // see https://github.com/ag-ui-protocol/ag-ui/issues/2874), so a remote
+          // producer aborted mid-stream is usually suspended in
+          // `await onChunk(...)`. Completing synchronously settled the run
+          // before that producer resumed and saw its aborted signal, so the
+          // caller observed a finished run while upstream had not stopped yet
+          // (#2288). One macrotask lets the pending microtasks run first, the
+          // producer's resumption included.
+          //
+          // Only complete() is deferred. The cancelled RUN_FINISHED above is
+          // still emitted synchronously, so the run's outcome is decided here
+          // and cannot be overtaken by the stream draining during that tick.
+          // In production the per-run handle also aborts the fetch, which tears
+          // the stream down regardless; this keeps the "producer stopped before
+          // the Observable settles" ordering the #2288 tests assert.
+          setTimeout(() => {
+            if (!subscriber.closed) subscriber.complete();
+          }, 0);
         },
         { once: true },
       );
@@ -1389,18 +1411,20 @@ export class MastraAgent extends AbstractAgent {
                       stopped = true;
                       return;
                     }
-                    if (handleChunk(chunk)) stopped = true;
+                    if (await handleChunk(chunk)) stopped = true;
                   },
                 });
               } catch (error) {
                 // A resumed call already streamed, and any output held behind
                 // it, must reach the client before the failure.
-                if (!abortController.signal.aborted) releaseDeferredReplay();
+                if (!abortController.signal.aborted) {
+                  await releaseDeferredReplay();
+                }
                 throw error;
               }
 
               if (!stopped) {
-                flush();
+                await flush();
                 await finishRun(async () => ({
                   traceId: await this.resolveTraceId(response),
                   usage: await this.resolveUsage(response, getUsage()),
@@ -2438,7 +2462,7 @@ export class MastraAgent extends AbstractAgent {
 
     // Emit the held call, its result, then the held output. Returns true when
     // a replayed chunk ended the stream, like handleChunk.
-    const releaseDeferredReplay = (): boolean => {
+    const releaseDeferredReplay = async (): Promise<boolean> => {
       if (!deferredReplay) return false;
       const { toolCallId, toolName, args, result, modelOutput } =
         deferredReplay;
@@ -2447,7 +2471,11 @@ export class MastraAgent extends AbstractAgent {
       callbacks.onToolResultPart?.({ toolCallId, result, modelOutput });
       const chunks = deferredChunks;
       deferredChunks = [];
-      return chunks.some((chunk) => processChunk(chunk));
+      // Sequential with short-circuit, matching the former `.some()`.
+      for (const chunk of chunks) {
+        if (await processChunk(chunk)) return true;
+      }
+      return false;
     };
 
     // Hold output chunks behind the call until one announces the id, then
@@ -2839,16 +2867,18 @@ export class MastraAgent extends AbstractAgent {
       }
     };
 
-    const handleChunk = (chunk: any): boolean => {
+    const handleChunk = async (chunk: any): Promise<boolean> => {
       if (deferredReplay) {
         const settled = settleDeferredReplay(chunk);
         if (settled === "hold") return false;
-        if (settled === "release" && releaseDeferredReplay()) return true;
+        if (settled === "release" && (await releaseDeferredReplay())) {
+          return true;
+        }
       }
       return processChunk(chunk);
     };
 
-    const processChunk = (chunk: any): boolean => {
+    const processChunk = async (chunk: any): Promise<boolean> => {
       // Observational Memory data parts arrive on fullStream as
       // `{ type: "data-om-*", data: {...} }` (no `payload`). Handle them before
       // the payload guard below so they map to activity when surfacing is on,
@@ -2988,7 +3018,9 @@ export class MastraAgent extends AbstractAgent {
             if (argsTextDelta != null) {
               workingMemoryStream.argsText += argsTextDelta;
               emitWorkingMemoryState(
-                parseStreamingWorkingMemoryUpdate(workingMemoryStream.argsText),
+                await parseStreamingWorkingMemoryUpdate(
+                  workingMemoryStream.argsText,
+                ),
               );
             }
             break;
@@ -3474,9 +3506,9 @@ export class MastraAgent extends AbstractAgent {
     return {
       handleChunk,
       releaseDeferredReplay,
-      flush: () => {
+      flush: async () => {
         // Still held at the end means no id came: release under the run's id.
-        releaseDeferredReplay();
+        await releaseDeferredReplay();
         flush();
         if (pendingRetryReason !== undefined) {
           const reason = pendingRetryReason;
@@ -3524,15 +3556,15 @@ export class MastraAgent extends AbstractAgent {
         // has ended a cancelled run, and an error reported itself through
         // onError.
         if (abortSignal.aborted) return "cancelled";
-        if (handleChunk(chunk)) return "error";
+        if (await handleChunk(chunk)) return "error";
       }
     } catch (error) {
       // A resumed call already streamed, and any output held behind it, must
       // reach the client before the failure.
-      if (!abortSignal.aborted) releaseDeferredReplay();
+      if (!abortSignal.aborted) await releaseDeferredReplay();
       throw error;
     }
-    flush();
+    await flush();
     return "completed";
   }
 
@@ -4131,11 +4163,11 @@ export class MastraAgent extends AbstractAgent {
                 stopped = true;
                 return;
               }
-              if (handleChunk(chunk)) stopped = true;
+              if (await handleChunk(chunk)) stopped = true;
             },
           });
           if (!stopped) {
-            flush();
+            await flush();
             await onRunFinished?.(async () => ({
               traceId: await this.resolveTraceId(response),
               usage: await this.resolveUsage(response, getUsage()),