DurableAgent: various compatibility fixes

.changeset/durable-agent-improvements.md

@@ -0,0 +1,5 @@
+---
+'@workflow/ai': patch
+---
+
+Add type helpers (`InferDurableAgentTools`, `InferDurableAgentUIMessage`), support `prepareStep` on `DurableAgent` constructor, fix `supportedUrls` causing `AI_DownloadError` for image URLs, and add telemetry span support for `experimental_telemetry`

packages/ai/package.json

@@ -72,9 +72,15 @@
     "workflow": "workspace:*"
   },
   "peerDependencies": {
+    "@opentelemetry/api": "^1.0.0",
     "ai": "^6",
     "workflow": "workspace:^"
   },
+  "peerDependenciesMeta": {
+    "@opentelemetry/api": {
+      "optional": true
+    }
+  },
   "dependencies": {
     "@ai-sdk/provider": "^3.0.0",
     "@workflow/serde": "workspace:^",

packages/ai/src/agent/do-stream-step.ts

@@ -21,6 +21,7 @@ import type {
   StreamTextTransform,
   TelemetrySettings,
 } from './durable-agent.js';
+import { recordSpan } from './telemetry.js';
 import type { CompatibleLanguageModel } from './types.js';
 
 export type FinishPart = Extract<LanguageModelV3StreamPart, { type: 'finish' }>;
@@ -163,7 +164,39 @@ export async function doStreamStep(
     }),
   };
 
-  const result = await model.doStream(callOptions);
+  const result = await recordSpan({
+    name: 'ai.streamText.doStream',
+    telemetry: options?.experimental_telemetry,
+    attributes: {
+      'ai.model.provider': model.provider,
+      'ai.model.id': model.modelId,
+      // gen_ai semantic convention attributes
+      'gen_ai.system': model.provider,
+      'gen_ai.request.model': model.modelId,
+      ...(options?.maxOutputTokens !== undefined && {
+        'gen_ai.request.max_tokens': options.maxOutputTokens,
+      }),
+      ...(options?.temperature !== undefined && {
+        'gen_ai.request.temperature': options.temperature,
+      }),
+      ...(options?.topP !== undefined && {
+        'gen_ai.request.top_p': options.topP,
+      }),
+      ...(options?.topK !== undefined && {
+        'gen_ai.request.top_k': options.topK,
+      }),
+      ...(options?.frequencyPenalty !== undefined && {
+        'gen_ai.request.frequency_penalty': options.frequencyPenalty,
+      }),
+      ...(options?.presencePenalty !== undefined && {
+        'gen_ai.request.presence_penalty': options.presencePenalty,
+      }),
+      ...(options?.stopSequences !== undefined && {
+        'gen_ai.request.stop_sequences': options.stopSequences,
+      }),
+    },
+    fn: () => model!.doStream(callOptions),
+  });
 
   let finish: FinishPart | undefined;
   const toolCalls: LanguageModelV3ToolCall[] = [];

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(),
@@ -277,6 +283,143 @@ describe('DurableAgent', () => {
       });
     });
 
