Fix token usage display and run cancellation

Author: Wistomize

apps/web/src/app/data-tasks/__tests__/live-run-state.test.ts

@@ -1259,6 +1259,29 @@ describe("live run state reducer", () => {
     expect(session.tokens.outputTokens).toBe(340);
   });
 
+  it("does not double count duplicate token_usage custom events", () => {
+    let run = createInitialLiveRun();
+    const event = {
+      type: "CUSTOM" as const,
+      name: "token_usage",
+      value: {
+        step_number: 1,
+        model: "qwen-plus",
+        input_tokens: 1200,
+        output_tokens: 340,
+      },
+    };
+
+    run = reduceLiveRunEvent(run, { type: "RUN_STARTED" });
+    run = reduceLiveRunEvent(run, event);
+    run = reduceLiveRunEvent(run, event);
+
+    const runUsage = deriveRunUsage(run);
+    expect(runUsage.tokens.inputTokens).toBe(1200);
+    expect(runUsage.tokens.outputTokens).toBe(340);
+    expect(run.tokenUsageEvents).toHaveLength(1);
+  });
+
   it("tracks token_usage model, cost, and step-level usage", () => {
     let run = createInitialLiveRun();
     run = reduceLiveRunEvent(run, { type: "RUN_STARTED" });

apps/web/src/app/data-tasks/__tests__/mastra-stream-hooks.test.ts

@@ -4,6 +4,10 @@ import {
   createMastraStreamNormalizerHooks,
   tokenUsageEventFromChunk,
 } from "../../../../../../packages/agent-runtime/src/stream/mastra-stream-hooks";
+import {
+  normalizeMastraFullStream,
+  wrapAgentForAgUi,
+} from "../../../../../../packages/agent-runtime/src/stream/mastra-stream-normalizer";
 
 describe("tokenUsageEventFromChunk", () => {
   it("extracts usage, tool_call_id, and tool_name from step output", () => {
@@ -78,3 +82,95 @@ describe("createMastraStreamNormalizerHooks token emission", () => {
     });
   });
 });
+
+describe("normalizeMastraFullStream", () => {
+  it("drops Mastra internal chunks before AG-UI processing", async () => {
+    const chunks = normalizeMastraFullStream(streamFrom([
+      { type: "abort", payload: { reason: "cancelled" } },
+      { type: "text-start", payload: { id: "msg-1" } },
+      { type: "finish", payload: {} },
+    ]));
+
+    await expect(collectStream(chunks)).resolves.toEqual([{ type: "finish", payload: {} }]);
+  });
+});
+
+describe("wrapAgentForAgUi abort signal passthrough", () => {
+  it("passes the run abort signal into stream options", async () => {
+    const runAbortController = new AbortController();
+    let receivedSignal: AbortSignal | undefined;
+    const agent = wrapAgentForAgUi(
+      {
+        async stream(_messages: unknown, options?: { abortSignal?: AbortSignal; marker?: string }) {
+          receivedSignal = options?.abortSignal;
+          return { fullStream: emptyStream() };
+        },
+      },
+      {},
+      { abortSignal: runAbortController.signal },
+    );
+
+    await agent.stream([], { marker: "kept" });
+
+    expect(receivedSignal).toBe(runAbortController.signal);
+  });
+
+  it("merges an existing stream abort signal with the run abort signal", async () => {
+    const upstreamAbortController = new AbortController();
+    const runAbortController = new AbortController();
+    let receivedSignal: AbortSignal | undefined;
+    const agent = wrapAgentForAgUi(
+      {
+        async stream(_messages: unknown, options?: { abortSignal?: AbortSignal }) {
+          receivedSignal = options?.abortSignal;
+          return { fullStream: emptyStream() };
+        },
+      },
+      {},
+      { abortSignal: runAbortController.signal },
+    );
+
+    await agent.stream([], { abortSignal: upstreamAbortController.signal });
+    runAbortController.abort(new Error("run cancelled"));
+
+    expect(receivedSignal).not.toBe(upstreamAbortController.signal);
+    expect(receivedSignal?.aborted).toBe(true);
+  });
+
+  it("passes the run abort signal into resumeStream options", async () => {
+    const runAbortController = new AbortController();
+    let receivedSignal: AbortSignal | undefined;
+    const agent = wrapAgentForAgUi(
+      {
+        async resumeStream(_resumeData: unknown, options?: { abortSignal?: AbortSignal }) {
+          receivedSignal = options?.abortSignal;
+          return { fullStream: emptyStream() };
+        },
+      },
+      {},
+      { abortSignal: runAbortController.signal },
+    );
+
+    await agent.resumeStream({ answer: "yes" }, {});
+
+    expect(receivedSignal).toBe(runAbortController.signal);
+  });
+});
+
+async function* emptyStream() {
+  yield { type: "finish", payload: {} };
+}
+
+async function* streamFrom(chunks: Array<{ payload?: unknown; type: string }>) {
+  for (const chunk of chunks) {
+    yield chunk;
+  }
+}
+
+async function collectStream<T>(stream: AsyncIterable<T>): Promise<T[]> {
+  const chunks: T[] = [];
+  for await (const chunk of stream) {
+    chunks.push(chunk);
+  }
+  return chunks;
+}

apps/web/src/app/data-tasks/live-run-state.ts

@@ -919,6 +919,10 @@ function reduceCustomEvent(state: LiveRun, event: AgUiLikeEvent): LiveRun {
         : {}),
       ...(stringValue(value?.model) ? { model: stringValue(value?.model) } : {}),
     };
+    const deltaKey = tokenUsageRecordKey(delta);
+    if (state.tokenUsageEvents.some((record) => tokenUsageRecordKey(record) === deltaKey)) {
+      return state;
+    }
     return {
       ...state,
       tokenUsage: mergeTokenUsageStats(state.tokenUsage ?? emptyTokenUsageStats(), delta),

apps/web/src/app/data-tasks/use-data-foundry-run.tsx

@@ -153,7 +153,7 @@ export function LiveRunProvider({ children }: { children: ReactNode }) {
     [isRestoringConversation],
   );
 
-  useEffect(() => {
+  useLayoutEffect(() => {
     const previous = prevRunStatusRef.current;
     const current = liveRun.runStatus;
     if (previous === "running" && current === "completed") {

packages/agent-runtime/src/index.ts

@@ -112,6 +112,7 @@ export {
 export {
   normalizeMastraFullStream,
   wrapAgentForAgUi,
+  type MastraAgentForAgUiOptions,
   type MastraStreamChunk,
   type MastraStreamNormalizerHooks
 } from "./stream/mastra-stream-normalizer.js";
@@ -441,6 +442,7 @@ export const createDataFoundry = async (
   const agentForAgUi = wrapAgentForAgUi(
     agent,
     createMastraStreamNormalizerHooks(input.emitter),
+    { ...(input.abortSignal ? { abortSignal: input.abortSignal } : {}) },
   );
   const mastra = input.taskStateRuntime
     ? new Mastra({

packages/agent-runtime/src/stream/mastra-stream-normalizer.ts

@@ -13,7 +13,12 @@ export type MastraStreamNormalizerHooks = {
   onQuarantine?: (chunk: MastraStreamChunk) => void;
 };
 
-const PAYLOADLESS_CHUNK_TYPES = new Set([
+export type MastraAgentForAgUiOptions = {
+  abortSignal?: AbortSignal | undefined;
+};
+
+const INTERNAL_CHUNK_TYPES = new Set([
+  "abort",
   "text-start",
   "text-end",
   "tool-call-input-streaming-start",
@@ -48,7 +53,7 @@ export async function* normalizeMastraFullStream(
       continue;
     }
 
-    if (type && PAYLOADLESS_CHUNK_TYPES.has(type) && chunk.payload === undefined) {
+    if (type && INTERNAL_CHUNK_TYPES.has(type)) {
       continue;
     }
 
@@ -85,15 +90,16 @@ export async function* normalizeMastraFullStream(
 /** Wrap a local Mastra Agent so only fullStream is normalized for AG-UI consumption. */
 export function wrapAgentForAgUi<TAgent extends object>(
   agent: TAgent,
-  hooks: MastraStreamNormalizerHooks = {}
+  hooks: MastraStreamNormalizerHooks = {},
+  options: MastraAgentForAgUiOptions = {}
 ): TAgent {
   return new Proxy(agent, {
     get(target, prop, receiver) {
       const value = Reflect.get(target, prop, receiver);
 
-      if (prop === "stream" && typeof value === "function") {
+      if ((prop === "stream" || prop === "resumeStream") && typeof value === "function") {
         return async (...args: unknown[]) => {
-          const response = await value.apply(target, args);
+          const response = await value.apply(target, withAbortSignal(args, options.abortSignal));
           if (!response || typeof response !== "object") {
             return response;
           }
@@ -122,6 +128,45 @@ export function wrapAgentForAgUi<TAgent extends object>(
   }) as TAgent;
 }
 
+const withAbortSignal = (args: unknown[], abortSignal?: AbortSignal | undefined): unknown[] => {
+  if (!abortSignal) {
+    return args;
+  }
+  const nextArgs = [...args];
+  const streamOptions = isRecord(nextArgs[1]) ? nextArgs[1] : {};
+  nextArgs[1] = {
+    ...streamOptions,
+    abortSignal: mergeAbortSignals(streamOptions.abortSignal, abortSignal)
+  };
+  return nextArgs;
+};
+
+const mergeAbortSignals = (
+  first: unknown,
+  second: AbortSignal
+): AbortSignal => {
+  if (!(first instanceof AbortSignal)) {
+    return second;
+  }
+  if (first === second) {
+    return second;
+  }
+  if (typeof AbortSignal.any === "function") {
+    return AbortSignal.any([first, second]);
+  }
+  if (first.aborted) {
+    return AbortSignal.abort(first.reason);
+  }
+  if (second.aborted) {
+    return AbortSignal.abort(second.reason);
+  }
+  const controller = new AbortController();
+  const abort = (signal: AbortSignal): void => controller.abort(signal.reason);
+  first.addEventListener("abort", () => abort(first), { once: true });
+  second.addEventListener("abort", () => abort(second), { once: true });
+  return controller.signal;
+};
+
 const stringifyStreamError = (error: unknown): string => {
   if (error instanceof Error) {
     return error.message;