feat: implement type inference for tool streaming

libs/sdk/src/react/stream.lgp.tsx

@@ -319,7 +319,13 @@ export function useStreamLGP<
         }
         case "on_tool_event": {
           if (existing) {
-            next.set(key, { ...existing, state: "running", data: data.data });
+            next.set(key, {
+              ...existing,
+              state: "running",
+              data: data.data,
+              result: undefined,
+              error: undefined,
+            });
           }
           break;
         }
@@ -329,13 +335,19 @@ export function useStreamLGP<
               ...existing,
               state: "completed",
               result: data.output,
+              error: undefined,
             });
           }
           break;
         }
         case "on_tool_error": {
           if (existing) {
-            next.set(key, { ...existing, state: "error", error: data.error });
+            next.set(key, {
+              ...existing,
+              state: "error",
+              error: data.error,
+              result: undefined,
+            });
           }
           break;
         }

libs/sdk/src/react/types.tsx

@@ -3,7 +3,7 @@ import type { Client } from "../client.js";
 
 import type { ThreadState } from "../schema.js";
 import type { Message } from "../types.messages.js";
-import type { StreamMode, ToolProgress } from "../types.stream.js";
+import type { GetToolProgressType, StreamMode } from "../types.stream.js";
 import type { Sequence } from "../ui/branching.js";
 import type {
   GetUpdateType,
@@ -137,7 +137,7 @@ export interface UseStream<
   /**
    * Progress of tool executions during streaming.
    */
-  toolProgress: ToolProgress[];
+  toolProgress: GetToolProgressType<Bag>[];
 
   /**
    * LangGraph SDK client used to send request and receive responses.

libs/sdk/src/tests/stream.test-d.ts

@@ -44,6 +44,7 @@ import type {
   InferSubagentState,
   InferSubagentNames,
   SubagentStateMap,
+  InferToolMapFromAgent,
 } from "../ui/types.js";
 import type { ResolveStreamOptions } from "../ui/stream/index.js";
 
@@ -693,3 +694,119 @@ describe("useStream type inference integration", () => {
     expectTypeOf(stream.values.preferences.language).toEqualTypeOf<string>();
   });
 });
+
+// ============================================================================
+// Type Tests: InferToolMapFromAgent (tool progress type inference)
+// ============================================================================
+
+describe("InferToolMapFromAgent", () => {
+  test("simple agent infers tool name and input; data/result are unknown for non-streaming tools", () => {
+    type Map = InferToolMapFromAgent<typeof simpleAgent>;
+
+    expectTypeOf<Map>().toHaveProperty("get_weather");
+    expectTypeOf<Map["get_weather"]["input"]>().toEqualTypeOf<{
+      location: string;
+    }>();
+    expectTypeOf<Map["get_weather"]>().toHaveProperty("data");
+    expectTypeOf<Map["get_weather"]>().toHaveProperty("result");
+  });
+
+  test("multi-tool agent infers all tool entries with correct input types", () => {
+    type Map = InferToolMapFromAgent<typeof multiToolAgent>;
+
+    expectTypeOf<Map>().toHaveProperty("get_weather");
+    expectTypeOf<Map>().toHaveProperty("search_web");
+    expectTypeOf<Map>().toHaveProperty("send_email");
+    expectTypeOf<Map["get_weather"]["input"]>().toEqualTypeOf<{
+      location: string;
+    }>();
+    expectTypeOf<Map["search_web"]["input"]>().toExtend<{
+      query: string;
+      maxResults?: number;
+    }>();
+    expectTypeOf<Map["send_email"]["input"]>().toEqualTypeOf<{
+      to: string;
+      subject: string;
+      body: string;
+    }>();
+  });
+
+  test("streaming tool (AsyncGenerator) infers typed data and result", () => {
+    type MockStreamingTool = {
+      name: "streaming_tool";
+      func: (arg: {
+        query: string;
+      }) => AsyncGenerator<{ progress: number }, string>;
+    };
+    type MockAgent = {
+      "~agentTypes": {
+        Response: unknown;
+        State: unknown;
+        Context: unknown;
+        Middleware: unknown;
+        Tools: readonly [MockStreamingTool];
+      };
+    };
+
+    type Map = InferToolMapFromAgent<MockAgent>;
+
+    expectTypeOf<Map>().toHaveProperty("streaming_tool");
+    expectTypeOf<Map["streaming_tool"]["data"]>().toExtend<
+      { progress: number } | undefined
+    >();
+    expectTypeOf<Map["streaming_tool"]["result"]>().toExtend<
+      string | undefined
+    >();
+  });
+
+  test("useStream with agent has toolProgress typed with literal tool names", () => {
+    const stream = useStream<typeof simpleAgent>({
+      assistantId: "agent",
+    });
+
+    expectTypeOf(stream).toHaveProperty("toolProgress");
+    const progress = stream.toolProgress[0];
+
+    if (progress) {
+      expectTypeOf(progress.name).toEqualTypeOf<"get_weather">();
+    }
+  });
+
+  test("useStream toolProgress narrows data and result by state", () => {
+    type MockStreamingTool = {
+      name: "live_search";
+      func: (arg: {
+        query: string;
+      }) => AsyncGenerator<{ progress: number; partial: string[] }, string>;
+    };
+    type MockAgent = {
+      "~agentTypes": {
+        Response: unknown;
+        State: undefined;
+        Context: unknown;
+        Middleware: readonly [];
+        Tools: readonly [MockStreamingTool];
+      };
+    };
+
+    const stream = useStream<MockAgent>({
+      assistantId: "agent",
+    });
+
+    const tp = stream.toolProgress[0];
+
+    if (tp) {
+      expectTypeOf(tp.name).toEqualTypeOf<"live_search">();
+    }
+
+    if (tp && tp.name === "live_search" && tp.state === "running") {
+      expectTypeOf(tp.data).toEqualTypeOf<
+        { progress: number; partial: string[] } | undefined
+      >();
+    }
+
+    if (tp && tp.name === "live_search" && tp.state === "completed") {
+      expectTypeOf(tp.result).toEqualTypeOf<string | undefined>();
+    }
+  });
+});

libs/sdk/src/types.stream.ts

@@ -1,5 +1,6 @@
 import type { Message } from "./types.messages.js";
 import type { Interrupt, Metadata, Config, ThreadTask } from "./schema.js";
+import { BagTemplate } from "./types.template.js";
 /**
 import type { SubgraphCheckpointsStreamEvent } from "./types.stream.subgraph.js";
  * Stream modes
@@ -278,19 +279,64 @@ export type ToolsStreamEvent = {
   data: ToolStreamEventData;
 };
 
-export type ToolProgress = {
+export type ToolTypes = { input?: unknown; data?: unknown; result?: unknown };
+
+export type ToolProgress<
+  TData = unknown,
+  TInput = unknown,
+  TResult = unknown,
+  TName extends string = string
+> = {
   toolCallId?: string;
-  name: string;
-  state: "starting" | "running" | "completed" | "error";
-  input?: unknown;
-  data?: unknown;
-  result?: unknown;
-  error?: unknown;
-};
+  name: TName;
+} & (
+  | {
+      state: "starting";
+      input?: TInput;
+      data?: undefined;
+      result?: undefined;
+      error?: undefined;
+    }
+  | {
+      state: "running";
+      data?: TData;
+      input?: TInput;
+      result?: undefined;
+      error?: undefined;
+    }
+  | {
+      state: "completed";
+      result?: TResult;
+      input?: TInput;
+      data?: TData;
+      error?: undefined;
+    }
+  | {
+      state: "error";
+      error?: Error | unknown;
+      input?: TInput;
+      data?: TData;
+      result?: undefined;
+    }
+);
 
 /** @internal */
 export type SubgraphToolsStreamEvent = AsSubgraph<ToolsStreamEvent>;
 
+export type DeriveToolProgress<T extends Record<string, ToolTypes>> = {
+  [K in keyof T & string]: ToolProgress<
+    T[K]["data"],
+    T[K]["input"],
+    T[K]["result"],
+    K
+  >;
+}[keyof T & string];
+
+export type GetToolProgressType<Bag extends BagTemplate> =
+  Bag["ToolMap"] extends Record<string, ToolTypes>
+    ? DeriveToolProgress<Bag["ToolMap"]>
+    : ToolProgress;
+
 type GetStreamModeMap<
   TStreamMode extends StreamMode | StreamMode[],
   TStateType = unknown,

libs/sdk/src/types.template.ts

@@ -12,4 +12,8 @@ export type BagTemplate = {
   CustomEventType?: unknown;
   UpdateType?: unknown;
   MetaType?: unknown;
+  ToolMap?: Record<
+    string,
+    { input?: unknown; data?: unknown; result?: unknown }
+  >;
 };

libs/sdk/src/ui/stream/base.ts

@@ -9,7 +9,7 @@
 
 import type { Client } from "../../client.js";
 import type { ThreadState, Interrupt } from "../../schema.js";
-import type { StreamMode, ToolProgress } from "../../types.stream.js";
+import type { GetToolProgressType, StreamMode } from "../../types.stream.js";
 import type { StreamEvent } from "../../types.js";
 import type { Message, DefaultToolCall } from "../../types.messages.js";
 import type { BagTemplate } from "../../types.template.js";
@@ -156,7 +156,7 @@ export interface BaseStream<
    * Progress of tool executions during streaming. Populated when stream mode includes "tools"
    * and tools yield or report progress.
    */
-  toolProgress: ToolProgress[];
+  toolProgress: GetToolProgressType<Bag>[];
 
   /**
    * LangGraph SDK client used to send requests and receive responses.

libs/sdk/src/ui/stream/index.ts

@@ -21,6 +21,7 @@ import type {
   AgentTypeConfigLike,
   DeepAgentTypeConfigLike,
   UseStreamOptions,
+  InferToolMapFromAgent,
 } from "../types.js";
 
 // Import for internal use
@@ -258,5 +259,5 @@ export type ResolveStreamOptions<
 export type InferBag<T, B extends BagTemplate = BagTemplate> = T extends {
   "~agentTypes": unknown;
 }
-  ? BagTemplate
+  ? Omit<BagTemplate, "ToolMap"> & { ToolMap: InferToolMapFromAgent<T> }
   : B;

libs/sdk/src/ui/types.ts

@@ -1305,3 +1305,44 @@ export type CustomSubmitOptions<
   SubmitOptions<StateType, ConfigurableType>,
   "optimisticValues" | "context" | "command" | "config"
 >;
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+/* eslint-disable @typescript-eslint/no-explicit-any */
+type ExtractAsyncGenTypes<T> = T extends AsyncGenerator<infer Y, infer R, any>
+  ? { data: Y | undefined; result: R | undefined }
+  : { data: unknown; result: unknown };
+
+type ExtractToolStreamTypes<T> = T extends { func: (...args: any[]) => infer R }
+  ? ExtractAsyncGenTypes<
+      Extract<R, AsyncGenerator<any, any, any>>
+    > extends infer G
+    ? [G] extends [never]
+      ? { data: unknown; result: unknown }
+      : G
+    : { data: unknown; result: unknown }
+  : { data: unknown; result: unknown };
+/* eslint-enable @typescript-eslint/no-explicit-any */
+
+type ToolMapEntryFromTool<T> = T extends { name: infer N }
+  ? N extends string
+    ? IsLiteralString<N> extends true
+      ? { input: InferToolInput<T> } & ExtractToolStreamTypes<T>
+      : never
+    : never
+  : never;
+
+/**
+ * Infer a tool map from an agent's tools array. Maps each tool name to { input, data, result } types.
+ */
+export type InferToolMapFromAgent<T> =
+  ExtractAgentConfig<T>["Tools"] extends readonly (infer Tool)[]
+    ? {
+        [K in Tool extends { name: infer N }
+          ? N extends string
+            ? IsLiteralString<N> extends true
+              ? N
+              : never
+            : never
+          : never]: ToolMapEntryFromTool<Extract<Tool, { name: K }>>;
+      }
+    : Record<string, { input?: unknown; data?: unknown; result?: unknown }>;