From 53d6a96c306c77421df00b1cbe4db02d1078cc55 Mon Sep 17 00:00:00 2001
From: kai-agent-free <kai@kdn.agency>
Date: Fri, 3 Apr 2026 12:43:54 +0000
Subject: [PATCH] fix: handle type-discriminator format in contentBlockFromData
 (#533)

When content blocks are stored by ORMs or spread ({...block}) instead
of serialized via toJSON(), the data has a flat shape with a type
discriminator (e.g. {type: 'toolUseBlock', name, toolUseId, input})
instead of the nested wrapper ({toolUse: {name, toolUseId, input}}).

This adds fallback handling in contentBlockFromData for ToolUseBlock,
TextBlock, and ToolResultBlock type discriminators, fixing the
'Unknown ContentBlockData type' error when recreating agents from
stored session data.

Fixes #533
---
 src/types/__tests__/messages.test.ts | 46 ++++++++++++++++++++++++++++
 src/types/messages.ts                | 20 ++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/src/types/__tests__/messages.test.ts b/src/types/__tests__/messages.test.ts
index 55c885556..ac9afa3e0 100644
--- a/src/types/__tests__/messages.test.ts
+++ b/src/types/__tests__/messages.test.ts
@@ -330,6 +330,52 @@ describe('Message.fromMessageData', () => {
     } as unknown as MessageData
     expect(() => Message.fromMessageData(messageData)).toThrow('Unknown ContentBlockData type')
   })
+  it('handles ToolUseBlock data with type discriminator instead of nested wrapper (issue #533)', () => {
+    // When content blocks are spread ({...block}) or stored by ORMs that flatten
+    // the nested structure, the data has {type: 'toolUseBlock', name, toolUseId, input}
+    // instead of {toolUse: {name, toolUseId, input}}
+    const messageData = {
+      role: 'assistant' as const,
+      content: [
+        { type: 'toolUseBlock', name: 'greet', toolUseId: 'tu_123', input: { name: 'Alice' } },
+      ],
+    } as unknown as MessageData
+    const message = Message.fromMessageData(messageData)
+    expect(message.content).toHaveLength(1)
+    expect(message.content[0].type).toBe('toolUseBlock')
+    const block = message.content[0] as ToolUseBlock
+    expect(block.name).toBe('greet')
+    expect(block.toolUseId).toBe('tu_123')
+    expect(block.input).toEqual({ name: 'Alice' })
+  })
+
+  it('handles TextBlock data with type discriminator instead of nested wrapper', () => {
+    const messageData = {
+      role: 'user' as const,
+      content: [
+        { type: 'textBlock', text: 'hello' },
+      ],
+    } as unknown as MessageData
+    const message = Message.fromMessageData(messageData)
+    expect(message.content).toHaveLength(1)
+    expect(message.content[0].type).toBe('textBlock')
+    expect((message.content[0] as TextBlock).text).toBe('hello')
+  })
+
+  it('handles ToolResultBlock data with type discriminator instead of nested wrapper', () => {
+    const messageData = {
+      role: 'user' as const,
+      content: [
+        { type: 'toolResultBlock', toolUseId: 'tu_123', content: [], status: 'success' },
+      ],
+    } as unknown as MessageData
+    const message = Message.fromMessageData(messageData)
+    expect(message.content).toHaveLength(1)
+    expect(message.content[0].type).toBe('toolResultBlock')
+    const block = message.content[0] as ToolResultBlock
+    expect(block.toolUseId).toBe('tu_123')
+  })
+
 })
 
 describe('systemPromptFromData', () => {
diff --git a/src/types/messages.ts b/src/types/messages.ts
index 33b7ad11b..e84549bc3 100644
--- a/src/types/messages.ts
+++ b/src/types/messages.ts
@@ -883,6 +883,26 @@ export function contentBlockFromData(data: ContentBlockData): ContentBlock {
   } else if ('citations' in data) {
     return CitationsBlock.fromJSON(data)
   } else {
+    // Fallback: handle data with a type discriminator instead of nested wrappers.
+    // This occurs when content blocks are spread ({...block}) or stored by ORMs
+    // that flatten the nested structure produced by toJSON().
+    const typed = data as Record<string, unknown>
+    if (typed.type === 'toolUseBlock' && typeof typed.name === 'string' && typeof typed.toolUseId === 'string') {
+      return new ToolUseBlock({
+        name: typed.name,
+        toolUseId: typed.toolUseId,
+        input: typed.input as JSONValue,
+        ...(typed.reasoningSignature !== undefined && { reasoningSignature: typed.reasoningSignature as string }),
+      })
+    } else if (typed.type === 'textBlock' && typeof typed.text === 'string') {
+      return new TextBlock(typed.text)
+    } else if (typed.type === 'toolResultBlock' && typeof typed.toolUseId === 'string') {
+      return new ToolResultBlock({
+        toolUseId: typed.toolUseId,
+        content: Array.isArray(typed.content) ? (typed.content as ToolResultContent[]) : [],
+        status: (typed.status as 'success' | 'error') ?? 'success',
+      })
+    }
     throw new Error('Unknown ContentBlockData type')
   }
 }