diff --git a/libs/sdk/src/react/stream.lgp.tsx b/libs/sdk/src/react/stream.lgp.tsx index 6d33428e6b..0adcb9aab9 100644 --- a/libs/sdk/src/react/stream.lgp.tsx +++ b/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; } diff --git a/libs/sdk/src/react/types.tsx b/libs/sdk/src/react/types.tsx index 28e3eede0f..53ab58da7a 100644 --- a/libs/sdk/src/react/types.tsx +++ b/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[]; /** * LangGraph SDK client used to send request and receive responses. diff --git a/libs/sdk/src/tests/stream.test-d.ts b/libs/sdk/src/tests/stream.test-d.ts index 8e8a69ba96..d4a4de92cd 100644 --- a/libs/sdk/src/tests/stream.test-d.ts +++ b/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(); }); }); + +// ============================================================================ +// 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; + + expectTypeOf().toHaveProperty("get_weather"); + expectTypeOf().toEqualTypeOf<{ + location: string; + }>(); + expectTypeOf().toHaveProperty("data"); + expectTypeOf().toHaveProperty("result"); + }); + + test("multi-tool agent infers all tool entries with correct input types", () => { + type Map = InferToolMapFromAgent; + + expectTypeOf().toHaveProperty("get_weather"); + expectTypeOf().toHaveProperty("search_web"); + expectTypeOf().toHaveProperty("send_email"); + expectTypeOf().toEqualTypeOf<{ + location: string; + }>(); + expectTypeOf().toExtend<{ + query: string; + maxResults?: number; + }>(); + expectTypeOf().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; + + expectTypeOf().toHaveProperty("streaming_tool"); + expectTypeOf().toExtend< + { progress: number } | undefined + >(); + expectTypeOf().toExtend< + string | undefined + >(); + }); + + test("useStream with agent has toolProgress typed with literal tool names", () => { + const stream = useStream({ + 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({ + 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(); + } + }); +}); diff --git a/libs/sdk/src/types.stream.ts b/libs/sdk/src/types.stream.ts index 74a581c6ad..339606759b 100644 --- a/libs/sdk/src/types.stream.ts +++ b/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; +export type DeriveToolProgress> = { + [K in keyof T & string]: ToolProgress< + T[K]["data"], + T[K]["input"], + T[K]["result"], + K + >; +}[keyof T & string]; + +export type GetToolProgressType = + Bag["ToolMap"] extends Record + ? DeriveToolProgress + : ToolProgress; + type GetStreamModeMap< TStreamMode extends StreamMode | StreamMode[], TStateType = unknown, diff --git a/libs/sdk/src/types.template.ts b/libs/sdk/src/types.template.ts index 068962a7ba..c055c6246d 100644 --- a/libs/sdk/src/types.template.ts +++ b/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 } + >; }; diff --git a/libs/sdk/src/ui/stream/base.ts b/libs/sdk/src/ui/stream/base.ts index 0d661f6d29..7c8eac4a82 100644 --- a/libs/sdk/src/ui/stream/base.ts +++ b/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[]; /** * LangGraph SDK client used to send requests and receive responses. diff --git a/libs/sdk/src/ui/stream/index.ts b/libs/sdk/src/ui/stream/index.ts index b08b451909..9912f1ee8d 100644 --- a/libs/sdk/src/ui/stream/index.ts +++ b/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 extends { "~agentTypes": unknown; } - ? BagTemplate + ? Omit & { ToolMap: InferToolMapFromAgent } : B; diff --git a/libs/sdk/src/ui/types.ts b/libs/sdk/src/ui/types.ts index dfde13047f..7db8105842 100644 --- a/libs/sdk/src/ui/types.ts +++ b/libs/sdk/src/ui/types.ts @@ -1305,3 +1305,44 @@ export type CustomSubmitOptions< SubmitOptions, "optimisticValues" | "context" | "command" | "config" >; + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +/* eslint-disable @typescript-eslint/no-explicit-any */ +type ExtractAsyncGenTypes = T extends AsyncGenerator + ? { data: Y | undefined; result: R | undefined } + : { data: unknown; result: unknown }; + +type ExtractToolStreamTypes = T extends { func: (...args: any[]) => infer R } + ? ExtractAsyncGenTypes< + Extract> + > 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 extends { name: infer N } + ? N extends string + ? IsLiteralString extends true + ? { input: InferToolInput } & ExtractToolStreamTypes + : 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 = + ExtractAgentConfig["Tools"] extends readonly (infer Tool)[] + ? { + [K in Tool extends { name: infer N } + ? N extends string + ? IsLiteralString extends true + ? N + : never + : never + : never]: ToolMapEntryFromTool>; + } + : Record;