Tools can also return DocumentBlock, ImageBlock, VideoBlock

README.md

@@ -149,6 +149,7 @@ const weatherTool = tool({
   }),
   callback: (input) => {
     // input is fully typed based on the Zod schema
+    // Tools can return strings, JSON values, or media blocks (DocumentBlock, ImageBlock, VideoBlock)
     return `The weather in ${input.location} is 72°F and sunny.`
   },
 })

src/models/bedrock.ts

@@ -558,6 +558,33 @@ export class BedrockModel extends Model<BedrockModelConfig> {
               return { text: content.text }
             case 'jsonBlock':
               return { json: content.json }
+            case 'documentBlock':
+              return {
+                document: {
+                  name: content.name,
+                  format: content.format as BedrockContentBlock.DocumentMember['document']['format'],
+                  source: this._formatDocumentSource(content.source),
+                  ...(content.citations && { citations: content.citations }),
+                  ...(content.context && { context: content.context }),
+                },
+              }
+            case 'imageBlock':
+              return {
+                image: {
+                  format: content.format as BedrockContentBlock.ImageMember['image']['format'],
+                  source: this._formatMediaSource(content.source),
+                },
+              }
+            case 'videoBlock':
+              return {
+                video: {
+                  format: content.format as BedrockContentBlock.VideoMember['video']['format'],
+                  source: this._formatMediaSource(content.source),
+                },
+              }
+            default:
+              // eslint-disable-next-line @typescript-eslint/no-explicit-any
+              return { text: `[Unsupported content type: ${(content as any).type}]` }
           }
         })
 

src/tools/__tests__/zod-tool.test.ts

@@ -98,6 +98,28 @@ describe('tool', () => {
         const result = await myTool.invoke({ count: 3 })
         expect(result).toBe(3)
       })