+    it('should pass through LanguageModelV3ToolResultOutput directly', async () => {
+      // Tool returns a pre-formatted content output (e.g., multimodal with images)
+      const contentOutput = {
+        type: 'content',
+        value: [
+          { type: 'text', text: 'Here is the image' },
+          {
+            type: 'file-data',
+            data: 'base64data',
+            mediaType: 'image/jpeg',
+          },
+        ],
+      };
+      const tools: ToolSet = {
+        visionTool: {
+          description: 'Returns multimodal content',
+          inputSchema: z.object({}),
+          execute: vi.fn().mockResolvedValue(contentOutput),
+        },
+      };
+
+      const mockModel = createMockModel();
+
+      const agent = new DurableAgent({
+        model: async () => mockModel,
+        tools,
+      });
+
+      const mockWritable = new WritableStream({
+        write: vi.fn(),
+        close: vi.fn(),
+      });
+
+      const mockMessages: LanguageModelV3Prompt = [
+        { role: 'user', content: [{ type: 'text', text: 'test' }] },
+      ];
+
+      const { streamTextIterator } = await import('./stream-text-iterator.js');
+      const mockIterator = {
+        next: vi
+          .fn()
+          .mockResolvedValueOnce({
+            done: false,
+            value: {
+              toolCalls: [
+                {
+                  toolCallId: 'vision-call-id',
+                  toolName: 'visionTool',
+                  input: '{}',
+                } as LanguageModelV3ToolCall,
+              ],
+              messages: mockMessages,
+            },
+          })
+          .mockResolvedValueOnce({ done: true, value: [] }),
+      };
+      vi.mocked(streamTextIterator).mockReturnValue(
+        mockIterator as unknown as MockIterator
+      );
+
+      await agent.stream({
+        messages: [{ role: 'user', content: 'test' }],
+        writable: mockWritable,
+      });
+
+      const toolResultsCall = mockIterator.next.mock.calls[1][0];
+      expect(toolResultsCall).toHaveLength(1);
+      expect(toolResultsCall[0]).toMatchObject({
+        type: 'tool-result',
+        toolCallId: 'vision-call-id',
+        toolName: 'visionTool',
+        output: contentOutput, // Passed through as-is, not wrapped in json
+      });
+    });
+
+    it('should pass through pre-formatted text output directly', async () => {
+      // Tool returns an already-formatted text output
+      const textOutput = { type: 'text', value: 'pre-formatted result' };
+      const tools: ToolSet = {
+        textTool: {
+          description: 'Returns pre-formatted text',
+          inputSchema: z.object({}),
+          execute: vi.fn().mockResolvedValue(textOutput),
+        },
+      };
+
+      const mockModel = createMockModel();
+
+      const agent = new DurableAgent({
+        model: async () => mockModel,
+        tools,
+      });
+
+      const mockWritable = new WritableStream({
+        write: vi.fn(),
+        close: vi.fn(),
+      });
+
+      const mockMessages: LanguageModelV3Prompt = [
+        { role: 'user', content: [{ type: 'text', text: 'test' }] },
+      ];
+
+      const { streamTextIterator } = await import('./stream-text-iterator.js');
+      const mockIterator = {
+        next: vi
+          .fn()
+          .mockResolvedValueOnce({
+            done: false,
+            value: {
+              toolCalls: [
+                {
+                  toolCallId: 'text-call-id',
+                  toolName: 'textTool',
+                  input: '{}',
+                } as LanguageModelV3ToolCall,
+              ],
+              messages: mockMessages,
+            },
+          })
+          .mockResolvedValueOnce({ done: true, value: [] }),
+      };
+      vi.mocked(streamTextIterator).mockReturnValue(
+        mockIterator as unknown as MockIterator
+      );
+
+      await agent.stream({
+        messages: [{ role: 'user', content: 'test' }],
+        writable: mockWritable,
+      });
+
+      const toolResultsCall = mockIterator.next.mock.calls[1][0];
+      expect(toolResultsCall[0]).toMatchObject({
+        type: 'tool-result',
+        output: textOutput, // Passed through, not re-wrapped
+      });
+    });
+
     it('should skip local execution for provider-executed tools', async () => {
       // This tool should NOT be called because the tool call is provider-executed
       const executeFn = vi.fn();
@@ -982,6 +1125,80 @@ describe('DurableAgent', () => {
       );
     });
 
+    it('should use prepareStep from the agent definition by default', async () => {
+      const mockModel = createMockModel();
+      const prepareStep: PrepareStepCallback = vi.fn().mockReturnValue({});
+
+      const agent = new DurableAgent({
+        model: async () => mockModel,
+        tools: {},
+        prepareStep,
+      });
+
+      const mockWritable = new WritableStream({
+        write: vi.fn(),
+        close: vi.fn(),
+      });
+
+      const { streamTextIterator } = await import('./stream-text-iterator.js');
+      const mockIterator = {
+        next: vi.fn().mockResolvedValueOnce({ done: true, value: [] }),
+      };
+      vi.mocked(streamTextIterator).mockReturnValue(
+        mockIterator as unknown as MockIterator
+      );
+
+      await agent.stream({
+        messages: [{ role: 'user', content: 'test' }],
+        writable: mockWritable,
+      });
+
+      expect(streamTextIterator).toHaveBeenCalledWith(
+        expect.objectContaining({
+          prepareStep,
+        })
+      );
+    });
+
+    it('should prefer a stream prepareStep over the agent definition', async () => {
+      const mockModel = createMockModel();
+      const agentPrepareStep: PrepareStepCallback = vi.fn().mockReturnValue({});
+      const streamPrepareStep: PrepareStepCallback = vi
+        .fn()
+        .mockReturnValue({});
+
+      const agent = new DurableAgent({
+        model: async () => mockModel,
+        tools: {},
+        prepareStep: agentPrepareStep,
+      });
+
+      const mockWritable = new WritableStream({
+        write: vi.fn(),
+        close: vi.fn(),
+      });
+
+      const { streamTextIterator } = await import('./stream-text-iterator.js');
+      const mockIterator = {
+        next: vi.fn().mockResolvedValueOnce({ done: true, value: [] }),
+      };
+      vi.mocked(streamTextIterator).mockReturnValue(
+        mockIterator as unknown as MockIterator
+      );
+
+      await agent.stream({
+        messages: [{ role: 'user', content: 'test' }],
+        writable: mockWritable,
+        prepareStep: streamPrepareStep,
+      });
+
+      expect(streamTextIterator).toHaveBeenCalledWith(
+        expect.objectContaining({
+          prepareStep: streamPrepareStep,
+        })
+      );
+    });
+
     it('should allow prepareStep to modify messages', async () => {
       const mockModel = createMockModel();