From 856f785210b59cad05ed9c2430ef18f482660bae Mon Sep 17 00:00:00 2001
From: nicoalbanese <49612682+nicoalbanese@users.noreply.github.com>
Date: Mon, 9 Mar 2026 11:46:40 +0000
Subject: [PATCH] feat: add type helpers for DurableAgent UI messages

---
 .../ai/src/agent/durable-agent-types.test.ts  | 40 +++++++++++++++++++
 packages/ai/src/agent/durable-agent.test.ts   |  8 +++-
 packages/ai/src/agent/durable-agent.ts        | 19 +++++++++
 3 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 packages/ai/src/agent/durable-agent-types.test.ts

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<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>>
+      >
+    >();
+  });
+});
diff --git a/packages/ai/src/agent/durable-agent.test.ts b/packages/ai/src/agent/durable-agent.test.ts
index 343e32a3c8..2052066d63 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 12839d4926..8bfd258ce6 100644
--- a/packages/ai/src/agent/durable-agent.ts
+++ b/packages/ai/src/agent/durable-agent.ts
@@ -9,6 +9,7 @@ import type {
 import {
   asSchema,
   type FinishReason,
+  type InferUITools,
   type LanguageModelResponseMetadata,
   type LanguageModelUsage,
   type ModelMessage,
@@ -35,6 +36,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.