diff --git a/packages/ai/src/agent/durable-agent-types.test.ts b/packages/ai/src/agent/durable-agent-types.test.ts new file mode 100644 index 0000000000..625278e277 --- /dev/null +++ b/packages/ai/src/agent/durable-agent-types.test.ts @@ -0,0 +1,40 @@ +import { type InferUITools, tool, type UIMessage } from 'ai'; +import { describe, expectTypeOf, it } from 'vitest'; +import { z } from 'zod'; +import type { + DurableAgent, + InferDurableAgentTools, + InferDurableAgentUIMessage, +} from './durable-agent.js'; + +const getWeather = tool({ + description: 'Get weather for a location', + inputSchema: z.object({ location: z.string() }), + execute: async ({ location }) => `Weather in ${location}`, +}); + +type WeatherAgent = DurableAgent<{ + getWeather: typeof getWeather; +}>; + +describe('InferDurableAgentTools', () => { + it('infers the tools from a durable agent', () => { + expectTypeOf>().toEqualTypeOf<{ + getWeather: typeof getWeather; + }>(); + }); +}); + +describe('InferDurableAgentUIMessage', () => { + it('infers the UI message type from a durable agent', () => { + expectTypeOf< + InferDurableAgentUIMessage + >().toEqualTypeOf< + UIMessage< + { threadId: string }, + never, + InferUITools> + > + >(); + }); +}); diff --git a/packages/ai/src/agent/durable-agent.test.ts b/packages/ai/src/agent/durable-agent.test.ts index 4adaf4fad7..ac5762549d 100644 --- a/packages/ai/src/agent/durable-agent.test.ts +++ b/packages/ai/src/agent/durable-agent.test.ts @@ -13,9 +13,15 @@ import type { } from '@ai-sdk/provider'; import type { StepResult, ToolSet } from 'ai'; import { describe, expect, it, vi } from 'vitest'; -import { FatalError } from 'workflow'; import { z } from 'zod'; +class FatalError extends Error { + constructor(message: string) { + super(message); + this.name = 'FatalError'; + } +} + // Mock the streamTextIterator vi.mock('./stream-text-iterator.js', () => ({ streamTextIterator: vi.fn(), diff --git a/packages/ai/src/agent/durable-agent.ts b/packages/ai/src/agent/durable-agent.ts index f558444205..b20d8fa1ca 100644 --- a/packages/ai/src/agent/durable-agent.ts +++ b/packages/ai/src/agent/durable-agent.ts @@ -10,6 +10,7 @@ import type { import { asSchema, type FinishReason, + type InferUITools, type LanguageModelResponseMetadata, type LanguageModelUsage, type ModelMessage, @@ -37,6 +38,24 @@ export type { CompatibleLanguageModel } from './types.js'; */ export { Output }; +/** + * Infer the type of the tools of a durable agent. + */ +export type InferDurableAgentTools = + DURABLE_AGENT extends DurableAgent ? TOOLS : never; + +/** + * Infer the UI message type of a durable agent. + */ +export type InferDurableAgentUIMessage< + DURABLE_AGENT, + MESSAGE_METADATA = unknown, +> = UIMessage< + MESSAGE_METADATA, + never, + InferUITools> +>; + /** * Output specification interface for structured outputs. * Use `Output.object({ schema })` or `Output.text()` to create an output specification.