+
+      it('handles DocumentBlock return', async () => {
+        const { DocumentBlock } = await import('../../types/media.js')
+
+        const docTool = tool({
+          name: 'create_document',
+          description: 'Creates a document',
+          inputSchema: z.object({ content: z.string() }),
+          callback: (input) => {
+            return new DocumentBlock({
+              name: 'RESULT',
+              format: 'md',
+              source: { bytes: new TextEncoder().encode(input.content) },
+            })
+          },
+        })
+
+        const result = await docTool.invoke({ content: 'Hello World!' })
+        expect(result.type).toBe('documentBlock')
+        expect(result.name).toBe('RESULT')
+        expect(result.format).toBe('md')
+      })
     })
 
     describe('validation', () => {

src/tools/function-tool.ts

@@ -5,6 +5,7 @@ import type { ToolSpec } from './types.js'
 import type { JSONSchema, JSONValue } from '../types/json.js'
 import { deepCopy } from '../types/json.js'
 import { JsonBlock, TextBlock, ToolResultBlock } from '../types/messages.js'
+import { DocumentBlock, ImageBlock, VideoBlock } from '../types/media.js'
 
 /**
  * Callback function for FunctionTool implementations.
@@ -218,6 +219,15 @@ export class FunctionTool extends Tool {
    */
   private _wrapInToolResult(value: unknown, toolUseId: string): ToolResultBlock {
     try {
+      // Handle DocumentBlock, ImageBlock, VideoBlock directly
+      if (value instanceof DocumentBlock || value instanceof ImageBlock || value instanceof VideoBlock) {
+        return new ToolResultBlock({
+          toolUseId,
+          status: 'success',
+          content: [value],
+        })
+      }
+
       // Handle null with special string representation as text content
       if (value === null) {
         return new ToolResultBlock({

src/tools/zod-tool.ts

@@ -4,6 +4,13 @@ import type { ToolSpec } from './types.js'
 import type { JSONSchema, JSONValue } from '../types/json.js'
 import { FunctionTool } from './function-tool.js'
 import { z, ZodVoid } from 'zod'
+import { DocumentBlock, ImageBlock, VideoBlock } from '../types/media.js'
+
+/**
+ * Valid return types for tool callbacks.
+ * Includes JSON-serializable values and media blocks.
+ */
+export type ToolReturnValue = JSONValue | DocumentBlock | ImageBlock | VideoBlock
 
 /**
  * Helper type to infer input type from Zod schema or default to never.
@@ -16,7 +23,7 @@ type ZodInferred<TInput> = TInput extends z.ZodType ? z.infer<TInput> : never
  * @typeParam TInput - Zod schema type for input validation
  * @typeParam TReturn - Return type of the callback function
  */
-export interface ToolConfig<TInput extends z.ZodType | undefined, TReturn extends JSONValue = JSONValue> {
+export interface ToolConfig<TInput extends z.ZodType | undefined, TReturn extends ToolReturnValue = JSONValue> {
   /** The name of the tool */
   name: string
 
@@ -46,7 +53,7 @@ export interface ToolConfig<TInput extends z.ZodType | undefined, TReturn extend
  * Internal implementation of a Zod-based tool.
  * Extends Tool abstract class and implements InvokableTool interface.
  */
-class ZodTool<TInput extends z.ZodType | undefined, TReturn extends JSONValue = JSONValue>
+class ZodTool<TInput extends z.ZodType | undefined, TReturn extends ToolReturnValue = JSONValue>
   extends Tool
   implements InvokableTool<ZodInferred<TInput>, TReturn>
 {
@@ -231,7 +238,7 @@ class ZodTool<TInput extends z.ZodType | undefined, TReturn extends JSONValue =
  * @param config - Tool configuration
  * @returns An InvokableTool that implements the Tool interface with invoke() method
  */
-export function tool<TInput extends z.ZodType | undefined, TReturn extends JSONValue = JSONValue>(
+export function tool<TInput extends z.ZodType | undefined, TReturn extends ToolReturnValue = JSONValue>(
   config: ToolConfig<TInput, TReturn>
 ): InvokableTool<ZodInferred<TInput>, TReturn> {
   return new ZodTool(config)

src/types/__tests__/messages.test.ts

@@ -176,6 +176,38 @@ describe('Message.fromMessageData', () => {
     expect(toolResultBlock.content[0]).toBeInstanceOf(JsonBlock)
   })
 
+  it('converts tool result block data to ToolResultBlock with document content', () => {
+    const messageData: MessageData = {
+      role: 'user',
+      content: [
+        {
+          toolResult: {
+            toolUseId: 'tooluse_HUxQGMooooooooooooooeV',
+            status: 'success',
+            content: [
+              {
+                document: {
+                  format: 'md',
+                  name: 'DOCUMENT',
+                  source: {
+                    bytes: Uint8Array.from(
+                      'CiMgUHJvZHVjdCBSZXF1aXJlbWVudHMgRG9jdW1lbnQKCiMjIE92ZXJ2aWV3ClRoaXMgZG9jdW1lbnQgb3V0bGluZXMgdGhlIHJlcXVpcmVtZW50cyBmb3IgYSBuZXcgZmVhdHVyZSBpbiBvdXIgYXBwbGljYXRpb24uCgojIyBSZXF1aXJlbWVudHMKMS4gVXNlciBhdXRoZW50aWNhdGlvbiBtdXN0IHN1cHBvcnQgT0F1dGggMi4wCjIuIEFQSSByZXNwb25zZSB0aW1lIG11c3QgYmUgdW5kZXIgMjAwbXMKMy4gU3lzdGVtIG11c3QgaGFuZGxlIDEwLDAwMCBjb25jdXJyZW50IHVzZXJzCjQuIERhdGEgbXVzdCBiZSBlbmNyeXB0ZWQgYXQgcmVzdCBhbmQgaW4gdHJhbnNpdAoKIyMgVGltZWxpbmUKLSBQaGFzZSAxOiBRMSAyMDI2Ci0gUGhhc2UgMjogUTIgMjAyNgo='
+                    ),
+                  },
+                },
+              },
+            ],
+          },
+        },
+      ],
+    }
+    const message = Message.fromMessageData(messageData)
+    expect(message.content).toHaveLength(1)
+    const toolResultBlock = message.content[0] as ToolResultBlock
+    expect(toolResultBlock.content).toHaveLength(1)
+    expect(toolResultBlock.content[0]).toBeInstanceOf(DocumentBlock)
+  })
+
   it('converts reasoning block data to ReasoningBlock', () => {
     const messageData: MessageData = {
       role: 'assistant',

src/types/messages.ts

@@ -192,9 +192,14 @@ export class ToolUseBlock implements ToolUseBlockData {
  *
  * This is a discriminated union where the object key determines the content format.
  */
-export type ToolResultContentData = TextBlockData | JsonBlockData
+export type ToolResultContentData =
+  | TextBlockData
+  | JsonBlockData
+  | { document: DocumentBlockData }
+  | { image: ImageBlockData }
+  | { video: VideoBlockData }
 
-export type ToolResultContent = TextBlock | JsonBlock
+export type ToolResultContent = TextBlock | JsonBlock | DocumentBlock | ImageBlock | VideoBlock
 
 /**
  * Data for a tool result block.
@@ -226,7 +231,7 @@ export interface ToolResultBlockData {
 /**
  * Tool result content block.
  */
-export class ToolResultBlock implements ToolResultBlockData {
+export class ToolResultBlock {
   /**
    * Discriminator for tool result content.
    */
@@ -596,6 +601,12 @@ export function contentBlockFromData(data: ContentBlockData): ContentBlock {
           return new TextBlock(contentItem.text)
         } else if ('json' in contentItem) {
           return new JsonBlock(contentItem)
+        } else if ('document' in contentItem) {
+          return new DocumentBlock(contentItem.document as DocumentBlockData)
+        } else if ('image' in contentItem) {
+          return new ImageBlock(contentItem.image as ImageBlockData)
+        } else if ('video' in contentItem) {
+          return new VideoBlock(contentItem.video as VideoBlockData)
         } else {
           throw new Error('Unknown ToolResultContentData type')
         }