Merge pull request #1395 from ag-ui-protocol/perf/structured-clone

fix: excessive cloning in runSubscribersWithMutation

Author: jpr5

sdks/typescript/packages/client/src/agent/__tests__/subscriber.test.ts

@@ -1,5 +1,5 @@
 import { AbstractAgent } from "../agent";
-import { AgentSubscriber } from "../subscriber";
+import { AgentSubscriber, runSubscribersWithMutation } from "../subscriber";
 import {
   BaseEvent,
   EventType,
@@ -1261,6 +1261,220 @@ describe("AgentSubscriber", () => {
     });
   });
 
+  describe("runSubscribersWithMutation isolation contract", () => {
+    const runWith = (
+      subscribers: AgentSubscriber[],
+      messages: Message[] = [{ id: "orig", role: "user", content: "hi" }],
+      state: Record<string, any> = {},
+    ) =>
+      runSubscribersWithMutation(subscribers, messages, state, (subscriber, msgs, st) =>
+        subscriber.onEvent?.({
+          messages: msgs,
+          state: st,
+          agent: {} as any,
+          input: {} as any,
+          event: { type: EventType.RUN_STARTED } as any,
+        }),
+      );
+
+    it("should detect changes when subscriber returns a new messages array with same content", async () => {
+      const cloningSubscriber: AgentSubscriber = {
+        onEvent: ({ messages }) => ({
+          messages: [...messages],
+        }),
+      };
+
+      const result = await runWith([cloningSubscriber]);
+
+      // Reference equality means a new array is detected as a change
+      // (trade-off: may cause extra onMessagesChanged callbacks, but avoids JSON.stringify cost)
+      expect(result.messages).toBeDefined();
+    });
+
+    it("should detect changes when subscriber returns a new state object with same content", async () => {
+      const cloningSubscriber: AgentSubscriber = {
+        onEvent: ({ state }) => ({
+          state: { ...state },
+        }),
+      };
+
+      const result = await runWith([cloningSubscriber]);
+
+      expect(result.state).toBeDefined();
+    });
+
+    it("should not report changes when subscriber returns the same messages reference", async () => {
+      const passThroughSubscriber: AgentSubscriber = {
+        onEvent: ({ messages }) => ({
+          messages: messages as Message[],
+        }),
+      };
+
+      const result = await runWith([passThroughSubscriber]);
+
+      // Same reference returned — detected as no-op, no clone needed
+      expect(result.messages).toBeUndefined();
+    });
+
+    it("should not report changes when subscriber returns the same state reference", async () => {
+      const passThroughSubscriber: AgentSubscriber = {
+        onEvent: ({ state }) => ({
+          state,
+        }),
+      };
+
+      const result = await runWith([passThroughSubscriber]);
+
+      // Same reference returned — detected as no-op, no clone needed
+      expect(result.state).toBeUndefined();
+    });
+
+    it("should not report in-place mutations that are not returned (production mode)", async () => {
+      // Stub to production mode: no freeze, so the push succeeds, but the mutation is
+      // not communicated via return value and therefore not detected.
+      vi.stubEnv("NODE_ENV", "production");
+      vi.stubEnv("VITEST_WORKER_ID", "");
+
+      try {
+        const inPlaceMutator: AgentSubscriber = {
+          onEvent: ({ messages }) => {
+            // In production mode inputs are not frozen, so push succeeds silently.
+            // But since nothing is returned, the change is not propagated — silent data loss.
+            // This documents the contract: subscribers MUST return mutations, not mutate in-place.
+            (messages as Message[]).push({
+              id: "injected",
+              role: "assistant",
+              content: "injected",
+            });
+            // Return void — no mutation communicated
+          },
+        };
+
+        const result = await runWith([inPlaceMutator]);
+
+        // The in-place push is NOT reflected in the result because it wasn't returned
+        expect(result.messages).toBeUndefined();
+      } finally {
+        vi.unstubAllEnvs();
+      }
+    });
+
+    it("should re-throw TypeError when subscriber mutates frozen inputs in dev/test mode", async () => {
+      // In dev/test mode, inputs are deep-frozen. Any in-place mutation attempt
+      // throws a TypeError that runSubscribersWithMutation re-throws so violations
+      // are visible rather than silently swallowed.
+      const inPlaceMutator: AgentSubscriber = {
+        onEvent: ({ messages }) => {
+          (messages as Message[]).push({ id: "bad", role: "assistant", content: "bad" });
+        },
+      };
+
+      await expect(runWith([inPlaceMutator])).rejects.toThrow(TypeError);
+    });
+
+    it("should log a freeze violation and continue in development mode (non-test)", async () => {
+      // In dev mode (NODE_ENV=development), freeze violations are logged with a specific
+      // message instead of re-throwing — the violating subscriber is skipped, and remaining
+      // subscribers continue to execute.
+      vi.stubEnv("NODE_ENV", "development");
+      vi.stubEnv("VITEST_WORKER_ID", "");
+
+      const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
+
+      try {
+        const inPlaceMutator: AgentSubscriber = {
+          onEvent: ({ messages }) => {
+            (messages as Message[]).push({ id: "bad", role: "assistant", content: "bad" });
+          },
+        };
+
+        // Should resolve (not reject) because dev mode logs instead of re-throwing
+        const result = await runWith([inPlaceMutator]);
+
+        expect(result.messages).toBeUndefined();
+        expect(consoleErrorSpy).toHaveBeenCalledWith(
+          expect.stringContaining("AG-UI:"),
+          expect.any(TypeError),
+        );
+      } finally {
+        consoleErrorSpy.mockRestore();
+        vi.unstubAllEnvs();
+      }
+    });
+
+    it("should return unfrozen messages even after freeze was applied internally", async () => {
+      // Verifies that the internal deepFreeze does not leak frozen references to callers.
+      // The returned AgentStateMutation.messages must be mutable.
+      const mutatingSubscriber: AgentSubscriber = {
+        onEvent: () => ({
+          messages: [{ id: "new", role: "assistant", content: "new message" }],
+        }),
+      };
+
+      // Add a second subscriber that does nothing, so the messages clone gets frozen
+      // on the next iteration — this is the scenario that could return a frozen value.
+      const noopSubscriber: AgentSubscriber = {
+        onEvent: () => undefined,
+      };
+
+      const result = await runWith([mutatingSubscriber, noopSubscriber]);
+
+      expect(result.messages).toBeDefined();
+      // Must not be frozen — callers should be able to mutate the returned value
+      expect(Object.isFrozen(result.messages)).toBe(false);
+    });
+
+    it("should return unfrozen state even after freeze was applied internally", async () => {
+      // Same as the messages variant above, but for state. The state clone at the
+      // return point must also be unfrozen (symmetric code path, separately tested).
+      const mutatingSubscriber: AgentSubscriber = {
+        onEvent: () => ({
+          state: { key: "value" },
+        }),
+      };
+
+      // Second subscriber triggers the freeze-on-next-iteration path for state.
+      const noopSubscriber: AgentSubscriber = {
+        onEvent: () => undefined,
+      };
+
+      const result = await runWith([mutatingSubscriber, noopSubscriber]);
+
+      expect(result.state).toBeDefined();
+      expect(Object.isFrozen(result.state)).toBe(false);
+    });
+
+    it("should give subscriber B a clone of subscriber A's mutation output", async () => {
+      // Verifies multi-subscriber isolation: the new cloning-on-output strategy
+      // must still ensure each subsequent subscriber sees its own defensive copy,
+      // not a direct reference to the previous subscriber's returned array.
+      let subscriberBReceivedMessages: ReadonlyArray<Readonly<Message>> | undefined;
+      let subscriberAReturnedMessages: Message[] | undefined;
+
+      const subscriberA: AgentSubscriber = {
+        onEvent: () => {
+          subscriberAReturnedMessages = [{ id: "a-msg", role: "assistant", content: "from A" }];
+          return { messages: subscriberAReturnedMessages };
+        },
+      };
+
+      const subscriberB: AgentSubscriber = {
+        onEvent: ({ messages }) => {
+          subscriberBReceivedMessages = messages;
+        },
+      };
+
+      await runWith([subscriberA, subscriberB]);
+
+      // B should see A's messages content
+      expect(subscriberBReceivedMessages).toBeDefined();
+      expect(subscriberBReceivedMessages).toHaveLength(1);
+      expect(subscriberBReceivedMessages![0].id).toBe("a-msg");
+      // But NOT the same reference — B received a defensive clone
+      expect(subscriberBReceivedMessages).not.toBe(subscriberAReturnedMessages);
+    });
+  });
+
   describe("EmptyError Bug Reproduction", () => {
     test("should demonstrate EmptyError with STEP_STARTED/STEP_FINISHED events that cause no mutations", async () => {
       const emptyAgent = new TestAgent();

sdks/typescript/packages/client/src/agent/subscriber.ts

@@ -41,8 +41,11 @@ export interface AgentStateMutation {
 }
 
 export interface AgentSubscriberParams {
-  messages: Message[];
-  state: State;
+  messages: ReadonlyArray<Readonly<Message>>;
+  // NOTE: State resolves to `any` at the type level (z.infer<typeof z.any()>), so Readonly<State>
+  // provides no compile-time mutation protection. Runtime enforcement via deepFreeze in
+  // dev/test mode is the only guard against in-place mutation of state.
+  state: Readonly<State>;
   agent: AbstractAgent;
   input: RunAgentInput;
 }
@@ -205,41 +208,66 @@ export interface AgentSubscriber {
   ): MaybePromise<void>;
 }
 
+function deepFreeze<T>(obj: T): T {
+  Object.freeze(obj);
+  if (obj !== null && typeof obj === "object") {
+    for (const value of Object.values(obj)) {
+      if (value !== null && typeof value === "object" && !Object.isFrozen(value)) {
+        deepFreeze(value);
+      }
+    }
+  }
+  return obj;
+}
+
 export async function runSubscribersWithMutation(
   subscribers: AgentSubscriber[],
   initialMessages: Message[],
   initialState: State,
   executor: (
     subscriber: AgentSubscriber,
-    messages: Message[],
-    state: State,
+    messages: ReadonlyArray<Readonly<Message>>,
+    state: Readonly<State>,
   ) => MaybePromise<AgentStateMutation | void>,
 ): Promise<AgentStateMutation> {
-  let messages: Message[] = initialMessages;
-  let state: State = initialState;
+  const isTestEnvironment =
+    process.env.NODE_ENV === "test" || Boolean(process.env.VITEST_WORKER_ID);
+  const isDev =
+    process.env.NODE_ENV === "development" ||
+    process.env.NODE_ENV === "test" ||
+    Boolean(process.env.VITEST_WORKER_ID);
+  const baselineMessages = structuredClone_(initialMessages);
+  const baselineState = structuredClone_(initialState);
+  let messages: Message[] = baselineMessages;
+  let state: State = baselineState;
 
   let stopPropagation: boolean | undefined = undefined;
 
   for (const subscriber of subscribers) {
     try {
-      const mutation = await executor(
-        subscriber,
-        structuredClone_(messages),
-        structuredClone_(state),
-      );
+      // Subscribers receive shared references and must not mutate them in-place.
+      // Mutations should only be communicated via the return value.
+      // In dev/test mode only: deep-freeze inputs so accidental in-place mutations surface
+      // as TypeErrors immediately. In production, enforcement is type-level only.
+      if (isDev) {
+        deepFreeze(messages);
+        deepFreeze(state);
+      }
+      const mutation = await executor(subscriber, messages, state);
 
       if (mutation === undefined) {
         // Nothing returned – keep going
         continue;
       }
 
-      // Merge messages/state so next subscriber sees latest view
-      if (mutation.messages !== undefined) {
-        messages = mutation.messages;
+      // Replace with a defensive copy of the subscriber's mutation,
+      // but skip if the subscriber returned the same reference (no-op).
+      if (mutation.messages !== undefined && mutation.messages !== messages) {
+        messages = structuredClone_(mutation.messages);
       }
 
-      if (mutation.state !== undefined) {
-        state = mutation.state;
+      if (mutation.state !== undefined && mutation.state !== state) {
+        state = structuredClone_(mutation.state);
       }
 
       stopPropagation = mutation.stopPropagation;
@@ -248,21 +276,37 @@ export async function runSubscribersWithMutation(
         break;
       }
     } catch (error) {
-      // Log subscriber errors but continue processing (silence during tests)
-      const isTestEnvironment =
-        process.env.NODE_ENV === "test" || process.env.VITEST_WORKER_ID !== undefined;
-
-      if (!isTestEnvironment) {
+      if (isDev && error instanceof TypeError) {
+        // Likely a freeze violation: subscriber attempted to mutate frozen inputs in-place.
+        // In test environments, re-throw so tests fail fast and the violation is visible.
+        // In development (non-test), log a specific message to distinguish freeze violations
+        // from ordinary subscriber errors.
+        if (isTestEnvironment) {
+          throw error;
+        }
+        console.error(
+          "AG-UI: Subscriber attempted to mutate frozen inputs in-place. " +
+            "Return mutations via AgentStateMutation instead of mutating directly.",
+          error,
+        );
+      } else if (!isTestEnvironment) {
         console.error("Subscriber error:", error);
       }
-      // Continue to next subscriber unless we want to stop propagation
+      // Skip this subscriber's mutation and continue
       continue;
     }
   }
 
+  // In dev/test mode, the canonical messages/state references may have been
+  // frozen in-place (for subscriber mutation detection). Clone them before
+  // returning so callers receive a mutable copy, not a frozen one.
   return {
-    ...(JSON.stringify(messages) !== JSON.stringify(initialMessages) ? { messages } : {}),
-    ...(JSON.stringify(state) !== JSON.stringify(initialState) ? { state } : {}),
+    ...(messages !== baselineMessages
+      ? { messages: isDev && Object.isFrozen(messages) ? structuredClone_(messages) : messages }
+      : {}),
+    ...(state !== baselineState
+      ? { state: isDev && Object.isFrozen(state) ? structuredClone_(state) : state }
+      : {}),
     ...(stopPropagation !== undefined ? { stopPropagation } : {}),
   };
 }

sdks/typescript/packages/client/src/utils.ts

@@ -8,6 +8,10 @@ export const structuredClone_ = <T>(obj: T): T => {
   try {
     return JSON.parse(JSON.stringify(obj));
   } catch (err) {
+    // Preserve array vs object type in shallow fallback
+    if (Array.isArray(obj)) {
+      return [...obj] as unknown as T;
+    }
     return { ...obj } as T;
   }
 };
@@ -23,7 +27,6 @@ export function randomUUID(): string {
 // Note: semver helpers were removed in favor of using
 // the external `compare-versions` library directly at call sites.
 
-
 /**
  * Parses a semantic version string into its numeric components.
  * Supports incomplete versions (e.g. "1", "1.2") by defaulting missing segments to zero.