feat: add type helpers for DurableAgent UI messages

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<InferDurableAgentTools<WeatherAgent>>().toEqualTypeOf<{
+      getWeather: typeof getWeather;
+    }>();
+  });
+});
+
+describe('InferDurableAgentUIMessage', () => {
+  it('infers the UI message type from a durable agent', () => {
+    expectTypeOf<
+      InferDurableAgentUIMessage<WeatherAgent, { threadId: string }>
+    >().toEqualTypeOf<
+      UIMessage<
+        { threadId: string },
+        never,
+        InferUITools<InferDurableAgentTools<WeatherAgent>>
+      >
+    >();
+  });
+});

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(),

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> =
+  DURABLE_AGENT extends DurableAgent<infer TOOLS> ? TOOLS : never;
+
+/**
+ * Infer the UI message type of a durable agent.
+ */
+export type InferDurableAgentUIMessage<
+  DURABLE_AGENT,
+  MESSAGE_METADATA = unknown,
+> = UIMessage<
+  MESSAGE_METADATA,
+  never,
+  InferUITools<InferDurableAgentTools<DURABLE_AGENT>>
+>;
+
 /**
  * Output specification interface for structured outputs.
  * Use `Output.object({ schema })` or `Output.text()` to create an output specification.