feat(core): add useToolArgsStatus hook (#3712)

okisdev · web-flow · commit af70d7fa9d6e · 2026-03-27T12:02:21.000-07:00
diff --git a/.changeset/tool-args-status-hook.md b/.changeset/tool-args-status-hook.md
@@ -0,0 +1,10 @@
+---
+"@assistant-ui/core": patch
+"@assistant-ui/react": patch
+"@assistant-ui/react-native": patch
+"@assistant-ui/react-ink": patch
+---
+
+feat: add useToolArgsStatus hook for per-prop streaming status
+
+Add a convenience hook that derives per-property streaming completion status from tool call args using structural partial JSON analysis.
diff --git a/apps/docs/content/docs/(docs)/guides/tool-ui.mdx b/apps/docs/content/docs/(docs)/guides/tool-ui.mdx
@@ -748,6 +748,35 @@ useAssistantToolUI({
   server.
 </Callout>
 
+## Per-Property Streaming Status
+
+When rendering a tool UI, you can track which arguments have finished streaming using `useToolArgsStatus`. This must be used inside a tool-call message part context.
+
+```tsx
+import { useToolArgsStatus } from "@assistant-ui/react";
+
+const WeatherUI = makeAssistantToolUI({
+  toolName: "weather",
+  render: ({ args }) => {
+    const { status, propStatus } = useToolArgsStatus<{
+      location: string;
+      unit: string;
+    }>();
+
+    return (
+      <div>
+        <span className={propStatus.location === "streaming" ? "animate-pulse" : ""}>
+          {args.location ?? "..."}
+        </span>
+        {status === "complete" && <WeatherChart data={args} />}
+      </div>
+    );
+  },
+});
+```
+
+`propStatus` maps each key to `"streaming"` | `"complete"` once the key appears in the partial JSON. Keys not yet present in the stream are absent from `propStatus`.
+
 ## Related Guides
 
 - [Tools Guide](/docs/guides/tools) - Learn how to create and use tools with AI models
diff --git a/packages/assistant-stream/src/utils.ts b/packages/assistant-stream/src/utils.ts
@@ -1,6 +1,7 @@
 export {
   parsePartialJsonObject,
   getPartialJsonObjectFieldState,
+  getPartialJsonObjectMeta,
 } from "./utils/json/parse-partial-json-object";
 export {
   type AsyncIterableStream,
diff --git a/packages/core/src/react/index.ts b/packages/core/src/react/index.ts
@@ -33,6 +33,10 @@ export {
   type AssistantInteractableProps,
 } from "./model-context/useAssistantInteractable";
 export { useInteractableState } from "./model-context/useInteractableState";
+export {
+  useToolArgsStatus,
+  type ToolArgsStatus,
+} from "./model-context/useToolArgsStatus";
 
 // client
 export { Tools } from "./client/Tools";
diff --git a/packages/core/src/react/model-context/useToolArgsStatus.ts b/packages/core/src/react/model-context/useToolArgsStatus.ts
@@ -0,0 +1,51 @@
+import { useMemo } from "react";
+import { useAuiState } from "@assistant-ui/store";
+import {
+  getPartialJsonObjectFieldState,
+  getPartialJsonObjectMeta,
+} from "assistant-stream/utils";
+
+type PropFieldStatus = "streaming" | "complete";
+
+export type ToolArgsStatus<
+  TArgs extends Record<string, unknown> = Record<string, unknown>,
+> = {
+  status: "running" | "complete" | "incomplete" | "requires-action";
+  propStatus: Partial<Record<keyof TArgs, PropFieldStatus>>;
+};
+
+export const useToolArgsStatus = <
+  TArgs extends Record<string, unknown> = Record<string, unknown>,
+>(): ToolArgsStatus<TArgs> => {
+  const part = useAuiState((s) => s.part);
+
+  return useMemo(() => {
+    const statusType = part.status.type;
+
+    if (part.type !== "tool-call") {
+      throw new Error(
+        "useToolArgsStatus can only be used inside tool-call message parts",
+      );
+    }
+
+    const isStreaming = statusType === "running";
+    const args = part.args as Record<string, unknown>;
+    const meta = getPartialJsonObjectMeta(args as Record<symbol, unknown>);
+    const propStatus: Partial<Record<string, PropFieldStatus>> = {};
+
+    for (const key of Object.keys(args)) {
+      if (meta) {
+        const fieldState = getPartialJsonObjectFieldState(args, [key]);
+        propStatus[key] =
+          fieldState === "complete" || !isStreaming ? "complete" : "streaming";
+      } else {
+        propStatus[key] = isStreaming ? "streaming" : "complete";
+      }
+    }
+
+    return {
+      status: statusType,
+      propStatus: propStatus as Partial<Record<keyof TArgs, PropFieldStatus>>,
+    };
+  }, [part]);
+};
diff --git a/packages/react-ink/src/index.ts b/packages/react-ink/src/index.ts
@@ -152,6 +152,8 @@ export {
   useAssistantInteractable,
   type AssistantInteractableProps,
   useInteractableState,
+  useToolArgsStatus,
+  type ToolArgsStatus,
 } from "@assistant-ui/core/react";
 export type {
   ModelContext,
diff --git a/packages/react-native/src/index.ts b/packages/react-native/src/index.ts
@@ -149,6 +149,8 @@ export {
   useAssistantInteractable,
   type AssistantInteractableProps,
   useInteractableState,
+  useToolArgsStatus,
+  type ToolArgsStatus,
 } from "@assistant-ui/core/react";
 
 export type {
diff --git a/packages/react/src/index.ts b/packages/react/src/index.ts
@@ -203,6 +203,8 @@ export {
   useAssistantInteractable,
   type AssistantInteractableProps,
   useInteractableState,
+  useToolArgsStatus,
+  type ToolArgsStatus,
 } from "@assistant-ui/core/react";
 
 // Core pass-through (unchanged)
