@assistant-ui/ai-sdk 0.0.3 → 0.0.5

package/dist/usage.js

@@ -1,5 +1,6 @@
-import { useAuiState } from "@assistant-ui/store";
 import { c } from "@assistant-ui/tap/react-shim/compiler-runtime";
+import "@assistant-ui/tap/react-shim";
+import { useAuiState } from "@assistant-ui/store";
 //#region src/usage.ts
 const USAGE_KEYS = [
 	"inputTokens",
@@ -93,13 +94,10 @@ function getThreadMessageTokenUsage(message) {
 	if (!metadata) return void 0;
 	const topLevelUsage = normalizeUsage(metadata.usage);
 	if (topLevelUsage) return withComputedTotal(topLevelUsage);
-	const legacyUsage = normalizeUsage(asRecord(metadata.custom)?.usage);
-	if (legacyUsage) return withComputedTotal(legacyUsage);
+	const customUsage = normalizeUsage(asRecord(metadata.custom)?.usage);
+	if (customUsage) return withComputedTotal(customUsage);
 	return usageFromSteps(metadata.steps);
 }
-function getLatestThreadTokenUsage(messages) {
-	return getThreadMessageTokenUsage(findLatestMessageWithUsage(messages));
-}
 function findLatestMessageWithUsage(messages) {
 	if (!messages) return void 0;
 	for (let idx = messages.length - 1; idx >= 0; idx -= 1) {
@@ -107,6 +105,14 @@ function findLatestMessageWithUsage(messages) {
 		if (getThreadMessageTokenUsage(message)) return message;
 	}
 }
+/**
+* Reads token usage from the newest assistant message that reports any.
+*
+* A route attaches usage through the AI SDK's `messageMetadata` option. Because
+* a thread message carries a fixed metadata shape, the converter moves every
+* other key the route returns into `metadata.custom`, which is where this hook
+* looks.
+*/
 function useThreadTokenUsage() {
 	const $ = c(2);
 	const msg = useAuiState(_temp);
@@ -122,6 +128,6 @@ function _temp(s) {
 	return findLatestMessageWithUsage(s.thread.messages);
 }
 //#endregion
-export { getLatestThreadTokenUsage, getThreadMessageTokenUsage, useThreadTokenUsage };
+export { getThreadMessageTokenUsage, useThreadTokenUsage };
 
 //# sourceMappingURL=usage.js.map

package/dist/usage.js.map

@@ -1 +1 @@
-{"version":3,"file":"usage.js","names":["useAuiState","ThreadTokenUsage","totalTokens","inputTokens","outputTokens","reasoningTokens","cachedInputTokens","TokenUsageExtractableMessage","role","metadata","UsageRecord","Record","USAGE_KEYS","const","asRecord","value","Array","isArray","undefined","asPositiveTokenCount","Number","isFinite","computeTotalTokens","usage","normalizeUsage","record","result","hasFields","key","count","outputTokenDetails","inputTokenDetails","cacheReadTokens","withComputedTotal","usageFromSteps","steps","sums","present","stepsWithUsage","stepsWithComputableTotal","step","stepTotal","getThreadMessageTokenUsage","message","topLevelUsage","legacyUsage","custom","getLatestThreadTokenUsage","messages","findLatestMessageWithUsage","idx","length","useThreadTokenUsage","$","_c","msg","_temp","t0","s","thread"],"sources":["../src/usage.ts"],"sourcesContent":["/// <reference types=\"@assistant-ui/core/react\" />\nimport { useAuiState } from \"@assistant-ui/store\";\n\nexport type ThreadTokenUsage = {\n  totalTokens?: number;\n  inputTokens?: number;\n  outputTokens?: number;\n  reasoningTokens?: number;\n  cachedInputTokens?: number;\n};\n\nexport interface TokenUsageExtractableMessage {\n  role?: string;\n  metadata?: unknown;\n}\n\ntype UsageRecord = Record<string, unknown>;\n\nconst USAGE_KEYS = [\n  \"inputTokens\",\n  \"outputTokens\",\n  \"reasoningTokens\",\n  \"cachedInputTokens\",\n  \"totalTokens\",\n] as const satisfies (keyof ThreadTokenUsage)[];\n\nfunction asRecord(value: unknown): UsageRecord | undefined {\n  if (!value || typeof value !== \"object\" || Array.isArray(value))\n    return undefined;\n  return value as UsageRecord;\n}\n\nfunction asPositiveTokenCount(value: unknown): number | undefined {\n  if (typeof value !== \"number\" || !Number.isFinite(value) || value < 0) {\n    return undefined;\n  }\n  return value;\n}\n\nfunction computeTotalTokens(usage: ThreadTokenUsage): number | undefined {\n  if (usage.totalTokens !== undefined) return usage.totalTokens;\n  if (usage.inputTokens !== undefined && usage.outputTokens !== undefined) {\n    return usage.inputTokens + usage.outputTokens;\n  }\n  return undefined;\n}\n\nfunction normalizeUsage(value: unknown): ThreadTokenUsage | undefined {\n  const record = asRecord(value);\n  if (!record) return undefined;\n\n  const result: ThreadTokenUsage = {};\n  let hasFields = false;\n  for (const key of USAGE_KEYS) {\n    const count = asPositiveTokenCount(record[key]);\n    if (count !== undefined) {\n      result[key] = count;\n      hasFields = true;\n    }\n  }\n  // AI SDK v7 moved these under token detail objects; v6 kept them top-level.\n  if (result.reasoningTokens === undefined) {\n    const count = asPositiveTokenCount(\n      asRecord(record.outputTokenDetails)?.reasoningTokens,\n    );\n    if (count !== undefined) {\n      result.reasoningTokens = count;\n      hasFields = true;\n    }\n  }\n  if (result.cachedInputTokens === undefined) {\n    const count = asPositiveTokenCount(\n      asRecord(record.inputTokenDetails)?.cacheReadTokens,\n    );\n    if (count !== undefined) {\n      result.cachedInputTokens = count;\n      hasFields = true;\n    }\n  }\n  return hasFields ? result : undefined;\n}\n\nfunction withComputedTotal(\n  usage: ThreadTokenUsage,\n): ThreadTokenUsage | undefined {\n  const totalTokens = computeTotalTokens(usage);\n  return { ...usage, ...(totalTokens !== undefined && { totalTokens }) };\n}\n\nfunction usageFromSteps(value: unknown): ThreadTokenUsage | undefined {\n  const steps = Array.isArray(value) ? value : [];\n\n  const sums: Record<string, number> = {};\n  const present: Record<string, boolean> = {};\n  let stepsWithUsage = 0;\n  let stepsWithComputableTotal = 0;\n\n  for (const step of steps) {\n    const usage = normalizeUsage(asRecord(step)?.usage);\n    if (!usage) continue;\n    stepsWithUsage++;\n\n    const stepTotal = computeTotalTokens(usage);\n    if (stepTotal !== undefined) {\n      sums.totalTokens = (sums.totalTokens ?? 0) + stepTotal;\n      stepsWithComputableTotal++;\n    }\n\n    for (const key of USAGE_KEYS) {\n      if (key === \"totalTokens\") continue;\n      if (usage[key] !== undefined) {\n        sums[key] = (sums[key] ?? 0) + usage[key];\n        present[key] = true;\n      }\n    }\n  }\n\n  if (stepsWithUsage === 0) return undefined;\n\n  const result: ThreadTokenUsage = {};\n  if (stepsWithComputableTotal === stepsWithUsage) {\n    result.totalTokens = sums.totalTokens!;\n  }\n  for (const key of USAGE_KEYS) {\n    if (key === \"totalTokens\") continue;\n    if (present[key]) {\n      result[key] = sums[key]!;\n    }\n  }\n  return result;\n}\n\nexport function getThreadMessageTokenUsage(\n  message: TokenUsageExtractableMessage | undefined,\n): ThreadTokenUsage | undefined {\n  if (!message || message.role !== \"assistant\") return undefined;\n\n  const metadata = asRecord(message.metadata);\n  if (!metadata) return undefined;\n\n  const topLevelUsage = normalizeUsage(metadata.usage);\n  if (topLevelUsage) return withComputedTotal(topLevelUsage);\n\n  const legacyUsage = normalizeUsage(asRecord(metadata.custom)?.usage);\n  if (legacyUsage) return withComputedTotal(legacyUsage);\n\n  return usageFromSteps(metadata.steps);\n}\n\nexport function getLatestThreadTokenUsage(\n  messages: readonly TokenUsageExtractableMessage[] | undefined,\n): ThreadTokenUsage | undefined {\n  return getThreadMessageTokenUsage(findLatestMessageWithUsage(messages));\n}\n\nfunction findLatestMessageWithUsage(\n  messages: readonly TokenUsageExtractableMessage[] | undefined,\n): TokenUsageExtractableMessage | undefined {\n  if (!messages) return undefined;\n\n  for (let idx = messages.length - 1; idx >= 0; idx -= 1) {\n    const message = messages[idx];\n    if (getThreadMessageTokenUsage(message)) {\n      return message;\n    }\n  }\n\n  return undefined;\n}\n\nexport function useThreadTokenUsage(): ThreadTokenUsage | undefined {\n  const msg = useAuiState((s) => findLatestMessageWithUsage(s.thread.messages));\n  return getThreadMessageTokenUsage(msg);\n}\n"],"mappings":";;;AAkBA,MAAMY,aAAa;CACjB;CACA;CACA;CACA;CACA;AAAa;AAGf,SAASE,SAASC,OAAyC;CACzD,IAAI,CAACA,SAAS,OAAOA,UAAU,YAAYC,MAAMC,QAAQF,KAAK,GAC5D,OAAOG,KAAAA;CACT,OAAOH;AACT;AAEA,SAASI,qBAAqBJ,OAAoC;CAChE,IAAI,OAAOA,UAAU,YAAY,CAACK,OAAOC,SAASN,KAAK,KAAKA,QAAQ,GAClE;CAEF,OAAOA;AACT;AAEA,SAASO,mBAAmBC,OAA6C;CACvE,IAAIA,MAAMrB,gBAAgBgB,KAAAA,GAAW,OAAOK,MAAMrB;CAClD,IAAIqB,MAAMpB,gBAAgBe,KAAAA,KAAaK,MAAMnB,iBAAiBc,KAAAA,GAC5D,OAAOK,MAAMpB,cAAcoB,MAAMnB;AAGrC;AAEA,SAASoB,eAAeT,OAA8C;CACpE,MAAMU,SAASX,SAASC,KAAK;CAC7B,IAAI,CAACU,QAAQ,OAAOP,KAAAA;CAEpB,MAAMQ,SAA2B,CAAC;CAClC,IAAIC,YAAY;CAChB,KAAK,MAAMC,OAAOhB,YAAY;EAC5B,MAAMiB,QAAQV,qBAAqBM,OAAOG,IAAI;EAC9C,IAAIC,UAAUX,KAAAA,GAAW;GACvBQ,OAAOE,OAAOC;GACdF,YAAY;EACd;CACF;CAEA,IAAID,OAAOrB,oBAAoBa,KAAAA,GAAW;EACxC,MAAMW,QAAQV,qBACZL,SAASW,OAAOK,kBAAkB,CAAC,EAAEzB,eACvC;EACA,IAAIwB,UAAUX,KAAAA,GAAW;GACvBQ,OAAOrB,kBAAkBwB;GACzBF,YAAY;EACd;CACF;CACA,IAAID,OAAOpB,sBAAsBY,KAAAA,GAAW;EAC1C,MAAMW,QAAQV,qBACZL,SAASW,OAAOM,iBAAiB,CAAC,EAAEC,eACtC;EACA,IAAIH,UAAUX,KAAAA,GAAW;GACvBQ,OAAOpB,oBAAoBuB;GAC3BF,YAAY;EACd;CACF;CACA,OAAOA,YAAYD,SAASR,KAAAA;AAC9B;AAEA,SAASe,kBACPV,OAC8B;CAC9B,MAAMrB,cAAcoB,mBAAmBC,KAAK;CAC5C,OAAO;EAAE,GAAGA;EAAO,GAAIrB,gBAAgBgB,KAAAA,KAAa,EAAEhB,YAAY;CAAG;AACvE;AAEA,SAASgC,eAAenB,OAA8C;CACpE,MAAMoB,QAAQnB,MAAMC,QAAQF,KAAK,IAAIA,QAAQ,CAAA;CAE7C,MAAMqB,OAA+B,CAAC;CACtC,MAAMC,UAAmC,CAAC;CAC1C,IAAIC,iBAAiB;CACrB,IAAIC,2BAA2B;CAE/B,KAAK,MAAMC,QAAQL,OAAO;EACxB,MAAMZ,QAAQC,eAAeV,SAAS0B,IAAI,CAAC,EAAEjB,KAAK;EAClD,IAAI,CAACA,OAAO;EACZe;EAEA,MAAMG,YAAYnB,mBAAmBC,KAAK;EAC1C,IAAIkB,cAAcvB,KAAAA,GAAW;GAC3BkB,KAAKlC,eAAekC,KAAKlC,eAAe,KAAKuC;GAC7CF;EACF;EAEA,KAAK,MAAMX,OAAOhB,YAAY;GAC5B,IAAIgB,QAAQ,eAAe;GAC3B,IAAIL,MAAMK,SAASV,KAAAA,GAAW;IAC5BkB,KAAKR,QAAQQ,KAAKR,QAAQ,KAAKL,MAAMK;IACrCS,QAAQT,OAAO;GACjB;EACF;CACF;CAEA,IAAIU,mBAAmB,GAAG,OAAOpB,KAAAA;CAEjC,MAAMQ,SAA2B,CAAC;CAClC,IAAIa,6BAA6BD,gBAC/BZ,OAAOxB,cAAckC,KAAKlC;CAE5B,KAAK,MAAM0B,OAAOhB,YAAY;EAC5B,IAAIgB,QAAQ,eAAe;EAC3B,IAAIS,QAAQT,MACVF,OAAOE,OAAOQ,KAAKR;CAEvB;CACA,OAAOF;AACT;AAEA,SAAgBgB,2BACdC,SAC8B;CAC9B,IAAI,CAACA,WAAWA,QAAQnC,SAAS,aAAa,OAAOU,KAAAA;CAErD,MAAMT,WAAWK,SAAS6B,QAAQlC,QAAQ;CAC1C,IAAI,CAACA,UAAU,OAAOS,KAAAA;CAEtB,MAAM0B,gBAAgBpB,eAAef,SAASc,KAAK;CACnD,IAAIqB,eAAe,OAAOX,kBAAkBW,aAAa;CAEzD,MAAMC,cAAcrB,eAAeV,SAASL,SAASqC,MAAM,CAAC,EAAEvB,KAAK;CACnE,IAAIsB,aAAa,OAAOZ,kBAAkBY,WAAW;CAErD,OAAOX,eAAezB,SAAS0B,KAAK;AACtC;AAEA,SAAgBY,0BACdC,UAC8B;CAC9B,OAAON,2BAA2BO,2BAA2BD,QAAQ,CAAC;AACxE;AAEA,SAASC,2BACPD,UAC0C;CAC1C,IAAI,CAACA,UAAU,OAAO9B,KAAAA;CAEtB,KAAK,IAAIgC,MAAMF,SAASG,SAAS,GAAGD,OAAO,GAAGA,OAAO,GAAG;EACtD,MAAMP,UAAUK,SAASE;EACzB,IAAIR,2BAA2BC,OAAO,GACpC,OAAOA;CAEX;AAGF;AAEA,SAAOS,sBAAA;CAAA,MAAAC,IAAAC,EAAA,CAAA;CACL,MAAAC,MAAYvD,YAAYwD,KAAoD;CAAE,IAAAC;CAAA,IAAAJ,EAAA,OAAAE,KAAA;EACvEE,KAAAf,2BAA2Ba,GAAG;EAACF,EAAA,KAAAE;EAAAF,EAAA,KAAAI;CAAA,OAAAA,KAAAJ,EAAA;CAAA,OAA/BI;AAA+B;AAFjC,SAAAD,MAAAE,GAAA;CAAA,OAC0BT,2BAA2BS,EAACC,OAAOX,QAAS;AAAC"}
+{"version":3,"file":"usage.js","names":["useMemo","useAuiState","ThreadTokenUsage","totalTokens","inputTokens","outputTokens","reasoningTokens","cachedInputTokens","TokenUsageExtractableMessage","role","metadata","UsageRecord","Record","USAGE_KEYS","const","asRecord","value","Array","isArray","undefined","asPositiveTokenCount","Number","isFinite","computeTotalTokens","usage","normalizeUsage","record","result","hasFields","key","count","outputTokenDetails","inputTokenDetails","cacheReadTokens","withComputedTotal","usageFromSteps","steps","sums","present","stepsWithUsage","stepsWithComputableTotal","step","stepTotal","getThreadMessageTokenUsage","message","topLevelUsage","customUsage","custom","findLatestMessageWithUsage","messages","idx","length","useThreadTokenUsage","$","_c","msg","_temp","t0","s","thread"],"sources":["../src/usage.ts"],"sourcesContent":["/// <reference types=\"@assistant-ui/core/react\" />\nimport { useMemo } from \"react\";\nimport { useAuiState } from \"@assistant-ui/store\";\n\nexport type ThreadTokenUsage = {\n  totalTokens?: number;\n  inputTokens?: number;\n  outputTokens?: number;\n  reasoningTokens?: number;\n  cachedInputTokens?: number;\n};\n\nexport interface TokenUsageExtractableMessage {\n  role?: string;\n  metadata?: unknown;\n}\n\ntype UsageRecord = Record<string, unknown>;\n\nconst USAGE_KEYS = [\n  \"inputTokens\",\n  \"outputTokens\",\n  \"reasoningTokens\",\n  \"cachedInputTokens\",\n  \"totalTokens\",\n] as const satisfies (keyof ThreadTokenUsage)[];\n\nfunction asRecord(value: unknown): UsageRecord | undefined {\n  if (!value || typeof value !== \"object\" || Array.isArray(value))\n    return undefined;\n  return value as UsageRecord;\n}\n\nfunction asPositiveTokenCount(value: unknown): number | undefined {\n  if (typeof value !== \"number\" || !Number.isFinite(value) || value < 0) {\n    return undefined;\n  }\n  return value;\n}\n\nfunction computeTotalTokens(usage: ThreadTokenUsage): number | undefined {\n  if (usage.totalTokens !== undefined) return usage.totalTokens;\n  if (usage.inputTokens !== undefined && usage.outputTokens !== undefined) {\n    return usage.inputTokens + usage.outputTokens;\n  }\n  return undefined;\n}\n\nfunction normalizeUsage(value: unknown): ThreadTokenUsage | undefined {\n  const record = asRecord(value);\n  if (!record) return undefined;\n\n  const result: ThreadTokenUsage = {};\n  let hasFields = false;\n  for (const key of USAGE_KEYS) {\n    const count = asPositiveTokenCount(record[key]);\n    if (count !== undefined) {\n      result[key] = count;\n      hasFields = true;\n    }\n  }\n  // AI SDK v7 moved these under token detail objects; v6 kept them top-level.\n  if (result.reasoningTokens === undefined) {\n    const count = asPositiveTokenCount(\n      asRecord(record.outputTokenDetails)?.reasoningTokens,\n    );\n    if (count !== undefined) {\n      result.reasoningTokens = count;\n      hasFields = true;\n    }\n  }\n  if (result.cachedInputTokens === undefined) {\n    const count = asPositiveTokenCount(\n      asRecord(record.inputTokenDetails)?.cacheReadTokens,\n    );\n    if (count !== undefined) {\n      result.cachedInputTokens = count;\n      hasFields = true;\n    }\n  }\n  return hasFields ? result : undefined;\n}\n\nfunction withComputedTotal(\n  usage: ThreadTokenUsage,\n): ThreadTokenUsage | undefined {\n  const totalTokens = computeTotalTokens(usage);\n  return { ...usage, ...(totalTokens !== undefined && { totalTokens }) };\n}\n\nfunction usageFromSteps(value: unknown): ThreadTokenUsage | undefined {\n  const steps = Array.isArray(value) ? value : [];\n\n  const sums: Record<string, number> = {};\n  const present: Record<string, boolean> = {};\n  let stepsWithUsage = 0;\n  let stepsWithComputableTotal = 0;\n\n  for (const step of steps) {\n    const usage = normalizeUsage(asRecord(step)?.usage);\n    if (!usage) continue;\n    stepsWithUsage++;\n\n    const stepTotal = computeTotalTokens(usage);\n    if (stepTotal !== undefined) {\n      sums.totalTokens = (sums.totalTokens ?? 0) + stepTotal;\n      stepsWithComputableTotal++;\n    }\n\n    for (const key of USAGE_KEYS) {\n      if (key === \"totalTokens\") continue;\n      if (usage[key] !== undefined) {\n        sums[key] = (sums[key] ?? 0) + usage[key];\n        present[key] = true;\n      }\n    }\n  }\n\n  if (stepsWithUsage === 0) return undefined;\n\n  const result: ThreadTokenUsage = {};\n  if (stepsWithComputableTotal === stepsWithUsage) {\n    result.totalTokens = sums.totalTokens!;\n  }\n  for (const key of USAGE_KEYS) {\n    if (key === \"totalTokens\") continue;\n    if (present[key]) {\n      result[key] = sums[key]!;\n    }\n  }\n  return result;\n}\n\nexport function getThreadMessageTokenUsage(\n  message: TokenUsageExtractableMessage | undefined,\n): ThreadTokenUsage | undefined {\n  if (!message || message.role !== \"assistant\") return undefined;\n\n  const metadata = asRecord(message.metadata);\n  if (!metadata) return undefined;\n\n  const topLevelUsage = normalizeUsage(metadata.usage);\n  if (topLevelUsage) return withComputedTotal(topLevelUsage);\n\n  const customUsage = normalizeUsage(asRecord(metadata.custom)?.usage);\n  if (customUsage) return withComputedTotal(customUsage);\n\n  return usageFromSteps(metadata.steps);\n}\n\nfunction findLatestMessageWithUsage(\n  messages: readonly TokenUsageExtractableMessage[] | undefined,\n): TokenUsageExtractableMessage | undefined {\n  if (!messages) return undefined;\n\n  for (let idx = messages.length - 1; idx >= 0; idx -= 1) {\n    const message = messages[idx];\n    if (getThreadMessageTokenUsage(message)) {\n      return message;\n    }\n  }\n\n  return undefined;\n}\n\n/**\n * Reads token usage from the newest assistant message that reports any.\n *\n * A route attaches usage through the AI SDK's `messageMetadata` option. Because\n * a thread message carries a fixed metadata shape, the converter moves every\n * other key the route returns into `metadata.custom`, which is where this hook\n * looks.\n */\nexport function useThreadTokenUsage(): ThreadTokenUsage | undefined {\n  const msg = useAuiState((s) => findLatestMessageWithUsage(s.thread.messages));\n  return useMemo(() => getThreadMessageTokenUsage(msg), [msg]);\n}\n"],"mappings":";;;;AAmBA,MAAMa,aAAa;CACjB;CACA;CACA;CACA;CACA;AAAa;AAGf,SAASE,SAASC,OAAyC;CACzD,IAAI,CAACA,SAAS,OAAOA,UAAU,YAAYC,MAAMC,QAAQF,KAAK,GAC5D,OAAOG,KAAAA;CACT,OAAOH;AACT;AAEA,SAASI,qBAAqBJ,OAAoC;CAChE,IAAI,OAAOA,UAAU,YAAY,CAACK,OAAOC,SAASN,KAAK,KAAKA,QAAQ,GAClE;CAEF,OAAOA;AACT;AAEA,SAASO,mBAAmBC,OAA6C;CACvE,IAAIA,MAAMrB,gBAAgBgB,KAAAA,GAAW,OAAOK,MAAMrB;CAClD,IAAIqB,MAAMpB,gBAAgBe,KAAAA,KAAaK,MAAMnB,iBAAiBc,KAAAA,GAC5D,OAAOK,MAAMpB,cAAcoB,MAAMnB;AAGrC;AAEA,SAASoB,eAAeT,OAA8C;CACpE,MAAMU,SAASX,SAASC,KAAK;CAC7B,IAAI,CAACU,QAAQ,OAAOP,KAAAA;CAEpB,MAAMQ,SAA2B,CAAC;CAClC,IAAIC,YAAY;CAChB,KAAK,MAAMC,OAAOhB,YAAY;EAC5B,MAAMiB,QAAQV,qBAAqBM,OAAOG,IAAI;EAC9C,IAAIC,UAAUX,KAAAA,GAAW;GACvBQ,OAAOE,OAAOC;GACdF,YAAY;EACd;CACF;CAEA,IAAID,OAAOrB,oBAAoBa,KAAAA,GAAW;EACxC,MAAMW,QAAQV,qBACZL,SAASW,OAAOK,kBAAkB,CAAC,EAAEzB,eACvC;EACA,IAAIwB,UAAUX,KAAAA,GAAW;GACvBQ,OAAOrB,kBAAkBwB;GACzBF,YAAY;EACd;CACF;CACA,IAAID,OAAOpB,sBAAsBY,KAAAA,GAAW;EAC1C,MAAMW,QAAQV,qBACZL,SAASW,OAAOM,iBAAiB,CAAC,EAAEC,eACtC;EACA,IAAIH,UAAUX,KAAAA,GAAW;GACvBQ,OAAOpB,oBAAoBuB;GAC3BF,YAAY;EACd;CACF;CACA,OAAOA,YAAYD,SAASR,KAAAA;AAC9B;AAEA,SAASe,kBACPV,OAC8B;CAC9B,MAAMrB,cAAcoB,mBAAmBC,KAAK;CAC5C,OAAO;EAAE,GAAGA;EAAO,GAAIrB,gBAAgBgB,KAAAA,KAAa,EAAEhB,YAAY;CAAG;AACvE;AAEA,SAASgC,eAAenB,OAA8C;CACpE,MAAMoB,QAAQnB,MAAMC,QAAQF,KAAK,IAAIA,QAAQ,CAAA;CAE7C,MAAMqB,OAA+B,CAAC;CACtC,MAAMC,UAAmC,CAAC;CAC1C,IAAIC,iBAAiB;CACrB,IAAIC,2BAA2B;CAE/B,KAAK,MAAMC,QAAQL,OAAO;EACxB,MAAMZ,QAAQC,eAAeV,SAAS0B,IAAI,CAAC,EAAEjB,KAAK;EAClD,IAAI,CAACA,OAAO;EACZe;EAEA,MAAMG,YAAYnB,mBAAmBC,KAAK;EAC1C,IAAIkB,cAAcvB,KAAAA,GAAW;GAC3BkB,KAAKlC,eAAekC,KAAKlC,eAAe,KAAKuC;GAC7CF;EACF;EAEA,KAAK,MAAMX,OAAOhB,YAAY;GAC5B,IAAIgB,QAAQ,eAAe;GAC3B,IAAIL,MAAMK,SAASV,KAAAA,GAAW;IAC5BkB,KAAKR,QAAQQ,KAAKR,QAAQ,KAAKL,MAAMK;IACrCS,QAAQT,OAAO;GACjB;EACF;CACF;CAEA,IAAIU,mBAAmB,GAAG,OAAOpB,KAAAA;CAEjC,MAAMQ,SAA2B,CAAC;CAClC,IAAIa,6BAA6BD,gBAC/BZ,OAAOxB,cAAckC,KAAKlC;CAE5B,KAAK,MAAM0B,OAAOhB,YAAY;EAC5B,IAAIgB,QAAQ,eAAe;EAC3B,IAAIS,QAAQT,MACVF,OAAOE,OAAOQ,KAAKR;CAEvB;CACA,OAAOF;AACT;AAEA,SAAgBgB,2BACdC,SAC8B;CAC9B,IAAI,CAACA,WAAWA,QAAQnC,SAAS,aAAa,OAAOU,KAAAA;CAErD,MAAMT,WAAWK,SAAS6B,QAAQlC,QAAQ;CAC1C,IAAI,CAACA,UAAU,OAAOS,KAAAA;CAEtB,MAAM0B,gBAAgBpB,eAAef,SAASc,KAAK;CACnD,IAAIqB,eAAe,OAAOX,kBAAkBW,aAAa;CAEzD,MAAMC,cAAcrB,eAAeV,SAASL,SAASqC,MAAM,CAAC,EAAEvB,KAAK;CACnE,IAAIsB,aAAa,OAAOZ,kBAAkBY,WAAW;CAErD,OAAOX,eAAezB,SAAS0B,KAAK;AACtC;AAEA,SAASY,2BACPC,UAC0C;CAC1C,IAAI,CAACA,UAAU,OAAO9B,KAAAA;CAEtB,KAAK,IAAI+B,MAAMD,SAASE,SAAS,GAAGD,OAAO,GAAGA,OAAO,GAAG;EACtD,MAAMN,UAAUK,SAASC;EACzB,IAAIP,2BAA2BC,OAAO,GACpC,OAAOA;CAEX;AAGF;;;;;;;;;AAUA,SAAOQ,sBAAA;CAAA,MAAAC,IAAAC,EAAA,CAAA;CACL,MAAAC,MAAYtD,YAAYuD,KAAoD;CAAE,IAAAC;CAAA,IAAAJ,EAAA,OAAAE,KAAA;EACzDE,KAAAd,2BAA2BY,GAAG;EAACF,EAAA,KAAAE;EAAAF,EAAA,KAAAI;CAAA,OAAAA,KAAAJ,EAAA;CAAA,OAA/BI;AAA+B;AAF/C,SAAAD,MAAAE,GAAA;CAAA,OAC0BV,2BAA2BU,EAACC,OAAOV,QAAS;AAAC"}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/ai-sdk",
-  "version": "0.0.3",
+  "version": "0.0.5",
   "description": "Vercel AI SDK adapter for assistant-ui",
   "keywords": [
     "ai-sdk",
@@ -46,14 +46,14 @@
   ],
   "sideEffects": false,
   "dependencies": {
-    "@ai-sdk/mcp": "^2.0.36",
-    "@ai-sdk/react": "^4.0.80",
-    "@assistant-ui/core": "^0.3.16",
-    "@assistant-ui/store": "^0.3.11",
-    "@assistant-ui/tap": "^0.9.15",
-    "ai": "^7.0.77",
+    "@ai-sdk/mcp": "^2.0.45",
+    "@ai-sdk/react": "^4.0.96",
+    "@assistant-ui/core": "^0.3.18",
+    "@assistant-ui/store": "^0.3.13",
+    "@assistant-ui/tap": "^0.9.17",
+    "ai": "^7.0.93",
     "assistant-cloud": "*",
-    "assistant-stream": "^0.3.40"
+    "assistant-stream": "^0.3.42"
   },
   "peerDependencies": {
     "@types/react": "*",
@@ -68,11 +68,11 @@
     }
   },
   "devDependencies": {
-    "@assistant-ui/x-buildutils": "0.0.24",
-    "@testing-library/react": "^16.3.2",
+    "@assistant-ui/x-buildutils": "0.0.26",
+    "@testing-library/react": "^16.3.3",
     "@types/json-schema": "^7.0.15",
     "@types/react": "^19.2.18",
-    "@types/react-dom": "^19.2.5",
+    "@types/react-dom": "^19.2.7",
     "jsdom": "^30.0.1",
     "react": "^19.2.8",
     "vitest": "^4.1.11"

package/src/converters/convertMessage.test.ts

@@ -26,6 +26,28 @@ describe("AISDKMessageConverter", () => {
     expect(converted[1]?.metadata.isOptimistic).toBe(true);
   });
 
+  it("keeps metadata outside the thread shape reachable under custom", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [{ type: "text", text: "yo" }],
+        metadata: {
+          usage: { inputTokens: 40, outputTokens: 2 },
+          modelId: "gpt-5.6-luna",
+          custom: { source: "route" },
+        },
+      },
+    ] as any);
+
+    expect(converted[0]?.metadata.custom).toEqual({
+      usage: { inputTokens: 40, outputTokens: 2 },
+      modelId: "gpt-5.6-luna",
+      source: "route",
+    });
+    expect(converted[0]?.metadata).not.toHaveProperty("usage");
+  });
+
   it("does not flag messages when no optimistic id is provided", () => {
     const converted = AISDKMessageConverter.toThreadMessages([
       { id: "a1", role: "assistant", parts: [{ type: "text", text: "yo" }] },
@@ -372,6 +394,108 @@ describe("AISDKMessageConverter", () => {
     });
   });
 
+  it("preserves producer-defined approval fields and gives prompt precedence", () => {
+    const descriptor = { scope: "account:deploy" };
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-deploy",
+            toolCallId: "tc-1",
+            state: "approval-responded",
+            input: { environment: "production" },
+            approval: {
+              id: "approval-1",
+              approved: true,
+              reason: "approved by operator",
+              prompt: "Deploy to production?",
+              descriptor,
+              requestReason: "Production access requires approval",
+              signature: "signed-approval",
+              futureField: "preserved",
+            },
+          },
+        ],
+      } as any,
+    ]);
+
+    const toolCall = converted[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(toolCall?.approval).toEqual({
+      id: "approval-1",
+      approved: true,
+      reason: "approved by operator",
+      prompt: "Deploy to production?",
+      descriptor,
+      requestReason: "Production access requires approval",
+      signature: "signed-approval",
+      futureField: "preserved",
+    });
+  });
+
+  it("drops fields the AI SDK cannot answer and uses requestReason as the prompt", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-deploy",
+            toolCallId: "tc-1",
+            state: "approval-requested",
+            input: {},
+            approval: {
+              id: "approval-1",
+              display: "select",
+              allowFreeform: true,
+              options: [{ id: "once", kind: "allow-once" }],
+              optionId: "once",
+              text: "an answer",
+              resolution: "cancelled",
+              requestReason: "kept",
+            },
+          },
+        ],
+      } as any,
+    ]);
+
+    const toolCall = converted[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(toolCall?.approval).toEqual({
+      id: "approval-1",
+      prompt: "kept",
+      resolution: "cancelled",
+      requestReason: "kept",
+    });
+  });
+
+  it("drops a resolution the core contract does not declare", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-deploy",
+            toolCallId: "tc-1",
+            state: "approval-requested",
+            input: {},
+            approval: { id: "approval-1", resolution: "whatever" },
+          },
+        ],
+      } as any,
+    ]);
+
+    const toolCall = converted[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(toolCall?.approval).toEqual({ id: "approval-1" });
+  });
+
   it("strips closing delimiters from streaming tool argsText", () => {
     const converted = AISDKMessageConverter.toThreadMessages([
       {

package/src/converters/convertMessage.ts

@@ -31,6 +31,30 @@ import {
 import { unwrapModelContentEnvelope } from "./modelContentEnvelope";
 
 type MessageMetadata = ThreadMessageLike["metadata"];
+
+const THREAD_METADATA_KEYS = new Set([
+  "unstable_state",
+  "unstable_annotations",
+  "unstable_data",
+  "steps",
+  "timing",
+  "submittedFeedback",
+  "isOptimistic",
+  "custom",
+]);
+
+const toThreadMetadata = (metadata: unknown): MessageMetadata => {
+  if (!metadata || typeof metadata !== "object") return undefined;
+  const result: Record<string, unknown> = {};
+  const extra: Record<string, unknown> = {};
+  for (const [key, value] of Object.entries(metadata)) {
+    (THREAD_METADATA_KEYS.has(key) ? result : extra)[key] = value;
+  }
+  if (Object.keys(extra).length > 0) {
+    result.custom = { ...extra, ...(result.custom as object | undefined) };
+  }
+  return result as MessageMetadata;
+};
 export type AISDKMessageConverterMetadata =
   useExternalMessageConverter.Metadata & {
     toolArgsKeyOrderCache?: Map<string, Map<string, string[]>>;
@@ -130,30 +154,50 @@ function extractMcpAppMetadata(
 
 function getToolApprovalAndInterrupt(
   part: {
-    approval?:
-      | {
-          id: string;
-          approved?: boolean;
-          reason?: string;
-          isAutomatic?: boolean;
-        }
-      | undefined;
+    approval?: Record<string, unknown> | undefined;
   },
   toolStatus: { type: string; payload?: unknown } | undefined,
 ): {
   approval?: NonNullable<ToolCallMessagePart["approval"]>;
   interrupt?: NonNullable<ToolCallMessagePart["interrupt"]>;
 } {
-  if (part.approval && typeof part.approval.id === "string") {
-    const { id, approved, reason, isAutomatic } = part.approval;
-    return {
-      approval: {
-        id,
-        ...(typeof approved === "boolean" && { approved }),
-        ...(typeof reason === "string" && { reason }),
-        ...(isAutomatic === true && { isAutomatic: true }),
-      },
-    };
+  if (part.approval) {
+    // The AI SDK sends only id, approved and reason back to the server, so a
+    // request shape promising any other answer would render controls whose
+    // response cannot travel.
+    const {
+      id,
+      prompt,
+      approved,
+      reason,
+      isAutomatic,
+      resolution,
+      display,
+      allowFreeform,
+      options,
+      optionId,
+      text,
+      ...additionalApprovalFields
+    } = part.approval;
+    const requestReason = additionalApprovalFields.requestReason;
+    if (typeof id === "string")
+      return {
+        approval: {
+          ...additionalApprovalFields,
+          id,
+          ...(typeof prompt === "string"
+            ? { prompt }
+            : typeof requestReason === "string"
+              ? { prompt: requestReason }
+              : {}),
+          ...(typeof approved === "boolean" && { approved }),
+          ...(typeof reason === "string" && { reason }),
+          ...(isAutomatic === true && { isAutomatic: true }),
+          ...((resolution === "cancelled" || resolution === "expired") && {
+            resolution,
+          }),
+        } as NonNullable<ToolCallMessagePart["approval"]>,
+      };
   }
 
   if (toolStatus?.type === "interrupt") {
@@ -427,7 +471,7 @@ export const AISDKMessageConverter = unstable_createMessageConverter(
                 status: { type: "complete" as const },
               };
             }),
-          metadata: message.metadata as MessageMetadata,
+          metadata: toThreadMetadata(message.metadata),
         };
 
       case "system":
@@ -442,7 +486,7 @@ export const AISDKMessageConverter = unstable_createMessageConverter(
           createdAt,
           content,
           metadata: {
-            ...(message.metadata as MessageMetadata),
+            ...toThreadMetadata(message.metadata),
             ...(timing && { timing }),
             ...(isOptimistic && { isOptimistic: true }),
           },

package/src/runtime/AISDKChat.integration.test.tsx

@@ -1,47 +1,68 @@
 // @vitest-environment jsdom
 
-import { useEffect } from "react";
-import { render, waitFor } from "@testing-library/react";
-import { StrictMode } from "react";
+import { StrictMode, type ReactNode } from "react";
+import { act, render, waitFor } from "@testing-library/react";
 import { AuiConfig, AuiProvider, useAui } from "@assistant-ui/store";
-import { flushTapSync } from "@assistant-ui/tap";
+import type { ChatTransport, UIMessage } from "ai";
 import { describe, expect, it } from "vitest";
 import { AISDKChat } from "./AISDKChat";
-import { createCancellableTransport } from "./__tests__/controlled-transport";
+import {
+  createCancellableTransport,
+  createStreamHarness,
+} from "./__tests__/controlled-transport";
 
 describe("AISDKChat React integration", () => {
-  it("does not treat React provider unmount as client destruction", async () => {
-    const { transport, getCancelCount, close } = createCancellableTransport();
-    let started = false;
-    let isRunning = () => false;
-
-    const SendOnMount = () => {
-      const aui = useAui();
-      isRunning = () => aui.thread.getState().isRunning;
-      useEffect(() => {
-        if (started) return;
-        started = true;
-        flushTapSync(() => aui.composer.setText("keep streaming"));
-        flushTapSync(() => aui.composer.send());
-      }, [aui]);
-      return null;
-    };
+  it("aborts the in-flight transport after a real unmount", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const { Probe, send, isRunning } = createStreamHarness();
 
     const view = render(
       <StrictMode>
         <AuiProvider config={AuiConfig({ threads: AISDKChat({ transport }) })}>
-          <SendOnMount />
+          <Probe />
         </AuiProvider>
       </StrictMode>,
     );
 
-    await waitFor(() => {
-      expect(isRunning()).toBe(true);
-    });
+    await act(async () => send());
+    await waitFor(() => expect(isRunning()).toBe(true));
+    // the Strict Mode double mount already ran a host cleanup by now
+    expect(getCancelCount()).toBe(0);
 
     view.unmount();
-    await new Promise((resolve) => setTimeout(resolve, 0));
+    await waitFor(() => expect(getCancelCount()).toBe(1));
+  });
+});
+
+describe("AISDKChat legacy useAui host integration", () => {
+  const LegacyProvider = ({
+    transport,
+    children,
+  }: {
+    transport: ChatTransport<UIMessage>;
+    children: ReactNode;
+  }) => {
+    const aui = useAui(AuiConfig({ threads: AISDKChat({ transport }) }));
+    return <AuiProvider value={aui}>{children}</AuiProvider>;
+  };
+
+  it("aborts the in-flight transport after a real unmount", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const { Probe, send, isRunning } = createStreamHarness();
+
+    const view = render(
+      <StrictMode>
+        <LegacyProvider transport={transport}>
+          <Probe />
+        </LegacyProvider>
+      </StrictMode>,
+    );
+
+    await act(async () => send());
+    await waitFor(() => expect(isRunning()).toBe(true));
     expect(getCancelCount()).toBe(0);
-    close();
+
+    view.unmount();
+    await waitFor(() => expect(getCancelCount()).toBe(1));
   });
 });

package/src/runtime/AISDKThreads.cloud.test.ts

@@ -95,7 +95,7 @@ describe("AISDKThreads cloud", () => {
     handle.destroy();
   });
 
-  it("stops an in-flight cloud chat when switching away", async () => {
+  it("keeps an in-flight cloud chat running across a switch and stops it on delete", async () => {
     const chat = createCancellableTransport();
     const handle = createAssistantClient(
       AuiConfig({
@@ -126,6 +126,16 @@ describe("AISDKThreads cloud", () => {
       await vi.waitFor(() => {
         expect(handle.getClient().threads.getState().mainThreadId).toBe("t2");
       });
+      expect(chat.getCancelCount()).toBe(0);
+      await vi.waitFor(() => {
+        expect(
+          handle.getClient().threads.item({ id: "t1" }).getState().isRunning,
+        ).toBe(true);
+      });
+
+      flushTapSync(() =>
+        handle.getClient().threads.item({ id: "t1" }).delete(),
+      );
       await vi.waitFor(() => {
         expect(chat.getCancelCount()).toBe(1);
       });

package/src/runtime/AISDKThreads.test.ts

@@ -546,3 +546,105 @@ describe("AISDKThreads", () => {
     }
   });
 });
+
+describe("AISDKThreads branch isolation", () => {
+  const completeRound = async (
+    handle: ReturnType<typeof createAssistantClient>,
+    emit: (...chunks: never[]) => void,
+    close: () => void,
+    question: string,
+    answer: string,
+  ) => {
+    const aui = handle.getClient();
+    flushTapSync(() => aui.composer.setText(question));
+    flushTapSync(() => aui.composer.send());
+    await vi.waitFor(() => {
+      expect(aui.thread.getState().messages.length).toBeGreaterThan(0);
+    });
+    emit(...(textReply(answer) as never[]));
+    close();
+    await vi.waitFor(() => {
+      expect(handle.getClient().thread.getState().isRunning).toBe(false);
+    });
+  };
+
+  it("keeps switching between populated threads free of cross-thread branches", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKThreads({ transport }) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    await completeRound(
+      handle,
+      emit as never,
+      close,
+      "thread a question",
+      "thread a answer",
+    );
+
+    flushTapSync(() => aui.threads.switchToNewThread());
+    await completeRound(
+      handle,
+      emit as never,
+      close,
+      "thread b question",
+      "thread b answer",
+    );
+
+    flushTapSync(() => aui.threads.switchToThread("main"));
+    await vi.waitFor(() => {
+      expect(threadText(handle as never)).toEqual([
+        "thread a question",
+        "thread a answer",
+      ]);
+    });
+    expect(aui.thread.message({ index: 0 }).getState().branchCount).toBe(1);
+    expect(aui.thread.message({ index: 1 }).getState().branchCount).toBe(1);
+
+    handle.destroy();
+  });
+
+  it("preserves intra-thread branches across a switch through an empty thread", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKThreads({ transport }) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    await completeRound(
+      handle,
+      emit as never,
+      close,
+      "branchy question",
+      "first answer",
+    );
+
+    flushTapSync(() => aui.thread.message({ index: 1 }).reload());
+    await vi.waitFor(() => {
+      expect(handle.getClient().thread.getState().isRunning).toBe(true);
+    });
+    emit(...(textReply("second answer") as never[]));
+    close();
+    await vi.waitFor(() => {
+      expect(handle.getClient().thread.getState().isRunning).toBe(false);
+      expect(aui.thread.message({ index: 1 }).getState().branchCount).toBe(2);
+    });
+
+    flushTapSync(() => aui.threads.switchToNewThread());
+    expect(handle.getClient().thread.getState().messages).toHaveLength(0);
+
+    flushTapSync(() => handle.getClient().threads.switchToThread("main"));
+    await vi.waitFor(() => {
+      expect(threadText(handle as never)).toEqual([
+        "branchy question",
+        "second answer",
+      ]);
+    });
+    expect(aui.thread.message({ index: 1 }).getState().branchCount).toBe(2);
+
+    handle.destroy();
+  });
+});