diff --git a/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.int.test.ts b/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.int.test.ts
index 8e2dea4a8..9d1deb35e 100644
--- a/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.int.test.ts
+++ b/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.int.test.ts
@@ -1,10 +1,10 @@
-import type { CoreMessage } from 'ai';
+import type { ModelMessage } from 'ai';
 import { describe, expect, it } from 'vitest';
 import { formatFollowUpMessageStepExecution } from './format-follow-up-message-step';
 
 describe('format-follow-up-message-step integration', () => {
   it('should format follow-up message with new issues and assumptions', async () => {
-    const mockConversationHistory: CoreMessage[] = [
+    const mockConversationHistory: ModelMessage[] = [
       {
         content:
           'How many stock Mountain-500 series bikes were sold online to NA customers using a ColonialVoice card?',
@@ -22,7 +22,7 @@ describe('format-follow-up-message-step integration', () => {
       {
         content: [
           {
-            args: {
+            input: {
               final_response:
                 "I found the answer to your specific question about Mountain-500 series bikes sold online to North American customers using ColonialVoice cards.\n\n**Result: 47 Mountain-500 series bikes were sold under these exact conditions.**\n\nHere's how I identified and filtered the data based on your requirements:\n\n- **Mountain-500 series bikes**: Found 10 different Mountain-500 variants in the database (5 black and 5 silver in various sizes: 40, 42, 44, 48, 52)\n- **Stock bikes**: Filtered for finished goods ready for sale (using the `finishedgoodsflag`)\n- **Online sales**: Used the online order flag to identify web-based purchases\n- **NA customers**: Included customers from United States, Canada, and Mexico\n- **ColonialVoice card**: Filtered specifically for this credit card type (which exists in the database alongside Vista, SuperiorCard, and Distinguish)\n\nThe analysis shows that despite the very specific filtering criteria, there was still a meaningful volume of 47 bikes sold through this particular channel combination. This suggests that ColonialVoice cardholders in North America represent an active customer segment for the Mountain-500 product line in online sales.",
             },
@@ -36,9 +36,7 @@ describe('format-follow-up-message-step integration', () => {
       {
         content: [
           {
-            result: {
-              success: true,
-            },
+            output: { success: true } as any,
             toolCallId: 'toolu_06WAfvCoQtpBoNdmNi17LKCe',
             toolName: 'doneTool',
             type: 'tool-result',
diff --git a/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.test.ts b/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.test.ts
index b72151f87..e07294a90 100644
--- a/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.test.ts
+++ b/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.test.ts
@@ -1,4 +1,4 @@
-import type { CoreMessage } from 'ai';
+import type { ModelMessage } from 'ai';
 import { describe, expect, test, vi } from 'vitest';
 import { formatFollowUpMessageStepExecution } from './format-follow-up-message-step';
 
@@ -38,7 +38,7 @@ vi.mock('../../../src/tools/post-processing/generate-update-message', () => ({
 
 describe('Format Follow-up Message Step Unit Tests', () => {
   test('should include chat history in context message when present', async () => {
-    const mockConversationHistory: CoreMessage[] = [
+    const mockConversationHistory: ModelMessage[] = [
       { role: 'user', content: 'Initial query about sales' },
       { role: 'assistant', content: 'Sales data analysis complete' },
       { role: 'user', content: 'Can you filter by last 6 months?' },
@@ -179,7 +179,7 @@ describe('Format Follow-up Message Step Unit Tests', () => {
   });
 
   test('should handle multiple major assumptions with conversation history', async () => {
-    const mockConversationHistory: CoreMessage[] = [
+    const mockConversationHistory: ModelMessage[] = [
       { role: 'user', content: 'Show me customer segments' },
       { role: 'assistant', content: 'Here are the segments' },
       { role: 'user', content: 'Filter by enterprise only' },
diff --git a/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.ts b/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.ts
index 1bafdb60a..71ebd5b3d 100644
--- a/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.ts
+++ b/packages/ai/src/steps/message-post-processing-steps/format-follow-up-message-step/format-follow-up-message-step.ts
@@ -4,11 +4,8 @@ import { wrapTraced } from 'braintrust';
 import { z } from 'zod';
 import { Sonnet4 } from '../../../llm/sonnet-4';
 import { postProcessingWorkflowOutputSchema } from '../schemas';
-// Import the schema from combine-parallel-results step
-import { combineParallelResultsOutputSchema } from './combine-parallel-results-step';
 
-// Input schema matches the output of combine-parallel-results step
-const inputSchema = combineParallelResultsOutputSchema;
+const inputSchema = postProcessingWorkflowOutputSchema;
 
 // Use the unified schema from the workflow
 export const formatFollowUpMessageOutputSchema = postProcessingWorkflowOutputSchema;
@@ -128,7 +125,7 @@ Generate a concise update message for the data team.`;
     schema: generateUpdateMessageOutputSchema,
     messages: systemAndUserMessages,
     temperature: 0,
-    maxTokens: 10000,
+    maxOutputTokens: 10000,
     providerOptions: {
       anthropic: {
         disableParallelToolCalls: true,